new center of gravity and energy tests

tests/test_spectra.py

@@ -0,0 +1,60 @@
+import numpy as np
+import pytest
+
+from scgenerator.math import abs2, wspace
+from scgenerator.physics.units import m_rads
+from scgenerator.spectra import Spectrum
+
+
+def test_center_gravity():
+    t = np.linspace(-1e-12, 1e-12, 1024)
+    w = wspace(t) + m_rads(800e-9)
+
+    spec = np.fft.fft(np.exp(-((t / 1e-13) ** 2)))
+    spec = Spectrum(spec, w, t)
+    assert spec.center_of_gravity == pytest.approx(0)
+
+    spec = np.fft.fft(
+        np.array([np.exp(-((t / 1e-13) ** 2)), np.exp(-(((t - 1e-15) / 1e-13) ** 2))])
+    )
+    spec = Spectrum(spec, w, t)
+    assert spec.center_of_gravity * 1e15 == pytest.approx((0, 1))
+
+    spec = np.fft.fft(
+        np.array(
+            [
+                [
+                    np.exp(-((t / 1e-13) ** 2)),
+                    np.exp(-(((t - 1e-15) / 1e-13) ** 2)),
+                    np.exp(-(((t + 2e-15) / 1e-13) ** 2)),
+                ],
+                [
+                    np.exp(-((t / 1e-13) ** 2)),
+                    np.exp(-(((t + 1e-15) / 1e-13) ** 2)),
+                    np.exp(-(((t - 2e-15) / 1e-13) ** 2)),
+                ],
+            ]
+        )
+    )
+    spec = Spectrum(spec, w, t)
+    assert spec.center_of_gravity[0] * 1e15 == pytest.approx((0, 1, -2))
+    assert spec.center_of_gravity[1] * 1e15 == pytest.approx((0, -1, 2))
+
+
+def test_energy():
+    t = np.linspace(-1e-12, 1e-12, 1024)
+    w = wspace(t) + m_rads(800e-9)
+
+    spec = np.random.rand(1024)
+    spec = Spectrum(spec, w, t)
+    assert spec.energy == np.trapz(abs2(np.fft.ifft(spec)), x=t)
+
+    spec = np.random.rand(2, 1024)
+    spec = Spectrum(spec, w, t)
+    assert np.all(spec.energy == np.trapz(abs2(np.fft.ifft(spec)), x=t, axis=-1))
+    assert spec.energy.shape == (2,)
+
+    spec = np.random.rand(3, 2, 1024)
+    spec = Spectrum(spec, w, t)
+    assert np.all(spec.energy == np.trapz(abs2(np.fft.ifft(spec)), x=t, axis=-1))
+    assert spec.energy.shape == (3, 2)