Working Erk43

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),
 ]

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):

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()

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(

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):
+class ERK43(RK4IPSD):
-    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
-
     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.current_error = compute_diff(new_coarse, new_fine)
-                self.next_h_factor = max(
+                if self.current_error > 0.0:
+                    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
+                else:
-                if self.current_error <= self.tolerated_error:
+                    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(RK4IPSD):
-class ERK54(ERK43):
     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
+                k1 = h * expD2 * k7
-                k2 = self.nl(A_I + 0.5 * h * k1)
+                k2 = h * self.nl(A_I + 0.5 * k1)
-                k3 = expD4p * self.nl(expD4m * (A_I + h / 16 * (3 * k1 + k2)))
+                k3 = h * expD4p * self.nl(expD4m * (A_I + 0.0625 * (3 * k1 + k2)))
-                k4 = self.nl(A_I + h / 4 * (k1 - k2 + 4 * k3))
+                k4 = h * self.nl(A_I + 0.25 * (k1 - k2 + 4 * k3))
-                k5 = expD4m * self.nl(expD4p * (A_I + 3 * h / 16 * (k1 + 3 * k4)))
+                k5 = h * expD4m * self.nl(expD4p * (A_I + 0.1875 * (k1 + 3 * k4)))
-                k6 = self.nl(expD2 * (A_I + h / 7 * (-2 * k1 + k2 + 12 * k3 - 12 * k4 + 8 * k5)))
+                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))
+                    expD2 * (A_I + 1 / 90 * (7 * k1 + 32 * k3 + 12 * k4 + 32 * k5)) + 7 / 90 * k6
-                    + 7 * h / 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.current_error = compute_diff(new_coarse, new_fine)
-                self.next_h_factor = max(
+                next_h_factor = max(
-                    0.5, min(2.0, (self.tolerated_error / self.current_error) ** 0.25)
+                    0.5, min(2.0, (self.tolerated_error / self.current_error) ** 0.2)
                 )
-                h_next_step = self.next_h_factor * h
+                h_next_step = next_h_factor * h
-                if self.current_error <= self.tolerated_error:
+                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())