From 25a201543f45a1420d857bf9cdd01b82cf480fcf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Sierro?= Date: Mon, 25 Oct 2021 09:25:37 +0200 Subject: [PATCH] Lots of progress, switching to tomli --- developement_help.md | 4 + requirements.txt | 3 +- setup.cfg | 2 +- src/scgenerator/errors.py | 4 + src/scgenerator/evaluator.py | 13 +- src/scgenerator/legacy.py | 11 +- src/scgenerator/math.py | 57 ++-- src/scgenerator/operators.py | 384 +++++++++++++++++++++++---- src/scgenerator/parameter.py | 23 +- src/scgenerator/physics/fiber.py | 2 +- src/scgenerator/physics/materials.py | 13 +- src/scgenerator/physics/simulate.py | 40 ++- src/scgenerator/plotting.py | 36 ++- src/scgenerator/scripts/__init__.py | 3 +- src/scgenerator/utils.py | 25 +- 15 files changed, 484 insertions(+), 136 deletions(-) diff --git a/developement_help.md b/developement_help.md index ceed26f..5dbb76c 100644 --- a/developement_help.md +++ b/developement_help.md @@ -5,3 +5,7 @@ - optional : add a default value - optional : add to valid_variable - optional : add to mandatory_parameters + +## complicated Rule conditions +- add the desired parameters to the init of the operator +- raise OperatorError if the conditions are not met \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index bb83987..a0d470c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,6 @@ numpy matplotlib scipy -toml +tomli +tomli-w tqdm \ No newline at end of file diff --git a/setup.cfg b/setup.cfg index 86b849d..6c4a4b8 100644 --- a/setup.cfg +++ b/setup.cfg @@ -23,7 +23,7 @@ install_requires = numba matplotlib scipy - toml + tomli tqdm diff --git a/src/scgenerator/errors.py b/src/scgenerator/errors.py index 3b1d790..93e69a3 100644 --- a/src/scgenerator/errors.py +++ b/src/scgenerator/errors.py @@ -42,3 +42,7 @@ class EvaluatorError(Exception): class NoDefaultError(EvaluatorError): pass + + +class OperatorError(EvaluatorError): + pass diff --git a/src/scgenerator/evaluator.py b/src/scgenerator/evaluator.py index db82c03..048e7a7 100644 --- a/src/scgenerator/evaluator.py +++ b/src/scgenerator/evaluator.py @@ -24,7 +24,7 @@ class Rule: targets = list(target) if isinstance(target, (list, tuple)) else [target] self.func = func if priorities is None: - priorities = [1] * len(targets) + priorities = [0] * len(targets) elif isinstance(priorities, (int, float, np.integer, np.floating)): priorities = [priorities] self.targets = dict(zip(targets, priorities)) @@ -293,7 +293,7 @@ class Evaluator: default_rules: list[Rule] = [ # Grid *Rule.deduce( - ["z_targets", "t", "time_window", "t_num", "dt", "w_c", "w0", "w", "l"], + ["z_targets", "t", "time_window", "t_num", "dt", "w_c", "w0", "w", "w_order", "l"], math.build_sim_grid, ["time_window", "t_num", "dt"], 2, @@ -403,11 +403,16 @@ default_rules: list[Rule] = [ Rule("raman_op", operators.Raman), Rule("raman_op", operators.NoRaman, priorities=-1), Rule("nonlinear_operator", operators.EnvelopeNonLinearOperator), - Rule("loss_op", operators.CustomConstantLoss, priorities=3), + Rule("loss_op", operators.CustomLoss, priorities=3), Rule("loss_op", operators.CapillaryLoss, priorities=2, conditions=dict(loss="capillary")), Rule("loss_op", operators.ConstantLoss, priorities=1), Rule("loss_op", operators.NoLoss, priorities=-1), + Rule("n_op", operators.ConstantRefractiveIndex), + Rule("n_op", operators.MarcatiliRefractiveIndex), + Rule("n_op", operators.MarcatiliAdjustedRefractiveIndex), + Rule("n_op", operators.HasanRefractiveIndex), Rule("dispersion_op", operators.ConstantPolyDispersion), + Rule("dispersion_op", operators.DirectDispersion), Rule("linear_operator", operators.LinearOperator), - Rule("conserved_quantity", operators.ConservedQuantity), + Rule("conserved_quantity", operators.conserved_quantity), ] diff --git a/src/scgenerator/legacy.py b/src/scgenerator/legacy.py index e2cd9f0..d33f3e5 100644 --- a/src/scgenerator/legacy.py +++ b/src/scgenerator/legacy.py @@ -5,7 +5,8 @@ from pathlib import Path from typing import Any, Set import numpy as np -import toml +import tomli +import tomli_w from .const import SPEC1_FN, SPEC1_FN_N, SPECN_FN1 from .parameter import Configuration, Parameters @@ -15,8 +16,8 @@ from .variationer import VariationDescriptor def load_config(path: os.PathLike) -> dict[str, Any]: - with open(path) as file: - d = toml.load(file) + with open(path, "rb") as file: + d = tomli.load(file) d.setdefault("variable", {}) return d @@ -40,8 +41,8 @@ def convert_sim_folder(path: os.PathLike): os.makedirs(new_root, exist_ok=True) _, configs = load_config_sequence(path) master_config = dict(name=path.name, Fiber=configs) - with open(new_root / "initial_config.toml", "w") as f: - toml.dump(master_config, f, encoder=toml.TomlNumpyEncoder()) + with open(new_root / "initial_config.toml", "wb") as f: + tomli_w.dump(Parameters.strip_params_dict(master_config), f) configuration = Configuration(path, final_output_path=new_root) pbar = PBars(configuration.total_num_steps, "Converting") diff --git a/src/scgenerator/math.py b/src/scgenerator/math.py index 68c0968..2fd8ebc 100644 --- a/src/scgenerator/math.py +++ b/src/scgenerator/math.py @@ -2,6 +2,7 @@ from typing import Union import numpy as np from scipy.special import jn_zeros + from .cache import np_cache pi = np.pi @@ -244,7 +245,6 @@ def build_sim_grid( length: float, z_num: int, wavelength: float, - interpolation_degree: int, time_window: float = None, t_num: int = None, dt: float = None, @@ -286,6 +286,8 @@ def build_sim_grid( pump angular frequency w : np.ndarray, shape (t_num, ) actual angualr frequency grid in rad/s + w_order : np.ndarray, shape (t_num, ) + result of w.argsort() l : np.ndarray, shape (t_num) wavelengths in m """ @@ -295,31 +297,34 @@ def build_sim_grid( dt = t[1] - t[0] t_num = len(t) z_targets = np.linspace(0, length, z_num) - w_c, w0, w = update_frequency_domain(t, wavelength, interpolation_degree) - l = 2 * pi * c / w - return z_targets, t, time_window, t_num, dt, w_c, w0, w, l - - -def update_frequency_domain( - t: np.ndarray, wavelength: float, deg: int -) -> tuple[np.ndarray, float, np.ndarray, np.ndarray]: - """updates the frequency grid - - Parameters - ---------- - t : np.ndarray - time array - wavelength : float - wavelength - deg : int - interpolation degree of the dispersion - - Returns - ------- - Tuple[np.ndarray, float, np.ndarray, np.ndarray] - w_c, w0, w - """ w_c = wspace(t) w0 = 2 * pi * c / wavelength w = w_c + w0 - return w_c, w0, w + w_order = np.argsort(w) + l = 2 * pi * c / w + return z_targets, t, time_window, t_num, dt, w_c, w0, w, w_order, l + + +def linear_extrapolation(y: np.ndarray) -> np.ndarray: + """extrapolates IN PLACE linearily on both side of the support + + Parameters + ---------- + y : np.ndarray + array + left_ind : int + first value we want to keep (extrapolate to the left of that) + right_ind : int + last value we want to keep (extrapolate to the right of that) + """ + out = y.copy() + order = np.argsort(y) + left_ind, *_, right_ind = np.nonzero(out[order])[0] + _linear_extrapolation_in_place(out[order], left_ind, right_ind) + return out + + +def _linear_extrapolation_in_place(y: np.ndarray, left_ind: int, right_ind: int): + y[:left_ind] = -(1 + np.arange(left_ind))[::-1] * (y[left_ind + 1] - y[left_ind]) + y[left_ind] + y[right_ind:] = np.arange(len(y) - right_ind) * (y[right_ind] - y[right_ind - 1]) + y[right_ind] + return y diff --git a/src/scgenerator/operators.py b/src/scgenerator/operators.py index 336da69..9257366 100644 --- a/src/scgenerator/operators.py +++ b/src/scgenerator/operators.py @@ -4,15 +4,19 @@ Nothing except the solver should depend on this file """ from __future__ import annotations -from abc import ABC, abstractmethod import dataclasses +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Any + import numpy as np from scipy.interpolate import interp1d from . import math +from .errors import OperatorError from .logger import get_logger - -from .physics import fiber, pulse +from .physics import fiber, materials, pulse, units +from .utils import load_material_dico class SpectrumDescriptor: @@ -57,10 +61,12 @@ class CurrentState: class Operator(ABC): def __repr__(self) -> str: value_pair_list = list(self.__dict__.items()) - if len(value_pair_list) > 1: - value_pair_str_list = [k + "=" + self.__value_repr(k, v) for k, v in value_pair_list] - else: + if len(value_pair_list) == 0: + value_pair_str_list = "" + elif len(value_pair_list) == 1: value_pair_str_list = [self.__value_repr(value_pair_list[0][0], value_pair_list[0][1])] + else: + value_pair_str_list = [k + "=" + self.__value_repr(k, v) for k, v in value_pair_list] return self.__class__.__name__ + "(" + ", ".join(value_pair_str_list) + ")" @@ -79,6 +85,238 @@ class NoOp: self.arr_const = np.zeros(t_num) +################################################## +###################### GAS ####################### +################################################## + + +class AbstractGas(ABC): + @abstractmethod + def pressure(self, state: CurrentState) -> float: + """returns the pressure at the current + + Parameters + ---------- + state : CurrentState + + Returns + ------- + float + pressure un bar + """ + + @abstractmethod + def number_density(self, state: CurrentState) -> float: + """returns the number density in 1/m^3 of at the current state + + Parameters + ---------- + state : CurrentState + + Returns + ------- + float + number density in 1/m^3 + """ + + @abstractmethod + def square_index(self, state: CurrentState) -> np.ndarray: + """returns the square of the material refractive index at the current state + + Parameters + ---------- + state : CurrentState + + Returns + ------- + np.ndarray + n^2 + """ + + +class ConstantGas(AbstractGas): + pressure_const: float + number_density_const: float + n_gas_2_const: np.ndarray + + def __init__( + self, + gas_name: str, + pressure: float, + temperature: float, + ideal_gas: bool, + wl_for_disp: np.ndarray, + ): + gas_dico = load_material_dico(gas_name) + self.pressure_const = pressure + if ideal_gas: + self.number_density_const = materials.number_density_van_der_waals( + pressure=pressure, temperature=temperature, material_dico=gas_dico + ) + else: + self.number_density_const = self.pressure_const / (units.kB * temperature) + self.n_gas_2_const = materials.n_gas_2( + wl_for_disp, gas_name, self.pressure_const, temperature, ideal_gas + ) + + def pressure(self, state: CurrentState = None) -> float: + return self.pressure_const + + def number_density(self, state: CurrentState = None) -> float: + return self.number_density_const + + def square_index(self, state: CurrentState = None) -> float: + return self.n_gas_2_const + + +class PressureGradientGas(AbstractGas): + name: str + p_in: float + p_out: float + temperature: float + gas_dico: dict[str, Any] + + def __init__( + self, + gas_name: str, + pressure_in: float, + pressure_out: float, + temperature: float, + ideal_gas: bool, + wl_for_disp: np.ndarray, + ): + self.name = gas_name + self.p_in = pressure_in + self.p_out = pressure_out + self.gas_dico = load_material_dico(gas_name) + self.temperature = temperature + self.ideal_gas = ideal_gas + self.wl_for_disp = wl_for_disp + + def pressure(self, state: CurrentState) -> float: + return materials.pressure_from_gradient(state.z_ratio, self.p_in, self.p_out) + + def number_density(self, state: CurrentState) -> float: + if self.ideal: + return self.pressure(state) / (units.kB * self.temperature) + else: + return materials.number_density_van_der_waals( + pressure=self.pressure(state), + temperature=self.temperature, + material_dico=self.gas_dico, + ) + + def square_index(self, state: CurrentState) -> np.ndarray: + return materials.fast_n_gas_2( + self.wl_for_disp, self.pressure(state), self.temperature, self.ideal_gas, self.gas_dico + ) + + +################################################## +################### DISPERSION ################### +################################################## + + +class AbstractRefractiveIndex(Operator): + @abstractmethod + def __call__(self, state: CurrentState) -> np.ndarray: + """returns the total/effective refractive index at this state + + Parameters + ---------- + state : CurrentState + + Returns + ------- + np.ndarray + refractive index + """ + + +class ConstantRefractiveIndex(AbstractRefractiveIndex): + n_eff_arr: np.ndarray + + def __init__(self, n_eff: np.ndarray): + self.n_eff_arr = n_eff + + def __call__(self, state: CurrentState = None) -> np.ndarray: + return self.n_eff_arr + + +class PCFRefractiveIndex(AbstractRefractiveIndex): + def __int__(self, wl_for_disp: np.ndarray, pitch: float, pitch_ratio: float): + self.n_eff_const = fiber.n_eff_pcf(wl_for_disp, pitch, pitch_ratio) + + +class MarcatiliRefractiveIndex(AbstractRefractiveIndex): + gas_op: AbstractGas + core_radius: float + wl_for_disp: np.ndarray + + def __init__(self, gas_op: ConstantGas, core_radius: float, wl_for_disp: np.ndarray): + self.gas_op = gas_op + self.core_radius = core_radius + self.wl_for_disp = wl_for_disp + + def __call__(self, state: CurrentState) -> np.ndarray: + return fiber.n_eff_marcatili( + self.wl_for_disp, self.gas_op.square_index(state), self.core_radius + ) + + +class MarcatiliAdjustedRefractiveIndex(MarcatiliRefractiveIndex): + def __call__(self, state: CurrentState) -> np.ndarray: + return fiber.n_eff_marcatili_adjusted( + self.wl_for_disp, self.gas_op.square_index(state), self.core_radius + ) + + +@dataclass(repr=False, eq=False) +class HasanRefractiveIndex(AbstractRefractiveIndex): + gas_op: ConstantGas + core_radius: float + capillary_num: int + capillary_nested: int + capillary_thickness: float + capillary_radius: float + capillary_resonance_strengths: list[float] + wl_for_disp: np.ndarray + + # def __init__( + # self, + # gas_op: ConstantGas, + # core_radius: float, + # capillary_num: int, + # capillary_nested: int, + # capillary_thickness: float, + # capillary_radius: float, + # capillary_resonance_strengths: list[float], + # wl_for_disp: np.ndarray, + # ): + # self.gas_op = gas_op + # self.core_radius = core_radius + # self.capillary_num = capillary_num + # self.capillary_nested = capillary_nested + # self.capillary_thickness = capillary_thickness + # self.capillary_radius = capillary_radius + # self.capillary_resonance_strengths = capillary_resonance_strengths + # self.wl_for_disp = wl_for_disp + + def __call__(self, state: CurrentState) -> np.ndarray: + return fiber.n_eff_hasan( + self.wl_for_disp, + self.gas_op.square_index(state), + self.core_radius, + self.capillary_num, + self.capillary_nested, + self.capillary_thickness, + fiber.capillary_spacing_hasan( + self.capillary_num, self.capillary_radius, self.core_radius + ), + self.capillary_resonance_strengths, + ) + + ################################################## ################### DISPERSION ################### ################################################## @@ -92,7 +330,6 @@ class AbstractDispersion(Operator): Parameters ---------- state : CurrentState - current state of the simulation Returns ------- @@ -118,6 +355,10 @@ class ConstantPolyDispersion(AbstractDispersion): w_c: np.ndarray, interpolation_degree: int, ): + if interpolation_degree < 1: + raise OperatorError( + f"interpolation degree of degree {interpolation_degree} incompatible" + ) self.coefs = fiber.dispersion_coefficients(w_for_disp, beta2_arr, w0, interpolation_degree) self.w_c = w_c self.w_power_fact = np.array( @@ -133,24 +374,66 @@ class ConstantDirectDispersion(AbstractDispersion): Direct dispersion for when the refractive index is known """ - disp_arr_const: np.ndarray + disp_arr: np.ndarray def __init__( self, w_for_disp: np.ndarray, w0: np.ndarray, t_num: int, - n_eff: np.ndarray, + n_op: ConstantRefractiveIndex, dispersion_ind: np.ndarray, + w_order: np.ndarray, ): - self.disp_arr_const = np.zeros(t_num) + self.disp_arr = np.zeros(t_num, dtype=complex) w0_ind = math.argclosest(w_for_disp, w0) - self.disp_arr_const[dispersion_ind] = fiber.fast_direct_dispersion( - w_for_disp, w0, n_eff, w0_ind + self.disp_arr[dispersion_ind] = fiber.fast_direct_dispersion( + w_for_disp, w0, n_op(), w0_ind )[2:-2] + left_ind, *_, right_ind = np.nonzero(self.disp_arr[w_order])[0] + self.disp_arr[w_order] = math._linear_extrapolation_in_place( + self.disp_arr[w_order], left_ind, right_ind + ) - def __call__(self, state: CurrentState = None): - return self.disp_arr_const + def __call__(self, state: CurrentState = None) -> np.ndarray: + return self.disp_arr + + +class DirectDispersion(AbstractDispersion): + def __new__( + cls, + w_for_disp: np.ndarray, + w0: np.ndarray, + t_num: int, + n_op: ConstantRefractiveIndex, + dispersion_ind: np.ndarray, + w_order: np.ndarray, + ): + if isinstance(n_op, ConstantRefractiveIndex): + return ConstantDirectDispersion(w_for_disp, w0, t_num, n_op, dispersion_ind, w_order) + return object.__new__(cls) + + def __init__( + self, + w_for_disp: np.ndarray, + w0: np.ndarray, + t_num: int, + n_op: ConstantRefractiveIndex, + dispersion_ind: np.ndarray, + w_order: np.ndarray, + ): + self.w_for_disp = w_for_disp + self.disp_ind = dispersion_ind + self.n_op = n_op + self.disp_arr = np.zeros(t_num) + self.w0 = w0 + self.w0_ind = math.argclosest(w_for_disp, w0) + + def __call__(self, state: CurrentState) -> np.ndarray: + self.disp_arr[self.disp_ind] = fiber.fast_direct_dispersion( + self.w_for_disp, self.w0, self.n_op(state), self.w0_ind + )[2:-2] + return self.disp_arr ################################################## @@ -166,7 +449,6 @@ class AbstractLoss(Operator): Parameters ---------- state : CurrentState - current state of the simulation Returns ------- @@ -203,18 +485,18 @@ class CapillaryLoss(ConstantLoss): self.arr = np.zeros(t_num) self.arr[dispersion_ind] = alpha[2:-2] - def __call__(self, state: CurrentState) -> np.ndarray: + def __call__(self, state: CurrentState = None) -> np.ndarray: return self.arr -class CustomConstantLoss(ConstantLoss): +class CustomLoss(ConstantLoss): def __init__(self, l: np.ndarray, loss_file: str): loss_data = np.load(loss_file) wl = loss_data["wavelength"] loss = loss_data["loss"] self.arr = interp1d(wl, loss, fill_value=0, bounds_error=False)(l) - def __call__(self, state: CurrentState) -> np.ndarray: + def __call__(self, state: CurrentState = None) -> np.ndarray: return self.arr @@ -234,7 +516,6 @@ class LinearOperator(Operator): Parameters ---------- state : CurrentState - current state of the simulation Returns ------- @@ -259,7 +540,6 @@ class AbstractRaman(Operator): Parameters ---------- state : CurrentState - current state of the simulation Returns ------- @@ -297,7 +577,6 @@ class AbstractSPM(Operator): Parameters ---------- state : CurrentState - current state of the simulation Returns ------- @@ -332,7 +611,6 @@ class AbstractSelfSteepening(Operator): Parameters ---------- state : CurrentState - current state of the simulation Returns ------- @@ -367,7 +645,6 @@ class AbstractGamma(Operator): Parameters ---------- state : CurrentState - current state of the simulation Returns ------- @@ -414,7 +691,6 @@ class NonLinearOperator(Operator): Parameters ---------- state : CurrentState - current state of the simulation Returns ------- @@ -450,39 +726,19 @@ class EnvelopeNonLinearOperator(NonLinearOperator): ################################################## -class ConservedQuantity(Operator): - @classmethod - def create( - cls, raman_op: AbstractGamma, gamma_op: AbstractGamma, loss_op: AbstractLoss, w: np.ndarray - ) -> ConservedQuantity: - logger = get_logger(__name__) - raman = not isinstance(raman_op, NoRaman) - loss = not isinstance(raman_op, NoLoss) - if raman and loss: - logger.debug("Conserved quantity : photon number with loss") - return PhotonNumberLoss(w, gamma_op, loss_op) - elif raman: - logger.debug("Conserved quantity : photon number without loss") - return PhotonNumberNoLoss(w, gamma_op) - elif loss: - logger.debug("Conserved quantity : energy with loss") - return EnergyLoss(w, loss_op) - else: - logger.debug("Conserved quantity : energy without loss") - return EnergyNoLoss(w) - +class AbstractConservedQuantity(Operator): @abstractmethod def __call__(self, state: CurrentState) -> float: pass -class NoConservedQuantity(ConservedQuantity): +class NoConservedQuantity(AbstractConservedQuantity): def __call__(self, state: CurrentState) -> float: return 0.0 -class PhotonNumberLoss(ConservedQuantity): - def __init__(self, w: np.ndarray, gamma_op: AbstractGamma, loss_op=AbstractLoss): +class PhotonNumberLoss(AbstractConservedQuantity): + def __init__(self, w: np.ndarray, gamma_op: AbstractGamma, loss_op: AbstractLoss): self.w = w self.dw = w[1] - w[0] self.gamma_op = gamma_op @@ -494,7 +750,7 @@ class PhotonNumberLoss(ConservedQuantity): ) -class PhotonNumberNoLoss(ConservedQuantity): +class PhotonNumberNoLoss(AbstractConservedQuantity): def __init__(self, w: np.ndarray, gamma_op: AbstractGamma): self.w = w self.dw = w[1] - w[0] @@ -504,7 +760,7 @@ class PhotonNumberNoLoss(ConservedQuantity): return pulse.photon_number(state.spec2, self.w, self.dw, self.gamma_op(state)) -class EnergyLoss(ConservedQuantity): +class EnergyLoss(AbstractConservedQuantity): def __init__(self, w: np.ndarray, loss_op: AbstractLoss): self.dw = w[1] - w[0] self.loss_op = loss_op @@ -515,9 +771,35 @@ class EnergyLoss(ConservedQuantity): ) -class EnergyNoLoss(ConservedQuantity): +class EnergyNoLoss(AbstractConservedQuantity): def __init__(self, w: np.ndarray): self.dw = w[1] - w[0] def __call__(self, state: CurrentState) -> float: return pulse.pulse_energy(math.abs2(state.C_to_A_factor * state.spectrum), self.dw) + + +def conserved_quantity( + adapt_step_size: bool, + raman_op: AbstractGamma, + gamma_op: AbstractGamma, + loss_op: AbstractLoss, + w: np.ndarray, +) -> AbstractConservedQuantity: + if not adapt_step_size: + return NoConservedQuantity() + logger = get_logger(__name__) + raman = not isinstance(raman_op, NoRaman) + loss = not isinstance(raman_op, NoLoss) + if raman and loss: + logger.debug("Conserved quantity : photon number with loss") + return PhotonNumberLoss(w, gamma_op, loss_op) + elif raman: + logger.debug("Conserved quantity : photon number without loss") + return PhotonNumberNoLoss(w, gamma_op) + elif loss: + logger.debug("Conserved quantity : energy with loss") + return EnergyLoss(w, loss_op) + else: + logger.debug("Conserved quantity : energy without loss") + return EnergyNoLoss(w) diff --git a/src/scgenerator/parameter.py b/src/scgenerator/parameter.py index 5871509..6c358c9 100644 --- a/src/scgenerator/parameter.py +++ b/src/scgenerator/parameter.py @@ -1,5 +1,6 @@ from __future__ import annotations + import datetime as datetime_module import enum import os @@ -16,7 +17,7 @@ from . import env, legacy, utils from .const import MANDATORY_PARAMETERS, PARAM_FN, VALID_VARIABLE, __version__ from .evaluator import Evaluator, pdict from .logger import get_logger -from .operators import LinearOperator, NonLinearOperator +from .operators import AbstractConservedQuantity, LinearOperator, NonLinearOperator from .utils import fiber_folder, update_path_name from .variationer import VariationDescriptor, Variationer @@ -261,7 +262,6 @@ class Parameters: name: str = Parameter(string, default="no name") prev_data_dir: str = Parameter(string) recovery_data_dir: str = Parameter(string) - previous_config_file: str = Parameter(string) output_path: Path = Parameter(type_checker(Path), default=Path("sc_data"), converter=Path) # # fiber @@ -334,7 +334,7 @@ class Parameters: tolerated_error: float = Parameter(in_range_excl(1e-15, 1e-3), default=1e-11) step_size: float = Parameter(non_negative(float, int), default=0) interpolation_range: tuple[float, float] = Parameter(float_pair) - interpolation_degree: int = Parameter(positive(int), default=8) + interpolation_degree: int = Parameter(non_negative(int)) prev_sim_dir: str = Parameter(string) recovery_last_stored: int = Parameter(non_negative(int), default=0) parallel: bool = Parameter(boolean, default=True) @@ -343,6 +343,9 @@ class Parameters: # computed linear_operator: LinearOperator = Parameter(type_checker(LinearOperator)) nonlinear_operator: NonLinearOperator = Parameter(type_checker(NonLinearOperator)) + conserved_quantity: AbstractConservedQuantity = Parameter( + type_checker(AbstractConservedQuantity) + ) field_0: np.ndarray = Parameter(type_checker(np.ndarray)) spec_0: np.ndarray = Parameter(type_checker(np.ndarray)) beta2: float = Parameter(type_checker(int, float)) @@ -372,7 +375,13 @@ class Parameters: self._evaluator.set(self._param_dico) def __repr__(self) -> str: - return "Parameter(" + ", ".join(f"{k}={v}" for k, v in self.dump_dict().items()) + ")" + return "Parameter(" + ", ".join(self.__repr_list__()) + ")" + + def __pformat__(self) -> str: + return "\n".join(["Parameter(", *list(self.__repr_list__()), ")"]) + + def __repr_list__(self) -> Iterator[str]: + yield from (f"{k}={v}" for k, v in self.dump_dict().items()) def dump_dict(self) -> dict[str, Any]: param = Parameters.strip_params_dict(self._param_dico) @@ -427,7 +436,7 @@ class Parameters: "nonlinear_op", "linear_op", } - types = (np.ndarray, float, int, str, list, tuple, dict) + types = (np.ndarray, float, int, str, list, tuple, dict, Path) out = {} for key, value in dico.items(): if key in forbiden_keys: @@ -436,8 +445,12 @@ class Parameters: continue if isinstance(value, dict): out[key] = Parameters.strip_params_dict(value) + elif isinstance(value, Path): + out[key] = str(value) elif isinstance(value, np.ndarray) and value.dtype == complex: continue + elif isinstance(value, np.ndarray): + out[key] = value.tolist() else: out[key] = value diff --git a/src/scgenerator/physics/fiber.py b/src/scgenerator/physics/fiber.py index 3997fd3..1789dec 100644 --- a/src/scgenerator/physics/fiber.py +++ b/src/scgenerator/physics/fiber.py @@ -1087,7 +1087,7 @@ def capillary_loss(wl: np.ndarray, he_mode: tuple[int, int], core_radius: float) np.ndarray loss in 1/m """ - chi_silica = mat.sellmeier(wl, utils.load_material_dico("silica")) + chi_silica = abs(mat.sellmeier(wl, utils.load_material_dico("silica"))) nu_n = 0.5 * (chi_silica + 2) / np.sqrt(chi_silica) return nu_n * (u_nm(*he_mode) * wl / pipi) ** 2 * core_radius ** -3 diff --git a/src/scgenerator/physics/materials.py b/src/scgenerator/physics/materials.py index 218b3ac..ef6611f 100644 --- a/src/scgenerator/physics/materials.py +++ b/src/scgenerator/physics/materials.py @@ -1,4 +1,4 @@ -from typing import Callable +from typing import Any, Callable import numpy as np import scipy.special @@ -18,6 +18,17 @@ def n_gas_2( """Returns the sqare of the index of refraction of the specified gas""" material_dico = utils.load_material_dico(gas_name) + n_gas_2 = fast_n_gas_2(wl_for_disp, pressure, temperature, ideal_gas, material_dico) + return n_gas_2 + + +def fast_n_gas_2( + wl_for_disp: np.ndarray, + pressure: float, + temperature: float, + ideal_gas: bool, + material_dico: dict[str, Any], +): if ideal_gas: n_gas_2 = sellmeier(wl_for_disp, material_dico, pressure, temperature) + 1 else: diff --git a/src/scgenerator/physics/simulate.py b/src/scgenerator/physics/simulate.py index 4a208c1..5b242e3 100644 --- a/src/scgenerator/physics/simulate.py +++ b/src/scgenerator/physics/simulate.py @@ -2,16 +2,21 @@ import multiprocessing import multiprocessing.connection import os from datetime import datetime -from pathlib import Path -from typing import Any, Generator, Type, Union, Optional from logging import Logger +from pathlib import Path +from typing import Any, Generator, Optional, Type, Union + import numpy as np from .. import utils +from ..errors import EvaluatorError from ..logger import get_logger +from ..operators import ( + AbstractConservedQuantity, + CurrentState, +) from ..parameter import Configuration, Parameters from ..pbar import PBars, ProgressBarActor, progress_worker -from ..operators import CurrentState, ConservedQuantity, NoConservedQuantity try: import ray @@ -35,7 +40,7 @@ class RK4IP: cons_qty: list[float] state: CurrentState - conserved_quantity_func: ConservedQuantity + conserved_quantity: AbstractConservedQuantity stored_spectra: list[np.ndarray] def __init__( @@ -82,22 +87,8 @@ class RK4IP: params.tolerated_error if self.params.adapt_step_size else self.params.step_size ) - self.set_cons_qty() self._setup_sim_parameters() - def set_cons_qty(self): - # Set up which quantity is conserved for adaptive step size - if self.params.adapt_step_size: - self.conserved_quantity_func = ConservedQuantity.create( - self.params.nonlinear_operator.raman_op, - self.params.nonlinear_operator.gamma_op, - self.params.linear_operator.loss_op, - self.params.w, - ) - else: - self.logger.debug(f"Using constant step size of {1e6*self.error_ok:.3f}μm") - self.conserved_quantity_func = NoConservedQuantity() - def _setup_sim_parameters(self): # making sure to keep only the z that we want self.z_stored = list(self.z_targets.copy()[0 : self.params.recovery_last_stored + 1]) @@ -106,7 +97,10 @@ class RK4IP: self.store_num = len(self.z_targets) # Setup initial values for every physical quantity that we want to track - C_to_A_factor = (self.params.A_eff_arr / self.params.A_eff_arr[0]) ** (1 / 4) + try: + C_to_A_factor = (self.params.A_eff_arr / self.params.A_eff_arr[0]) ** (1 / 4) + except EvaluatorError: + C_to_A_factor = 1.0 z = self.z_targets.pop(0) # Initial step size if self.params.adapt_step_size: @@ -124,7 +118,7 @@ class RK4IP: self.state.spectrum.copy() ] self.cons_qty = [ - self.conserved_quantity_func(self.state), + self.params.conserved_quantity(self.state), 0, ] self.size_fac = 2 ** (1 / 5) @@ -267,7 +261,7 @@ class RK4IP: new_state = self.state.replace(expD * (A_I + k1 / 6 + k2 / 3 + k3 / 3) + k4 / 6) if self.params.adapt_step_size: - self.cons_qty[step] = self.conserved_quantity_func(new_state) + self.cons_qty[step] = self.params.conserved_quantity(new_state) curr_p_change = np.abs(self.cons_qty[step - 1] - self.cons_qty[step]) cons_qty_change_ok = self.error_ok * self.cons_qty[step - 1] @@ -531,7 +525,7 @@ class MultiProcSimulations(Simulations, priority=1): super().run() def new_sim(self, params: Parameters): - self.queue.put((params,), block=True, timeout=None) + self.queue.put(params.dump_dict(), block=True, timeout=None) def finish(self): """0 means finished""" @@ -556,7 +550,7 @@ class MultiProcSimulations(Simulations, priority=1): if raw_data == 0: queue.task_done() return - (params,) = raw_data + params = Parameters(**raw_data) MutliProcRK4IP( params, p_queue, diff --git a/src/scgenerator/plotting.py b/src/scgenerator/plotting.py index 1df1851..4b1822b 100644 --- a/src/scgenerator/plotting.py +++ b/src/scgenerator/plotting.py @@ -174,7 +174,6 @@ def create_zoom_axis( # copy the plot content if plot: lines = axis.get_lines() - ymin, ymax = 0, 0 for line in lines: xdata = line.get_xdata() xdata, ind, _ = sort_axis(xdata, (*xlim, units.s)) @@ -1071,9 +1070,36 @@ def partial_plot(root: os.PathLike): spec_list = sorted( path.glob(SPEC1_FN.format("*")), key=lambda el: int(re.search("[0-9]+", el.name)[0]) ) + params = Parameters.load(path / "params.toml") + params.z_targets = params.z_targets[: len(spec_list)] raw_values = np.array([load_spectrum(s) for s in spec_list]) - specs = units.to_log2D(math.abs2(np.fft.fftshift(raw_values, axes=-1))) - fields = math.abs2(np.fft.ifft(raw_values)) - left.imshow(specs, origin="lower", aspect="auto", vmin=-60, interpolation="nearest") - right.imshow(fields, origin="lower", aspect="auto", interpolation="nearest") + + wl, z, values = transform_2D_propagation( + raw_values, + PlotRange( + 0.5 * params.interpolation_range[0] * 1e9, + 1.1 * params.interpolation_range[1] * 1e9, + "nm", + ), + params, + log="2D", + ) + left.imshow( + values, + origin="lower", + aspect="auto", + vmin=-60, + interpolation="nearest", + extent=get_extent(wl, z), + ) + + t, z, values = transform_2D_propagation( + np.fft.ifft(raw_values), + PlotRange(-10, 10, "ps"), + params, + log=False, + ) + right.imshow( + values, origin="lower", aspect="auto", interpolation="nearest", extent=get_extent(t, z) + ) return left, right diff --git a/src/scgenerator/scripts/__init__.py b/src/scgenerator/scripts/__init__.py index af28fc2..eadb2f1 100644 --- a/src/scgenerator/scripts/__init__.py +++ b/src/scgenerator/scripts/__init__.py @@ -13,7 +13,8 @@ from ..parameter import Configuration, Parameters from ..physics import fiber, units from ..plotting import plot_setup from ..spectra import SimulationSeries -from ..utils import _open_config, auto_crop, save_toml, simulations_list, translate_parameters +from ..utils import _open_config, auto_crop, save_toml, simulations_list +from ..legacy import translate_parameters def fingerprint(params: Parameters): diff --git a/src/scgenerator/utils.py b/src/scgenerator/utils.py index 151edf1..a90bae2 100644 --- a/src/scgenerator/utils.py +++ b/src/scgenerator/utils.py @@ -18,7 +18,8 @@ from typing import Any, Callable, MutableMapping, Sequence, TypeVar, Set import numpy as np import pkg_resources as pkg -import toml +import tomli +import tomli_w from .const import PARAM_FN, PARAM_SEPARATOR, SPEC1_FN, Z_FN, ROOT_PARAMETERS from .logger import get_logger @@ -47,8 +48,8 @@ class Paths: def get(cls, key): if key not in cls.paths: if os.path.exists("paths.toml"): - with open("paths.toml") as file: - paths_dico = toml.load(file) + with open("paths.toml", "rb") as file: + paths_dico = tomli.load(file) for k, v in paths_dico.items(): cls.paths[k] = v if key not in cls.paths: @@ -182,18 +183,18 @@ def load_toml(descr: os.PathLike) -> dict[str, Any]: descr = str(descr) if ":" in descr: path, entry = descr.split(":", 1) - with open(path) as file: - return toml.load(file)[entry] + with open(path, "rb") as file: + return tomli.load(file)[entry] else: - with open(descr) as file: - return toml.load(file) + with open(descr, "rb") as file: + return tomli.load(file) def save_toml(path: os.PathLike, dico): """saves a dictionary into a toml file""" path = conform_toml_path(path) - with open(path, mode="w") as file: - toml.dump(dico, file) + with open(path, mode="wb") as file: + tomli_w.dump(dico, file) return dico @@ -247,7 +248,7 @@ def load_material_dico(name: str) -> dict[str, Any]: ---------- material_dico : dict """ - return toml.loads(Paths.gets("materials"))[name] + return tomli.loads(Paths.gets("materials"))[name] def save_data(data: np.ndarray, data_dir: Path, file_name: str): @@ -478,8 +479,8 @@ def save_parameters( os.makedirs(file_path.parent, exist_ok=True) # save toml of the simulation - with open(file_path, "w") as file: - toml.dump(params, file, encoder=toml.TomlNumpyEncoder()) + with open(file_path, "wb") as file: + tomli_w.dump(params, file) return file_path