Some more cleanups, especially tests

tests/test_integrator.py

@@ -0,0 +1,87 @@
+import matplotlib.pyplot as plt
+import numpy as np
+import pytest
+
+import scgenerator as sc
+import scgenerator.operators as op
+
+
+def test_rk43_absorbtion_only():
+    n = 129
+    w_c = np.linspace(-5, 5, n)
+    spec0 = np.exp(-(w_c**2))
+
+    lin = op.envelope_linear_operator(
+        op.constant_quantity(np.zeros(n)),
+        op.constant_quantity(np.ones(n) * np.log(2)),
+    )
+    non_lin = op.no_op_freq(n)
+
+    res = sc.integrate(spec0, 1.0, lin, non_lin, targets=[1.0])
+    assert np.max(sc.abs2(res.spectra[-1])) == pytest.approx(0.5)
+
+
+def test_rk43_soliton(plot=False):
+    """
+    create a N=3 soliton and test that the spectrum at after one oscillation goes back to the same
+    maximum value
+    """
+    n = 1024
+    l0 = 835e-9
+    w0 = sc.units.m(l0)
+    b2 = sc.fiber.D_to_beta2(sc.units.D_ps_nm_km(24), l0)
+    gamma = 0.08
+    t0_fwhm = 50e-15
+    p0 = 1.26e3
+    t0 = sc.pulse.width_to_t0(t0_fwhm, "sech")
+    soliton_num = 3
+    p0 = soliton_num**2 * np.abs(b2) / (gamma * t0**2)
+
+    disp_len = t0**2 / np.abs(b2)
+    end = disp_len * 0.5 * np.pi
+    targets = np.linspace(0, end, 64)
+
+    t = np.linspace(-200e-15, 200e-15, n)
+    w_c = np.pi * 2 * np.fft.fftfreq(n, t[1] - t[0])
+    field0 = sc.pulse.sech_pulse(t, t0, p0)
+    spec0 = np.fft.fft(field0)
+    no_op = op.no_op_freq(n)
+
+    lin = op.envelope_linear_operator(
+        op.constant_polynomial_dispersion([b2], w_c),
+        op.constant_quantity(np.zeros(n)),
+    )
+    non_lin = op.envelope_nonlinear_operator(
+        op.constant_quantity(np.ones(n) * gamma),
+        op.constant_quantity(np.zeros(n)),
+        op.envelope_spm(0),
+        no_op,
+    )
+
+    res = sc.integrate(spec0, end, lin, non_lin, targets=targets, atol=1e-10, rtol=1e-9)
+    if plot:
+        x, y, z = sc.plotting.transform_2D_propagation(
+            res.spectra,
+            sc.PlotRange(500, 1300, "nm"),
+            w_c + w0,
+            targets,
+        )
+        plt.imshow(z, extent=sc.plotting.get_extent(x, y), origin="lower", aspect="auto", vmin=-40)
+        plt.show()
+
+        plt.plot(sc.abs2(spec0))
+        plt.plot(sc.abs2(res.spectra[-1]))
+        plt.yscale("log")
+        plt.show()
+
+    assert sc.abs2(spec0).max() == pytest.approx(sc.abs2(res.spectra[-1]).max(), rel=0.01)
+
+
+def benchmark():
+    for _ in range(50):
+        test_rk43_soliton()
+
+
+if __name__ == "__main__":
+    test_rk43_soliton()
+    benchmark()