Huge plotting revamp

src/scgenerator/__init__.py

@@ -9,7 +9,7 @@ from .io import Paths, load_toml, load_params
 from .math import abs2, argclosest, span
 from .physics import fiber, materials, pulse, simulate, units
 from .physics.simulate import RK4IP, new_simulation, resume_simulations
-from .plotting import plot_avg, plot_results_1D, plot_results_2D, plot_spectrogram
+from .plotting import mean_values_plot, single_position_plot, propagation_plot, plot_spectrogram
 from .spectra import Pulse, Spectrum
 from .utils.parameter import BareParams, BareConfig
 from . import utils, io, initialize, math

src/scgenerator/initialize.py

@@ -14,7 +14,7 @@ from .errors import *
 from .logger import get_logger
 from .math import power_fact
 from .physics import fiber, pulse, units
-from .utils import count_variations, override_config, pretty_format_value, required_simulations
+from .utils import override_config, required_simulations
 from .utils.parameter import BareConfig, BareParams, hc_model_specific_parameters
 
 
@@ -320,8 +320,10 @@ class ParamSequence:
         config_dict : Union[Dict[str, Any], os.PathLike, BareConfig]
             Can be either a dictionary, a path to a config toml file or BareConfig obj
         """
-        if isinstance(config_dict, BareConfig):
+        if isinstance(config_dict, Config):
             self.config = config_dict
+        elif isinstance(config_dict, BareConfig):
+            self.config = Config.from_bare(config_dict)
         else:
             if not isinstance(config_dict, Mapping):
                 config_dict = io.load_toml(config_dict)
@@ -329,7 +331,7 @@ class ParamSequence:
         self.name = self.config.name
         self.logger = get_logger(__name__)
 
-        self.update_num_sim(count_variations(self.config))
+        self.update_num_sim()
 
     def __iter__(self) -> Iterator[Tuple[List[Tuple[str, Any]], Params]]:
         """iterates through all possible parameters, yielding a config as well as a flattened
@@ -343,14 +345,18 @@ class ParamSequence:
     def __repr__(self) -> str:
         return f"dispatcher generated from config {self.name}"
 
-    def update_num_sim(self, num_sim):
+    def update_num_sim(self):
+        num_sim = self.count_variations()
         self.num_sim = num_sim
         self.num_steps = self.num_sim * self.config.z_num
         self.single_sim = self.num_sim == 1
 
+    def count_variations(self) -> int:
+        return count_variations(self.config)
+
 
 class ContinuationParamSequence(ParamSequence):
-    def __init__(self, prev_sim_dir: os.PathLike, new_config_dict: Dict[str, Any]):
+    def __init__(self, prev_sim_dir: os.PathLike, new_config: BareConfig):
         """Parameter sequence that builds on a previous simulation but with a new configuration
         It is recommended that only the fiber and the number of points stored may be changed and
         changing other parameters could results in unexpected behaviors. The new config doesn't have to
@@ -364,30 +370,19 @@ class ContinuationParamSequence(ParamSequence):
             new config
         """
         self.prev_sim_dir = Path(prev_sim_dir)
-        init_config = io.load_config(self.prev_sim_dir / "initial_config.toml")
+        self.bare_configs = io.load_config_sequence(new_config.previous_config_file)
-
+        self.bare_configs.append(new_config)
-        new_variable_keys = set(new_config_dict.get("variable", {}).keys())
+        self.bare_configs[0] = Config.from_bare(self.bare_configs[0])
-        new_config = utils.override_config(new_config_dict, init_config)
+        final_config = utils.final_config_from_sequence(*self.bare_configs)
-        super().__init__(new_config)
+        super().__init__(final_config)
-        additional_sims_factor = int(
-            np.prod(
-                [
-                    len(init_config.variable[k])
-                    for k in (new_variable_keys & init_config.variable.keys())
-                ]
-            )
-        )
-        self.update_num_sim(self.num_sim * additional_sims_factor)
 
     def __iter__(self) -> Iterator[Tuple[List[Tuple[str, Any]], Params]]:
         """iterates through all possible parameters, yielding a config as well as a flattened
         computed parameters set each time"""
-        for variable_list, bare_params in required_simulations(self.config):
+        for variable_list, bare_params in required_simulations(*self.bare_configs):
-            variable_list.insert(1, ("prev_data_dir", None))
+            prev_data_dir = self.find_prev_data_dirs(variable_list)[0]
-            for prev_data_dir in self.find_prev_data_dirs(variable_list):
+            bare_params.prev_data_dir = str(prev_data_dir.resolve())
-                variable_list[1] = ("prev_data_dir", str(prev_data_dir.name))
+            yield variable_list, Params.from_bare(bare_params)
-                bare_params.prev_data_dir = str(prev_data_dir.resolve())
-                yield variable_list, Params.from_bare(bare_params)
 
     def find_prev_data_dirs(self, new_variable_list: List[Tuple[str, Any]]) -> List[Path]:
         """finds the previous simulation data that this new config should start from
@@ -419,6 +414,17 @@ class ContinuationParamSequence(ParamSequence):
 
         return path_dic[max_in_common]
 
+    def count_variations(self) -> int:
+        return count_variations(*self.bare_configs)
+
+
+def count_variations(*bare_configs: BareConfig) -> int:
+    sim_num = 1
+    for conf in bare_configs:
+        for l in conf.variable.values():
+            sim_num *= len(l)
+    return sim_num * (bare_configs[0].repeat or 1)
+
 
 class RecoveryParamSequence(ParamSequence):
     def __init__(self, config_dict, task_id):
@@ -506,7 +512,7 @@ class RecoveryParamSequence(ParamSequence):
         return path_dic[max_in_common]
 
 
-def validate_config_sequence(*configs: os.PathLike) -> Tuple[Config, int]:
+def validate_config_sequence(*configs: os.PathLike) -> tuple[str, int]:
     """validates a sequence of configs where all but the first one may have
     parameters missing
 
@@ -517,19 +523,17 @@ def validate_config_sequence(*configs: os.PathLike) -> Tuple[Config, int]:
 
     Returns
     -------
-    Dict[str, Any]
+    int
-        the final config as would be simulated, but of course missing input fields in the middle
+        total number of simulations
     """
+
     previous = None
-    variables = set()
     for config in configs:
         if (p := Path(config)).is_dir():
             config = p / "initial_config.toml"
-        dico = io.load_toml(config)
+        new_conf = io.load_config(config)
-        previous = Config.from_bare(override_config(dico, previous))
+        previous = Config.from_bare(override_config(new_conf, previous))
-        variables |= {(k, tuple(v)) for k, v in previous.variable.items()}
+    return previous.name, count_variations(*configs)
-        variables.add(("repeat", range(previous.repeat)))
-    return previous, int(np.product([len(v) for _, v in variables if len(v) > 0]))
 
 
 def wspace(t, t_num=0):

src/scgenerator/io.py

@@ -162,6 +162,36 @@ def load_config(path: os.PathLike) -> BareConfig:
     return BareConfig(**config)
 
 
+def load_config_sequence(*config_paths: os.PathLike) -> list[BareConfig]:
+    """Loads a sequence of
+
+    Parameters
+    ----------
+    config_paths : os.PathLike
+        either one path (the last config containing previous_config_file parameter)
+        or a list of config path in the order they have to be simulated
+
+    Returns
+    -------
+    list[BareConfig]
+        all loaded configs
+    """
+    if config_paths[0] is None:
+        return []
+    all_configs = [load_config(config_paths[0])]
+    if len(config_paths) == 1:
+        while True:
+            if all_configs[0].previous_config_file is not None:
+                all_configs.insert(0, load_config(all_configs[0].previous_config_file))
+            else:
+                break
+    else:
+        for i, path in enumerate(config_paths[1:]):
+            all_configs.append(load_config(path))
+            all_configs[i + 1].previous_config_file = config_paths[i]
+    return all_configs
+
+
 def load_material_dico(name: str) -> dict[str, Any]:
     """loads a material dictionary
     Parameters

src/scgenerator/physics/simulate.py

@@ -679,18 +679,11 @@ def run_simulation_sequence(
     method=None,
     prev_sim_dir: os.PathLike = None,
 ):
-    config_files = list(config_files)
+    configs = io.load_config_sequence(*config_files)
-    if len(config_files) == 1:
-        while True:
-            conf = io.load_toml(config_files[0])
-            if (prev := conf.get("previous_config_file")) is not None:
-                config_files.insert(0, prev)
-            else:
-                break
 
     prev = prev_sim_dir
-    for config_file in config_files:
+    for config in configs:
-        sim = new_simulation(config_file, prev, method)
+        sim = new_simulation(config, prev, method)
         sim.run()
         prev = sim.sim_dir
     path_trees = io.build_path_trees(sim.sim_dir)
@@ -703,25 +696,21 @@ def run_simulation_sequence(
 
 
 def new_simulation(
-    config: Union[dict, os.PathLike],
+    config: utils.BareConfig,
     prev_sim_dir=None,
     method: Type[Simulations] = None,
 ) -> Simulations:
-    if isinstance(config, dict):
-        config_dict = config
-    else:
-        config_dict = io.load_toml(config)
     logger = get_logger(__name__)
 
     if prev_sim_dir is not None:
-        config_dict["prev_sim_dir"] = str(prev_sim_dir)
+        config.prev_sim_dir = str(prev_sim_dir)
 
     task_id = random.randint(1e9, 1e12)
 
     if prev_sim_dir is None:
-        param_seq = initialize.ParamSequence(config_dict)
+        param_seq = initialize.ParamSequence(config)
     else:
-        param_seq = initialize.ContinuationParamSequence(prev_sim_dir, config_dict)
+        param_seq = initialize.ContinuationParamSequence(prev_sim_dir, config)
 
     logger.info(f"running {param_seq.name}")
 

src/scgenerator/physics/units.py

@@ -159,6 +159,11 @@ def D_ps_nm_km(D: _T) -> _T:
     return 1e-6 * D
 
 
+@unit("OTHER", r"a.u.")
+def unity(x: _T) -> _T:
+    return x
+
+
 @unit("TEMPERATURE", r"Temperature (K)")
 def K(t: _T) -> _T:
     return t
@@ -229,7 +234,7 @@ def standardize_dictionary(dico):
     return dico
 
 
-def sort_axis(axis, plt_range: PlotRange):
+def sort_axis(axis, plt_range: PlotRange) -> tuple[np.ndarray, np.ndarray, tuple[float, float]]:
     """
     given an axis, returns this axis cropped according to the given range, converted and sorted
 

src/scgenerator/scripts/__init__.py

@@ -14,6 +14,7 @@ from ..initialize import ParamSequence
 from ..physics import units, fiber
 from ..spectra import Pulse
 from ..utils import pretty_format_value, pretty_format_from_file_name, auto_crop
+from ..plotting import plot_setup
 from .. import env, math
 
 
@@ -33,20 +34,22 @@ def plot_all(sim_dir: Path, limits: list[str], **opts):
     limits = [
         tuple(func(el) for func, el in zip([float, float, str], lim.split(","))) for lim in limits
     ]
-    print(limits)
     with tqdm(total=len(dir_list) * len(limits)) as bar:
         for p in dir_list:
             pulse = Pulse(p)
             for left, right, unit in limits:
+                path, fig, ax = plot_setup(
+                    pulse.path.parent / f"{pulse.path.name}_{left:.1f}_{right:.1f}_{unit}"
+                )
                 pulse.plot_2D(
                     left,
                     right,
                     unit,
-                    file_name=p.parent
+                    ax,
-                    / f"{pretty_format_from_file_name(p.name)} {left} {right} {unit}",
                     **opts,
                 )
                 bar.update()
+                fig.savefig(path, bbox_inches="tight")
             plt.close("all")
 
 

src/scgenerator/scripts/slurm_submit.py

@@ -127,22 +127,20 @@ def main():
         )
 
     if args.command == "merge":
-        final_config = load_config(Path(args.configs[0]) / "initial_config.toml")
+        final_name = load_config(Path(args.configs[0]) / "initial_config.toml").name
         sim_num = "many"
         args.nodes = 1
         args.cpus_per_node = 1
     else:
         config_paths = args.configs
-        final_config, sim_num = validate_config_sequence(*config_paths)
+        final_name, sim_num = validate_config_sequence(*config_paths)
 
         args.nodes, args.cpus_per_node = distribute(sim_num, args.nodes, args.cpus_per_node)
 
-    submit_path = Path(
+    submit_path = Path("submit " + final_name + "-" + format(datetime.now(), "%Y%m%d%H%M") + ".sh")
-        "submit " + final_config.name + "-" + format(datetime.now(), "%Y%m%d%H%M") + ".sh"
-    )
     tmp_path = Path("submit tmp.sh")
 
-    job_name = f"supercontinuum {final_config.name}"
+    job_name = f"supercontinuum {final_name}"
     submit_sh = template.format(
         job_name=job_name, configs_list=" ".join(f'"{c}"' for c in args.configs), **vars(args)
     )

src/scgenerator/spectra.py

@@ -1,22 +1,17 @@
 import os
 from collections.abc import Sequence
 from pathlib import Path
-from re import UNICODE
+from typing import Callable, Dict, Iterable, Union
-from typing import Callable, Dict, Iterable, Optional, Union
-from matplotlib.pyplot import subplot
-from dataclasses import replace
 
+import matplotlib.pyplot as plt
 import numpy as np
-from numpy.lib import utils
-from numpy.lib.arraysetops import isin
-from tqdm.std import Bar
 
 from . import initialize, io, math
-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 .physics import pulse, units
-from .utils.parameter import BareParams, validator_and
+from .plotting import mean_values_plot, propagation_plot, single_position_plot
+from .utils.parameter import BareParams
 
 
 class Spectrum(np.ndarray):
@@ -255,7 +250,7 @@ class Pulse(Sequence):
     def _to_time_amp(self, spectrum):
         return np.fft.ifft(spectrum)
 
-    def all_spectra(self, ind) -> Spectrum:
+    def all_spectra(self, ind=None) -> Spectrum:
         """
         loads the data already simulated.
         defauft shape is (z_targets, n, nt)
@@ -318,35 +313,38 @@ class Pulse(Sequence):
         left: float,
         right: float,
         unit: Union[Callable[[float], float], str],
+        ax: plt.Axes,
         z_pos: Union[int, Iterable[int]] = None,
         sim_ind: int = 0,
         **kwargs,
     ):
         plt_range, vals = self.retrieve_plot_values(left, right, unit, z_pos, sim_ind)
-        return plot_results_2D(vals, plt_range, self.params, **kwargs)
+        return propagation_plot(vals, plt_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,
     ):
         plt_range, vals = self.retrieve_plot_values(left, right, unit, z_pos, sim_ind)
-        return plot_results_1D(vals, plt_range, self.params, **kwargs)
+        return single_position_plot(vals, plt_range, self.params, ax, **kwargs)
 
-    def plot_avg(
+    def plot_mean(
         self,
         left: float,
         right: float,
         unit: Union[Callable[[float], float], str],
+        ax: plt.Axes,
         z_pos: int,
         **kwargs,
     ):
         plt_range, vals = self.retrieve_plot_values(left, right, unit, z_pos, slice(None))
-        return plot_avg(vals, plt_range, self.params, **kwargs)
+        return mean_values_plot(vals, plt_range, self.params, ax, **kwargs)
 
     def retrieve_plot_values(self, left, right, unit, z_pos, sim_ind):
         plt_range = units.PlotRange(left, right, unit)

src/scgenerator/utils/__init__.py

@@ -182,13 +182,6 @@ def progress_worker(
             pbars[0].update()
 
 
-def count_variations(config: BareConfig) -> int:
-    """returns (sim_num, variable_params_num) where sim_num is the total number of simulations required and
-    variable_params_num is the number of distinct parameters that will vary."""
-    sim_num = np.prod([len(l) for l in config.variable.values()]) * config.repeat
-    return int(sim_num)
-
-
 def format_variable_list(l: List[Tuple[str, Any]]):
     joints = 2 * PARAM_SEPARATOR
     str_list = []
@@ -229,7 +222,7 @@ def pretty_format_from_file_name(name: str) -> str:
     return PARAM_SEPARATOR.join(out)
 
 
-def variable_iterator(config: BareConfig) -> Iterator[Tuple[List[Tuple[str, Any]], BareParams]]:
+def variable_iterator(config: BareConfig) -> Iterator[Tuple[List[Tuple[str, Any]], dict[str, Any]]]:
     """given a config with "variable" parameters, iterates through every possible combination,
     yielding a a list of (parameter_name, value) tuples and a full config dictionary.
 
@@ -240,10 +233,10 @@ def variable_iterator(config: BareConfig) -> Iterator[Tuple[List[Tuple[str, Any]
 
     Yields
     -------
-    Iterator[Tuple[List[Tuple[str, Any]], BareParams]]
+    Iterator[Tuple[List[Tuple[str, Any]], dict[str, Any]]]
         variable_list : a list of (name, value) tuple of parameter name and value that are variable.
 
-        params : a BareParams obj for one simulation
+        params : a dict[str, Any] to be fed to Params
     """
     possible_keys = []
     possible_ranges = []
@@ -264,10 +257,12 @@ def variable_iterator(config: BareConfig) -> Iterator[Tuple[List[Tuple[str, Any]
         param_dict = asdict(config)
         param_dict.pop("variable")
         param_dict.update(indiv_config)
-        yield variable_list, BareParams(**param_dict)
+        yield variable_list, param_dict
 
 
-def required_simulations(config: BareConfig) -> Iterator[Tuple[List[Tuple[str, Any]], BareParams]]:
+def required_simulations(
+    *configs: BareConfig,
+) -> Iterator[Tuple[List[Tuple[str, Any]], BareParams]]:
     """takes the output of `scgenerator.utils.variable_iterator` which is a new dict per different
     parameter set and iterates through every single necessary simulation
 
@@ -281,22 +276,49 @@ def required_simulations(config: BareConfig) -> Iterator[Tuple[List[Tuple[str, A
         dict : a config dictionary for one simulation
     """
     i = 0  # unique sim id
-    for variable_only, bare_params in variable_iterator(config):
+    for data in itertools.product(*[variable_iterator(config) for config in configs]):
-        for j in range(config.repeat):
+        all_variable_only, all_params_dict = list(zip(*data))
+        params_dict = all_params_dict[0]
+        for p in all_params_dict[1:]:
+            params_dict.update({k: v for k, v in p.items() if v is not None})
+        variable_only = reduce_all_variable(all_variable_only)
+        for j in range(configs[0].repeat or 1):
             variable_ind = [("id", i)] + variable_only + [("num", j)]
             i += 1
-            yield variable_ind, bare_params
+            yield variable_ind, BareParams(**params_dict)
 
 
-def override_config(new: Dict[str, Any], old: BareConfig = None) -> BareConfig:
+def reduce_all_variable(all_variable: list[list[tuple[str, Any]]]) -> list[tuple[str, Any]]:
+    out = []
+    for n, variable_list in enumerate(all_variable):
+        out += [("fiber", "ABCDEFGHIJKLMNOPQRSTUVWXYZ"[n % 26] * (n // 26 + 1)), *variable_list]
+    return out
+
+
+def override_config(new: BareConfig, old: BareConfig = None) -> BareConfig:
     """makes sure all the parameters set in new are there, leaves untouched parameters in old"""
+    new_dict = asdict(new)
     if old is None:
-        return BareConfig(**new)
+        return BareConfig(**new_dict)
     variable = deepcopy(old.variable)
-    variable.update(new.pop("variable", {}))  # add new variable
+    new_dict = {k: v for k, v in new_dict.items() if v is not None}
-    for k in new:
+
-        variable.pop(k, None)  # remove old ones
+    for k, v in new_dict.pop("variable", {}).items():
-    return replace(old, variable=variable, **{k: None for k in variable}, **new)
+        variable[k] = v
+    for k in variable:
+        new_dict[k] = None
+    return replace(old, variable=variable, **new_dict)
+
+
+def final_config_from_sequence(*configs: BareConfig) -> BareConfig:
+    if len(configs) == 0:
+        raise ValueError("Must provide at least one config")
+    if len(configs) == 1:
+        return configs[0]
+    elif len(configs) == 2:
+        return override_config(*configs[::-1])
+    else:
+        return override_config(configs[-1], final_config_from_sequence(*configs[:-1]))
 
 
 def auto_crop(x: np.ndarray, y: np.ndarray, rel_thr: float = 0.01) -> np.ndarray:

testing/test_new_iterator.py

@@ -0,0 +1,14 @@
+import scgenerator as sc
+from pathlib import Path
+
+p = Path("/Users/benoitsierro/Nextcloud/PhD/Supercontinuum/PCF Simulations/PPP")
+
+configs = [
+    sc.io.load_config(p / c)
+    for c in ("PM1550.toml", "PMHNLF_appended.toml", "PM2000_appended.toml")
+]
+
+for variable, params in sc.utils.required_simulations(*configs):
+    print(variable)
+
+# sc.initialize.ContinuationParamSequence(configs[-1])