diff --git a/src/scgenerator/evaluator.py b/src/scgenerator/evaluator.py index a34449e..ccccd6e 100644 --- a/src/scgenerator/evaluator.py +++ b/src/scgenerator/evaluator.py @@ -378,7 +378,6 @@ default_rules: list[Rule] = [ Rule("n_op", operators.HasanRefractiveIndex), Rule("raman_op", operators.NoRaman, priorities=-1), Rule("loss_op", operators.NoLoss, priorities=-1), - Rule("plasma_op", operators.NoPlasma, priorities=-1), # solvers Rule("integrator", solver.ConstantStepIntegrator, conditions=dict(adapt_step_size=False)), Rule( @@ -453,5 +452,7 @@ full_field_rules = default_rules + [ "linear_operator", operators.FullFieldLinearOperator, ), + Rule("plasma_op", operators.Plasma, conditions=dict(photoionization=True)), + Rule("plasma_op", operators.NoPlasma, priorities=-1), Rule("nonlinear_operator", operators.FullFieldNonLinearOperator), ] diff --git a/src/scgenerator/operators.py b/src/scgenerator/operators.py index 22dd39f..111afac 100644 --- a/src/scgenerator/operators.py +++ b/src/scgenerator/operators.py @@ -13,7 +13,7 @@ from scipy.interpolate import interp1d from . import math from .logger import get_logger -from .physics import fiber, materials, pulse, units +from .physics import fiber, materials, pulse, units, plasma from .utils import load_material_dico @@ -198,6 +198,9 @@ class NoOpFreq(Operator): class AbstractGas(ABC): + gas_name: str + material_dico: dict[str, Any] + @abstractmethod def pressure(self, state: CurrentState) -> float: """returns the pressure at the current @@ -254,11 +257,12 @@ class ConstantGas(AbstractGas): ideal_gas: bool, wl_for_disp: np.ndarray, ): - gas_dico = load_material_dico(gas_name) + self.material_dico = load_material_dico(gas_name) + self.gas_name = 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 + pressure=pressure, temperature=temperature, material_dico=self.material_dico ) else: self.number_density_const = self.pressure_const / (units.kB * temperature) @@ -277,11 +281,9 @@ class ConstantGas(AbstractGas): class PressureGradientGas(AbstractGas): - name: str p_in: float p_out: float temperature: float - gas_dico: dict[str, Any] def __init__( self, @@ -292,10 +294,10 @@ class PressureGradientGas(AbstractGas): ideal_gas: bool, wl_for_disp: np.ndarray, ): - self.name = gas_name + self.gas_name = gas_name self.p_in = pressure_in self.p_out = pressure_out - self.gas_dico = load_material_dico(gas_name) + self.material_dico = load_material_dico(gas_name) self.temperature = temperature self.ideal_gas = ideal_gas self.wl_for_disp = wl_for_disp @@ -310,12 +312,16 @@ class PressureGradientGas(AbstractGas): return materials.number_density_van_der_waals( pressure=self.pressure(state), temperature=self.temperature, - material_dico=self.gas_dico, + material_dico=self.material_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 + self.wl_for_disp, + self.pressure(state), + self.temperature, + self.ideal_gas, + self.material_dico, ) @@ -793,7 +799,20 @@ class ConstantGamma(AbstractGamma): class Plasma(Operator): - pass + mat_plasma: plasma.Plasma + gas_op: AbstractGas + + def __init__(self, dt: float, gas_op: AbstractGas): + self.gas_op = gas_op + self.mat_plasma = plasma.Plasma( + dt, + self.gas_op.material_dico["ionization_energy"], + self.gas_op.material_dico["atomic_number"], + ) + + def __call__(self, state: CurrentState) -> np.ndarray: + N0 = self.gas_op.number_density(state) + return self.mat_plasma(state.field, N0) class NoPlasma(NoOpTime, Plasma): diff --git a/src/scgenerator/parameter.py b/src/scgenerator/parameter.py index 9b3b8f1..c1c8815 100644 --- a/src/scgenerator/parameter.py +++ b/src/scgenerator/parameter.py @@ -422,6 +422,13 @@ class Parameters: def __repr_list__(self) -> Iterator[str]: yield from (f"{k}={v}" for k, v in self.dump_dict().items()) + def __getstate__(self) -> dict[str, Any]: + return self.dump_dict(add_metadata=False) + + def __setstate__(self, dumped_dict: dict[str, Any]): + self._param_dico = dumped_dict + self.__post_init__() + def dump_dict(self, compute=True, add_metadata=True) -> dict[str, Any]: if compute: self.compute_in_place() diff --git a/src/scgenerator/physics/simulate.py b/src/scgenerator/physics/simulate.py index f1d057e..8f99cdf 100644 --- a/src/scgenerator/physics/simulate.py +++ b/src/scgenerator/physics/simulate.py @@ -164,12 +164,11 @@ class RK4IP: state = self.init_state.copy() yield len(self.stored_spectra) - 1, state if self.params.adapt_step_size: - integrator = solver.ConservedQuantityIntegrator( + integrator = solver.RK4IPSD( state, self.params.linear_operator, self.params.nonlinear_operator, self.params.tolerated_error, - self.params.conserved_quantity, ) else: integrator = solver.ConstantStepIntegrator( diff --git a/src/scgenerator/solver.py b/src/scgenerator/solver.py index 83408ab..beba2dc 100644 --- a/src/scgenerator/solver.py +++ b/src/scgenerator/solver.py @@ -4,9 +4,9 @@ import logging from abc import abstractmethod from typing import Iterator, Type +import numba import numpy as np -from . import math from .logger import get_logger from .operators import ( AbstractConservedQuantity, @@ -196,7 +196,7 @@ class RK4IPSD(Integrator): self.nonlinear_operator(new_fine_inter_state), ) new_coarse = self.take_step(h, self.state.spectrum, lin, nonlin) - self.current_error = self.compute_diff(new_coarse, new_fine) + self.current_error = compute_diff(new_coarse, new_fine) if self.current_error > 2 * self.tolerated_error: h_next_step = h * 0.5 @@ -218,42 +218,14 @@ class RK4IPSD(Integrator): ) -> np.ndarray: return rk4ip_step(self.nonlinear_operator, self.state, spec, h, lin, nonlin) - def compute_diff(self, coarse_spec: np.ndarray, fine_spec: np.ndarray) -> float: - return np.sqrt(math.abs2(coarse_spec - fine_spec).sum() / math.abs2(fine_spec).sum()) - def values(self) -> dict[str, float]: return dict( - step=self.state.step, z=self.state.z, local_error=self.current_error, - next_h_factor=self.next_h_factor, ) -class ERK43(Integrator): - state: CurrentState - linear_operator: LinearOperator - nonlinear_operator: NonLinearOperator - tolerated_error: float - dt: float - current_error: float - next_h_factor = 1.0 - - def __init__( - self, - init_state: CurrentState, - linear_operator: LinearOperator, - nonlinear_operator: NonLinearOperator, - tolerated_error: float, - dw: float, - ): - self.state = init_state - self.linear_operator = linear_operator - self.nonlinear_operator = nonlinear_operator - self.dw = dw - self.tolerated_error = tolerated_error - self.current_error = 0.0 - +class ERK43(RK4IPSD): def __iter__(self) -> Iterator[CurrentState]: h_next_step = self.state.current_step_size k5 = self.nonlinear_operator(self.state) @@ -276,13 +248,17 @@ class ERK43(Integrator): new_coarse = r + h / 30 * (2 * k4 + 3 * tmp_k5) - self.current_error = np.sqrt(self.dw * math.abs2(new_fine - new_coarse).sum()) - self.next_h_factor = max( - 0.5, min(2.0, (self.tolerated_error / self.current_error) ** 0.25) - ) - h_next_step = self.next_h_factor * h - if self.current_error <= self.tolerated_error: + self.current_error = compute_diff(new_coarse, new_fine) + if self.current_error > 0.0: + next_h_factor = max( + 0.5, min(2.0, (self.tolerated_error / self.current_error) ** 0.25) + ) + else: + next_h_factor = 2.0 + h_next_step = next_h_factor * h + if self.current_error <= 2 * self.tolerated_error: break + h_next_step = min(0.9, next_h_factor) * h self.logger.info( f"step {self.state.step} rejected : {h=}, {self.current_error=}, {h_next_step=}" ) @@ -291,16 +267,8 @@ class ERK43(Integrator): self.state.spectrum = new_fine yield self.accept_step(self.state, h, h_next_step) - def values(self) -> dict[str, float]: - return dict( - step=self.state.step, - z=self.state.z, - local_error=self.current_error, - next_h_factor=self.next_h_factor, - ) - -class ERK54(ERK43): +class ERK54(RK4IPSD): def __iter__(self) -> Iterator[CurrentState]: self.logger.info("using ERK54") h_next_step = self.state.current_step_size @@ -315,29 +283,31 @@ class ERK54(ERK43): expD4m = 1 / expD4p A_I = expD2 * self.state.spectrum - k1 = expD2 * k7 - k2 = self.nl(A_I + 0.5 * h * k1) - k3 = expD4p * self.nl(expD4m * (A_I + h / 16 * (3 * k1 + k2))) - k4 = self.nl(A_I + h / 4 * (k1 - k2 + 4 * k3)) - k5 = expD4m * self.nl(expD4p * (A_I + 3 * h / 16 * (k1 + 3 * k4))) - k6 = self.nl(expD2 * (A_I + h / 7 * (-2 * k1 + k2 + 12 * k3 - 12 * k4 + 8 * k5))) + k1 = h * expD2 * k7 + k2 = h * self.nl(A_I + 0.5 * k1) + k3 = h * expD4p * self.nl(expD4m * (A_I + 0.0625 * (3 * k1 + k2))) + k4 = h * self.nl(A_I + 0.25 * (k1 - k2 + 4 * k3)) + k5 = h * expD4m * self.nl(expD4p * (A_I + 0.1875 * (k1 + 3 * k4))) + k6 = h * self.nl( + expD2 * (A_I + 1 / 7 * (-2 * k1 + k2 + 12 * k3 - 12 * k4 + 8 * k5)) + ) new_fine = ( - expD2 * (A_I + h / 90 * (7 * k1 + 32 * k3 + 12 * k4 + 32 * k5)) - + 7 * h / 90 * k6 + expD2 * (A_I + 1 / 90 * (7 * k1 + 32 * k3 + 12 * k4 + 32 * k5)) + 7 / 90 * k6 ) tmp_k7 = self.nl(new_fine) new_coarse = ( - expD2 * (A_I + h / 42 * (3 * k1 + 16 * k3 + 4 * k4 + 16 * k5)) + h / 14 * k7 + expD2 * (A_I + 1 / 42 * (3 * k1 + 16 * k3 + 4 * k4 + 16 * k5)) + h / 14 * k7 ) - self.current_error = np.sqrt(self.dw * math.abs2(new_fine - new_coarse).sum()) - self.next_h_factor = max( - 0.5, min(2.0, (self.tolerated_error / self.current_error) ** 0.25) + self.current_error = compute_diff(new_coarse, new_fine) + next_h_factor = max( + 0.5, min(2.0, (self.tolerated_error / self.current_error) ** 0.2) ) - h_next_step = self.next_h_factor * h - if self.current_error <= self.tolerated_error: + h_next_step = next_h_factor * h + if self.current_error <= 2 * self.tolerated_error: break + h_next_step = min(0.9, next_h_factor) * h self.logger.info( f"step {self.state.step} rejected : {h=}, {self.current_error=}, {h_next_step=}" ) @@ -385,3 +355,10 @@ def rk4ip_step( k4 = h * nonlinear_operator(init_state.replace(expD * (A_I + k3))) return expD * (A_I + k1 / 6 + k2 / 3 + k3 / 3) + k4 / 6 + + +@numba.jit(nopython=True) +def compute_diff(coarse_spec: np.ndarray, fine_spec: np.ndarray) -> float: + diff = coarse_spec - fine_spec + diff2 = diff.imag ** 2 + diff.real ** 2 + return np.sqrt(diff2.sum() / (fine_spec.real ** 2 + fine_spec.imag ** 2).sum())