From 4491fc0248f8e898830983a2a056eaa7a61cc734 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Sierro?= <benoitsierro@bluewin.ch>
Date: Thu, 31 Mar 2022 15:19:34 +0200
Subject: [PATCH] working on pressure gradients

---
 src/scgenerator/evaluator.py |  2 ++
 src/scgenerator/operators.py | 30 +++++++++++++++++++++++++-----
 src/scgenerator/parameter.py |  6 +++---
 3 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/src/scgenerator/evaluator.py b/src/scgenerator/evaluator.py
index e617046..8ccdcc7 100644
--- a/src/scgenerator/evaluator.py
+++ b/src/scgenerator/evaluator.py
@@ -386,6 +386,7 @@ default_rules: list[Rule] = [
     Rule("n_op", operators.MarcatiliAdjustedRefractiveIndex),
     Rule("n_op", operators.HasanRefractiveIndex),
     Rule("gas_op", operators.ConstantGas),
+    Rule("gas_op", operators.PressureGradientGas),
     Rule("loss_op", operators.NoLoss, priorities=-1),
     Rule("conserved_quantity", operators.NoConservedQuantity, priorities=-1),
 ]
@@ -416,6 +417,7 @@ envelope_rules = default_rules + [
     Rule("gamma_op", operators.ConstantGamma, priorities=1),
     Rule("gamma_op", operators.ConstantScalarGamma),
     Rule("gamma_op", operators.NoGamma, priorities=-1),
+    Rule("gamma_op", operators.VariableScalarGamma, priorities=2),
     Rule("ss_op", operators.SelfSteepening),
     Rule("ss_op", operators.NoSelfSteepening, priorities=-1),
     Rule("spm_op", operators.NoEnvelopeSPM, priorities=-1),
diff --git a/src/scgenerator/operators.py b/src/scgenerator/operators.py
index f773c40..c8506aa 100644
--- a/src/scgenerator/operators.py
+++ b/src/scgenerator/operators.py
@@ -165,7 +165,7 @@ class ValueTracker(ABC):
         return self.__class__.__name__ + "(" + ", ".join(value_pair_str_list) + ")"
 
     def __value_repr(self, k: str, v) -> str:
-        if k.endswith("_const"):
+        if k.endswith("_const") and isinstance(v, (list, np.ndarray, tuple)):
             return repr(v[0])
         return repr(v)
 
@@ -198,10 +198,13 @@ class NoOpFreq(Operator):
 ##################################################
 
 
-class AbstractGas(ABC):
+class AbstractGas(Operator):
     gas_name: str
     material_dico: dict[str, Any]
 
+    def __call__(self, state: CurrentState) -> np.ndarray:
+        return self.square_index(state)
+
     @abstractmethod
     def pressure(self, state: CurrentState) -> float:
         """returns the pressure at the current
@@ -307,7 +310,7 @@ class PressureGradientGas(AbstractGas):
         return materials.pressure_from_gradient(state.z_ratio, self.p_in, self.p_out)
 
     def number_density(self, state: CurrentState) -> float:
-        if self.ideal:
+        if self.ideal_gas:
             return self.pressure(state) / (units.kB * self.temperature)
         else:
             return materials.number_density_van_der_waals(
@@ -505,7 +508,7 @@ class DirectDispersion(AbstractDispersion):
         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.disp_arr = np.zeros(t_num, dtype=complex)
         self.w0 = w0
         self.w0_ind = math.argclosest(w_for_disp, w0)
 
@@ -782,7 +785,7 @@ class AbstractGamma(Operator):
 
 
 class ConstantScalarGamma(AbstractGamma):
-    def __init__(self, gamma: np.ndarray, t_num: int):
+    def __init__(self, gamma: float, t_num: int):
         self.arr_const = gamma * np.ones(t_num)
 
     def __call__(self, state: CurrentState) -> np.ndarray:
@@ -802,6 +805,23 @@ class ConstantGamma(AbstractGamma):
         return self.arr
 
 
+class VariableScalarGamma(AbstractGamma):
+    def __init__(
+        self, gas_op: AbstractGas, temperature: float, w0: float, A_eff: float, t_num: int
+    ):
+        self.gas_op = gas_op
+        self.temperature = temperature
+        self.w0 = w0
+        self.A_eff = A_eff
+        self.arr = np.ones(t_num)
+
+    def __call__(self, state: CurrentState) -> np.ndarray:
+        n2 = materials.non_linear_refractive_index(
+            self.gas_op.material_dico, self.gas_op.pressure(state), self.temperature
+        )
+        return self.arr * fiber.gamma_parameter(n2, self.w0, self.A_eff)
+
+
 ##################################################
 ##################### PLASMA #####################
 ##################################################
diff --git a/src/scgenerator/parameter.py b/src/scgenerator/parameter.py
index b5ba335..d0d1cf4 100644
--- a/src/scgenerator/parameter.py
+++ b/src/scgenerator/parameter.py
@@ -334,9 +334,9 @@ class Parameters:
 
     # gas
     gas_name: str = Parameter(string, converter=str.lower, default="vacuum")
-    pressure: Union[float, Iterable[float]] = Parameter(
-        validator_or(non_negative(float, int), num_list), display_info=(1e-5, "bar")
-    )
+    pressure: float = Parameter(non_negative(float, int), display_info=(1e-5, "bar"))
+    pressure_in: float = Parameter(non_negative(float, int), display_info=(1e-5, "bar"))
+    pressure_out: float = Parameter(non_negative(float, int), display_info=(1e-5, "bar"))
     temperature: float = Parameter(positive(float, int), display_info=(1, "K"), default=300)
     plasma_density: float = Parameter(non_negative(float, int), default=0)