good progress on plasma, no luck for bug

src/scgenerator/evaluator.py

@@ -387,9 +387,9 @@ default_rules: list[Rule] = [
     Rule("gamma_op", operators.NoGamma, priorities=-1),
     Rule("ss_op", operators.SelfSteepening),
     Rule("ss_op", operators.NoSelfSteepening, priorities=-1),
-    Rule("spm_op", operators.SPM),
+    Rule("spm_op", operators.EnvelopeSPM),
     Rule("spm_op", operators.NoSPM, priorities=-1),
-    Rule("raman_op", operators.Raman),
+    Rule("raman_op", operators.EnvelopeRaman),
     Rule("raman_op", operators.NoRaman, priorities=-1),
     Rule("nonlinear_operator", operators.EnvelopeNonLinearOperator),
     Rule("loss_op", operators.CustomLoss, priorities=3),
@@ -402,6 +402,6 @@ default_rules: list[Rule] = [
     Rule("n_op", operators.HasanRefractiveIndex),
     Rule("dispersion_op", operators.ConstantPolyDispersion),
     Rule("dispersion_op", operators.DirectDispersion),
-    Rule("linear_operator", operators.LinearOperator),
+    Rule("linear_operator", operators.EnvelopeLinearOperator),
     Rule("conserved_quantity", operators.conserved_quantity),
 ]

src/scgenerator/operators.py

@@ -321,8 +321,7 @@ class ConstantPolyDispersion(AbstractDispersion):
     evaluating on the envelope
     """
 
-    coefs: np.ndarray
+    disp_arr: np.ndarray
-    w_c: np.ndarray
 
     def __init__(
         self,
@@ -336,14 +335,14 @@ class ConstantPolyDispersion(AbstractDispersion):
             raise OperatorError(
                 f"interpolation degree of degree {interpolation_degree} incompatible"
             )
-        self.coefs = fiber.dispersion_coefficients(w_for_disp, beta2_arr, w0, interpolation_degree)
+        coefs = fiber.dispersion_coefficients(w_for_disp, beta2_arr, w0, interpolation_degree)
-        self.w_c = w_c
+        w_power_fact = np.array(
-        self.w_power_fact = np.array(
             [math.power_fact(w_c, k) for k in range(2, interpolation_degree + 3)]
         )
+        self.disp_arr = fiber.fast_poly_dispersion_op(w_c, coefs, w_power_fact)
 
     def __call__(self, state: CurrentState = None) -> np.ndarray:
-        return fiber.fast_poly_dispersion_op(self.w_c, self.coefs, self.w_power_fact)
+        return self.disp_arr
 
 
 class ConstantDirectDispersion(AbstractDispersion):
@@ -413,6 +412,45 @@ class DirectDispersion(AbstractDispersion):
         return self.disp_arr
 
 
+##################################################
+################## WAVE VECTOR ###################
+##################################################
+
+
+class AbstractWaveVector(Operator):
+    @abstractmethod
+    def __call__(self, state: CurrentState) -> np.ndarray:
+        """returns the wave vector beta in the frequency domain
+
+        Parameters
+        ----------
+        state : CurrentState
+
+        Returns
+        -------
+        np.ndarray
+            wave vector
+        """
+
+
+class ConstantWaveVector(AbstractWaveVector):
+    beta_arr: np.ndarray
+
+    def __init__(
+        self,
+        n_op: ConstantRefractiveIndex,
+        w_for_disp: np.ndarray,
+        t_num: int,
+        dispersion_ind: np.ndarray,
+    ):
+
+        self.beta_arr = np.zeros(t_num, dtype=float)
+        self.beta_arr[dispersion_ind] = fiber.beta(w_for_disp, n_op())[2:-2]
+
+    def __call__(self, state: CurrentState) -> np.ndarray:
+        return self.beta_arr
+
+
 ##################################################
 ###################### LOSS ######################
 ##################################################
@@ -477,31 +515,6 @@ class CustomLoss(ConstantLoss):
         return self.arr
 
 
-##################################################
-##################### LINEAR #####################
-##################################################
-
-
-class LinearOperator(Operator):
-    def __init__(self, dispersion_op: AbstractDispersion, loss_op: AbstractLoss):
-        self.dispersion_op = dispersion_op
-        self.loss_op = loss_op
-
-    def __call__(self, state: CurrentState) -> np.ndarray:
-        """returns the linear operator to be multiplied by the spectrum in the frequency domain
-
-        Parameters
-        ----------
-        state : CurrentState
-
-        Returns
-        -------
-        np.ndarray
-            linear component
-        """
-        return self.dispersion_op(state) - self.loss_op(state) / 2
-
-
 ##################################################
 ###################### RAMAN #####################
 ##################################################
@@ -530,7 +543,7 @@ class NoRaman(NoOp, AbstractRaman):
         return self.arr_const
 
 
-class Raman(AbstractRaman):
+class EnvelopeRaman(AbstractRaman):
     def __init__(self, raman_type: str, t: np.ndarray):
         self.hr_w = fiber.delayed_raman_w(t, raman_type)
         self.f_r = 0.245 if raman_type == "agrawal" else 0.18
@@ -539,6 +552,16 @@ class Raman(AbstractRaman):
         return self.f_r * np.fft.ifft(self.hr_w * np.fft.fft(state.field2))
 
 
+class FullFieldRaman(AbstractRaman):
+    def __init__(self, raman_type: str, t: np.ndarray, chi3: float):
+        self.hr_w = fiber.delayed_raman_w(t, raman_type)
+        self.f_r = 0.245 if raman_type == "agrawal" else 0.18
+        self.multiplier = units.epsilon0 * chi3 * self.f_r
+
+    def __call__(self, state: CurrentState) -> np.ndarray:
+        return self.multiplier * np.fft.ifft(np.fft.fft(state.field2) * self.hr_w)
+
+
 ##################################################
 ####################### SPM ######################
 ##################################################
@@ -567,7 +590,15 @@ class NoSPM(NoOp, AbstractSPM):
         return self.arr_const
 
 
-class SPM(AbstractSPM):
+class EnvelopeSPM(AbstractSPM):
+    def __init__(self, raman_op: AbstractRaman):
+        self.fraction = 1 - raman_op.f_r
+
+    def __call__(self, state: CurrentState) -> np.ndarray:
+        return self.fraction * state.field2
+
+
+class FullFieldSPM(AbstractSPM):
     def __init__(self, raman_op: AbstractRaman):
         self.fraction = 1 - raman_op.f_r
 
@@ -605,7 +636,7 @@ class SelfSteepening(AbstractSelfSteepening):
     def __init__(self, w_c: np.ndarray, w0: float):
         self.arr = w_c / w0
 
-    def __call__(self, state: CurrentState) -> np.ndarray:
+    def __call__(self, state: CurrentState = None) -> np.ndarray:
         return self.arr
 
 
@@ -620,6 +651,7 @@ class AbstractGamma(Operator):
         """returns the gamma component
 
         Parameters
+
         ----------
         state : CurrentState
 
@@ -655,47 +687,14 @@ class ConstantGamma(AbstractGamma):
 ##################### PLASMA #####################
 ##################################################
 
-##################################################
+
-#################### NONLINEAR ###################
+class Plasma(Operator):
-##################################################
+    pass
 
 
-class NonLinearOperator(Operator):
+class NoPlasma(NoOp, Plasma):
-    @abstractmethod
     def __call__(self, state: CurrentState) -> np.ndarray:
-        """returns the nonlinear operator applied on the spectrum in the frequency domain
+        return self.arr_const
-
-        Parameters
-        ----------
-        state : CurrentState
-
-        Returns
-        -------
-        np.ndarray
-            nonlinear component
-        """
-
-
-class EnvelopeNonLinearOperator(NonLinearOperator):
-    def __init__(
-        self,
-        gamma_op: AbstractGamma,
-        ss_op: AbstractSelfSteepening,
-        spm_op: AbstractSPM,
-        raman_op: AbstractRaman,
-    ):
-        self.gamma_op = gamma_op
-        self.ss_op = ss_op
-        self.spm_op = spm_op
-        self.raman_op = raman_op
-
-    def __call__(self, state: CurrentState) -> np.ndarray:
-        return (
-            -1j
-            * self.gamma_op(state)
-            * (1 + self.ss_op(state))
-            * np.fft.fft(state.field * (self.spm_op(state) + self.raman_op(state)))
-        )
 
 
 ##################################################
@@ -780,3 +779,108 @@ def conserved_quantity(
     else:
         logger.debug("Conserved quantity : energy without loss")
         return EnergyNoLoss(w)
+
+
+##################################################
+##################### LINEAR #####################
+##################################################
+
+
+class EnvelopeLinearOperator(Operator):
+    def __init__(self, dispersion_op: AbstractDispersion, loss_op: AbstractLoss):
+        self.dispersion_op = dispersion_op
+        self.loss_op = loss_op
+
+    def __call__(self, state: CurrentState) -> np.ndarray:
+        """returns the linear operator to be multiplied by the spectrum in the frequency domain
+
+        Parameters
+        ----------
+        state : CurrentState
+
+        Returns
+        -------
+        np.ndarray
+            linear component
+        """
+        return self.dispersion_op(state) - self.loss_op(state) / 2
+
+
+class FullFieldLinearOperator(Operator):
+    def __init__(
+        self,
+        wave_vector: AbstractWaveVector,
+        loss_op: AbstractLoss,
+        reference_velocity: float,
+        w: np.ndarray,
+    ):
+        self.delay = w / reference_velocity
+        self.wave_vector = wave_vector
+        self.loss_op = loss_op
+
+    def __call__(self, state: CurrentState) -> np.ndarray:
+        """returns the linear operator to be multiplied by the spectrum in the frequency domain
+
+        Parameters
+        ----------
+        state : CurrentState
+
+        Returns
+        -------
+        np.ndarray
+            linear component
+        """
+        return 1j * (self.wave_vector(state) - self.delay) - self.loss_op(state) / 2
+
+
+##################################################
+#################### NONLINEAR ###################
+##################################################
+
+
+class NonLinearOperator(Operator):
+    @abstractmethod
+    def __call__(self, state: CurrentState) -> np.ndarray:
+        """returns the nonlinear operator applied on the spectrum in the frequency domain
+
+        Parameters
+        ----------
+        state : CurrentState
+
+        Returns
+        -------
+        np.ndarray
+            nonlinear component
+        """
+
+
+class EnvelopeNonLinearOperator(NonLinearOperator):
+    def __init__(
+        self,
+        gamma_op: AbstractGamma,
+        ss_op: AbstractSelfSteepening,
+        spm_op: AbstractSPM,
+        raman_op: AbstractRaman,
+    ):
+        self.gamma_op = gamma_op
+        self.ss_op = ss_op
+        self.spm_op = spm_op
+        self.raman_op = raman_op
+
+    def __call__(self, state: CurrentState) -> np.ndarray:
+        return (
+            -1j
+            * self.gamma_op(state)
+            * (1 + self.ss_op(state))
+            * np.fft.fft(state.field * (self.spm_op(state) + self.raman_op(state)))
+        )
+
+
+class FullFieldNonLinearOperator(NonLinearOperator):
+    def __init__(self, raman_op: AbstractRaman, spm_op: AbstractSPM, plasma_op: Plasma):
+        self.raman_op = raman_op
+        self.spm_op = spm_op
+        self.plasma_op = plasma_op
+
+    def __call__(self, state: CurrentState) -> np.ndarray:
+        return self.spm_op(state) + self.raman_op(state) + self.plasma_op(state)

src/scgenerator/parameter.py

@@ -18,7 +18,7 @@ from .const import MANDATORY_PARAMETERS, PARAM_FN, VALID_VARIABLE, __version__
 from .errors import EvaluatorError
 from .evaluator import Evaluator
 from .logger import get_logger
-from .operators import AbstractConservedQuantity, LinearOperator, NonLinearOperator
+from .operators import AbstractConservedQuantity, EnvelopeLinearOperator, NonLinearOperator
 from .utils import fiber_folder, update_path_name
 from .variationer import VariationDescriptor, Variationer
 
@@ -367,7 +367,7 @@ class Parameters:
     worker_num: int = Parameter(positive(int))
 
     # computed
-    linear_operator: LinearOperator = Parameter(type_checker(LinearOperator))
+    linear_operator: EnvelopeLinearOperator = Parameter(type_checker(EnvelopeLinearOperator))
     nonlinear_operator: NonLinearOperator = Parameter(type_checker(NonLinearOperator))
     conserved_quantity: AbstractConservedQuantity = Parameter(
         type_checker(AbstractConservedQuantity)

src/scgenerator/physics/materials.py

@@ -8,7 +8,7 @@ from .. import utils
 from ..cache import np_cache
 from ..logger import get_logger
 from . import units
-from .units import NA, c, e, hbar, kB, me
+from .units import NA, c, e, hbar, kB, me, epsilon0
 
 
 @np_cache
@@ -233,63 +233,35 @@ def gas_n2(gas_name: str, pressure: float, temperature: float) -> float:
     return non_linear_refractive_index(utils.load_material_dico(gas_name), pressure, temperature)
 
 
-def adiabadicity(w: np.ndarray, I: float, field: np.ndarray) -> np.ndarray:
+def gas_chi3(gas_name: str, wavelength: float, pressure: float, temperature: float) -> float:
-    return w * np.sqrt(2 * me * I) / (e * np.abs(field))
+    """returns the chi3 of a particular material
+
+    Parameters
+    ----------
+    gas_name : str
+        [description]
+    pressure : float
+        [description]
+    temperature : float
+        [description]
+
+    Returns
+    -------
+    float
+        [description]
+    """
+    mat_dico = utils.load_material_dico(gas_name)
+    chi = sellmeier(wavelength, mat_dico, pressure=pressure, temperature=temperature)
+    return n2_to_chi3(
+        non_linear_refractive_index(mat_dico, pressure, temperature), np.sqrt(chi + 1)
+    )
 
 
-def free_electron_density(
+def n2_to_chi3(n2: float, n0: float) -> float:
-    t: np.ndarray, field: np.ndarray, N0: float, rate_func: Callable[[np.ndarray], np.ndarray]
+    return n2 / 3.0 * 4 * epsilon0 * n0 * c ** 2
-) -> np.ndarray:
-    return N0 * (1 - np.exp(-cumulative_trapezoid(rate_func(field), t, initial=0)))
 
 
-def ionization_rate_ADK(
+def chi3_to_n2(chi3: float, n0: float) -> float:
-    ionization_energy: float, atomic_number
+    return 3.0 * chi3 / (4.0 * epsilon0 * c * n0 ** 2)
-) -> Callable[[np.ndarray], np.ndarray]:
-    Z = -(atomic_number - 1) * e
-
-    omega_p = ionization_energy / hbar
-    nstar = Z * np.sqrt(2.1787e-18 / ionization_energy)
-
-    def omega_t(field):
-        return e * np.abs(field) / np.sqrt(2 * me * ionization_energy)
-
-    Cnstar = 2 ** (2 * nstar) / (scipy.special.gamma(nstar + 1) ** 2)
-    omega_pC = omega_p * Cnstar
-
-    def rate(field: np.ndarray) -> np.ndarray:
-        opt4 = 4 * omega_p / omega_t(field)
-        return omega_pC * opt4 ** (2 * nstar - 1) * np.exp(-opt4 / 3)
-
-    return rate
 
 
-class Plasma:
-    def __init__(self, t: np.ndarray, ionization_energy: float, atomic_number: int):
-        self.t = t
-        self.Ip = ionization_energy
-        self.atomic_number = atomic_number
-        self.rate = ionization_rate_ADK(self.Ip, self.atomic_number)
-
-    def __call__(self, field: np.ndarray, N0: float) -> np.ndarray:
-        """returns the number density of free electrons as function of time
-
-        Parameters
-        ----------
-        field : np.ndarray
-            electric field in V/m
-        N0 : float
-            total number density of matter
-
-        Returns
-        -------
-        np.ndarray
-            number density of free electrons as function of time
-        """
-        Ne = free_electron_density(self.t, field, N0, self.rate)
-        return cumulative_trapezoid(
-            np.gradient(Ne, self.t) * self.Ip / field
-            + e ** 2 / me * cumulative_trapezoid(Ne * field, self.t, initial=0),
-            self.t,
-            initial=0,
-        )

src/scgenerator/physics/plasma.py

@@ -0,0 +1,69 @@
+import numpy as np
+import scipy.special
+from scipy.integrate import cumulative_trapezoid
+
+from .units import e, hbar, me
+
+
+class IonizationRate:
+    def __call__(self, field: np.ndarray) -> np.ndarray:
+        ...
+
+
+class IonizationRateADK(IonizationRate):
+    def __init__(self, ionization_energy: float, atomic_number: int):
+        self.Z = -(atomic_number - 1) * e
+        self.ionization_energy = ionization_energy
+        self.omega_p = ionization_energy / hbar
+
+        self.nstar = self.Z * np.sqrt(2.1787e-18 / ionization_energy)
+        Cnstar = 2 ** (2 * self.nstar) / (scipy.special.gamma(self.nstar + 1) ** 2)
+        self.omega_pC = self.omega_p * Cnstar
+
+    def omega_t(self, field):
+        return e * np.abs(field) / np.sqrt(2 * me * self.ionization_energy)
+
+    def __call__(self, field: np.ndarray) -> np.ndarray:
+        opt4 = 4 * self.omega_p / self.omega_t(field)
+        return self.omega_pC * opt4 ** (2 * self.nstar - 1) * np.exp(-opt4 / 3)
+
+
+class Plasma:
+    def __init__(self, t: np.ndarray, ionization_energy: float, atomic_number: int):
+        self.t = t
+        self.Ip = ionization_energy
+        self.atomic_number = atomic_number
+        self.rate = IonizationRateADK(self.Ip, self.atomic_number)
+
+    def __call__(self, field: np.ndarray, N0: float) -> np.ndarray:
+        """returns the number density of free electrons as function of time
+
+        Parameters
+        ----------
+        field : np.ndarray
+            electric field in V/m
+        N0 : float
+            total number density of matter
+
+        Returns
+        -------
+        np.ndarray
+            number density of free electrons as function of time
+        """
+        Ne = free_electron_density(self.t, field, N0, self.rate)
+        return cumulative_trapezoid(
+            np.gradient(Ne, self.t) * self.Ip / field
+            + e ** 2 / me * cumulative_trapezoid(Ne * field, self.t, initial=0),
+            self.t,
+            initial=0,
+        )
+
+
+def adiabadicity(w: np.ndarray, I: float, field: np.ndarray) -> np.ndarray:
+    return w * np.sqrt(2 * me * I) / (e * np.abs(field))
+
+
+def free_electron_density(
+    t: np.ndarray, field: np.ndarray, N0: float, rate: IonizationRate
+) -> np.ndarray:
+    return N0 * (1 - np.exp(-cumulative_trapezoid(rate(field), t, initial=0)))

src/scgenerator/physics/simulate.py

@@ -10,7 +10,7 @@ import numpy as np
 
 from .. import utils
 from ..logger import get_logger
-from ..operators import AbstractConservedQuantity, CurrentState
+from ..operators import CurrentState
 from ..parameter import Configuration, Parameters
 from ..pbar import PBars, ProgressBarActor, progress_worker
 
@@ -36,7 +36,6 @@ class RK4IP:
     cons_qty: list[float]
 
     state: CurrentState
-    conserved_quantity: AbstractConservedQuantity
     stored_spectra: list[np.ndarray]
 
     def __init__(
@@ -229,7 +228,7 @@ class RK4IP:
                 store = True
                 self.state.h = self.z_targets[0] - self.state.z
 
-    def take_step(self, step: int) -> tuple[float, float, np.ndarray]:
+    def take_step(self, step: int) -> float:
         """computes a new spectrum, whilst adjusting step size if required, until the error estimation
         validates the new spectrum. Saves the result in the internal state attribute