working direct dispersion

README.md

@@ -4,24 +4,40 @@ It is recommended to import scgenerator in the following manner :
 # How to run a set of simulations
 create a config file
 
-run `sc.parallel_simulations(config_file)` or `sc.simulate(config_file)`
+run `sc.run_simulation(config_file)`
 
 # How to analyse a simulation
 
-load data with the load_sim_data method
+Load completed simulations with the SimulationSeries class. The path argument should be that of the last fiber in the series (of the form `xx fiber A...` where `xx` is an integer)
-spectra, params = load_sim_data("varyTechNoise100kW_sim_data")
+
-    to plot
+The SimulationSeries class has basic plotting functions available. See the class documentation for more info
-        plot_results_2D(spectra[0], (600, 1450, nm), params)
+
+```
+series = sc.SimulationSeries(path_to_data_folder)
+fig, ax = plt.subplots()
+series.plot_2D(800, 1600, "nm", ax)
+plt.show()
+```
 
 # Environment variables
 
-SCGENERATOR_PBAR_POLICY : "none", "file", "print", "both", optional
+`SCGENERATOR_PBAR_POLICY` : "none", "file", "print", "both", optional
     whether progress should be printed to a file ("file"), to the standard output ("print") or both, default : print
 
+`SCGENERATOR_LOG_FILE_LEVEL` : "debug", "info", "warning", "error", "critical", optional
+    level of logging printed in $PWD/scgenerator.log
+
+`SCGENERATOR_LOG_PRINT_LEVEL` : "debug", "info", "warning", "error", "critical", optional
+    level of logging printed in the cli.
+
 # Configuration
 
 You can load parameters by simply passing the path to a toml file to the appropriate simulation function. Each possible key of this dictionary is described below. Every value must be given in standard SI units (m, s, W, J, ...)
 the root of the file has information concerning the whole simulation : name, grid information, input pulse, ...
+Then, there must be a `[[Fiber]]` array with at least one fiber with fiber-specific parameters.
+
+Parameters can be variable (either in the root or in one fiber). if at most one single `[variable]` dict is specified by section, all the possible combinations of those variable parameters are considered. Another possibility is to specify a `[[variable]]` array where the length of each set of parameter is the same so they're coupled to each other. 
+
 Examples : 
 
 ```
@@ -38,11 +54,14 @@ n2 = 2.2e-20
 [variable]
 n2 = [2.1e-20, 2.4e-20, 2.6e-20]
 ```
-NOT ALLOWED
+NOT ALLOWED. You cannot specified the same parameter twice.
 
 Here is an example of a configuration file
 
 ```
+# these be applied to the whole simulation fiber PM1550_1 only
+# fiber parameters specified here would apply to the whole simulation as well
+# unless overridden in one of the individual fiber
 name = "Test/Compound 1"
 
 field_file = "Toptica/init_field.npz"
@@ -57,11 +76,9 @@ z_num = 32
 mean_power = 200e-3
 repeat = 3
 
-# will be applied to the whole simulation fiber PM1550_1 only
+
-# fiber parameters specified here would apply to the whole simulation as well
+[[variable]]
-# unless overridden in one of the individual fiber
+spm = [true, false]
-[variable] 
-behaviors = [["spm", "ss", "raman"], ["spm", "ss"]]
 raman_type = ["agrawal", "stolen"]
 
 [[Fiber]]
@@ -82,8 +99,7 @@ dispersion_file = "PM2000D/Dispersion/PM2000D_1 extrapolated 0 4.npz"
 input_transmission = [0.9, 0.95]
 ```
 
-
+this means that only `(spm=true, raman_type="agrawal")` and `(spm=false, raman_type="stolen")` are considered and not `(spm=false, raman_type="agrawal")` for example. In the end, 12 simulations are ran with this configuration.
-note : internally, another structure with a flattened dictionary is used
 
 
 

make_new_dispersion.py

@@ -0,0 +1,37 @@
+import numpy as np
+import scgenerator as sc
+import matplotlib.pyplot as plt
+
+
+def main():
+    params = sc.Configuration(
+        "/Users/benoitsierro/tests/test_sc/Travers2019 Fig2 dev/initial_config.toml"
+    ).first
+    evalu = sc.Evaluator.default()
+    evalu.set(**params.prepare_for_dump())
+    n_eff = evalu.compute("n_eff")
+    wl = evalu.compute("wl_for_disp")
+    w = evalu.compute("w_for_disp")
+    print(w.max(), w.min())
+    disp_inf = evalu.compute("dispersion_ind")
+    # quit()
+
+    params.interpolation_degree = 6
+    params.compute()
+    current_disp = params.linear_operator.dispersion_op
+    beta = n_eff * w / 3e8
+    beta1 = np.gradient(beta, w)
+    ind = sc.argclosest(w, params.w0)
+    disp = -1j * (beta - beta1[ind] * (w - params.w0) - beta[ind])
+    disp2 = sc.fiber.fast_direct_dispersion(w, params.w0, n_eff, ind)
+
+    plt.plot(params.l * 1e9, current_disp(None).imag)
+    plt.plot(wl * 1e9, disp.imag)
+    plt.plot(wl * 1e9, disp2.imag)
+    plt.xlim(100, 3000)
+    plt.ylim(-1000, 4000)
+    plt.show()
+
+
+if __name__ == "__main__":
+    main()

src/scgenerator/__init__.py

@@ -19,3 +19,4 @@ from .plotting import (
 from .spectra import SimulationSeries, Spectrum
 from .utils import Paths, _open_config, open_single_config
 from .variationer import DescriptorDict, VariationDescriptor, Variationer, VariationSpecsError
+from .evaluator import Evaluator

src/scgenerator/evaluator.py

@@ -131,9 +131,24 @@ class Evaluator:
                     self.rules[t].append(r)
                     self.rules[t].sort(key=lambda el: el.targets[t], reverse=True)
 
-    def set(self, **params: Any):
+    def set(self, dico: dict[str, Any] = None, **params: Any):
-        self.params.update(params)
+        """sets the internal set of parameters
-        for k in params:
+
+        Parameters
+        ----------
+        dico : dict, optional
+            if given, replace current dict of parameters with this one
+            (not a copy of it), by default None
+        params : Any
+            if dico is None, update the internal dict of parameters with params
+        """
+        if dico is None:
+            dico = params
+            self.params.update(dico)
+        else:
+            self.reset()
+            self.params = dico
+        for k in dico:
             self.eval_stats[k].priority = np.inf
 
     def reset(self):
@@ -175,7 +190,9 @@ class Evaluator:
                 self.__curent_lookup.append(target)
 
             if len(self.rules[target]) == 0:
-                error = EvaluatorError(f"no rule for {target}")
+                error = EvaluatorError(
+                    f"no rule for {target}, trying to evaluate {self.__curent_lookup!r}"
+                )
             else:
                 error = None
 
@@ -312,13 +329,9 @@ default_rules: list[Rule] = [
     Rule("L_NL", pulse.L_NL),
     Rule("L_sol", pulse.L_sol),
     # Fiber Dispersion
-    Rule("wl_for_disp", fiber.lambda_for_dispersion),
+    Rule(["wl_for_disp", "dispersion_ind"], fiber.lambda_for_dispersion),
     Rule("w_for_disp", units.m, ["wl_for_disp"]),
-    Rule(
+    Rule("beta2_coefficients", fiber.dispersion_coefficients),
-        "beta2_coefficients",
-        fiber.dispersion_coefficients,
-        ["wl_for_disp", "beta2_arr", "w0", "interpolation_range", "interpolation_degree"],
-    ),
     Rule("beta2_arr", fiber.beta2),
     Rule("beta2_arr", fiber.dispersion_from_coefficients),
     Rule("beta2", lambda beta2_coefficients: beta2_coefficients[0]),

src/scgenerator/operators.py

@@ -75,8 +75,8 @@ class Operator(ABC):
 
 
 class NoOp:
-    def __init__(self, w: np.ndarray):
+    def __init__(self, t_num: int):
-        self.arr_const = np.zeros_like(w)
+        self.arr_const = np.zeros(t_num)
 
 
 ##################################################
@@ -112,23 +112,45 @@ class ConstantPolyDispersion(AbstractDispersion):
 
     def __init__(
         self,
-        wl_for_disp: np.ndarray,
+        w_for_disp: np.ndarray,
         beta2_arr: np.ndarray,
         w0: float,
         w_c: np.ndarray,
-        interpolation_range: tuple[float, float],
         interpolation_degree: int,
     ):
-        self.coefs = fiber.dispersion_coefficients(
+        self.coefs = fiber.dispersion_coefficients(w_for_disp, beta2_arr, w0, interpolation_degree)
-            wl_for_disp, beta2_arr, w0, interpolation_range, interpolation_degree
-        )
         self.w_c = w_c
         self.w_power_fact = np.array(
             [math.power_fact(w_c, k) for k in range(2, interpolation_degree + 3)]
         )
 
-    def __call__(self, state: CurrentState) -> np.ndarray:
+    def __call__(self, state: CurrentState = None) -> np.ndarray:
-        return fiber.fast_dispersion_op(self.w_c, self.coefs, self.w_power_fact)
+        return fiber.fast_poly_dispersion_op(self.w_c, self.coefs, self.w_power_fact)
+
+
+class ConstantDirectDispersion(AbstractDispersion):
+    """
+    Direct dispersion for when the refractive index is known
+    """
+
+    disp_arr_const: np.ndarray
+
+    def __init__(
+        self,
+        w_for_disp: np.ndarray,
+        w0: np.ndarray,
+        t_num: int,
+        n_eff: np.ndarray,
+        dispersion_ind: np.ndarray,
+    ):
+        self.disp_arr_const = np.zeros(t_num)
+        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
+        )[2:-2]
+
+    def __call__(self, state: CurrentState = None):
+        return self.disp_arr_const
 
 
 ##################################################
@@ -286,8 +308,8 @@ class AbstractGamma(Operator):
 
 
 class ConstantScalarGamma(AbstractGamma):
-    def __init__(self, gamma: np.ndarray, w: np.ndarray):
+    def __init__(self, gamma: np.ndarray, t_num: int):
-        self.arr_const = gamma * np.ones_like(w)
+        self.arr_const = gamma * np.ones(t_num)
 
     def __call__(self, state: CurrentState) -> np.ndarray:
         return self.arr_const
@@ -373,8 +395,8 @@ class AbstractLoss(Operator):
 class ConstantLoss(AbstractLoss):
     arr_const: np.ndarray
 
-    def __init__(self, alpha: float, w: np.ndarray):
+    def __init__(self, alpha: float, t_num: int):
-        self.arr_const = alpha * np.ones_like(w)
+        self.arr_const = alpha * np.ones(t_num)
 
     def __call__(self, state: CurrentState = None) -> np.ndarray:
         return self.arr_const
@@ -388,15 +410,15 @@ class NoLoss(ConstantLoss):
 class CapillaryLoss(ConstantLoss):
     def __init__(
         self,
-        l: np.ndarray,
+        wl_for_disp: np.ndarray,
+        dispersion_ind: np.ndarray,
+        t_num: int,
         core_radius: float,
-        interpolation_range: tuple[float, float],
         he_mode: tuple[int, int],
     ):
-        mask = (l < interpolation_range[1]) & (l > interpolation_range[0])
+        alpha = fiber.capillary_loss(wl_for_disp, he_mode, core_radius)
-        alpha = fiber.capillary_loss(l[mask], he_mode, core_radius)
+        self.arr = np.zeros(t_num)
-        self.arr = np.zeros_like(l)
+        self.arr[dispersion_ind] = alpha[2:-2]
-        self.arr[mask] = alpha
 
     def __call__(self, state: CurrentState) -> np.ndarray:
         return self.arr

src/scgenerator/physics/fiber.py

@@ -5,11 +5,11 @@ from numpy import e
 from numpy.fft import fft
 from numpy.polynomial.chebyshev import Chebyshev, cheb2poly
 from scipy.interpolate import interp1d
+from sympy import re
 
 from .. import utils
 from ..cache import np_cache
-from ..logger import get_logger
+from ..math import argclosest, u_nm
-from ..math import argclosest, power_fact, u_nm
 from . import materials as mat
 from . import units
 from .units import c, pi
@@ -18,14 +18,28 @@ pipi = 2 * pi
 T = TypeVar("T")
 
 
-def lambda_for_dispersion(interpolation_range: tuple[float, float]) -> np.ndarray:
+def lambda_for_dispersion(
+    l: np.ndarray, interpolation_range: tuple[float, float]
+) -> tuple[np.ndarray, np.ndarray]:
     """Returns a wl vector for dispersion calculation
 
     Returns
     -------
-    array of wl values
+    np.ndarray
+        subset of l in the interpolation range with two extra values on each side
+        to accomodate for taking gradients
+    np.ndarray
+        indices of the original l where the values are valid (i.e. without the two extra on each side)
     """
-    return np.arange(interpolation_range[0] - 2e-9, interpolation_range[1] + 3e-9, 1e-9)
+    su = np.where((l >= interpolation_range[0]) & (l <= interpolation_range[1]))[0]
+    ind_above_cond = su >= len(l) // 2
+    ind_above = su[ind_above_cond]
+    ind_below = su[~ind_above_cond]
+    fu = np.concatenate((ind_below, (ind_below + 2)[-2:], (ind_above - 2)[:2], ind_above))
+    fs = fu[np.argsort(l[fu])[::-1]]
+    l_out = l[fs]
+    ind_out = fs[2:-2]
+    return l_out, ind_out
 
 
 def is_dynamic_dispersion(pressure=None):
@@ -107,9 +121,8 @@ def plasma_dispersion(wl_for_disp, number_density, simple=False):
         w_pl = number_density * e2_me_e0
         return -(w_pl ** 2) / (c * w ** 2)
 
-    beta = w / c * np.sqrt(1 - number_density * e2_me_e0 / w ** 2)
+    beta2_arr = beta2(w, np.sqrt(1 - number_density * e2_me_e0 / w ** 2))
-    beta2 = np.gradient(np.gradient(beta, w), w)
+    return beta2_arr
-    return beta2
 
 
 def n_eff_marcatili(wl_for_disp, n_gas_2, core_radius, he_mode=(1, 1)):
@@ -561,6 +574,14 @@ def HCPF_ZDW(
     return l[zdw_ind]
 
 
+def beta(w_for_disp: np.ndarray, n_eff: np.ndarray) -> np.ndarray:
+    return n_eff * w_for_disp / c
+
+
+def beta1(w_for_disp: np.ndarray, n_eff: np.ndarray) -> np.ndarray:
+    return np.gradient(beta(w_for_disp, n_eff), w_for_disp)
+
+
 def beta2(w_for_disp: np.ndarray, n_eff: np.ndarray) -> np.ndarray:
     """computes the dispersion parameter beta2 according to the effective refractive index of the fiber and the frequency range
 
@@ -575,7 +596,7 @@ def beta2(w_for_disp: np.ndarray, n_eff: np.ndarray) -> np.ndarray:
     -------
     beta2 : ndarray, shape (n, )
     """
-    return np.gradient(np.gradient(n_eff * w_for_disp / c, w_for_disp), w_for_disp)
+    return np.gradient(np.gradient(beta(w_for_disp, n_eff), w_for_disp), w_for_disp)
 
 
 def HCPCF_dispersion(
@@ -849,11 +870,10 @@ def load_custom_loss(l: np.ndarray, loss_file: str) -> np.ndarray:
 
 @np_cache
 def dispersion_coefficients(
-    wl_for_disp: np.ndarray,
+    w_for_disp: np.ndarray,
     beta2_arr: np.ndarray,
     w0: float,
-    interpolation_range=None,
+    interpolation_degree: int,
-    interpolation_degree=8,
 ):
     """Computes the taylor expansion of beta2 to be used in dispersion_op
 
@@ -865,8 +885,6 @@ def dispersion_coefficients(
         beta2 as function of wl_for_disp
     w0 : float
         pump angular frequency
-    interpolation_range : slice-like
-        index-style specifying wl range over which to fit to get beta2 coefficients
     interpolation_degree : int
         degree of polynomial fit. Will return deg+1 coefficients
 
@@ -875,29 +893,13 @@ def dispersion_coefficients(
         beta2_coef : 1D array
             Taylor coefficients in decreasing order
     """
-    logger = get_logger()
-    if interpolation_range is None:
-        r = slice(2, -2)
-    else:
-        # 2 discrete gradients are computed before getting to
-        # beta2, so we need to make sure coefficients are not affected
-        # by edge effects
-        r = (wl_for_disp >= interpolation_range[0]) & (wl_for_disp <= interpolation_range[1])
-    logger.debug(
-        f"interpolating dispersion between {wl_for_disp[r].min()*1e9:.1f}nm and {wl_for_disp[r].max()*1e9:.1f}nm"
-    )
 
     # we get the beta2 Taylor coeffiecients by making a fit around w0
-    w_c = units.m(wl_for_disp) - w0
+    w_c = w_for_disp - w0
-    interp = interp1d(w_c[r], beta2_arr[r])
-    w_c = np.linspace(w_c[r].min(), w_c[r].max(), len(w_c[r]))
 
-    # import matplotlib.pyplot as plt
+    w_c = w_c[2:-2]
-
+    beta2_arr = beta2_arr[2:-2]
-    # ax = plt.gca()
+    fit = Chebyshev.fit(w_c, beta2_arr, interpolation_degree)
-    # ax.plot(w_c, interp(w_c) * 1e28)
-
-    fit = Chebyshev.fit(w_c, interp(w_c), interpolation_degree)
     poly_coef = cheb2poly(fit.convert().coef)
     beta2_coef = poly_coef * np.cumprod([1] + list(range(1, interpolation_degree + 1)))
 
@@ -981,7 +983,7 @@ def delayed_raman_w(t: np.ndarray, raman_type: str) -> np.ndarray:
     return fft(delayed_raman_t(t, raman_type)) * (t[1] - t[0])
 
 
-def fast_dispersion_op(w_c, beta_arr, power_fact_arr, where=slice(None)):
+def fast_poly_dispersion_op(w_c, beta_arr, power_fact_arr, where=slice(None)):
     """
     dispersive operator
 
@@ -1015,39 +1017,32 @@ def _fast_disp_loop(dispersion: np.ndarray, beta_arr, power_fact_arr):
     return dispersion
 
 
-def dispersion_op(w_c, beta2_coefficients, where=None):
+def direct_dispersion(w: np.ndarray, w0: float, n_eff: np.ndarray) -> np.ndarray:
-    """
+    """returns the dispersive operator in direct form (without polynomial interpolation)
-    dispersive operator
+    i.e. -1j * (beta(w) - beta1 * (w - w0) - beta0)
 
     Parameters
     ----------
-    w_c : 1d array
+    w : np.ndarray
-        angular frequencies centered around 0
+        angular frequency array
-    beta2_coefficients : 1d array
+    w0 : float
-        beta coefficients returned by scgenerator.physics.fiber.dispersion_coefficients
+        center frequency
-    where : indices over which to apply the operatory, otherwise 0
+    n_eff : np.ndarray
+        effectiv refractive index
 
     Returns
     -------
-    disp_arr : dispersive component as an array of len = len(w_c)
+    np.ndarray
+        dispersive operator
     """
-
+    w0_ind = argclosest(w, w0)
-    dispersion = np.zeros_like(w_c)
+    return fast_direct_dispersion(w, w0, n_eff, w0_ind)
-
-    for k, beta in reversed(list(enumerate(beta2_coefficients))):
-        dispersion = dispersion + beta * power_fact(w_c, k + 2)
-
-    out = np.zeros_like(dispersion)
-    out[where] = dispersion[where]
-
-    return -1j * out
 
 
-def _get_radius(radius_param, wl_for_disp=None):
+def fast_direct_dispersion(w: np.ndarray, w0: float, n_eff: np.ndarray, w0_ind: int) -> np.ndarray:
-    if isinstance(radius_param, tuple) and wl_for_disp is not None:
+    beta_arr = beta(w, n_eff)
-        return effective_core_radius(wl_for_disp, *radius_param)
+    beta1_arr = np.gradient(beta_arr, w)
-    else:
+    return -1j * (beta_arr - beta1_arr[w0_ind] * (w - w0) - beta_arr[w0_ind])
-        return radius_param
 
 
 def effective_core_radius(wl_for_disp, core_radius, s=0.08, h=200e-9):