working on simpler integrating sequence

examples/test_integrator.py

@@ -0,0 +1,151 @@
+from collections import defaultdict
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pytest
+from scipy.interpolate import interp1d
+
+import scgenerator as sc
+import scgenerator.operators as op
+import scgenerator.solver as so
+
+
+def test_rk43_absorbtion_only():
+    n = 129
+    end = 1.0
+    h = 2**-3
+    w_c = np.linspace(-5, 5, n)
+    ind = np.argsort(w_c)
+    spec0 = np.exp(-(w_c**2))
+    init_state = op.SimulationState(spec0, end, h)
+
+    lin = op.envelope_linear_operator(
+        op.no_op_freq(n),
+        op.constant_array_operator(np.ones(n) * np.log(2)),
+    )
+    non_lin = op.no_op_freq(n)
+
+    judge = so.adaptive_judge(1e-6, 4)
+    stepper = so.ERK43(lin, non_lin)
+
+    for state in so.integrate(stepper, init_state, h, judge, max_step_size=0.125):
+        if state.z >= end:
+            break
+    assert np.max(state.spectrum2) == pytest.approx(0.5)
+
+
+def test_rk43_soliton(plot=False):
+    n = 1024
+    b2 = sc.fiber.D_to_beta2(sc.units.D_ps_nm_km(24), 835e-9)
+    gamma = 0.08
+    t0_fwhm = 50e-15
+    p0 = 1.26e3
+    t0 = sc.pulse.width_to_t0(t0_fwhm, "sech")
+    print(np.sqrt(t0**2 / np.abs(b2) * gamma * p0))
+
+    disp_len = t0**2 / np.abs(b2)
+    end = disp_len * 0.5 * np.pi
+    targets = np.linspace(0, end, 32)
+    print(end)
+
+    h = 2**-6
+    t = np.linspace(-200e-15, 200e-15, n)
+    w_c = np.pi * 2 * np.fft.fftfreq(n, t[1] - t[0])
+    ind = np.argsort(w_c)
+    field0 = sc.pulse.sech_pulse(t, t0, p0)
+    init_state = op.SimulationState(np.fft.fft(field0), end, h)
+    no_op = op.no_op_freq(n)
+
+    lin = op.envelope_linear_operator(
+        op.constant_polynomial_dispersion([b2], w_c),
+        no_op,
+        # op.constant_array_operator(np.ones(n) * np.log(2)),
+    )
+    non_lin = op.envelope_nonlinear_operator(
+        op.constant_array_operator(np.ones(n) * gamma),
+        no_op,
+        op.envelope_spm(0),
+        no_op,
+    )
+
+    # new_state = init_state.copy()
+    # plt.plot(t, init_state.field2)
+    # new_state.set_spectrum(non_lin(init_state))
+    # plt.plot(t, new_state.field2)
+    # new_state.set_spectrum(lin(init_state))
+    # plt.plot(t, new_state.field2)
+    # print(new_state.spectrum2.max())
+    # plt.show()
+    # return
+
+    judge = so.adaptive_judge(1e-6, 4)
+    stepper = so.ERKIP43Stepper(lin, non_lin)
+
+    # stepper.set_state(init_state)
+    # state, error = stepper.take_step(init_state, 1e-3, True)
+    # print(error)
+    # plt.plot(t, stepper.fine.field2)
+    # plt.plot(t, stepper.coarse.field2)
+    # plt.show()
+    # return
+
+    target = 0
+    stats = defaultdict(list)
+    saved = []
+    zs = []
+    for state in so.integrate(stepper, init_state, h, judge, max_step_size=0.125):
+        # print(f"z = {state.z*100:.2f}")
+        saved.append(state.spectrum2[ind])
+        zs.append(state.z)
+        for k, v in state.stats.items():
+            stats[k].append(v)
+        if state.z > end:
+            break
+    print(len(saved))
+    if plot:
+        interp = interp1d(zs, saved, axis=0)
+        plt.imshow(sc.units.to_log(interp(targets)), origin="lower", aspect="auto", vmin=-40)
+        plt.show()
+
+        plt.plot(stats["z"][1:], np.diff(stats["z"]))
+        plt.show()
+
+
+def test_simple_euler():
+    n = 129
+    end = 1.0
+    h = 2**-3
+    w_c = np.linspace(-5, 5, n)
+    ind = np.argsort(w_c)
+    spec0 = np.exp(-(w_c**2))
+    init_state = op.SimulationState(spec0, end, h)
+
+    lin = op.envelope_linear_operator(
+        op.no_op_freq(n),
+        op.constant_array_operator(np.ones(n) * np.log(2)),
+    )
+    euler = so.ConstantEuler(lin)
+
+    target = 0
+    end = 1.0
+    h = 2**-6
+    for state in so.integrate(euler, init_state, h):
+        if state.z >= target:
+            target += 0.125
+            plt.plot(w_c[ind], state.spectrum2[ind], label=f"z={state.z:.3f}")
+        if target > end:
+            print(np.max(state.spectrum2))
+            break
+    plt.title(f"{h = }")
+    plt.legend()
+    plt.show()
+
+
+def benchmark():
+    for _ in range(50):
+        test_rk43_soliton()
+
+
+if __name__ == "__main__":
+    test_rk43_soliton()
+    benchmark()