Small improvements

src/scgenerator/__init__.py

@@ -1,27 +1,9 @@
 # # flake8: noqa
-# from scgenerator import math, operators
-# from scgenerator.evaluator import Evaluator
-# from scgenerator.helpers import *
-# from scgenerator.math import abs2, argclosest, normalized, span, tspace, wspace
+from scgenerator import math, operators, plotting
+from scgenerator.math import abs2, argclosest, normalized, span, tspace, wspace
 from scgenerator.parameter import FileConfiguration, Parameters
-# from scgenerator.physics import fiber, materials, plasma, pulse, simulate, units
-# from scgenerator.physics.simulate import RK4IP, parallel_RK4IP, run_simulation
-# from scgenerator.physics.units import PlotRange
-# from scgenerator.plotting import (
-#     get_extent,
-#     mean_values_plot,
-#     plot_spectrogram,
-#     propagation_plot,
-#     single_position_plot,
-#     transform_1D_values,
-#     transform_2D_propagation,
-#     transform_mean_values,
-# )
-# from scgenerator.spectra import SimulationSeries, Spectrum
+from scgenerator.physics import fiber, materials, plasma, pulse, units
+from scgenerator.physics.units import PlotRange
 from scgenerator.utils import Paths, _open_config, open_single_config, simulations_list
-# from scgenerator.variationer import (
-#     DescriptorDict,
-#     VariationDescriptor,
-#     Variationer,
-#     VariationSpecsError,
-# )
+from scgenerator.solver import solve43, integrate
+

src/scgenerator/evaluator.py

@@ -403,7 +403,7 @@ default_rules: list[Rule] = [
     # Raman
     Rule(["hr_w", "raman_fraction"], fiber.delayed_raman_w),
     Rule("raman_fraction", fiber.raman_fraction),
-    Rule("raman-fraction", lambda:0, priorities=-1),
+    Rule("raman_fraction", lambda:0, priorities=-1),
     # loss
     Rule("alpha_arr", fiber.scalar_loss),
     Rule("alpha_arr", fiber.safe_capillary_loss, conditions=dict(loss="capillary")),

src/scgenerator/math.py

@@ -5,7 +5,7 @@ collection of purely mathematical function
 import math
 from dataclasses import dataclass
 from functools import cache
-from typing import Sequence, Union
+from typing import Sequence
 
 import numba
 import numpy as np
@@ -25,13 +25,19 @@ def span(*vec: np.ndarray) -> tuple[float, float]:
     """returns the min and max of whatever array-like is given. can accept many args"""
     out = (np.inf, -np.inf)
     if len(vec) == 0 or len(vec[0]) == 0:
-        raise ValueError(f"did not provide any value to span")
+        raise ValueError("did not provide any value to span")
     for x in vec:
         x = np.atleast_1d(x)
         out = (min(np.min(x), out[0]), max(np.max(x), out[1]))
     return out
 
 
+def total_extent(*vec: np.ndarray) -> float:
+    """measure the distance between the min and max value of all given arrays"""
+    left, right = span(*vec)
+    return right - left
+
+
 def argclosest(array: np.ndarray, target: float | int | Sequence[float | int]) -> int | np.ndarray:
     """
     returns the index/indices corresponding to the closest matches of target in array

src/scgenerator/operators.py

@@ -335,6 +335,8 @@ def ionization(
     def operate(field: np.ndarray, z: float) -> np.ndarray:
         N0 = number_density(z)
         plasma_info = plasma_obj(field, N0)
+
+
         # state.stats["ionization_fraction"] = plasma_info.electron_density[-1] / N0
         # state.stats["electron_density"] = plasma_info.electron_density[-1]
         return plasma_info.polarization

src/scgenerator/physics/fiber.py

@@ -34,7 +34,7 @@ def lambda_for_envelope_dispersion(
     if l[su].min() > 1.01 * interpolation_range[0]:
         raise ValueError(
             f"lower range of {1e9*interpolation_range[0]:.1f}nm is not reached by the grid. "
-            "try a finer grid"
+            f"Minimum of grid is {1e9*l[su].min():.1f}nm. Try a finer grid"
         )
 
     ind_above_cond = su >= len(l) // 2

src/scgenerator/physics/pulse.py

@@ -189,9 +189,6 @@ def modify_field_ratio(
     elif peak_power is not None:
         ratio *= np.sqrt(peak_power / math.abs2(pre_field_0).max())
 
-    if intensity_noise is not None:
-        d_int, _ = technical_noise(intensity_noise, noise_correlation)
-        ratio *= np.sqrt(d_int)
     return ratio
 
 
@@ -219,6 +216,10 @@ def c_to_a_factor(A_eff_arr: np.ndarray) -> np.ndarray:
     return (A_eff_arr / A_eff_arr[0]) ** (1 / 4)
 
 
+def a_to_c_factor(A_eff_arr: np.ndarray) -> np.ndarray:
+    return (A_eff_arr / A_eff_arr[0]) ** (-1 / 4)
+
+
 def spectrum_factor_envelope(dt: float) -> float:
     return dt / np.sqrt(2 * np.pi)
 
@@ -497,10 +498,16 @@ def finalize_pulse(
     dt: float,
     additional_noise_factor: float,
     input_transmission: float,
+    intensity_noise: float | None,
 ) -> np.ndarray:
     if quantum_noise:
         pre_field_0 = pre_field_0 + shot_noise(w, time_window, dt, additional_noise_factor)
-    return np.sqrt(input_transmission) * pre_field_0
+
+    ratio = 1
+    if intensity_noise is not None:
+        d_int, _ = technical_noise(intensity_noise, 0)
+        ratio *= np.sqrt(d_int)
+    return np.sqrt(input_transmission) * pre_field_0 * ratio
 
 
 def mean_phase(spectra):
@@ -1094,7 +1101,7 @@ def measure_properties(spectra, t, compress=True, return_limits=False, debug="")
     field = np.mean(fields, axis=0)
     ideal_field = math.abs2(fftshift(ifft(np.sqrt(np.mean(math.abs2(spectra), axis=0)))))
 
-    # Isolate whole central lobe of bof mean and ideal field
+    # Isolate whole central lobe of both mean and ideal field
     lobe_lim, fwhm_lim, _, big_spline = find_lobe_limits(t, field, debug)
     lobe_lim_i, _, _, big_spline_i = find_lobe_limits(t, ideal_field, debug)
 
@@ -1107,7 +1114,7 @@ def measure_properties(spectra, t, compress=True, return_limits=False, debug="")
 
     # Compute mean coherence
     mean_g12 = avg_g12(spectra)
-    fwhm_abs = math.span(fwhm_lim)
+    fwhm_abs = np.max(fwhm_lim) - np.min(fwhm_lim)
 
     # To compute amplitude and fwhm fluctuations, we need to measure every single peak
     P0 = []
@@ -1119,7 +1126,7 @@ def measure_properties(spectra, t, compress=True, return_limits=False, debug="")
         lobe_lim, fwhm_lim, _, big_spline = find_lobe_limits(t, f, debug)
         all_lims.append((lobe_lim, fwhm_lim))
         P0.append(big_spline(lobe_lim[2]))
-        fwhm.append(math.span(fwhm_lim))
+        fwhm.append(np.max(fwhm_lim) - np.min(fwhm_lim))
         t_offset.append(lobe_lim[2])
         energies.append(np.trapz(fields, t))
     fwhm_var = np.std(fwhm) / np.mean(fwhm)
@@ -1165,7 +1172,7 @@ def measure_field(t: np.ndarray, field: np.ndarray) -> Tuple[float, float, float
     else:
         intensity = field
     _, fwhm_lim, _, _ = find_lobe_limits(t, intensity)
-    fwhm = math.span(fwhm_lim)
+    fwhm = math.total_extent(fwhm_lim)
     peak_power = intensity.max()
     energy = np.trapz(intensity, t)
     return fwhm, peak_power, energy

src/scgenerator/physics/units.py

@@ -254,14 +254,17 @@ def bar(p: _T) -> _T:
 def bar_inv(p):
     return p * 1e-5
 
+
 @unit("PRESSURE", "Pressure (mbar)")
-def mbar(p: _T)->_T:
+def mbar(p: _T) -> _T:
     return 1e2 * p
 
+
 @mbar.inverse
 def mbar_inv(p):
     return 1e-2 * p
 
+
 @unit("OTHER", r"$\beta_2$ (fs$^2$/cm)")
 def beta2_fs_cm(b2: _T) -> _T:
     return 1e-28 * b2
@@ -312,6 +315,11 @@ def C_inv(t_K):
     return t_K - 272.15
 
 
+@unit("OTHER", r"a.u")
+def no_unit(x: _T) -> _T:
+    return x
+
+
 def get_unit(unit: Union[str, Callable]) -> Callable[[float], float]:
     if isinstance(unit, str):
         return units_map[unit]
@@ -377,7 +385,7 @@ def to_WL(spectrum: np.ndarray, lambda_: np.ndarray) -> np.ndarray:
     np.ndarray, shape (n, )
         intensity spectrum correctly scaled
     """
-    m = 2 * pi * c / (lambda_ ** 2) * spectrum
+    m = 2 * pi * c / (lambda_**2) * spectrum
     return m
 
 

src/scgenerator/solver.py

@@ -6,7 +6,6 @@ from typing import Any, Iterator, Sequence
 
 import numba
 import numpy as np
-
 from scgenerator.math import abs2
 from scgenerator.operators import SpecOperator
 from scgenerator.utils import TimedMessage
@@ -89,6 +88,7 @@ def solve43(
     rtol: float,
     safety: float,
     h_const: float | None = None,
+    targets: Sequence[float] | None = None,
 ) -> Iterator[tuple[np.ndarray, dict[str, Any]]]:
     """
     Solve the GNLSE using an embedded Runge-Kutta of order 4(3) in the interaction picture.
@@ -128,13 +128,18 @@ def solve43(
     k5 = nonlinear(spec, 0)
     z = 0
     stats = {}
+    rejected = []
+    if targets is not None:
+        targets = list(sorted(set(targets)))
+        if targets[0] == 0:
+            targets.pop(0)
 
-    step_ind = 1
+    step_ind = 0
     msg = TimedMessage(2)
     running = True
     last_error = 0
     error = 0
-    rejected = []
+    store_next = False
 
     def stats():
         return dict(z=z, rejected=rejected.copy(), error=error, h=h)
@@ -174,10 +179,13 @@ def solve43(
             step_ind += 1
             last_error = error
 
-            yield fine, stats()
+            if targets is None or store_next:
+                if targets is not None:
+                    targets.pop(0)
+                yield fine, stats()
 
             rejected.clear()
-            if z > z_max:
+            if z >= z_max:
                 return
             if const_step_size:
                 continue
@@ -186,6 +194,13 @@ def solve43(
             print(f"{z = :.3f} rejected step {step_ind} with {h = :.2g}, {error = :.2g}")
 
         h = h * next_h_factor
+
+        if targets is not None and z + h > targets[0]:
+            h = targets[0] - z
+            store_next = True
+        else:
+            store_next = False
+
         if msg.ready():
             print(f"step {step_ind}, {z = :.3f}, {error = :g}, {h = :.3g}")
 
@@ -198,6 +213,7 @@ def integrate(
     atol: float = 1e-6,
     rtol: float = 1e-6,
     safety: float = 0.9,
+    targets: Sequence[float] | None = None,
 ) -> SimulationResult:
     spec0 = initial_spectrum.copy()
     all_spectra = []
@@ -206,7 +222,7 @@ def integrate(
     with warnings.catch_warnings():
         warnings.filterwarnings("error")
         for i, (spec, new_stat) in enumerate(
-            solve43(spec0, linear, nonlinear, length, atol, rtol, safety)
+            solve43(spec0, linear, nonlinear, length, atol, rtol, safety, targets=targets)
         ):
             if msg.ready():
                 print(f"step {i}, z = {new_stat['z']*100:.2f}cm")

src/scgenerator/transform.py

@@ -1,38 +1,48 @@
 import numpy as np
-
 import scgenerator.math as math
 import scgenerator.physics.units as units
 
 
+def normalize_range(
+    axis: np.ndarray, _range: tuple | units.PlotRange | None, num: int
+) -> tuple[units.PlotRange, np.ndarray]:
+    if _range is None:
+        _range = units.PlotRange(axis.min(), axis.max(), units.no_unit)
+    elif not isinstance(_range, units.PlotRange):
+        _range = units.PlotRange(*_range)
+    new_axis = np.linspace(_range[0], _range[1], num)
+    return _range, new_axis
+
+
 def prop_2d(
     values: np.ndarray,
     h_axis: np.ndarray,
     v_axis: np.ndarray,
-    horizontal_range: tuple | units.PlotRange | None,
-    vertical_range: tuple | units.PlotRange | None,
+    h_range: tuple | units.PlotRange | None = None,
+    v_range: tuple | units.PlotRange | None = None,
     h_num: int = 1024,
     v_num: int = 1024,
     z_lim: tuple[float, float] | None = None,
 ):
     if values.ndim != 2:
         raise TypeError("prop_2d can only transform 2d data")
-    
-    if horizontal_range is None:
-        horizontal_range = units.PlotRange(h_axis.min(), h_axis.max(), units.no_unit)
-    elif not isinstance(horizontal_range, units.PlotRange):
-        horizontal_range = units.PlotRange(*horizontal_range)
-
-    if vertical_range is None:
-        vertical_range = units.PlotRange(h_axis.min(), h_axis.max(), units.no_unit)
-    elif not isinstance(vertical_range, units.PlotRange):
-        vertical_range = units.PlotRange(*vertical_range)
-
-    if np.iscomplex(values):
+    if np.iscomplexobj(values):
         values = math.abs2(values)
 
-    horizontal = np.linspace(horizontal_range[0], horizontal_range[1], h_num)
-    vertical = np.linspace(vertical_range[0], vertical_range[1], v_num)
+    horizontal_range, horizontal = normalize_range(h_axis, h_range, h_num)
+    vertical_range, vertical = normalize_range(v_axis, v_range, v_num)
+
+    values = math.interp_2d(
+        h_axis, v_axis, values, horizontal_range.unit(horizontal), vertical_range.unit(vertical)
+    )
+
+    if horizontal_range.must_correct_wl:
+        values = np.apply_along_axis(
+            lambda x: units.to_WL(x, horizontal_range.unit.to.m(horizontal)), 1, values
+        )
+    elif vertical_range.must_correct_wl:
+        values = np.apply_along_axis(
+            lambda x: units.to_WL(x, vertical_range.unit.to.m(vertical)), 0, values
+        )
 
-    values = math.interp_2d(h_axis, v_axis, values, horizontal_range.unit(horizontal), vertical_range.unit(vertical))
     return horizontal, vertical, values
-

tests/Optica_PM2000D/test_Optica_PM2000D.py

@@ -1,3 +1,9 @@
+"""
+May 2023
+
+Testing the new solver / operators structure
+using parameters from the 2022 Optica paper
+"""
 import matplotlib.pyplot as plt
 import numpy as np
 from scipy.interpolate import interp1d
@@ -19,11 +25,12 @@ def main():
     #
     # plt.plot(params.w, params.linear_operator(0).imag)
     # plt.show()
-    
 
     breakpoint()
 
-    res = sol.integrate(params.spec_0, params.length, params.linear_operator, params.nonlinear_operator)
+    res = sol.integrate(
+        params.spec_0, params.length, params.linear_operator, params.nonlinear_operator
+    )
 
     new_z = np.linspace(0, params.length, 256)
 

tests/Travers/Travers.toml

@@ -0,0 +1,15 @@
+name = "Travers"
+repetition_rate = 1e3
+wavelength = 800e-9
+tolerated_error = 1e-6
+z_num = 128
+t_num = 4096
+dt = 50e-18
+core_radius = 125e-6
+gas_name = "helium"
+pressure = 400e2
+energy = 0.4e-3
+width = 10e-15
+length = 3
+full_field = true
+shape = "sech"

tests/Travers/Travers2019.py

@@ -0,0 +1,69 @@
+"""
+Testing the the solver / operator mechanism with
+parameters from the 2019 Travers paper
+"""
+
+import matplotlib.pyplot as plt
+import numpy as np
+from scipy.interpolate import interp1d
+
+import scgenerator as sc
+import scgenerator.math as math
+import scgenerator.physics.units as units
+import scgenerator.plotting as plot
+import scgenerator.solver as sol
+
+
+def main():
+    params = sc.Parameters(**sc.open_single_config("./tests/Travers/Travers.toml"))
+    # print(params.nonlinear_operator)
+    # print(params.compute("dispersion_op"))
+    # print(params.linear_operator)
+    # print(params.spec_0)
+    # print(params.compute("gamma_op"))
+    #
+    # plt.plot(params.w, params.linear_operator(0).imag)
+    # plt.show()
+
+    breakpoint()
+
+    res = sol.integrate(
+        params.spec_0, params.length, params.linear_operator, params.nonlinear_operator
+    )
+
+    new_z = np.linspace(0, params.length, 256)
+
+    specs2 = math.abs2(res.spectra)
+    specs2 = units.to_WL(specs2, params.l)
+    x = params.l
+    # x = units.THz.inv(w)
+    # new_x = np.linspace(100, 2200, 1024)
+    new_x = np.linspace(100e-9, 1200e-9, 1024)
+    solution = interp1d(res.z, specs2, axis=0)(new_z)
+    solution = interp1d(x, solution)(new_x)
+    solution = units.to_log2D(solution)
+
+    plt.imshow(
+        solution,
+        origin="lower",
+        aspect="auto",
+        extent=plot.get_extent(1e9 * new_x, new_z * 1e2),
+        vmin=-30,
+    )
+    plt.show()
+
+    fields = np.fft.irfft(res.spectra)
+    solution = math.abs2(fields)
+    solution = interp1d(res.z, solution, axis=0)(new_z)
+    solution.T[:] /= solution.max(axis=1)
+    plt.imshow(
+        solution,
+        origin="lower",
+        aspect="auto",
+        extent=plot.get_extent(params.t * 1e15, new_z * 1e2),
+    )
+    plt.show()
+
+
+if __name__ == "__main__":
+    main()