diff --git a/src/scgenerator/__init__.py b/src/scgenerator/__init__.py index 4a9a18d..f53314d 100644 --- a/src/scgenerator/__init__.py +++ b/src/scgenerator/__init__.py @@ -1,15 +1,16 @@ +from . import initialize, io, math, utils from .initialize import ( - ParamSequence, - RecoveryParamSequence, + Config, ContinuationParamSequence, Params, - Config, + ParamSequence, + RecoveryParamSequence, ) -from .io import Paths, load_toml, load_params +from .io import Paths, load_params, load_toml from .math import abs2, argclosest, span from .physics import fiber, materials, pulse, simulate, units from .physics.simulate import RK4IP, new_simulation, resume_simulations -from .plotting import mean_values_plot, single_position_plot, propagation_plot, plot_spectrogram +from .physics.units import PlotRange +from .plotting import mean_values_plot, plot_spectrogram, propagation_plot, single_position_plot from .spectra import Pulse, Spectrum -from .utils.parameter import BareParams, BareConfig -from . import utils, io, initialize, math +from .utils.parameter import BareConfig, BareParams diff --git a/src/scgenerator/defaults.py b/src/scgenerator/defaults.py index 5a6097e..c568efd 100644 --- a/src/scgenerator/defaults.py +++ b/src/scgenerator/defaults.py @@ -8,6 +8,7 @@ default_parameters = dict( fit_parameters=(0.08, 200e-9), model="custom", length=1, + n2=2.2e-20, capillary_resonance_strengths=[], capillary_nested=0, gas_name="vacuum", diff --git a/src/scgenerator/initialize.py b/src/scgenerator/initialize.py index 1d61d26..1376500 100644 --- a/src/scgenerator/initialize.py +++ b/src/scgenerator/initialize.py @@ -97,8 +97,16 @@ class Params(BareParams): ) temp_gamma = None - if self.effective_mode_diameter is not None: - self.A_eff = (self.effective_mode_diameter / 2) ** 2 * pi + if self.A_eff_file is not None: + self.A_eff_arr = fiber.compute_custom_A_eff(self) + elif self.A_eff is not None: + self.A_eff_arr = np.ones(self.t_num) * self.A_eff + elif self.effective_mode_diameter is not None: + self.A_eff_arr = np.ones(self.t_num) * (self.effective_mode_diameter / 2) ** 2 * pi + else: + self.A_eff_arr = np.ones(self.t_num) * self.n2 * self.w0 / (299792458.0 * self.gamma) + self.A_eff = self.A_eff_arr[0] + if self.beta is not None: self.beta = np.array(self.beta) self.dynamic_dispersion = False @@ -112,8 +120,11 @@ class Params(BareParams): temp_gamma = temp_gamma(0) if self.gamma is None: - self.gamma = temp_gamma + self.gamma_arr = temp_gamma + self.gamma = temp_gamma[0] logger.info(f"using computed \u0263 = {self.gamma:.2e} W/m\u00B2") + else: + self.gamma_arr = np.ones(self.t_num) * self.gamma # Raman response if "raman" in self.behaviors: @@ -155,14 +166,15 @@ class Config(BareConfig): self.simulation_consistency() def fiber_consistency(self): - if self.contains("beta"): - if not (self.contains("A_eff") or self.contains("effective_mode_diameter")): - self.get("gamma", specified_parameters=["beta"]) - self.setdefault("model", "custom") - elif self.contains("dispersion_file"): - if not (self.contains("A_eff") or self.contains("effective_mode_diameter")): - self.get("gamma", specified_parameters=["dispersion_file"]) + if self.contains("dispersion_file") or self.contains("beta"): + if not ( + self.contains("A_eff") + or self.contains("A_eff_file") + or self.contains("effective_mode_diameter") + ): + self.get("gamma", specified_parameters=["custom fiber model"]) + self.get("n2", specified_parameters=["custom fiber model"]) self.setdefault("model", "custom") else: diff --git a/src/scgenerator/physics/fiber.py b/src/scgenerator/physics/fiber.py index 995cf92..7e2384f 100644 --- a/src/scgenerator/physics/fiber.py +++ b/src/scgenerator/physics/fiber.py @@ -1,4 +1,4 @@ -from typing import Any, Dict, Iterable, List, Literal, Optional, Tuple, Union +from typing import Any, Dict, Iterable, List, Literal, Optional, Tuple, Union, TypeVar import numpy as np from numpy.ma import core @@ -18,7 +18,9 @@ from . import materials as mat from . import units from .units import c, pi + pipi = 2 * pi +T = TypeVar("T") def lambda_for_dispersion(left: float, right: float) -> np.ndarray: @@ -102,6 +104,14 @@ def D_to_beta2(D, λ): return -(λ ** 2) / (pipi * c) * D +def A_to_C(A: np.ndarray, A_eff_arr: np.ndarray) -> np.ndarray: + return (A_eff_arr / A_eff_arr[0]) ** (-1 / 4) * A + + +def C_to_A(C: np.ndarray, A_eff_arr: np.ndarray) -> np.ndarray: + return (A_eff_arr / A_eff_arr[0]) ** (1 / 4) * C + + def plasma_dispersion(lambda_, number_density, simple=False): """computes dispersion (beta2) for constant plasma @@ -559,8 +569,12 @@ def dynamic_HCPCF_dispersion( return beta2_func, gamma_func -def gamma_parameter(n2, w0, A_eff): - return n2 * w0 / (A_eff * c) +def gamma_parameter(n2: float, w0: float, A_eff: T) -> T: + if isinstance(A_eff, np.ndarray): + A_eff_term = np.sqrt(A_eff * A_eff[0]) + else: + A_eff_term = A_eff + return n2 * w0 / (A_eff_term * c) @np_cache @@ -656,6 +670,15 @@ def PCF_dispersion(lambda_, pitch, ratio_d, w0=None, n2=None, A_eff=None): return beta2, gamma +def compute_custom_A_eff(params: BareParams) -> np.ndarray: + data = np.load(params.A_eff_file) + A_eff = data["A_eff"] + wl = data["wavelength"] + return interp1d( + wl, A_eff, fill_value=(A_eff[wl.argmin()], A_eff[wl.argmax()]), bounds_error=False + )(params.l) + + def compute_loss(params: BareParams) -> Optional[np.ndarray]: if params.loss_file is not None: loss_data = np.load(params.loss_file) @@ -671,7 +694,7 @@ def compute_loss(params: BareParams) -> Optional[np.ndarray]: return None -def compute_dispersion(params: BareParams): +def compute_dispersion(params: BareParams) -> tuple[np.ndarray, np.ndarray, tuple[float, float]]: """dispatch function depending on what type of fiber is used Parameters @@ -762,12 +785,14 @@ def compute_dispersion(params: BareParams): ) if gamma is None: - if params.A_eff is not None: - gamma = gamma_parameter(params.n2, params.w0, params.A_eff) + if params.A_eff_arr is not None: + gamma_arr = gamma_parameter(params.n2, params.w0, params.A_eff_arr) else: - gamma = 0 + gamma_arr = np.zeros(params.t_num) + else: + gamma_arr = np.ones(params.t_num) * gamma - return beta2_coef, gamma, interp_range + return beta2_coef, gamma_arr, interp_range @np_cache @@ -954,7 +979,7 @@ def create_non_linear_op(behaviors, w_c, w0, gamma, raman_type="stolen", f_r=0, ss_part = w_c / w0 if "ss" in behaviors else 0 - if isinstance(gamma, (float, int)): + if isinstance(gamma, (float, int, np.ndarray)): def N_func(spectrum: np.ndarray, r=0) -> np.ndarray: field = ifft(spectrum) diff --git a/src/scgenerator/physics/pulse.py b/src/scgenerator/physics/pulse.py index 8106282..8f1299c 100644 --- a/src/scgenerator/physics/pulse.py +++ b/src/scgenerator/physics/pulse.py @@ -14,6 +14,7 @@ import os from pathlib import Path from typing import Literal, Tuple, TypeVar from collections import namedtuple +from dataclasses import dataclass, astuple import matplotlib.pyplot as plt import numpy as np @@ -49,10 +50,54 @@ P0T0_to_E0_fac = dict( ) """relates the total energy (amplitue^2) to the t0 parameter of the amplitude and the peak intensity (peak_amplitude^2)""" -PulseProperties = namedtuple( - "PulseProperties", - "quality mean_coherence fwhm_noise mean_fwhm peak_rin energy_rin timing_jitter", -) + +@dataclass +class PulseProperties: + quality: float + mean_coherence: float + fwhm_noise: float + mean_fwhm: float + peak_rin: float + energy_rin: float + timing_jitter: float + + @classmethod + def header(cls, delimiter: str = ",", quotechar: str = "") -> str: + return delimiter.join( + quotechar + f + quotechar + for f in [ + "quality", + "mean_coherence", + "fwhm_noise", + "mean_fwhm", + "peak_rin", + "energy_rin", + "timing_jitter", + ] + ) + + @classmethod + def save_all( + cls, + destination: os.PathLike, + *props: "PulseProperties", + delimiter: str = ",", + quotechar: str = "", + ): + out = np.zeros((len(props), 7)) + for i, p in enumerate(props): + out[i] = astuple(p) + np.savetxt( + destination, + out, + header=cls.header(delimiter=delimiter, quotechar=quotechar), + delimiter=delimiter, + ) + + @classmethod + def load(cls, path: os.PathLike, delimiter: str = ",") -> list["PulseProperties"]: + arr = np.loadtxt(path, delimiter=delimiter) + return [cls(*a) for a in arr] def initial_field(t, shape, t0, peak_power): @@ -954,6 +999,25 @@ def measure_properties(spectra, t, compress=True, return_limits=False, debug="") return pp +def rin_curve(spectra: np.ndarray) -> np.ndarray: + """computes the rin curve, i.e. the rin at every single point + + Parameters + ---------- + spectra : np.ndarray, shape (n, nt) + a collection of n spectra from which to compute the RIN + + Returns + ------- + rin_curve : np.ndarray + RIN curve + """ + spec2 = abs2(spectra) + # return np.std(spec, axis=0) / np.mean(spec, axis=0) + m = np.mean(spec2, axis=0) + return np.sqrt(np.mean((spec2 - m) ** 2)) / m + + def measure_field(t: np.ndarray, field: np.ndarray) -> Tuple[float, float, float]: """returns fwhm, peak_power, energy""" intensity = abs2(field) diff --git a/src/scgenerator/physics/simulate.py b/src/scgenerator/physics/simulate.py index 5c2837d..9bf8175 100644 --- a/src/scgenerator/physics/simulate.py +++ b/src/scgenerator/physics/simulate.py @@ -13,6 +13,7 @@ from ..errors import IncompleteDataFolderError from ..logger import get_logger from . import pulse from .fiber import create_non_linear_op, fast_dispersion_op +from scgenerator.physics import fiber try: import ray @@ -72,7 +73,8 @@ class RK4IP: self.z_targets = params.z_targets self.z_final = params.length self.beta = params.beta_func if params.beta_func is not None else params.beta - self.gamma = params.gamma_func if params.gamma_func is not None else params.gamma + self.gamma = params.gamma_func if params.gamma_func is not None else params.gamma_arr + self.C_to_A_factor = (params.A_eff_arr / params.A_eff_arr[0]) ** (-1 / 4) self.behaviors = params.behaviors self.raman_type = params.raman_type self.hr_w = params.hr_w @@ -135,7 +137,7 @@ class RK4IP: self.z = self.z_targets.pop(0) # Setup initial values for every physical quantity that we want to track - self.current_spectrum = self.spec_0.copy() + self.current_spectrum = self.spec_0.copy() / self.C_to_A_factor self.stored_spectra = self.starting_num * [None] + [self.current_spectrum.copy()] self.cons_qty = [ self.conserved_quantity_func(self.current_spectrum, 0), @@ -160,7 +162,7 @@ class RK4IP: num : int index of the z postition """ - self._save_data(self.current_spectrum, f"spectrum_{num}") + self._save_data(self.C_to_A_factor * self.current_spectrum, f"spectrum_{num}") self._save_data(self.cons_qty, f"cons_qty") self.step_saved() diff --git a/src/scgenerator/physics/units.py b/src/scgenerator/physics/units.py index 6e040e2..d49de0e 100644 --- a/src/scgenerator/physics/units.py +++ b/src/scgenerator/physics/units.py @@ -2,14 +2,10 @@ # For example, nm(X) means "I give the number X in nm, figure out the ang. freq." # to be used especially when giving plotting ranges : (400, 1400, nm), (-4, 8, ps), ... -import re -from threading import settrace from typing import Callable, TypeVar, Union from dataclasses import dataclass -from matplotlib import pyplot as plt -from numpy.lib.arraysetops import isin -from ..utils.parameter import Parameter, type_checker +from ..utils.parameter import Parameter, boolean, type_checker import numpy as np from numpy import pi @@ -190,6 +186,7 @@ class PlotRange: left: float = Parameter(type_checker(int, float)) right: float = Parameter(type_checker(int, float)) unit: Callable[[float], float] = Parameter(is_unit, converter=get_unit) + conserved_quantity: bool = Parameter(boolean, default=True) def __str__(self): return f"{self.left:.1f}-{self.right:.1f} {self.unit.__name__}" diff --git a/src/scgenerator/plotting.py b/src/scgenerator/plotting.py index 3d08f76..ce4a5a2 100644 --- a/src/scgenerator/plotting.py +++ b/src/scgenerator/plotting.py @@ -535,7 +535,7 @@ def transform_mean_values( if log is not True and isinstance(log, (int, float, np.integer, np.floating)): ref = log else: - ref = mean_values.max() + ref = float(mean_values.max()) values = apply_log(values, ref) mean_values = apply_log(mean_values, ref) return new_axis, mean_values, values @@ -875,7 +875,7 @@ def uniform_axis( new_axis = tmp_axis values = values[:, ind] else: - if plt_range.unit.type == "WL": + if plt_range.unit.type == "WL" and plt_range.conserved_quantity: values[:, ind] = np.apply_along_axis(units.to_WL, 1, values[:, ind], tmp_axis) new_axis = np.linspace(tmp_axis.min(), tmp_axis.max(), len(tmp_axis)) values = np.array([interp1d(tmp_axis, v[ind])(new_axis) for v in values]) diff --git a/src/scgenerator/utils/parameter.py b/src/scgenerator/utils/parameter.py index b852a78..dda2cf3 100644 --- a/src/scgenerator/utils/parameter.py +++ b/src/scgenerator/utils/parameter.py @@ -252,6 +252,7 @@ valid_variable = { "dispersion_file", "field_file", "loss_file", + "A_eff_file", "beta", "gamma", "pitch", @@ -325,6 +326,7 @@ class BareParams: loss_file: str = Parameter(string) effective_mode_diameter: float = Parameter(positive(float, int)) A_eff: float = Parameter(non_negative(float, int)) + A_eff_file: str = Parameter(string) pitch: float = Parameter(in_range_excl(0, 1e-3)) pitch_ratio: float = Parameter(in_range_excl(0, 1)) core_radius: float = Parameter(in_range_excl(0, 1e-3)) @@ -386,6 +388,8 @@ class BareParams: field_0: np.ndarray = Parameter(type_checker(np.ndarray)) spec_0: np.ndarray = Parameter(type_checker(np.ndarray)) alpha: np.ndarray = Parameter(type_checker(np.ndarray)) + gamma_arr: np.ndarray = Parameter(type_checker(np.ndarray)) + A_eff_arr: np.ndarray = Parameter(type_checker(np.ndarray)) w: np.ndarray = Parameter(type_checker(np.ndarray)) l: np.ndarray = Parameter(type_checker(np.ndarray)) w_c: np.ndarray = Parameter(type_checker(np.ndarray)) @@ -434,6 +438,8 @@ class BareParams: "z_targets", "l", "alpha", + "gamma_arr", + "A_eff_arr", ] types = (np.ndarray, float, int, str, list, tuple, dict) out = {}