spectra working

src/scgenerator/__init__.py

@@ -3,7 +3,7 @@ from .math import abs2, argclosest, span
 from .physics import fiber, materials, pulse, simulate, units
 from .physics.simulate import RK4IP, parallel_RK4IP, run_simulation
 from .plotting import mean_values_plot, plot_spectrogram, propagation_plot, single_position_plot
-from .spectra import Pulse, Spectrum
+from .spectra import Pulse, Spectrum, SimulationSeries
 from ._utils import Paths, open_config, parameter
 from ._utils.parameter import Configuration, Parameters
 from ._utils.utils import PlotRange

src/scgenerator/spectra.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import os
 from collections.abc import Sequence
 from pathlib import Path
-from typing import Any, Callable, Dict, Iterable, Optional, Union
+from typing import Any, Callable, Dict, Iterable, Iterator, Optional, Union
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -30,6 +30,9 @@ class SimulationSeries:
     total_length: float
     total_num_steps: int
     previous: SimulationSeries = None
+    fiber_lengths: list[tuple[str, float]]
+    fiber_positions: list[tuple[str, float]]
+    z_inds: np.ndarray
 
     class Config:
         arbitrary_types_allowed = True
@@ -37,15 +40,25 @@ class SimulationSeries:
     def __init__(self, path: os.PathLike):
         self.path = Path(path)
         self.params = Parameters.load(self.path / PARAM_FN)
+        self.params.compute(["name", "t", "l", "w_c", "w0", "z_targets"])
+        self.t = self.params.t
+        self.w = self.params.w
         if self.params.prev_data_dir is not None:
             self.previous = SimulationSeries(self.params.prev_data_dir)
         self.total_length = self.accumulate_params("length")
         self.total_num_steps = self.accumulate_params("z_num")
-
+        self.z_inds = np.arange(len(self.params.z_targets))
-    def fiber_map(self):
+        self.z = self.params.z_targets
-        lengths = self.all_params("length")
+        if self.previous is not None:
-        return [
+            print(f"{self.params.z_targets=}")
-            (this[0], following[1]) for this, following in zip(lengths, [(None, 0.0)] + lengths)
+            self.z += self.previous.params.z_targets[-1]
+            self.params.z_targets = np.concatenate((self.previous.z, self.params.z_targets))
+            self.z_inds += self.previous.z_inds[-1] + 1
+        print(f"{self.z=}")
+        self.fiber_lengths = self.all_params("length")
+        self.fiber_positions = [
+            (this[0], following[1])
+            for this, following in zip(self.fiber_lengths, [(None, 0.0)] + self.fiber_lengths)
         ]
 
     def all_params(self, key: str) -> list[tuple[str, Any]]:
@@ -79,8 +92,135 @@ class SimulationSeries:
         """
         return sum(el[1] for el in self.all_params(key))
 
+    def spectra(
+        self, z_descr: Union[float, int, None] = None, sim_ind: Optional[int] = 0
+    ) -> Spectrum:
+        if z_descr is None:
+            out = [self.spectra(i, sim_ind) for i in range(self.total_num_steps)]
+        else:
+            if isinstance(z_descr, (float, np.floating)):
+                if self.z[0] <= z_descr <= self.z[-1]:
+                    z_ind = self.z_inds[np.argmin(np.abs(self.z - z_descr))]
+                elif 0 <= z_descr < self.z[0]:
+                    return self.previous.spectra(z_descr, sim_ind)
+                else:
+                    raise ValueError(
+                        f"cannot match z={z_descr} with max length of {self.total_length}"
+                    )
+            else:
+                z_ind = z_descr
+
+            if z_ind < self.z_inds[0]:
+                return self.previous.spectra(z_ind, sim_ind)
+            if sim_ind is None:
+                out = [self._load_1(z_ind, i) for i in range(self.params.repeat)]
+            else:
+                out = self._load_1(z_ind)
+        return Spectrum(out, self.params)
+
+    def fields(
+        self, z_descr: Union[float, int, None] = None, sim_ind: Optional[int] = 0
+    ) -> Spectrum:
+        return np.fft.ifft(self.spectra(z_descr, sim_ind))
+
+    # Plotting
+
+    def plot_2D(
+        self,
+        left: float,
+        right: float,
+        unit: Union[Callable[[float], float], str],
+        ax: plt.Axes,
+        sim_ind: int = 0,
+        **kwargs,
+    ):
+        plot_range = PlotRange(left, right, unit)
+        vals = self.retrieve_plot_values(plot_range, None, sim_ind)
+        return propagation_plot(vals, plot_range, self.params, ax, **kwargs)
+
+    def plot_1D(
+        self,
+        left: float,
+        right: float,
+        unit: Union[Callable[[float], float], str],
+        ax: plt.Axes,
+        z_pos: int,
+        sim_ind: int = 0,
+        **kwargs,
+    ):
+        plot_range = PlotRange(left, right, unit)
+        vals = self.retrieve_plot_values(plot_range, z_pos, sim_ind)
+        return single_position_plot(vals, plot_range, self.params, ax, **kwargs)
+
+    def plot_mean(
+        self,
+        left: float,
+        right: float,
+        unit: Union[Callable[[float], float], str],
+        ax: plt.Axes,
+        z_pos: int,
+        **kwargs,
+    ):
+        plot_range = PlotRange(left, right, unit)
+        vals = self.retrieve_plot_values(plot_range, z_pos, None)
+        return mean_values_plot(vals, plot_range, self.params, ax, **kwargs)
+
+    def retrieve_plot_values(
+        self, plot_range: PlotRange, z_pos: Optional[Union[int, float]], sim_ind: Optional[int]
+    ):
+
+        if plot_range.unit.type == "TIME":
+            return self.fields(z_pos, sim_ind)
+        else:
+            return self.spectra(z_pos, sim_ind)
+
+    def rin_propagation(
+        self, left: float, right: float, unit: str
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """returns the RIN as function of unit and z
+
+        Parameters
+        ----------
+        left : float
+            left limit in unit
+        right : float
+            right limit in unit
+        unit : str
+            unit descriptor
+
+        Returns
+        -------
+        x : np.ndarray, shape (nt,)
+            x axis
+        y : np.ndarray, shape (z_num, )
+            y axis
+        rin_prop : np.ndarray, shape (z_num, nt)
+            RIN
+        """
+        spectra = []
+        for spec in np.moveaxis(self.spectra(None, None), 1, 0):
+            x, z, tmp = transform_2D_propagation(spec, (left, right, unit), self.params, False)
+            spectra.append(tmp)
+        return x, z, pulse.rin_curve(np.moveaxis(spectra, 0, 1))
+
+    # Private
+
     def _load_1(self, z_ind: int, sim_ind=0) -> np.ndarray:
-        return load_spectrum(self.path / SPEC1_FN_N.format(z_ind, sim_ind))
+        """loads a spectrum file
+
+        Parameters
+        ----------
+        z_ind : int
+            z_index relative to the entire simulation
+        sim_ind : int, optional
+            simulation index, used when repeated simulations with same parameters are ran, by default 0
+
+        Returns
+        -------
+        np.ndarray
+            loaded spectrum file
+        """
+        return load_spectrum(self.path / SPEC1_FN_N.format(z_ind - self.z_inds[0], sim_ind))
 
     def _all_params(self, key: str, l: list) -> list:
         l.append((self.params.name, getattr(self.params, key)))
@@ -88,6 +228,12 @@ class SimulationSeries:
             return self.previous._all_params(key, l)
         return l
 
+    # Magic methods
+
+    def __iter__(self) -> Iterator[Spectrum]:
+        for i in range(self.total_num_steps):
+            yield self.spectra(i, None)
+
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}(path={self.path}, previous={self.previous!r})"
 
@@ -127,18 +273,6 @@ class Spectrum(np.ndarray):
     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) -> 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)