working on simpler integrating sequence

examples/chang2011.py

@@ -0,0 +1,50 @@
+from collections import defaultdict
+
+import matplotlib.pyplot as plt
+import numpy as np
+from scipy.interpolate import interp1d
+from tqdm import tqdm
+
+import scgenerator as sc
+from scgenerator import solver as so
+
+PARAMS = dict(
+    wavelength=800e-9,
+    width=30e-15,
+    energy=2.5e-6,
+    core_radius=10e-6,
+    length=10e-2,
+    gas_name="argon",
+    pressure=4e5,
+    t_num=4096,
+    dt=0.1e-15,
+    photoionization=True,
+    full_field=True,
+    model="marcatili",
+)
+
+
+params = sc.Parameters(**PARAMS)
+init_state = so.SimulationState(params.spec_0, params.length, 5e-6, converter=params.ifft)
+stepper = so.ERKIP43Stepper(params.linear_operator, params.nonlinear_operator)
+solution = []
+stats = defaultdict(list)
+for state in so.integrate(stepper, init_state, step_judge=so.adaptive_judge(1e-6, 4)):
+    solution.append(state.spectrum2)
+    for k, v in state.stats.items():
+        stats[k].append(v)
+    if state.z > params.length:
+        break
+quit()
+interp = interp1d(stats["z"], solution, axis=0)
+z = np.linspace(0, params.length, 128)
+plt.imshow(
+    sc.units.to_log(interp(z)),
+    vmin=-50,
+    extent=sc.get_extent(sc.units.THz_inv(params.w), z),
+    origin="lower",
+    aspect="auto",
+)
+plt.figure()
+plt.plot(stats["z"][1:], stats["electron_density"])
+plt.show()

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()

examples/test_rate.py

@@ -0,0 +1,19 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
+import scgenerator as sc
+
+
+def field(t: np.ndarray, fwhm=10e-15) -> np.ndarray:
+    t0 = sc.pulse.width_to_t0(fwhm, "gaussian")
+    return sc.pulse.initial_full_field(t, "gaussian", 1e-4, t0, 2e14, sc.units.nm(800), 1)
+
+
+rate = sc.plasma.create_ion_rate_func(sc.materials.Gas("argon").ionization_energy)
+
+fig, (top, mid, bot) = plt.subplots(3, 1)
+t = np.linspace(-10e-15, 10e-15, 1024)
+E = field(t)
+top.plot(t * 1e15, field(t))
+mid.plot(t * 1e15, rate(np.abs(E)))
+plt.show()