From 514dc4d99fedf5066e57294253fa86addab6ced5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Sierro?= Date: Wed, 28 Jun 2023 14:04:53 +0200 Subject: [PATCH] Small improvements --- src/scgenerator/__init__.py | 30 ++------- src/scgenerator/evaluator.py | 2 +- src/scgenerator/math.py | 10 ++- src/scgenerator/operators.py | 2 + src/scgenerator/physics/fiber.py | 2 +- src/scgenerator/physics/pulse.py | 23 ++++--- src/scgenerator/physics/units.py | 12 +++- src/scgenerator/solver.py | 28 +++++++-- src/scgenerator/transform.py | 48 ++++++++------ tests/Optica_PM2000D/test_Optica_PM2000D.py | 11 +++- tests/Travers/Travers.toml | 15 +++++ tests/Travers/Travers2019.py | 69 +++++++++++++++++++++ 12 files changed, 187 insertions(+), 65 deletions(-) create mode 100755 tests/Travers/Travers.toml create mode 100644 tests/Travers/Travers2019.py diff --git a/src/scgenerator/__init__.py b/src/scgenerator/__init__.py index c4fb4c6..cbae4c9 100644 --- a/src/scgenerator/__init__.py +++ b/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 + diff --git a/src/scgenerator/evaluator.py b/src/scgenerator/evaluator.py index 5cca77b..7f9867e 100644 --- a/src/scgenerator/evaluator.py +++ b/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")), diff --git a/src/scgenerator/math.py b/src/scgenerator/math.py index a083830..70dd6fb 100644 --- a/src/scgenerator/math.py +++ b/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 diff --git a/src/scgenerator/operators.py b/src/scgenerator/operators.py index 2f936a8..a986daa 100755 --- a/src/scgenerator/operators.py +++ b/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 diff --git a/src/scgenerator/physics/fiber.py b/src/scgenerator/physics/fiber.py index de50703..96af0b0 100644 --- a/src/scgenerator/physics/fiber.py +++ b/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 diff --git a/src/scgenerator/physics/pulse.py b/src/scgenerator/physics/pulse.py index a5d3fad..d3b9167 100644 --- a/src/scgenerator/physics/pulse.py +++ b/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 diff --git a/src/scgenerator/physics/units.py b/src/scgenerator/physics/units.py index bc8c0d2..b1ebd02 100644 --- a/src/scgenerator/physics/units.py +++ b/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 diff --git a/src/scgenerator/solver.py b/src/scgenerator/solver.py index 7df11b8..3a82dce 100644 --- a/src/scgenerator/solver.py +++ b/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") diff --git a/src/scgenerator/transform.py b/src/scgenerator/transform.py index 7515282..1ca8c45 100644 --- a/src/scgenerator/transform.py +++ b/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 - diff --git a/tests/Optica_PM2000D/test_Optica_PM2000D.py b/tests/Optica_PM2000D/test_Optica_PM2000D.py index aff1f8f..0c44c0b 100644 --- a/tests/Optica_PM2000D/test_Optica_PM2000D.py +++ b/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) diff --git a/tests/Travers/Travers.toml b/tests/Travers/Travers.toml new file mode 100755 index 0000000..7de44be --- /dev/null +++ b/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" diff --git a/tests/Travers/Travers2019.py b/tests/Travers/Travers2019.py new file mode 100644 index 0000000..58c3ee4 --- /dev/null +++ b/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()