new center of gravity and energy tests

src/scgenerator/physics/pulse.py

@@ -387,8 +387,8 @@ def soliton_length(dispersion_length: float):
     return pi / 2 * dispersion_length
 
 
-def center_of_gravitiy(t: np.ndarray, intensity: np.ndarray):
-    return np.sum(intensity * t) / np.sum(intensity)
+def center_of_gravity(t: np.ndarray, intensity: np.ndarray):
+    return np.sum(intensity * t, axis=-1) / np.sum(intensity, axis=-1)
 
 
 def adjust_custom_field(

src/scgenerator/spectra.py

@@ -136,6 +136,10 @@ class Spectrum(np.ndarray):
     def center_of_gravity(self):
         return pulse.center_of_gravity(self.t, self.time_int)
 
+    @property
+    def energy(self) -> np.ndarray:
+        return np.trapz(self.time_int, x=self.t, axis=-1)
+
     def mask_wl(self, pos: float, width: float) -> Spectrum:
         """
         Filters the spectrum with a bandpass centered at `pos` of FWHM `width`.

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)