added l to parameters, some plotting stuff

src/scgenerator/initialize.py

@@ -635,7 +635,8 @@ def build_sim_grid(
     t_num = len(t)
     z_targets = np.linspace(0, length, z_num)
     w_c, w0, w, w_power_fact = update_frequency_domain(t, wavelength)
-    return z_targets, t, time_window, t_num, dt, w_c, w0, w, w_power_fact
+    l = units.To.m(w)
+    return z_targets, t, time_window, t_num, dt, w_c, w0, w, w_power_fact, l
 
 
 def build_sim_grid_in_place(params: BareParams):
@@ -650,6 +651,7 @@ def build_sim_grid_in_place(params: BareParams):
         params.w0,
         params.w,
         params.w_power_fact,
+        params.l,
     ) = build_sim_grid(
         params.length, params.z_num, params.wavelength, params.time_window, params.t_num, params.dt
     )

src/scgenerator/physics/units.py

@@ -2,7 +2,9 @@
 # For example, nm(X) means "I give the number X in nm, figure out the ang. freq."
 # to be used especially when giving plotting ranges : (400, 1400, nm), (-4, 8, ps), ...
 
-from typing import Callable, Union
+import re
+from threading import settrace
+from typing import Callable, TypeVar, Union
 from dataclasses import dataclass
 from ..utils.parameter import Parameter, type_checker
 import numpy as np
@@ -21,204 +23,147 @@ prefix = dict(P=1e12, G=1e9, M=1e6, k=1e3, d=1e-1, c=1e-2, m=1e-3, u=1e-6, n=1e-
 
 """
 Below are common units. You can define your own unit function
-this function must have a few porperties:
+provided you decorate it with @unit and provide at least a type and a label
-inv : function
+types are "WL", "FREQ", "AFREQ", "TIME", "PRESSURE", "TEMPERATURE", "OTHER"
-    inverse of the function. example :
-        um(1) -> 883651567308853.2
-        um.inv(883651567308853.2) -> 1.0
-label : str
-    label to be displayed on plot
-type : ("WL", "FREQ", "AFREQ", "TIME", "OTHER")
 """
+_T = TypeVar("_T")
 
 
-def m(l):
+class From:
+    pass
+
+
+class To:
+    pass
+
+
+units_map = dict()
+
+
+def unit(tpe: str, label: str, inv: Callable = None):
+    def unit_maker(func):
+        nonlocal inv
+        name = func.__name__
+        if inv is None:
+            inv = func
+        setattr(From, name, func.__call__)
+        setattr(To, name, inv.__call__)
+        func.type = tpe
+        func.label = label
+        func.inv = inv
+        if name in units_map:
+            raise NameError(f"Two unit functions with the same name {name!r} were defined")
+        units_map[name] = func
+        return func
+
+    return unit_maker
+
+
+@unit("WL", r"Wavelength $\lambda$ (m)")
+def m(l: _T) -> _T:
     return 2 * pi * c / l
 
 
-m.inv = m
+@unit("WL", r"Wavelength $\lambda$ (nm)")
-m.label = r"Wavelength $\lambda$ (m)"
+def nm(l: _T) -> _T:
-m.type = "WL"
-
-
-def nm(l):
     return 2 * pi * c / (l * 1e-9)
 
 
-nm.inv = nm
+@unit("WL", r"Wavelength $\lambda$ (μm)")
-nm.label = r"Wavelength $\lambda$ (nm)"
+def um(l: _T) -> _T:
-nm.type = "WL"
-
-
-def um(l):
     return 2 * pi * c / (l * 1e-6)
 
 
-um.inv = um
+@unit("FREQ", r"Frequency $f$ (THz)", lambda w: w / (1e12 * 2 * pi))
-um.label = r"Wavelength $\lambda$ ($\mathrm{\mu}$m)"
+def THz(f: _T) -> _T:
-um.type = "WL"
-
-
-def THz(f):
     return 1e12 * 2 * pi * f
 
 
-THz.inv = lambda w: w / (1e12 * 2 * pi)
+@unit("FREQ", r"Frequency $f$ (PHz)", lambda w: w / (1e15 * 2 * pi))
-THz.label = r"Frequency $f$ (THz)"
+def PHz(f: _T) -> _T:
-THz.type = "FREQ"
-
-
-def PHz(f):
     return 1e15 * 2 * pi * f
 
 
-PHz.inv = lambda w: w / (1e15 * 2 * pi)
+@unit("AFREQ", r"Angular frequency $\omega$ ($\frac{\mathrm{rad}}{\mathrm{s}}$)")
-PHz.label = r"Frequency $f$ (PHz)"
+def rad_s(w: _T) -> _T:
-PHz.type = "FREQ"
-
-
-def rad_s(w):
     return w
 
 
-rad_s.inv = rad_s
+@unit(
-rad_s.label = r"Angular frequency $\omega$ ($\frac{\mathrm{rad}}{\mathrm{s}}$)"
+    "AFREQ", r"Angular frequency $\omega$ ($\frac{\mathrm{Prad}}{\mathrm{s}}$)", lambda w: 1e-15 * w
-rad_s.type = "AFREQ"
+)
-
+def Prad_s(w: _T) -> _T:
-
-def Prad_s(w):
     return w * 1e15
 
 
-Prad_s.inv = lambda w: 1e-15 * w
+@unit("TIME", r"relative time ${\tau}/{\tau_\mathrm{0, FWHM}}$")
-Prad_s.label = r"Angular frequency $\omega$ ($\frac{\mathrm{Prad}}{\mathrm{s}}$)"
+def rel_time(t: _T) -> _T:
-Prad_s.type = "AFREQ"
-
-
-def rel_time(t):
     return t
 
 
-rel_time.inv = rel_time
+@unit("FREQ", r"relative angular freq. $(\omega - \omega_0)/\Delta\omega_0$")
-rel_time.label = r"relative time ${\tau}/{\tau_\mathrm{0, FWHM}}$"
+def rel_freq(f: _T) -> _T:
-rel_time.type = "TIME"
-
-
-def rel_freq(f):
     return f
 
 
-rel_freq.inv = rel_freq
+@unit("TIME", r"Time $t$ (s)")
-rel_freq.label = r"relative angular freq. $(\omega - \omega_0)/\Delta\omega_0$"
+def s(t: _T) -> _T:
-rel_freq.type = "FREQ"
-
-
-def s(t):
     return t
 
 
-s.inv = s
+@unit("TIME", r"Time $t$ (us)", lambda t: t * 1e6)
-s.label = r"Time $t$ (s)"
+def us(t: _T) -> _T:
-s.type = "TIME"
-
-
-def us(t):
     return t * 1e-6
 
 
-us.inv = lambda t: t * 1e6
+@unit("TIME", r"Time $t$ (ns)", lambda t: t * 1e9)
-us.label = r"Time $t$ (us)"
+def ns(t: _T) -> _T:
-us.type = "TIME"
-
-
-def ns(t):
     return t * 1e-9
 
 
-ns.inv = lambda t: t * 1e9
+@unit("TIME", r"Time $t$ (ps)", lambda t: t * 1e12)
-ns.label = r"Time $t$ (ns)"
+def ps(t: _T) -> _T:
-ns.type = "TIME"
-
-
-def ps(t):
     return t * 1e-12
 
 
-ps.inv = lambda t: t * 1e12
+@unit("TIME", r"Time $t$ (fs)", lambda t: t * 1e15)
-ps.label = r"Time $t$ (ps)"
+def fs(t: _T) -> _T:
-ps.type = "TIME"
-
-
-def fs(t):
     return t * 1e-15
 
 
-fs.inv = lambda t: t * 1e15
+@unit("WL", "inverse")
-fs.label = r"Time $t$ (fs)"
+def inv(x: _T) -> _T:
-fs.type = "TIME"
-
-
-def inv(x):
     return 1 / x
 
 
-inv.inv = inv
+@unit("PRESSURE", "Pressure (bar)", lambda p: 1e-5 * p)
-inv.label = "inverse"
+def bar(p: _T) -> _T:
-inv.type = "WL"
-
-
-def bar(p):
     return 1e5 * p
 
 
-bar.inv = lambda p: 1e-5 * p
+@unit("OTHER", r"$\beta_2$ (fs$^2$/cm)", lambda b2: 1e28 * b2)
-bar.label = "Pressure (bar)"
+def beta2_fs_cm(b2: _T) -> _T:
-bar.type = "PRESSURE"
-
-
-def beta2_fs_cm(b2):
     return 1e-28 * b2
 
 
-beta2_fs_cm.inv = lambda b2: 1e28 * b2
+@unit("OTHER", r"$\beta_2$ (ps$^2$/km)", lambda b2: 1e27 * b2)
-beta2_fs_cm.label = r"$\beta_2$ (fs$^2$/cm)"
+def beta2_ps_km(b2: _T) -> _T:
-beta2_fs_cm.type = "OTHER"
-
-
-def beta2_ps_km(b2):
     return 1e-27 * b2
 
 
-beta2_ps_km.inv = lambda b2: 1e27 * b2
+@unit("OTHER", r"$D$ (ps/(nm km))", lambda D: 1e6 * D)
-beta2_ps_km.label = r"$\beta_2$ (ps$^2$/km)"
+def D_ps_nm_km(D: _T) -> _T:
-beta2_ps_km.type = "OTHER"
-
-
-def D_ps_nm_km(D):
     return 1e-6 * D
 
 
-D_ps_nm_km.inv = lambda D: 1e6 * D
+@unit("TEMPERATURE", r"Temperature (K)")
-D_ps_nm_km.label = r"$D$ (ps/(nm km))"
+def K(t: _T) -> _T:
-D_ps_nm_km.type = "OTHER"
+    return t
 
 
-units_map = dict(
+@unit("TEMPERATURE", r"Temperature (°C)", lambda t_K: t_K - 272.15)
-    nm=nm,
+def C(t_C: _T) -> _T:
-    um=um,
+    return t_C + 272.15
-    m=m,
-    THz=THz,
-    PHz=PHz,
-    rad_s=rad_s,
-    Prad_s=Prad_s,
-    rel_freq=rel_freq,
-    rel_time=rel_time,
-    s=s,
-    us=us,
-    ns=ns,
-    ps=ps,
-    fs=fs,
-)
 
 
 def get_unit(unit: Union[str, Callable]) -> Callable[[float], float]:
@@ -325,13 +270,22 @@ def sort_axis(axis, plt_range: PlotRange):
     return out_ax, indices, (out_ax[0], out_ax[-1])
 
 
-def to_WL(spectrum, frep, lambda_):
+def to_WL(spectrum: np.ndarray, lambda_: np.ndarray) -> np.ndarray:
+    """rescales the spectrum because of uneven binning when going from freq to wl
+
+    Parameters
+    ----------
+    spectrum : np.ndarray, shape (n, )
+        intensity spectrum
+    lambda_ : np.ndarray, shape (n, )
+        wavelength in m
+
+    Returns
+    -------
+    np.ndarray, shape (n, )
+        intensity spectrum correctly scaled
     """
-    When a spectrogram is displayed as function of wl instead of frequency, we
+    m = 2 * pi * c / (lambda_ ** 2) * spectrum
-    need to adjust the amplitude of each bin for the integral over the whole frequency
-    range to match.
-    """
-    m = 2 * pi * c / (lambda_ ** 2) * frep * spectrum
     return m
 
 

src/scgenerator/plotting.py

@@ -421,9 +421,7 @@ def plot_spectrogram(
     new_f, ind_f, _ = units.sort_axis(params.w, f_range)
     values = spec[ind_t][:, ind_f]
     if f_range[2].type == "WL":
-        values = np.apply_along_axis(
+        values = np.apply_along_axis(units.to_WL, 1, values, units.m(f_range[2].inv(new_f)))
-            units.to_WL, 1, values, params.repetition_rate, units.m(f_range[2].inv(new_f))
-        )
         values = np.apply_along_axis(make_uniform_1D, 1, values, new_f)
 
     if time_axis == 0:
@@ -528,7 +526,7 @@ def plot_results_2D(
     # make uniform if converting to wavelength
     if plt_range.unit.type == "WL":
         if is_spectrum:
-            values = np.apply_along_axis(units.to_WL, 1, values, params.repetition_rate, x_axis)
+            values = np.apply_along_axis(units.to_WL, 1, values, x_axis)
         values = np.array(
             [make_uniform_1D(v, x_axis, n=len(x_axis), method="linear") for v in values]
         )
@@ -648,7 +646,7 @@ def plot_results_1D(
     # make uniform if converting to wavelength
     if plt_range.unit.type == "WL":
         if is_spectrum:
-            values = units.to_WL(values, params.repetition_rate, units.m.inv(params.w[ind]))
+            values = units.to_WL(values, units.m.inv(params.w[ind]))
 
     # change the resolution
     if isinstance(spacing, float):
@@ -810,8 +808,8 @@ def plot_avg(
     values *= yscaling
     mean_values = np.mean(values, axis=0)
     if plt_range.unit.type == "WL" and renormalize:
-        values = np.apply_along_axis(units.to_WL, 1, values, params.repetition_rate, x_axis)
+        values = np.apply_along_axis(units.to_WL, 1, values, x_axis)
-        mean_values = units.to_WL(mean_values, params.repetition_rate, x_axis)
+        mean_values = units.to_WL(mean_values, x_axis)
 
     # change the resolution
     if isinstance(spacing, float):
@@ -962,7 +960,7 @@ def prepare_plot_1D(values, plt_range, x_axis, yscaling=1, spacing=1, frep=80e6)
     values = values[:, ind]
 
     if plt_range.unit.type == "WL":
-        values = np.apply_along_axis(units.to_WL, -1, values, frep, x_axis)
+        values = np.apply_along_axis(units.to_WL, -1, values, x_axis)
 
     if isinstance(spacing, float):
         new_x_axis = np.linspace(*span(x_axis), int(len(x_axis) / spacing))

src/scgenerator/spectra.py

@@ -3,24 +3,30 @@ from collections.abc import Sequence
 from pathlib import Path
 from re import UNICODE
 from typing import Callable, Dict, Iterable, Optional, Union
+from matplotlib.pyplot import subplot
+from dataclasses import replace
 
 import numpy as np
+from tqdm.std import Bar
 
 from . import initialize, io, math
-from .physics import units
+from .physics import units, pulse
 from .const import SPECN_FN
 from .logger import get_logger
 from .plotting import plot_avg, plot_results_1D, plot_results_2D
+from .utils.parameter import BareParams
 
 
 class Spectrum(np.ndarray):
-    def __new__(cls, input_array, wl, frep=1):
+    params: BareParams
+
+    def __new__(cls, input_array, params: BareParams):
         # Input array is an already formed ndarray instance
         # We first cast to be our class type
         obj = np.asarray(input_array).view(cls)
         # add the new attribute to the created instance
-        obj.frep = frep
+        obj.params = params
-        obj.wl = wl
+
         # Finally, we must return the newly created object:
         return obj
 
@@ -28,8 +34,91 @@ class Spectrum(np.ndarray):
         # see InfoArray.__array_finalize__ for comments
         if obj is None:
             return
-        self.frep = getattr(obj, "frep", None)
+        self.params = getattr(obj, "params", None)
-        self.wl = getattr(obj, "wl", None)
+
+    def __getitem__(self, key) -> "Spectrum":
+        return super().__getitem__(key)
+
+    def energy(self) -> Union[np.ndarray, float]:
+        if self.ndim == 1:
+            m = np.argwhere(self.params.l > 0)[:, 0]
+            m = np.array(sorted(m, key=lambda el: self.params.l[el]))
+            return np.trapz(self.wl_int[m], self.params.l[m])
+        else:
+            return np.array([s.energy() for s in self])
+
+    def crop_wl(self, left: float, right: float) -> tuple[np.ndarray, np.ndarray]:
+        cond = (self.params.l >= left) & (self.params.l <= right)
+        return cond
+
+    @property
+    def wl_int(self):
+        return units.to_WL(math.abs2(self), self.params.l)
+
+    @property
+    def freq_int(self):
+        return math.abs2(self)
+
+    @property
+    def afreq_int(self):
+        return math.abs2(self)
+
+    @property
+    def time_int(self):
+        return math.abs2(np.fft.ifft(self))
+
+    def amplitude(self, unit):
+        if unit.type in ["WL", "FREQ", "AFREQ"]:
+            x_axis = unit.inv(self.params.w)
+        else:
+            x_axis = unit.inv(self.params.t)
+
+        order = np.argsort(x_axis)
+        func = dict(
+            WL=self.wl_amp,
+            FREQ=self.freq_amp,
+            AFREQ=self.afreq_amp,
+            TIME=self.time_amp,
+        )[unit.type]
+
+        for spec in self:
+            yield x_axis[order], func(spec)[:, order]
+
+    @property
+    def wl_amp(self):
+        return (
+            np.sqrt(
+                units.to_WL(
+                    math.abs2(self),
+                    self.params.l,
+                )
+            )
+            * self
+            / np.abs(self)
+        )
+
+    @property
+    def freq_amp(self):
+        return self
+
+    @property
+    def afreq_amp(self):
+        return self
+
+    @property
+    def time_amp(self):
+        return np.fft.ifft(self)
+
+    @property
+    def wl_max(self):
+        if self.ndim == 1:
+            return self.params.l[np.argmax(self.wl_int, axis=-1)]
+        return np.array([s.wl_max for s in self])
+
+    def mask_wl(self, pos: float, width: float) -> "Spectrum":
+        return self * np.exp(
+            -(((self.params.l - pos) / (pulse.fwhm_to_T0_fac["gaussian"] * width)) ** 2)
+        )
 
 
 class Pulse(Sequence):
@@ -112,7 +201,7 @@ class Pulse(Sequence):
             yield x_axis[order], func(spec)[:, order]
 
     def _to_wl_int(self, spectrum):
-        return units.to_WL(math.abs2(spectrum), spectrum.frep, spectrum.wl)
+        return units.to_WL(math.abs2(spectrum), spectrum.wl)
 
     def _to_freq_int(self, spectrum):
         return math.abs2(spectrum)
@@ -145,7 +234,6 @@ class Pulse(Sequence):
             np.sqrt(
                 units.to_WL(
                     math.abs2(spectrum),
-                    spectrum.frep,
                     spectrum.wl,
                 )
             )
@@ -162,7 +250,7 @@ class Pulse(Sequence):
     def _to_time_amp(self, spectrum):
         return np.fft.ifft(spectrum)
 
-    def all_spectra(self, ind=None):
+    def all_spectra(self, ind=None) -> Spectrum:
         """
         loads the data already simulated.
         defauft shape is (z_targets, n, nt)
@@ -194,7 +282,6 @@ class Pulse(Sequence):
         spectra = []
         for i in ind:
             spectra.append(self._load1(i))
-        spectra = np.array(spectra)
 
         self.logger.debug(f"all spectra from {self.path} successfully loaded")
         if len(ind) == 1:
@@ -212,7 +299,7 @@ class Pulse(Sequence):
             return self.cache[i]
         spec = np.load(self.path / SPECN_FN.format(i))
         spec = np.atleast_2d(spec)
-        spec = Spectrum(spec, self.wl, self.params.repetition_rate)
+        spec = Spectrum(spec, self.params)
         self.cache[i] = spec
         return spec
 

src/scgenerator/utils/parameter.py

@@ -384,6 +384,7 @@ class BareParams:
     field_0: np.ndarray = Parameter(type_checker(np.ndarray))
     spec_0: np.ndarray = Parameter(type_checker(np.ndarray))
     w: np.ndarray = Parameter(type_checker(np.ndarray))
+    l: np.ndarray = Parameter(type_checker(np.ndarray))
     w_c: np.ndarray = Parameter(type_checker(np.ndarray))
     w0: float = Parameter(positive(float))
     w_power_fact: np.ndarray = Parameter(validator_list(type_checker(np.ndarray)))