Dudley2006 matched

.gitignore

@@ -2,6 +2,7 @@
 .idea
 .conda-env
 /play.py
+/examples/*.zip
 
 pyrightconfig.json
 

examples/dudley2006.py

@@ -0,0 +1,135 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from plotapp import PlotApp
+
+import scgenerator as sc
+
+params = sc.Parameters(
+    wavelength=835e-9,
+    width=50e-15,
+    peak_power=10e3,
+    repetition_rate=20e6,
+    shape="sech",
+    # fiber
+    length=15e-2,
+    raman_type="measured",
+    beta2_coefficients=[
+        -11.830e-27,
+        8.1038e-41,
+        -9.5205e-56,
+        2.0737e-70,
+        -5.3943e-85,
+        1.3486e-99,
+        -2.5495e-114,
+        3.0524e-129,
+        -1.7140e-144,
+    ],
+    gamma=0.11,
+    # simulation
+    tolerated_error=1e-8,
+    wavelength_window=[300e-9, 1500e-9],
+    t_num=8192,
+    quantum_noise=False,
+    z_num=128,
+)
+
+
+def compute_manual():
+    spec0 = params.compute("spec_0")
+    w_c, w0, gamma = params.compute("w_c", "w0", "gamma")
+    p = params.compile()
+    print(p.dt)
+    beta_op = sc.operators.constant_polynomial_dispersion(
+        params.beta2_coefficients, w_c, params.compute("dispersion_ind")
+    )
+    linear = sc.operators.envelope_linear_operator(
+        beta_op,
+        # sc.operators.constant_quantity(0),
+        sc.operators.constant_quantity(0),
+    )
+    # with PlotApp() as app:
+    #     o = np.argsort(p.l)
+    #     sax = app["spectrum"]
+    #     sax.set_line_data("beta2", p.l[o], beta_op(0)[o].imag)
+    # return
+
+    # spm = sc.operators.envelope_spm(0)
+    # plt.plot(p.l, sc.abs2(p.spec_0))
+    # plt.plot(p.l, sc.abs2(np.fft.fft(spm(np.fft.ifft(p.spec_0), 0))), ls=":")
+    # plt.xlim(100e-9, 1500e-9)
+    # plt.show()
+    # return
+
+    # nonlinear = sc.operators.envelope_nonlinear_operator(
+    #     gamma_op=sc.operators.constant_quantity(params.gamma),
+    #     ss_op=sc.operators.constant_quantity(w_c / w0),
+    #     spm_op=sc.operators.envelope_spm(0),
+    #     raman_op=sc.operators.no_op_time(params.t_num),
+    # )
+
+    hr_w = params.compute("hr_w")
+
+    def nonlinear(spec, z):
+        field = np.fft.ifft(spec)
+        field2 = sc.abs2(field)
+        fr = 0.18
+        return (
+            -1j
+            * gamma
+            * (1 + w_c / w0)
+            * np.fft.fft(field * ((1 - fr) * field2 + fr * np.fft.ifft(hr_w * np.fft.fft(field2))))
+        )
+
+    above0 = p.l > 0
+    linear_arr = np.zeros(p.t_num, dtype=complex)
+    w_power_fact = np.array(
+        [sc.math.power_fact(w_c[above0], k) for k in range(2, len(p.beta2_coefficients) + 2)]
+    )
+    for i, wn in reversed(list(enumerate(w_power_fact))):
+        print(i, p.beta2_coefficients[i])
+        linear_arr[above0] += p.beta2_coefficients[i] * sc.math.power_fact(w_c[above0], i + 2)
+    linear_arr *= -1j
+
+    def linear(_):
+        return linear_arr
+
+    prop = sc.propagation("examples/dudley_manual.zip", params)
+    z = []
+    for i, (spec, stat) in enumerate(
+        sc.solve43(spec0, linear, nonlinear, params.length, 1e-6, 1e-6, 0.9, h_const=20e-6)
+    ):
+        if i % 20:
+            continue
+        print(stat, end="\r")
+        prop.append(spec)
+        z.append(stat["z"])
+
+    # with PlotApp(i=range(len(z))) as app:
+    #     o = np.argsort(p.l)
+    #     sax = app["spectrum"]
+    #     tax = app["field"]
+
+    #     @app.update
+    #     def draw(i):
+    #         sax.set_line_data("spectrum", p.l[o], sc.abs2(prop[i][o]))
+    #         tax.set_line_data("field", p.t, sc.abs2(np.fft.ifft(prop[i])))
+
+
+def compute_auto():
+    sc.compute(params, True, "examples/dudley2006")
+
+
+def plot():
+    prop = sc.propagation("examples/dudley_manual.zip")
+    newprop = sc.propagation("examples/dudley2006.zip")
+    wl_r = sc.PlotRange(450, 2500, "nm")
+    fig, (top, bot) = plt.subplots(2, 2)
+    z = np.linspace(0, prop.parameters.length, len(prop))
+    sc.plotting.summary_plot(prop[:], z, axes=top, wl_range=wl_r)
+    sc.plotting.summary_plot(newprop[:], newprop.parameters.z_targets, axes=bot, wl_range=wl_r)
+    plt.show()
+
+
+if __name__ == "__main__":
+    # compute_auto()
+    plot()

src/scgenerator/helpers.py

@@ -1,6 +1,7 @@
 """
 series of helper functions
 """
+import os
 import warnings
 from pathlib import Path
 
@@ -162,8 +163,14 @@ def extend_axis(axis: np.ndarray) -> np.ndarray:
     return axis
 
 
-def compute(parameters: Parameters, overwrite: bool = False) -> Propagation:
+def compute(
+    parameters: Parameters, overwrite: bool = False, output: os.PathLike | None = None
+) -> Propagation:
+    if output is None:
         name = Path(parameters.compute("name")).stem + ".zip"
+    else:
+        path = Path(output)
+        name = path.parent / (path.stem + ".zip")
 
     prop_params = parameters.compile()
     prop = propagation(name, prop_params, bundle_data=True, overwrite=overwrite)
@@ -183,8 +190,6 @@ def compute(parameters: Parameters, overwrite: bool = False) -> Propagation:
             )
         ):
             pbar.update()
-            plt.plot(prop_params.t, abs2(prop_params.ifft(spec)))
-            plt.show()
             prop.append(spec)
 
     return prop

src/scgenerator/operators.py

@@ -194,6 +194,7 @@ def constant_polynomial_dispersion(
     w_power_fact = np.array(
         [math.power_fact(w_c, k) for k in range(2, len(beta2_coefficients) + 2)]
     )
+    disp_arr = np.zeros(len(w_c), dtype=complex)
     disp_arr = fiber.fast_poly_dispersion_op(w_c, beta2_coefficients, w_power_fact)
 
     return constant_quantity(disp_arr)

src/scgenerator/plotting.py

@@ -1156,9 +1156,11 @@ def summary_plot(
     specs: Spectrum,
     z: Sequence[float] | None = None,
     wl_range: PlotRange | None = None,
-    t_range: PlotRange | None = None,
+    time_range: PlotRange | None = None,
     db_min: float = -50.0,
     axes: tuple[Axes, Axes] | None = None,
+    wl_db="1D",
+    time_db=False,
 ):
     wl_int = specs.wl_int
     time_int = specs.time_int
@@ -1166,23 +1168,24 @@ def summary_plot(
     if z is None:
         z = np.arange(specs.shape[0])
     elif len(z) > specs.shape[0]:
-        z = z[:specs.shape[0]]
+        z = z[: specs.shape[0]]
+    z = np.asarray(z)
 
     if wl_range is None:
         imin, imax = math.span_above(wl_int, wl_int.max() * 1e-6)
         wl_range = PlotRange(specs.wl_disp[imin] * 1e9, specs.wl_disp[imax] * 1e9, "nm")
 
-    if t_range is None:
+    if time_range is None:
         imin, imax = math.span_above(time_int, time_int.max() * 1e-6)
-        t_range = PlotRange(specs.t[imin] * 1e15, specs.t[imax] * 1e15, "fs")
+        time_range = PlotRange(specs.t[imin] * 1e15, specs.t[imax] * 1e15, "fs")
 
     if axes is None:
         _, (left, right) = plt.subplots(1, 2)
     else:
         left, right = axes
 
-    x, y, values = transform_2D_propagation(wl_int, wl_range, specs.wl_disp, z, log="1D")
+    x, y, values = transform_2D_propagation(wl_int, wl_range, specs.wl_disp, z, log=wl_db)
     left.imshow(values, extent=get_extent(x, y), origin="lower", aspect="auto", vmin=db_min)
 
-    x, y, values = transform_2D_propagation(time_int, t_range, specs.t, z, log=False)
+    x, y, values = transform_2D_propagation(time_int, time_range, specs.t, z, log=time_db)
     right.imshow(values, extent=get_extent(x, y), origin="lower", aspect="auto", vmin=db_min)

src/scgenerator/solver.py

@@ -172,6 +172,7 @@ def solve43(
             return
         targets = list(sorted(set(targets)))
         z = targets[0]
+        if not const_step_size:
             h = min(h, (targets[1] - targets[0]) / 2)
         targets.pop(0)