From 3f3913526d0d47ab08b6d0dfc85a796f9f80d2b7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Sierro?= <benoit.sierro@bluewin.ch>
Date: Wed, 4 Oct 2023 10:45:27 +0200
Subject: [PATCH] new center of gravity and energy tests

---
 src/scgenerator/physics/pulse.py |  4 +--
 src/scgenerator/spectra.py       |  4 +++
 tests/test_spectra.py            | 60 ++++++++++++++++++++++++++++++++
 3 files changed, 66 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_spectra.py

diff --git a/src/scgenerator/physics/pulse.py b/src/scgenerator/physics/pulse.py
index 482e078..6b4136e 100644
--- a/src/scgenerator/physics/pulse.py
+++ b/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(
diff --git a/src/scgenerator/spectra.py b/src/scgenerator/spectra.py
index 88570ae..4abce9d 100644
--- a/src/scgenerator/spectra.py
+++ b/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`.
diff --git a/tests/test_spectra.py b/tests/test_spectra.py
new file mode 100644
index 0000000..40aa6db
--- /dev/null
+++ b/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)