From f56cbb674a84d34d8b63f72df0ab2617c83a2001 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Sierro?= Date: Thu, 23 Sep 2021 10:28:21 +0200 Subject: [PATCH 01/11] work in progress --- src/scgenerator/utils/__init__.py | 11 -- src/scgenerator/utils/parameter.py | 197 --------------------------- src/scgenerator/utils/utils.py | 12 ++ src/scgenerator/utils/variationer.py | 164 ++++++++++++++++++++++ 4 files changed, 176 insertions(+), 208 deletions(-) create mode 100644 src/scgenerator/utils/utils.py create mode 100644 src/scgenerator/utils/variationer.py diff --git a/src/scgenerator/utils/__init__.py b/src/scgenerator/utils/__init__.py index 9ca5ac3..8062bd3 100644 --- a/src/scgenerator/utils/__init__.py +++ b/src/scgenerator/utils/__init__.py @@ -238,17 +238,6 @@ def update_appended_params(source: Path, destination: Path, z: Sequence): save_toml(destination, params) -def to_62(i: int) -> str: - arr = [] - if i == 0: - return "0" - i = abs(i) - while i: - i, value = divmod(i, 62) - arr.append(str_printable[value]) - return "".join(reversed(arr)) - - def build_path_trees(sim_dir: Path) -> list[PathTree]: sim_dir = sim_dir.resolve() path_branches: list[tuple[Path, ...]] = [] diff --git a/src/scgenerator/utils/parameter.py b/src/scgenerator/utils/parameter.py index 6822aa7..17fe9f4 100644 --- a/src/scgenerator/utils/parameter.py +++ b/src/scgenerator/utils/parameter.py @@ -1022,71 +1022,6 @@ class Configuration: return param -class DataPather: - def __init__(self, dl: list[dict[str, Any]]): - self.dict_list = dl - - def vary_list_iterator( - self, index: int - ) -> Generator[tuple[tuple[tuple[int, ...]], list[list[tuple[str, Any]]]], None, None]: - """iterates through every possible combination of a list of dict of lists - - Parameters - ---------- - index : int - up to where in the stored dict_list to go - - Yields - ------- - list[list[tuple[str, Any]]] - list of list of (key, value) pairs - - Example - ------- - - self.dict_list = [{a:[56, 57], b:["?", "!"]}, {c:[0, -1]}] -> - [ - [[(a, 56), (b, "?")], [(c, 0)]], - [[(a, 56), (b, "?")], [(c, 1)]], - [[(a, 56), (b, "!")], [(c, 0)]], - [[(a, 56), (b, "!")], [(c, 1)]], - [[(a, 57), (b, "?")], [(c, 0)]], - [[(a, 57), (b, "?")], [(c, 1)]], - [[(a, 57), (b, "!")], [(c, 0)]], - [[(a, 57), (b, "!")], [(c, 1)]], - ] - """ - if index < 0: - index = len(self.dict_list) - index - d_tem_list = [el for d in self.dict_list[: index + 1] for el in d.items()] - dict_pos = np.cumsum([0] + [len(d) for d in self.dict_list[: index + 1]]) - ranges = [range(len(l)) for _, l in d_tem_list] - - for r in itertools.product(*ranges): - flat = [(d_tem_list[i][0], d_tem_list[i][1][j]) for i, j in enumerate(r)] - pos = tuple(r) - out = [flat[left:right] for left, right in zip(dict_pos[:-1], dict_pos[1:])] - pos = tuple(pos[left:right] for left, right in zip(dict_pos[:-1], dict_pos[1:])) - yield pos, out - - def all_vary_list(self, index): - for sim_index, l in self.vary_list_iterator(index): - unique_vary: list[tuple[str, Any]] = [] - for ll in l[: index + 1]: - for pname, pval in ll: - for i, (pn, _) in enumerate(unique_vary): - if pn == pname: - del unique_vary[i] - break - unique_vary.append((pname, pval)) - yield sim_index, format_variable_list( - reduce_all_variable(l[:index]), add_iden=True - ), format_variable_list(reduce_all_variable(l), add_iden=True), unique_vary - - def __repr__(self): - return f"DataPather([{', '.join(repr(d) for d in self.dict_list)}])" - - @dataclass(frozen=True) class PlotRange: left: float = Parameter(type_checker(int, float)) @@ -1202,70 +1137,6 @@ def _mock_function(num_args: int, num_returns: int) -> Callable: return out_func -def format_variable_list(l: list[tuple[str, Any]], add_iden=False) -> str: - """formats a variable list into a str such that each simulation has a unique - directory name. A u_XXX unique identifier and b_XXX (ignoring repeat simulations) - branch identifier are added at the beginning. - - Parameters - ---------- - l : list[tuple[str, Any]] - list of variable parameters - add_iden : bool - add unique simulation and parameter-set identifiers - - Returns - ------- - str - directory name - """ - str_list = [] - for p_name, p_value in l: - ps = p_name.replace("/", "").replace(PARAM_SEPARATOR, "") - vs = format_value(p_name, p_value).replace("/", "").replace(PARAM_SEPARATOR, "") - str_list.append(ps + PARAM_SEPARATOR + vs) - tmp_name = PARAM_SEPARATOR.join(str_list) - if not add_iden: - return tmp_name - unique_id = unique_identifier(l) - branch_id = branch_identifier(l) - return unique_id + PARAM_SEPARATOR + branch_id + PARAM_SEPARATOR + tmp_name - - -def branch_identifier(l): - branch_id = "b_" + utils.to_62(hash(str([el for el in l if el[0] != "num"]))) - return branch_id - - -def unique_identifier(l): - unique_id = "u_" + utils.to_62(hash(str(l))) - return unique_id - - -def format_value(name: str, value) -> str: - if value is True or value is False: - return str(value) - elif isinstance(value, (float, int)): - try: - return getattr(Parameters, name).display(value) - except AttributeError: - return format(value, ".9g") - elif isinstance(value, (list, tuple, np.ndarray)): - return "-".join([str(v) for v in value]) - elif isinstance(value, str): - p = Path(value) - if p.exists(): - return p.stem - return str(value) - - -def pretty_format_value(name: str, value) -> str: - try: - return getattr(Parameters, name).display(value) - except AttributeError: - return name + PARAM_SEPARATOR + str(value) - - def pretty_format_from_sim_name(name: str) -> str: """formats a pretty version of a simulation directory @@ -1289,74 +1160,6 @@ def pretty_format_from_sim_name(name: str) -> str: return PARAM_SEPARATOR.join(out) -def variable_iterator( - config: dict[str, Any], first: bool -) -> Generator[tuple[list[tuple[str, Any]], dict[str, Any]], None, None]: - """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. - - Parameters - ---------- - config : BareConfig - initial config obj - first : int - whether it is the first fiber or not (only the first fiber get a sim number) - - Yields - ------- - 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 dict[str, Any] to be fed to Parameters - """ - possible_keys = [] - possible_ranges = [] - - for key, values in config.get("variable", {}).items(): - possible_keys.append(key) - possible_ranges.append(range(len(values))) - - combinations = itertools.product(*possible_ranges) - - master_index = 0 - repeat = config.get("repeat", 1) if first else 1 - for combination in combinations: - indiv_config = {} - variable_list = [] - for i, key in enumerate(possible_keys): - parameter_value = config["variable"][key][combination[i]] - indiv_config[key] = parameter_value - variable_list.append((key, parameter_value)) - param_dict = deepcopy(config) - param_dict.pop("variable") - param_dict.update(indiv_config) - for repeat_index in range(repeat): - # variable_ind = [("id", master_index)] + variable_list - variable_ind = variable_list - if first: - variable_ind += [("num", repeat_index)] - yield variable_ind, param_dict - master_index += 1 - - -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 strip_vary_list(all_variable: T) -> T: - if len(all_variable) == 0: - return all_variable - elif isinstance(all_variable[0], Sequence) and ( - len(all_variable[0]) == 0 or not isinstance(all_variable[0][0], str) - ): - return [strip_vary_list(el) for el in all_variable] - else: - return [el for el in all_variable if el[0] != "num"] - - default_rules: list[Rule] = [ # Grid *Rule.deduce( diff --git a/src/scgenerator/utils/utils.py b/src/scgenerator/utils/utils.py new file mode 100644 index 0000000..b81744d --- /dev/null +++ b/src/scgenerator/utils/utils.py @@ -0,0 +1,12 @@ +from string import printable as str_printable + + +def to_62(i: int) -> str: + arr = [] + if i == 0: + return "0" + i = abs(i) + while i: + i, value = divmod(i, 62) + arr.append(str_printable[value]) + return "".join(reversed(arr)) diff --git a/src/scgenerator/utils/variationer.py b/src/scgenerator/utils/variationer.py new file mode 100644 index 0000000..cfed7e4 --- /dev/null +++ b/src/scgenerator/utils/variationer.py @@ -0,0 +1,164 @@ +from pydantic import BaseModel, validator +from typing import Union, Iterable, Generator, Any +from collections.abc import Sequence, MutableMapping +import itertools +from ..const import PARAM_SEPARATOR +from . import utils +import numpy as np +from pathlib import Path + + +def format_value(name: str, value) -> str: + if value is True or value is False: + return str(value) + elif isinstance(value, (float, int)): + try: + return getattr(Parameters, name).display(value) + except AttributeError: + return format(value, ".9g") + elif isinstance(value, (list, tuple, np.ndarray)): + return "-".join([str(v) for v in value]) + elif isinstance(value, str): + p = Path(value) + if p.exists(): + return p.stem + return str(value) + + +def pretty_format_value(name: str, value) -> str: + try: + return getattr(Parameters, name).display(value) + except AttributeError: + return name + PARAM_SEPARATOR + str(value) + + +class HashableBaseModel(BaseModel): + """Pydantic BaseModel that's immutable and can be hashed""" + + def __hash__(self) -> int: + return hash(type(self)) + sum(hash(v) for v in self.__dict__.values()) + + class Config: + allow_mutation = False + + +class VariationSpecsError(ValueError): + pass + + +class Variationer: + """ + manages possible combinations of values given dicts of lists + + Example + ------- + `var = Variationer([dict(a=[1, 2]), [dict(b=["000", "111"], c=["a", "-1"])]])` + + """ + + all_indices: list[list[int]] + all_dicts: list[list[dict[str, list]]] + + def __init__(self, variables: Iterable[Union[list[MutableMapping], MutableMapping]]): + self.all_indices = [] + self.all_dicts = [] + for i, el in enumerate(variables): + if not isinstance(el, Sequence): + el = [{k: v} for k, v in el.items()] + else: + el = list(el) + self.append(el) + + def append(self, var_list: list[dict[str, list]]): + num_vars = [] + for d in var_list: + values = list(d.values()) + len_to_test = len(values[0]) + if not all(len(v) == len_to_test for v in values[1:]): + raise VariationSpecsError( + f"variable items should all have the same number of parameters" + ) + num_vars.append(len_to_test) + if len(num_vars) == 0: + num_vars = [1] + self.all_indices.append(num_vars) + self.all_dicts.append(var_list) + + def iterate(self, index: int = -1) -> Generator["SimulationDescriptor", None, None]: + if index < 0: + index = len(self.all_indices) + index + 1 + flattened_indices = sum(self.all_indices[:index], []) + index_positions = np.cumsum([0] + [len(i) for i in self.all_indices[:index]]) + ranges = [range(i) for i in flattened_indices] + for r in itertools.product(*ranges): + out: list[list[tuple[str, Any]]] = [] + for i, (start, end) in enumerate(zip(index_positions[:-1], index_positions[1:])): + out.append([]) + for value_index, var_d in zip(r[start:end], self.all_dicts[i]): + for k, v in var_d.items(): + out[-1].append((k, v[value_index])) + yield SimulationDescriptor(raw_descr=out) + + +class SimulationDescriptor(HashableBaseModel): + raw_descr: tuple[tuple[tuple[str, Any], ...], ...] + separator: str = "fiber" + + def __str__(self) -> str: + return self.descriptor(add_identifier=False) + + def descriptor(self, add_identifier=False) -> str: + """formats a variable list into a str such that each simulation has a unique + directory name. A u_XXX unique identifier and b_XXX (ignoring repeat simulations) + branch identifier can added at the beginning. + + Parameters + ---------- + add_identifier : bool + add unique simulation and parameter-set identifiers + + Returns + ------- + str + simulation descriptor + """ + str_list = [] + + for p_name, p_value in self.flat: + ps = p_name.replace("/", "").replace("\\", "").replace(PARAM_SEPARATOR, "") + vs = format_value(p_name, p_value).replace("/", "").replace(PARAM_SEPARATOR, "") + str_list.append(ps + PARAM_SEPARATOR + vs) + tmp_name = PARAM_SEPARATOR.join(str_list) + if not add_identifier: + return tmp_name + return ( + self.identifier + PARAM_SEPARATOR + self.branch.identifier + PARAM_SEPARATOR + tmp_name + ) + + @property + def flat(self) -> list[tuple[str, Any]]: + out = [] + for n, variable_list in enumerate(self.raw_descr): + out += [ + (self.separator, "ABCDEFGHIJKLMNOPQRSTUVWXYZ"[n % 26] * (n // 26 + 1)), + *variable_list, + ] + return out + + @property + def branch(self) -> "BranchDescriptor": + return SimulationDescriptor(raw_descr=self.raw_descr, separator=self.separator) + + @property + def identifier(self) -> str: + return "u_" + utils.to_62(hash(str(self.flat))) + + +class BranchDescriptor(SimulationDescriptor): + @property + def identifier(self) -> str: + return "b_" + utils.to_62(hash(str(self.flat))) + + @validator("raw_descr") + def validate_raw_descr(cls, v): + return tuple(tuple(el for el in variable if el[0] != "num") for variable in v) From ef48711aa6db5029a6dc3e374534e26c1fefab89 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Sierro?= Date: Thu, 23 Sep 2021 10:40:17 +0200 Subject: [PATCH 02/11] added docstrings --- src/scgenerator/utils/variationer.py | 32 +++++++++++++++++++++++++++- 1 file changed, 31 insertions(+), 1 deletion(-) diff --git a/src/scgenerator/utils/variationer.py b/src/scgenerator/utils/variationer.py index cfed7e4..2876957 100644 --- a/src/scgenerator/utils/variationer.py +++ b/src/scgenerator/utils/variationer.py @@ -52,7 +52,15 @@ class Variationer: Example ------- - `var = Variationer([dict(a=[1, 2]), [dict(b=["000", "111"], c=["a", "-1"])]])` + `>> var = Variationer([dict(a=[1, 2]), [dict(b=["000", "111"], c=["a", "-1"])]]) + list(v.raw_descr for v in var.iterate()) + + [ + ((("a", 1),), (("b", "000"), ("c", "a"))), + ((("a", 1),), (("b", "111"), ("c", "-1"))), + ((("a", 2),), (("b", "000"), ("c", "a"))), + ((("a", 2),), (("b", "111"), ("c", "-1"))), + ]` """ @@ -70,6 +78,28 @@ class Variationer: self.append(el) def append(self, var_list: list[dict[str, list]]): + """append a list of variable parameter sets + each call to append creates a new group of parameters + + Parameters + ---------- + var_list : list[dict[str, list]] + each dict in the list is treated as an independent parameter + this means that if for one dict, len > 1, the lists of possible values + must be the same length + + Example + ------- + `>> append([dict(wavelength=[800e-9, 900e-9], power=[1e3, 2e3]), dict(length=[3e-2, 3.5e-2, 4e-2])])` + + means that for every parameter variations, wavelength=800e-9 will always occur when power=1e3 and + vice versa, while length is free to vary independently + + Raises + ------ + VariationSpecsError + raised when possible values lists in a same dict are not the same length + """ num_vars = [] for d in var_list: values = list(d.values()) From 695ac3bd730d8de9a5247bda40cd3f8c32c41de1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Sierro?= Date: Mon, 27 Sep 2021 13:21:02 +0200 Subject: [PATCH 03/11] Sim starts, merge not --- src/scgenerator/__init__.py | 3 +- src/scgenerator/physics/simulate.py | 20 +- src/scgenerator/plotting.py | 3 +- src/scgenerator/scripts/__init__.py | 5 +- src/scgenerator/spectra.py | 3 +- src/scgenerator/utils/__init__.py | 50 +++-- src/scgenerator/utils/parameter.py | 293 +++++++-------------------- src/scgenerator/utils/utils.py | 134 ++++++++++++ src/scgenerator/utils/variationer.py | 143 +++++++------ 9 files changed, 335 insertions(+), 319 deletions(-) diff --git a/src/scgenerator/__init__.py b/src/scgenerator/__init__.py index 4c9e73d..997e495 100644 --- a/src/scgenerator/__init__.py +++ b/src/scgenerator/__init__.py @@ -5,4 +5,5 @@ 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 .utils import Paths, open_config, parameter -from .utils.parameter import Configuration, Parameters, PlotRange +from .utils.parameter import Configuration, Parameters +from .utils.utils import PlotRange diff --git a/src/scgenerator/physics/simulate.py b/src/scgenerator/physics/simulate.py index 4d34743..477d9c5 100644 --- a/src/scgenerator/physics/simulate.py +++ b/src/scgenerator/physics/simulate.py @@ -11,7 +11,7 @@ from send2trash import send2trash from .. import env, utils from ..logger import get_logger -from ..utils.parameter import Configuration, Parameters, format_variable_list +from ..utils.parameter import Configuration, Parameters from . import pulse from .fiber import create_non_linear_op, fast_dispersion_op @@ -466,14 +466,14 @@ class Simulations: self.configuration = configuration - self.name = self.configuration.final_path - self.sim_dir = self.configuration.final_sim_dir + self.name = self.configuration.name + self.sim_dir = self.configuration.final_path self.configuration.save_parameters() self.sim_jobs_per_node = 1 def finished_and_complete(self): - for sim in self.configuration.all_configs_dict.values(): + for sim in self.configuration.all_configs.values(): if ( self.configuration.sim_status(sim.output_path)[0] != self.configuration.State.COMPLETE @@ -487,7 +487,7 @@ class Simulations: def _run_available(self): for variable, params in self.configuration: - v_list_str = format_variable_list(variable, add_iden=True) + v_list_str = variable.formatted_descriptor(True) utils.save_parameters(params.prepare_for_dump(), Path(params.output_path)) self.new_sim(v_list_str, params) @@ -526,7 +526,9 @@ class SequencialSimulations(Simulations, priority=0): def __init__(self, configuration: Configuration, task_id): super().__init__(configuration, task_id=task_id) self.pbars = utils.PBars( - self.configuration.total_num_steps, "Simulating " + self.configuration.final_path, 1 + self.configuration.total_num_steps, + "Simulating " + self.configuration.final_path.name, + 1, ) self.configuration.skip_callback = lambda num: self.pbars.update(0, num) @@ -569,7 +571,7 @@ class MultiProcSimulations(Simulations, priority=1): self.p_worker = multiprocessing.Process( target=utils.progress_worker, args=( - self.configuration.final_path, + self.configuration.final_path.name, self.sim_jobs_per_node, self.configuration.total_num_steps, self.progress_queue, @@ -716,7 +718,7 @@ def run_simulation( sim = new_simulation(config, method) sim.run() - path_trees = utils.build_path_trees(config.sim_dirs[-1]) + path_trees = utils.build_path_trees(config.fiber_paths[-1]) final_name = env.get(env.OUTPUT_PATH) if final_name is None: @@ -724,7 +726,7 @@ def run_simulation( utils.merge(final_name, path_trees) try: - send2trash(config.sim_dirs) + send2trash(config.fiber_paths) except (PermissionError, OSError): get_logger(__name__).error("Could not send temporary directories to trash") diff --git a/src/scgenerator/plotting.py b/src/scgenerator/plotting.py index fbaa903..712bf12 100644 --- a/src/scgenerator/plotting.py +++ b/src/scgenerator/plotting.py @@ -14,7 +14,8 @@ from .const import PARAM_SEPARATOR from .defaults import default_plotting as defaults from .math import abs2, span from .physics import pulse, units -from .utils.parameter import Parameters, PlotRange, sort_axis +from .utils.parameter import Parameters +from .utils.utils import PlotRange, sort_axis RangeType = tuple[float, float, Union[str, Callable]] NO_LIM = object() diff --git a/src/scgenerator/scripts/__init__.py b/src/scgenerator/scripts/__init__.py index ede8211..029268f 100644 --- a/src/scgenerator/scripts/__init__.py +++ b/src/scgenerator/scripts/__init__.py @@ -16,9 +16,8 @@ from ..utils import auto_crop, open_config, save_toml, translate_parameters from ..utils.parameter import ( Configuration, Parameters, - pretty_format_from_sim_name, - pretty_format_value, ) +from ..utils.variationer import VariationDescriptor def fingerprint(params: Parameters): @@ -46,7 +45,7 @@ def plot_all(sim_dir: Path, limits: list[str], show=False, **opts): path, fig, ax = plot_setup( pulse.path.parent / ( - pretty_format_from_sim_name(pulse.path.name) + pulse.path.name + PARAM_SEPARATOR + f"{left:.1f}{PARAM_SEPARATOR}{right:.1f}{PARAM_SEPARATOR}{unit}" ) diff --git a/src/scgenerator/spectra.py b/src/scgenerator/spectra.py index 1ec018c..fdd6b78 100644 --- a/src/scgenerator/spectra.py +++ b/src/scgenerator/spectra.py @@ -16,7 +16,8 @@ from .plotting import ( single_position_plot, transform_2D_propagation, ) -from .utils.parameter import Parameters, PlotRange +from .utils.parameter import Parameters +from .utils.utils import PlotRange from .utils import load_spectrum diff --git a/src/scgenerator/utils/__init__.py b/src/scgenerator/utils/__init__.py index 8062bd3..fdcbcea 100644 --- a/src/scgenerator/utils/__init__.py +++ b/src/scgenerator/utils/__init__.py @@ -161,29 +161,35 @@ def save_toml(path: os.PathLike, dico): return dico -def load_config_sequence(final_config_path: os.PathLike) -> tuple[list[dict[str, Any]], str]: - loaded_config = open_config(final_config_path) - final_name = loaded_config.get("name") - fiber_list = loaded_config.pop("Fiber") - configs = [] - if fiber_list is not None: - master_variable = loaded_config.get("variable", {}) - for i, params in enumerate(fiber_list): - params.setdefault("variable", master_variable if i == 0 else {}) - if i == 0: - params["variable"] |= master_variable - configs.append(loaded_config | params) - else: - configs.append(loaded_config) - while "previous_config_file" in configs[0]: - configs.insert(0, open_config(configs[0]["previous_config_file"])) - configs[0].setdefault("variable", {}) - for pre, nex in zip(configs[:-1], configs[1:]): - variable = nex.pop("variable", {}) - nex.update({k: v for k, v in pre.items() if k not in nex}) - nex["variable"] = variable +def load_config_sequence(path: os.PathLike) -> tuple[Path, list[dict[str, Any]]]: + """loads a configuration file - return configs, final_name + Parameters + ---------- + path : os.PathLike + path to the config toml file + + Returns + ------- + final_path : Path + output name of the simulation + list[dict[str, Any]] + one config per fiber + + """ + loaded_config = open_config(path) + + fiber_list: list[dict[str, Any]] = loaded_config.pop("Fiber") + if len(fiber_list) == 0: + raise ValueError(f"No fiber in config {path}") + final_path = loaded_config.get("name") + configs = [] + for i, params in enumerate(fiber_list): + params.setdefault("variable", {}) + configs.append(loaded_config | params) + configs[0]["variable"] = loaded_config.get("variable", {}) | configs[0]["variable"] + + return Path(final_path), configs def save_parameters( diff --git a/src/scgenerator/utils/parameter.py b/src/scgenerator/utils/parameter.py index 17fe9f4..fa43c58 100644 --- a/src/scgenerator/utils/parameter.py +++ b/src/scgenerator/utils/parameter.py @@ -11,14 +11,18 @@ from dataclasses import asdict, dataclass, fields from functools import cache, lru_cache from pathlib import Path from typing import Any, Callable, Generator, Iterable, Literal, Optional, Sequence, TypeVar, Union + import numpy as np from numpy.lib import isin +from scgenerator.utils import ensure_folder, variationer from .. import math, utils from ..const import PARAM_FN, PARAM_SEPARATOR, __version__ from ..errors import EvaluatorError, NoDefaultError from ..logger import get_logger from ..physics import fiber, materials, pulse, units +from ..utils.variationer import VariationDescriptor, Variationer +from .utils import func_rewrite, _mock_function, get_arg_names T = TypeVar("T") @@ -256,7 +260,7 @@ class Parameter: ---------- tpe : type type of the paramter - validators : Callable[[str, Any], None] + validator : Callable[[str, Any], None] signature : validator(name, value) must raise a ValueError when value doesn't fit the criteria checked by validator. name is passed to validator to be included in the error message @@ -290,7 +294,6 @@ class Parameter: if isinstance(value, Parameter): defaut = None if self.default is None else copy(self.default) instance.__dict__[self.name] = defaut - # instance.__dict__[self.name] = None else: if value is not None: if self.converter is not None: @@ -768,9 +771,11 @@ class Configuration: obj with the output path of the simulation saved in its output_path attribute. """ - master_configs: list[dict[str, Any]] - sim_dirs: list[Path] + fiber_configs: list[dict[str, Any]] + master_config: dict[str, Any] + fiber_paths: list[Path] num_sim: int + num_fibers: int repeat: int z_num: int total_num_steps: int @@ -778,19 +783,17 @@ class Configuration: parallel: bool overwrite: bool final_path: str - all_configs_dict: dict[tuple[tuple[int, ...], ...], "Configuration.__SimConfig"] - all_configs_list: list[list["Configuration.__SimConfig"]] + all_configs: dict[tuple[tuple[int, ...], ...], "Configuration.__SimConfig"] @dataclass(frozen=True) class __SimConfig: - vary_list: list[tuple[str, Any]] + descriptor: VariationDescriptor config: dict[str, Any] output_path: Path - index: tuple[tuple[int, ...], ...] @property def sim_num(self) -> int: - return len(self.index) + return len(self.descriptor.index) class State(enum.Enum): COMPLETE = enum.auto() @@ -810,48 +813,48 @@ class Configuration: ): self.logger = get_logger(__name__) - self.master_configs, self.final_path = utils.load_config_sequence(final_config_path) - if self.final_path is None: - self.final_path = Parameters.name.default - self.name = Path(self.final_path).name + self.overwrite = overwrite + self.final_path, self.fiber_configs = utils.load_config_sequence(final_config_path) + self.final_path = utils.ensure_folder( + self.final_path, mkdir=False, prevent_overwrite=not self.overwrite + ) + self.master_config = self.fiber_configs[0] + self.name = self.final_path.name self.z_num = 0 self.total_num_steps = 0 - self.sim_dirs = [] - self.overwrite = overwrite + self.fiber_paths = [] + self.all_configs = {} self.skip_callback = skip_callback - self.worker_num = self.master_configs[0].get("worker_num", max(1, os.cpu_count() // 2)) - self.repeat = self.master_configs[0].get("repeat", 1) + self.worker_num = self.master_config.get("worker_num", max(1, os.cpu_count() // 2)) + self.repeat = self.master_config.get("repeat", 1) + self.variationer = Variationer() - names = set() - for i, config in enumerate(self.master_configs): + fiber_names = set() + self.num_fibers = 0 + for i, config in enumerate(self.fiber_configs): + config.setdefault("name", Parameters.name.default) self.z_num += config["z_num"] - config.setdefault("name", f"{Parameters.name.default} {i}") - given_name = config["name"] - fn_i = 0 - while config["name"] in names: - config["name"] = given_name + f"_{fn_i}" - fn_i += 1 - names.add(config["name"]) - - self.sim_dirs.append( + fiber_names.add(config["name"]) + self.variationer.append(config.pop("variable")) + self.fiber_paths.append( utils.ensure_folder( - Path("_".join(["_", self.name, Path(config["name"]).name, "_"])), + self.fiber_path(i, config), mkdir=False, prevent_overwrite=not self.overwrite, ) ) self.__validate_variable(config) - self.__compute_sim_dirs() - [Evaluator.evaluate_default(c[0].config, True) for c in self.all_configs_list] - self.num_sim = len(self.all_configs_list[-1]) + self.num_fibers += 1 + Evaluator.evaluate_default(config, True) + self.num_sim = self.variationer.var_num() self.total_num_steps = sum( - config["z_num"] * len(self.all_configs_list[i]) - for i, config in enumerate(self.master_configs) + config["z_num"] * self.variationer.var_num(i) + for i, config in enumerate(self.fiber_configs) ) - self.final_sim_dir = utils.ensure_folder( - Path(self.master_configs[-1]["name"]), mkdir=False, prevent_overwrite=not self.overwrite - ) - self.parallel = self.master_configs[0].get("parallel", Parameters.parallel.default) + self.parallel = self.master_config.get("parallel", Parameters.parallel.default) + + def fiber_path(self, i: int, full_config: dict[str, Any]) -> Path: + return self.final_path / PARAM_SEPARATOR.join([format(i), self.name, full_config["name"]]) def __validate_variable(self, config: dict[str, Any]): for k, v in config.get("variable", {}).items(): @@ -862,76 +865,62 @@ class Configuration: if len(v) == 0: raise ValueError(f"variable parameter {k!r} must not be empty") - def __compute_sim_dirs(self): - self.all_configs_dict = {} - self.all_configs_list = [] - self.master_configs[0]["variable"]["num"] = list( - range(self.master_configs[0].get("repeat", 1)) - ) - dp = DataPather([c["variable"] for c in self.master_configs]) - for i, conf in enumerate(self.master_configs): - self.all_configs_list.append([]) - for sim_index, prev_path, this_path, this_vary in dp.all_vary_list(i): - this_conf = conf.copy() + def __iter__(self) -> Generator[tuple[VariationDescriptor, Parameters], None, None]: + for i in range(self.num_fibers): + for sim_config in self.iterate_single_fiber(i): + if i > 0: - prev_path = utils.ensure_folder( - self.sim_dirs[i - 1] / prev_path, not self.overwrite, False + sim_config.config["prev_data_dir"] = str( + self.fiber_paths[i - 1] / sim_config.descriptor[:i].formatted_descriptor() ) - this_conf["prev_data_dir"] = str(prev_path) - - this_path = utils.ensure_folder( - self.sim_dirs[i] / this_path, not self.overwrite, False - ) - this_conf.pop("variable") - conf_to_use = {k: v for k, v in this_vary if k != "num"} | this_conf - self.all_configs_dict[sim_index] = self.__SimConfig( - this_vary, conf_to_use, this_path, sim_index - ) - self.all_configs_list[i].append(self.all_configs_dict[sim_index]) - - def __iter__(self) -> Generator[tuple[list[tuple[str, Any]], Parameters], None, None]: - for i, sim_config_list in enumerate(self.all_configs_list): - for sim_config, params in self.__iter_1_sim(sim_config_list): + params = Parameters(**sim_config.config) + params.compute() fiber_map = [] for j in range(i + 1): - this_conf = self.all_configs_dict[sim_config.index[: j + 1]].config + this_conf = self.all_configs[sim_config.descriptor.index[: j + 1]].config if j > 0: - prev_conf = self.all_configs_dict[sim_config.index[:j]].config + prev_conf = self.all_configs[sim_config.descriptor.index[:j]].config length = prev_conf["length"] + fiber_map[j - 1][0] else: length = 0.0 fiber_map.append((length, this_conf["name"])) - params.output_path = str(sim_config.output_path) params.fiber_map = fiber_map - yield sim_config.vary_list, params + yield sim_config.descriptor, params - def __iter_1_sim( - self, configs: list["Configuration.__SimConfig"] - ) -> Generator[tuple["Configuration.__SimConfig", Parameters], None, None]: + def iterate_single_fiber( + self, index: int + ) -> Generator["Configuration.__SimConfig", None, None]: """iterates through the parameters of only one fiber. It takes care of recovering partially completed simulations, skipping complete ones and waiting for the previous fiber to finish Parameters ---------- - configs : list[__SimConfig] - list of configuration obj + index : int + which fiber to iterate over Yields ------- __SimConfig configuration obj - Parameters - computed Parameters obj """ - sim_dict: dict[Path, Configuration.__SimConfig] = {s.output_path: s for s in configs} + sim_dict: dict[Path, self.__SimConfig] = {} + for descr in self.variationer.iterate(index): + cfg = descr.update_config(self.fiber_configs[index]) + p = ensure_folder( + self.fiber_paths[index] / descr.formatted_descriptor(), + not self.overwrite, + False, + ) + cfg["output_path"] = str(p) + sim_config = self.__SimConfig(descr, cfg, p) + sim_dict[p] = sim_config + self.all_configs[sim_config.descriptor.index] = sim_config while len(sim_dict) > 0: for data_dir, sim_config in sim_dict.items(): task, config_dict = self.__decide(sim_config) if task == self.Action.RUN: sim_dict.pop(data_dir) - p = Parameters(**config_dict) - p.compute() - yield sim_config, p + yield sim_config if "recovery_last_stored" in config_dict and self.skip_callback is not None: self.skip_callback(config_dict["recovery_last_stored"]) break @@ -956,7 +945,7 @@ class Configuration: Returns ------- - str : {'run', 'wait', 'skip'} + str : Configuration.Action what to do config_dict : dict[str, Any] config dictionary. The only key possibly modified is 'prev_data_dir', which @@ -1012,7 +1001,7 @@ class Configuration: raise ValueError(f"Too many spectra in {data_dir}") def save_parameters(self): - for config, sim_dir in zip(self.master_configs, self.sim_dirs): + for config, sim_dir in zip(self.fiber_configs, self.fiber_paths): os.makedirs(sim_dir, exist_ok=True) utils.save_toml(sim_dir / f"initial_config.toml", config) @@ -1022,144 +1011,6 @@ class Configuration: return param -@dataclass(frozen=True) -class PlotRange: - left: float = Parameter(type_checker(int, float)) - right: float = Parameter(type_checker(int, float)) - unit: Callable[[float], float] = Parameter(units.is_unit, converter=units.get_unit) - conserved_quantity: bool = Parameter(boolean, default=True) - - def __str__(self): - return f"{self.left:.1f}-{self.right:.1f} {self.unit.__name__}" - - def sort_axis(self, axis: np.ndarray) -> tuple[np.ndarray, np.ndarray, tuple[float, float]]: - return sort_axis(axis, self) - - def __iter__(self): - yield self.left - yield self.right - yield self.unit.__name__ - - -def sort_axis( - axis: np.ndarray, 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 - - Parameters - ---------- - axis : 1D array containing the original axis (usual the w or t array) - plt_range : tupple (min, max, conversion_function) used to crop the axis - - Returns - ------- - cropped : the axis cropped, converted and sorted - indices : indices to use to slice and sort other array in the same fashion - extent : tupple with min and max of cropped - - Example - ------- - w = np.append(np.linspace(0, -10, 20), np.linspace(0, 10, 20)) - t = np.linspace(-10, 10, 400) - W, T = np.meshgrid(w, t) - y = np.exp(-W**2 - T**2) - - # Define ranges - rw = (-4, 4, s) - rt = (-2, 6, s) - - w, cw = sort_axis(w, rw) - t, ct = sort_axis(t, rt) - - # slice y according to the given ranges - y = y[ct][:, cw] - """ - if isinstance(plt_range, tuple): - plt_range = PlotRange(*plt_range) - r = np.array((plt_range.left, plt_range.right), dtype="float") - - indices = np.arange(len(axis))[ - (axis <= np.max(plt_range.unit(r))) & (axis >= np.min(plt_range.unit(r))) - ] - cropped = axis[indices] - order = np.argsort(plt_range.unit.inv(cropped)) - indices = indices[order] - cropped = cropped[order] - out_ax = plt_range.unit.inv(cropped) - - return out_ax, indices, (out_ax[0], out_ax[-1]) - - -def get_arg_names(func: Callable) -> list[str]: - spec = inspect.getfullargspec(func) - args = spec.args - if spec.defaults is not None and len(spec.defaults) > 0: - args = args[: -len(spec.defaults)] - return args - - -def validate_arg_names(names: list[str]): - for n in names: - if re.match(r"^[^\s\-'\(\)\"\d][^\(\)\-\s'\"]*$", n) is None: - raise ValueError(f"{n} is an invalid parameter name") - - -def func_rewrite(func: Callable, kwarg_names: list[str], arg_names: list[str] = None) -> Callable: - if arg_names is None: - arg_names = get_arg_names(func) - else: - validate_arg_names(arg_names) - validate_arg_names(kwarg_names) - sign_arg_str = ", ".join(arg_names + kwarg_names) - call_arg_str = ", ".join(arg_names + [f"{s}={s}" for s in kwarg_names]) - tmp_name = f"{func.__name__}_0" - func_str = f"def {tmp_name}({sign_arg_str}):\n return __func__({call_arg_str})" - scope = dict(__func__=func) - exec(func_str, scope) - out_func = scope[tmp_name] - out_func.__module__ = "evaluator" - return out_func - - -@cache -def _mock_function(num_args: int, num_returns: int) -> Callable: - if not isinstance(num_args, int) and isinstance(num_returns, int): - raise TypeError(f"num_args and num_returns must be int") - arg_str = ", ".join("a" * (n + 1) for n in range(num_args)) - return_str = ", ".join("True" for _ in range(num_returns)) - func_name = f"__mock_{num_args}_{num_returns}" - func_str = f"def {func_name}({arg_str}):\n return {return_str}" - scope = {} - exec(func_str, scope) - out_func = scope[func_name] - out_func.__module__ = "evaluator" - return out_func - - -def pretty_format_from_sim_name(name: str) -> str: - """formats a pretty version of a simulation directory - - Parameters - ---------- - name : str - name of the simulation (directory name) - - Returns - ------- - str - prettier name - """ - s = name.split(PARAM_SEPARATOR) - out = [] - for key, value in zip(s[::2], s[1::2]): - try: - out += [key.replace("_", " "), getattr(Parameters, key).display(float(value))] - except (AttributeError, ValueError): - out.append(key + PARAM_SEPARATOR + value) - return PARAM_SEPARATOR.join(out) - - default_rules: list[Rule] = [ # Grid *Rule.deduce( diff --git a/src/scgenerator/utils/utils.py b/src/scgenerator/utils/utils.py index b81744d..a18aacb 100644 --- a/src/scgenerator/utils/utils.py +++ b/src/scgenerator/utils/utils.py @@ -1,4 +1,23 @@ +import inspect +import re +from functools import cache from string import printable as str_printable +from typing import Callable + +import numpy as np +from pydantic import BaseModel + +from ..physics.units import get_unit + + +class HashableBaseModel(BaseModel): + """Pydantic BaseModel that's immutable and can be hashed""" + + def __hash__(self) -> int: + return hash(type(self)) + sum(hash(v) for v in self.__dict__.values()) + + class Config: + allow_mutation = False def to_62(i: int) -> str: @@ -10,3 +29,118 @@ def to_62(i: int) -> str: i, value = divmod(i, 62) arr.append(str_printable[value]) return "".join(reversed(arr)) + + +class PlotRange(HashableBaseModel): + left: float + right: float + unit: Callable[[float], float] + conserved_quantity: bool = True + + def __init__(self, left, right, unit, **kwargs): + super().__init__(left=left, right=right, unit=get_unit(unit), **kwargs) + + def __str__(self): + return f"{self.left:.1f}-{self.right:.1f} {self.unit.__name__}" + + def sort_axis(self, axis: np.ndarray) -> tuple[np.ndarray, np.ndarray, tuple[float, float]]: + return sort_axis(axis, self) + + def __iter__(self): + yield self.left + yield self.right + yield self.unit.__name__ + + +def sort_axis( + axis: np.ndarray, 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 + + Parameters + ---------- + axis : 1D array containing the original axis (usual the w or t array) + plt_range : tupple (min, max, conversion_function) used to crop the axis + + Returns + ------- + cropped : the axis cropped, converted and sorted + indices : indices to use to slice and sort other array in the same fashion + extent : tupple with min and max of cropped + + Example + ------- + w = np.append(np.linspace(0, -10, 20), np.linspace(0, 10, 20)) + t = np.linspace(-10, 10, 400) + W, T = np.meshgrid(w, t) + y = np.exp(-W**2 - T**2) + + # Define ranges + rw = (-4, 4, s) + rt = (-2, 6, s) + + w, cw = sort_axis(w, rw) + t, ct = sort_axis(t, rt) + + # slice y according to the given ranges + y = y[ct][:, cw] + """ + if isinstance(plt_range, tuple): + plt_range = PlotRange(*plt_range) + r = np.array((plt_range.left, plt_range.right), dtype="float") + + indices = np.arange(len(axis))[ + (axis <= np.max(plt_range.unit(r))) & (axis >= np.min(plt_range.unit(r))) + ] + cropped = axis[indices] + order = np.argsort(plt_range.unit.inv(cropped)) + indices = indices[order] + cropped = cropped[order] + out_ax = plt_range.unit.inv(cropped) + + return out_ax, indices, (out_ax[0], out_ax[-1]) + + +def get_arg_names(func: Callable) -> list[str]: + spec = inspect.getfullargspec(func) + args = spec.args + if spec.defaults is not None and len(spec.defaults) > 0: + args = args[: -len(spec.defaults)] + return args + + +def validate_arg_names(names: list[str]): + for n in names: + if re.match(r"^[^\s\-'\(\)\"\d][^\(\)\-\s'\"]*$", n) is None: + raise ValueError(f"{n} is an invalid parameter name") + + +def func_rewrite(func: Callable, kwarg_names: list[str], arg_names: list[str] = None) -> Callable: + if arg_names is None: + arg_names = get_arg_names(func) + else: + validate_arg_names(arg_names) + validate_arg_names(kwarg_names) + sign_arg_str = ", ".join(arg_names + kwarg_names) + call_arg_str = ", ".join(arg_names + [f"{s}={s}" for s in kwarg_names]) + tmp_name = f"{func.__name__}_0" + func_str = f"def {tmp_name}({sign_arg_str}):\n return __func__({call_arg_str})" + scope = dict(__func__=func) + exec(func_str, scope) + out_func = scope[tmp_name] + out_func.__module__ = "evaluator" + return out_func + + +@cache +def _mock_function(num_args: int, num_returns: int) -> Callable: + arg_str = ", ".join("a" * (n + 1) for n in range(num_args)) + return_str = ", ".join("True" for _ in range(num_returns)) + func_name = f"__mock_{num_args}_{num_returns}" + func_str = f"def {func_name}({arg_str}):\n return {return_str}" + scope = {} + exec(func_str, scope) + out_func = scope[func_name] + out_func.__module__ = "evaluator" + return out_func diff --git a/src/scgenerator/utils/variationer.py b/src/scgenerator/utils/variationer.py index 2876957..d53777d 100644 --- a/src/scgenerator/utils/variationer.py +++ b/src/scgenerator/utils/variationer.py @@ -1,45 +1,14 @@ -from pydantic import BaseModel, validator -from typing import Union, Iterable, Generator, Any -from collections.abc import Sequence, MutableMapping +from math import prod import itertools +from collections.abc import MutableMapping, Sequence +from pathlib import Path +from typing import Any, Callable, Generator, Iterable, Union + +import numpy as np +from pydantic import validator + from ..const import PARAM_SEPARATOR from . import utils -import numpy as np -from pathlib import Path - - -def format_value(name: str, value) -> str: - if value is True or value is False: - return str(value) - elif isinstance(value, (float, int)): - try: - return getattr(Parameters, name).display(value) - except AttributeError: - return format(value, ".9g") - elif isinstance(value, (list, tuple, np.ndarray)): - return "-".join([str(v) for v in value]) - elif isinstance(value, str): - p = Path(value) - if p.exists(): - return p.stem - return str(value) - - -def pretty_format_value(name: str, value) -> str: - try: - return getattr(Parameters, name).display(value) - except AttributeError: - return name + PARAM_SEPARATOR + str(value) - - -class HashableBaseModel(BaseModel): - """Pydantic BaseModel that's immutable and can be hashed""" - - def __hash__(self) -> int: - return hash(type(self)) + sum(hash(v) for v in self.__dict__.values()) - - class Config: - allow_mutation = False class VariationSpecsError(ValueError): @@ -67,23 +36,20 @@ class Variationer: all_indices: list[list[int]] all_dicts: list[list[dict[str, list]]] - def __init__(self, variables: Iterable[Union[list[MutableMapping], MutableMapping]]): + def __init__(self, variables: Iterable[Union[list[MutableMapping], MutableMapping]] = None): self.all_indices = [] self.all_dicts = [] - for i, el in enumerate(variables): - if not isinstance(el, Sequence): - el = [{k: v} for k, v in el.items()] - else: - el = list(el) - self.append(el) + if variables is not None: + for i, el in enumerate(variables): + self.append(el) - def append(self, var_list: list[dict[str, list]]): + def append(self, var_list: Union[list[MutableMapping], MutableMapping]): """append a list of variable parameter sets each call to append creates a new group of parameters Parameters ---------- - var_list : list[dict[str, list]] + var_list : Union[list[MutableMapping], MutableMapping] each dict in the list is treated as an independent parameter this means that if for one dict, len > 1, the lists of possible values must be the same length @@ -100,6 +66,10 @@ class Variationer: VariationSpecsError raised when possible values lists in a same dict are not the same length """ + if not isinstance(var_list, Sequence): + var_list = [{k: v} for k, v in var_list.items()] + else: + var_list = list(var_list) num_vars = [] for d in var_list: values = list(d.values()) @@ -114,30 +84,43 @@ class Variationer: self.all_indices.append(num_vars) self.all_dicts.append(var_list) - def iterate(self, index: int = -1) -> Generator["SimulationDescriptor", None, None]: - if index < 0: - index = len(self.all_indices) + index + 1 - flattened_indices = sum(self.all_indices[:index], []) - index_positions = np.cumsum([0] + [len(i) for i in self.all_indices[:index]]) + def iterate(self, index: int = -1) -> Generator["VariationDescriptor", None, None]: + index = self.__index(index) + flattened_indices = sum(self.all_indices[: index + 1], []) + index_positions = np.cumsum([0] + [len(i) for i in self.all_indices[: index + 1]]) ranges = [range(i) for i in flattened_indices] for r in itertools.product(*ranges): out: list[list[tuple[str, Any]]] = [] + indicies: list[list[int]] = [] for i, (start, end) in enumerate(zip(index_positions[:-1], index_positions[1:])): out.append([]) + indicies.append([]) for value_index, var_d in zip(r[start:end], self.all_dicts[i]): for k, v in var_d.items(): out[-1].append((k, v[value_index])) - yield SimulationDescriptor(raw_descr=out) + indicies[-1].append(value_index) + yield VariationDescriptor(raw_descr=out, index=indicies) + + def __index(self, index: int) -> int: + if index < 0: + index = len(self.all_indices) + index + return index + + def var_num(self, index: int = -1) -> int: + index = self.__index(index) + return max(1, prod(prod(el) for el in self.all_indices[: index + 1])) -class SimulationDescriptor(HashableBaseModel): +class VariationDescriptor(utils.HashableBaseModel): raw_descr: tuple[tuple[tuple[str, Any], ...], ...] + index: tuple[tuple[int, ...], ...] separator: str = "fiber" + _format_registry: dict[str, Callable[..., str]] = {} def __str__(self) -> str: - return self.descriptor(add_identifier=False) + return self.formatted_descriptor(add_identifier=False) - def descriptor(self, add_identifier=False) -> str: + def formatted_descriptor(self, add_identifier=False) -> str: """formats a variable list into a str such that each simulation has a unique directory name. A u_XXX unique identifier and b_XXX (ignoring repeat simulations) branch identifier can added at the beginning. @@ -156,7 +139,7 @@ class SimulationDescriptor(HashableBaseModel): for p_name, p_value in self.flat: ps = p_name.replace("/", "").replace("\\", "").replace(PARAM_SEPARATOR, "") - vs = format_value(p_name, p_value).replace("/", "").replace(PARAM_SEPARATOR, "") + vs = self.format_value(p_name, p_value).replace("/", "").replace(PARAM_SEPARATOR, "") str_list.append(ps + PARAM_SEPARATOR + vs) tmp_name = PARAM_SEPARATOR.join(str_list) if not add_identifier: @@ -165,6 +148,34 @@ class SimulationDescriptor(HashableBaseModel): self.identifier + PARAM_SEPARATOR + self.branch.identifier + PARAM_SEPARATOR + tmp_name ) + @classmethod + def register_formatter(cls, p_name: str, func: Callable[..., str]): + cls._format_registry[p_name] = func + + def format_value(self, name: str, value) -> str: + if value is True or value is False: + return str(value) + elif isinstance(value, (float, int)): + try: + return self._format_registry[name](value) + except KeyError: + return format(value, ".9g") + elif isinstance(value, (list, tuple, np.ndarray)): + return "-".join([str(v) for v in value]) + elif isinstance(value, str): + p = Path(value) + if p.exists(): + return p.stem + return str(value) + + def __getitem__(self, key) -> "VariationDescriptor": + return VariationDescriptor( + raw_descr=self.raw_descr[key], index=self.index[key], separator=self.separator + ) + + def update_config(self, cfg: dict[str, Any]): + return cfg | {k: v for k, v in self.raw_descr[-1]} + @property def flat(self) -> list[tuple[str, Any]]: out = [] @@ -177,17 +188,27 @@ class SimulationDescriptor(HashableBaseModel): @property def branch(self) -> "BranchDescriptor": - return SimulationDescriptor(raw_descr=self.raw_descr, separator=self.separator) + for i in reversed(range(len(self.raw_descr))): + for j in reversed(range(len(self.raw_descr[i]))): + if self.raw_descr[i][j][0] == "num": + del self.raw_descr[i][j] + return VariationDescriptor( + raw_descr=self.raw_descr, index=self.index, separator=self.separator + ) @property def identifier(self) -> str: return "u_" + utils.to_62(hash(str(self.flat))) -class BranchDescriptor(SimulationDescriptor): +class BranchDescriptor(VariationDescriptor): + __ids: dict[int, int] = {} + @property def identifier(self) -> str: - return "b_" + utils.to_62(hash(str(self.flat))) + branch_id = hash(str(self.flat)) + self.__ids.setdefault(branch_id, len(self.__ids)) + return str(self.__ids[branch_id]) @validator("raw_descr") def validate_raw_descr(cls, v): From 262a5b9701d54c52e6831de0eeb1b7a57aa44d78 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Sierro?= Date: Tue, 28 Sep 2021 13:28:00 +0200 Subject: [PATCH 04/11] partway trough big data structure revamp --- play.py | 19 -- src/scgenerator/__init__.py | 9 +- src/scgenerator/{utils => _utils}/__init__.py | 8 +- src/scgenerator/{utils => _utils}/cache.py | 0 .../{utils => _utils}/parameter.py | 67 ++--- src/scgenerator/{utils => _utils}/utils.py | 57 ++++ .../{utils => _utils}/variationer.py | 38 ++- src/scgenerator/cli/cli.py | 5 +- src/scgenerator/const.py | 3 +- src/scgenerator/env.py | 4 +- src/scgenerator/math.py | 2 +- src/scgenerator/physics/__init__.py | 4 +- src/scgenerator/physics/fiber.py | 4 +- src/scgenerator/physics/materials.py | 2 +- src/scgenerator/physics/pulse.py | 2 - src/scgenerator/physics/simulate.py | 18 +- src/scgenerator/physics/units.py | 1 - src/scgenerator/plotting.py | 4 +- src/scgenerator/scripts/__init__.py | 5 +- src/scgenerator/scripts/slurm_submit.py | 4 +- src/scgenerator/spectra.py | 258 ++++++++++-------- testing/long_tests/test_recovery_param_seq.py | 35 --- testing/test_initialize.py | 216 --------------- testing/test_pulse.py | 41 --- testing/test_utils.py | 67 ----- testing/test_variationer.py | 54 ++++ 26 files changed, 349 insertions(+), 578 deletions(-) rename src/scgenerator/{utils => _utils}/__init__.py (98%) rename src/scgenerator/{utils => _utils}/cache.py (100%) rename src/scgenerator/{utils => _utils}/parameter.py (94%) rename src/scgenerator/{utils => _utils}/utils.py (69%) rename src/scgenerator/{utils => _utils}/variationer.py (87%) delete mode 100644 testing/long_tests/test_recovery_param_seq.py delete mode 100644 testing/test_initialize.py delete mode 100644 testing/test_pulse.py delete mode 100644 testing/test_utils.py create mode 100644 testing/test_variationer.py diff --git a/play.py b/play.py index 9f687f3..d0cc8a8 100644 --- a/play.py +++ b/play.py @@ -4,22 +4,3 @@ import scgenerator as sc import matplotlib.pyplot as plt from pathlib import Path from pprint import pprint - - -def _main(): - print(os.getcwd()) - for v_list, params in sc.Configuration("PM1550+PM2000D+PM1550/Pos30000.toml"): - print(params.fiber_map) - - -def main(): - drr = os.getcwd() - os.chdir("/Users/benoitsierro/Nextcloud/PhD/Supercontinuum/PCF Simulations") - try: - _main() - finally: - os.chdir(drr) - - -if __name__ == "__main__": - main() diff --git a/src/scgenerator/__init__.py b/src/scgenerator/__init__.py index 997e495..6cb8fcb 100644 --- a/src/scgenerator/__init__.py +++ b/src/scgenerator/__init__.py @@ -1,9 +1,10 @@ -from . import math, utils +from . import math 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 .utils import Paths, open_config, parameter -from .utils.parameter import Configuration, Parameters -from .utils.utils import PlotRange +from ._utils import Paths, open_config, parameter +from ._utils.parameter import Configuration, Parameters +from ._utils.utils import PlotRange +from ._utils.variationer import Variationer, VariationDescriptor, VariationSpecsError diff --git a/src/scgenerator/utils/__init__.py b/src/scgenerator/_utils/__init__.py similarity index 98% rename from src/scgenerator/utils/__init__.py rename to src/scgenerator/_utils/__init__.py index fdcbcea..03c8859 100644 --- a/src/scgenerator/utils/__init__.py +++ b/src/scgenerator/_utils/__init__.py @@ -25,7 +25,7 @@ import pkg_resources as pkg import toml from tqdm import tqdm -from ..const import PARAM_FN, PARAM_SEPARATOR, SPEC1_FN, SPECN_FN, Z_FN, __version__ +from ..const import PARAM_FN, PARAM_SEPARATOR, SPEC1_FN, SPECN_FN1, Z_FN, __version__ from ..env import pbar_policy from ..logger import get_logger @@ -143,7 +143,8 @@ def resolve_loadfile_arg(dico: dict[str, Any]) -> dict[str, Any]: return dico -def load_toml(descr: str) -> dict[str, Any]: +def load_toml(descr: os.PathLike) -> dict[str, Any]: + descr = str(descr) if ":" in descr: path, entry = descr.split(":", 1) with open(path) as file: @@ -188,6 +189,7 @@ def load_config_sequence(path: os.PathLike) -> tuple[Path, list[dict[str, Any]]] params.setdefault("variable", {}) configs.append(loaded_config | params) configs[0]["variable"] = loaded_config.get("variable", {}) | configs[0]["variable"] + configs[0]["variable"]["num"] = list(range(configs[0].get("repeat", 1))) return Path(final_path), configs @@ -341,7 +343,7 @@ def merge_path_tree( for i, (z, merged_spectra) in enumerate(merge_spectra(path_tree)): z_arr.append(z) - spec_out_name = SPECN_FN.format(i) + spec_out_name = SPECN_FN1.format(i) np.save(destination / spec_out_name, merged_spectra) if z_callback is not None: z_callback(i) diff --git a/src/scgenerator/utils/cache.py b/src/scgenerator/_utils/cache.py similarity index 100% rename from src/scgenerator/utils/cache.py rename to src/scgenerator/_utils/cache.py diff --git a/src/scgenerator/utils/parameter.py b/src/scgenerator/_utils/parameter.py similarity index 94% rename from src/scgenerator/utils/parameter.py rename to src/scgenerator/_utils/parameter.py index fa43c58..f7d7c1a 100644 --- a/src/scgenerator/utils/parameter.py +++ b/src/scgenerator/_utils/parameter.py @@ -14,14 +14,15 @@ from typing import Any, Callable, Generator, Iterable, Literal, Optional, Sequen import numpy as np from numpy.lib import isin -from scgenerator.utils import ensure_folder, variationer -from .. import math, utils +from .. import math from ..const import PARAM_FN, PARAM_SEPARATOR, __version__ from ..errors import EvaluatorError, NoDefaultError from ..logger import get_logger from ..physics import fiber, materials, pulse, units -from ..utils.variationer import VariationDescriptor, Variationer +from .._utils.variationer import VariationDescriptor, Variationer +from .. import _utils as utils +from .. import env from .utils import func_rewrite, _mock_function, get_arg_names T = TypeVar("T") @@ -312,13 +313,6 @@ class Parameter: return f"{num_str} {unit}" -def fiber_map_converter(d: dict[str, str]) -> list[tuple[float, str]]: - if isinstance(d, dict): - return [(float(k), v) for k, v in d.items()] - else: - return [(float(k), v) for k, v in d] - - @dataclass class Parameters: """ @@ -432,15 +426,13 @@ class Parameters: const_qty: np.ndarray = Parameter(type_checker(np.ndarray)) beta_func: Callable[[float], list[float]] = Parameter(func_validator) gamma_func: Callable[[float], float] = Parameter(func_validator) - fiber_map: list[tuple[float, str]] = Parameter( - validator_list(type_checker(tuple)), converter=fiber_map_converter - ) + + num: int = Parameter(non_negative(int)) datetime: datetime_module.datetime = Parameter(type_checker(datetime_module.datetime)) version: str = Parameter(string) def prepare_for_dump(self) -> dict[str, Any]: param = asdict(self) - param["fiber_map"] = [(str(z), n) for z, n in param.get("fiber_map", [])] param = Parameters.strip_params_dict(param) param["datetime"] = datetime_module.datetime.now() param["version"] = __version__ @@ -816,7 +808,9 @@ class Configuration: self.overwrite = overwrite self.final_path, self.fiber_configs = utils.load_config_sequence(final_config_path) self.final_path = utils.ensure_folder( - self.final_path, mkdir=False, prevent_overwrite=not self.overwrite + Path(env.get(env.OUTPUT_PATH, self.final_path)), + mkdir=False, + prevent_overwrite=not self.overwrite, ) self.master_config = self.fiber_configs[0] self.name = self.final_path.name @@ -868,23 +862,8 @@ class Configuration: def __iter__(self) -> Generator[tuple[VariationDescriptor, Parameters], None, None]: for i in range(self.num_fibers): for sim_config in self.iterate_single_fiber(i): - - if i > 0: - sim_config.config["prev_data_dir"] = str( - self.fiber_paths[i - 1] / sim_config.descriptor[:i].formatted_descriptor() - ) params = Parameters(**sim_config.config) params.compute() - fiber_map = [] - for j in range(i + 1): - this_conf = self.all_configs[sim_config.descriptor.index[: j + 1]].config - if j > 0: - prev_conf = self.all_configs[sim_config.descriptor.index[:j]].config - length = prev_conf["length"] + fiber_map[j - 1][0] - else: - length = 0.0 - fiber_map.append((length, this_conf["name"])) - params.fiber_map = fiber_map yield sim_config.descriptor, params def iterate_single_fiber( @@ -903,18 +882,21 @@ class Configuration: __SimConfig configuration obj """ - sim_dict: dict[Path, self.__SimConfig] = {} - for descr in self.variationer.iterate(index): - cfg = descr.update_config(self.fiber_configs[index]) - p = ensure_folder( - self.fiber_paths[index] / descr.formatted_descriptor(), + sim_dict: dict[Path, Configuration.__SimConfig] = {} + for descriptor in self.variationer.iterate(index): + cfg = descriptor.update_config(self.fiber_configs[index]) + if index > 0: + cfg["prev_data_dir"] = str( + self.fiber_paths[index - 1] / descriptor[:index].formatted_descriptor(True) + ) + p = utils.ensure_folder( + self.fiber_paths[index] / descriptor.formatted_descriptor(True), not self.overwrite, False, ) cfg["output_path"] = str(p) - sim_config = self.__SimConfig(descr, cfg, p) - sim_dict[p] = sim_config - self.all_configs[sim_config.descriptor.index] = sim_config + sim_config = self.__SimConfig(descriptor, cfg, p) + sim_dict[p] = self.all_configs[sim_config.descriptor.index] = sim_config while len(sim_dict) > 0: for data_dir, sim_config in sim_dict.items(): task, config_dict = self.__decide(sim_config) @@ -1001,9 +983,12 @@ class Configuration: raise ValueError(f"Too many spectra in {data_dir}") def save_parameters(self): - for config, sim_dir in zip(self.fiber_configs, self.fiber_paths): - os.makedirs(sim_dir, exist_ok=True) - utils.save_toml(sim_dir / f"initial_config.toml", config) + os.makedirs(self.final_path, exist_ok=True) + cfgs = [ + cfg | dict(variable=self.variationer.all_dicts[i]) + for i, cfg in enumerate(self.fiber_configs) + ] + utils.save_toml(self.final_path / f"initial_config.toml", dict(name=self.name, Fiber=cfgs)) @property def first(self) -> Parameters: diff --git a/src/scgenerator/utils/utils.py b/src/scgenerator/_utils/utils.py similarity index 69% rename from src/scgenerator/utils/utils.py rename to src/scgenerator/_utils/utils.py index a18aacb..44eadf3 100644 --- a/src/scgenerator/utils/utils.py +++ b/src/scgenerator/_utils/utils.py @@ -1,12 +1,17 @@ import inspect +import os import re +from collections import defaultdict from functools import cache +from pathlib import Path from string import printable as str_printable from typing import Callable import numpy as np from pydantic import BaseModel +from .._utils import load_toml, save_toml +from ..const import PARAM_FN, Z_FN from ..physics.units import get_unit @@ -144,3 +149,55 @@ def _mock_function(num_args: int, num_returns: int) -> Callable: out_func = scope[func_name] out_func.__module__ = "evaluator" return out_func + + +def combine_simulations(path: Path, dest: Path = None): + """combines raw simulations into one folder per branch + + Parameters + ---------- + path : Path + source of the simulations (must contain u_xx directories) + dest : Path, optional + if given, moves the simulations to dest, by default None + """ + paths: dict[str, list[Path]] = defaultdict(list) + if dest is None: + dest = path + + for p in path.glob("u_*b_*"): + if p.is_dir(): + paths[p.name.split()[1]].append(p) + for l in paths.values(): + l.sort(key=lambda el: re.search(r"(?<=num )[0-9]+", el.name)[0]) + for pulses in paths.values(): + new_path = dest / update_path(pulses[0].name) + os.makedirs(new_path, exist_ok=True) + for num, pulse in enumerate(pulses): + params_ok = False + for file in pulse.glob("*"): + if file.name == PARAM_FN: + if not params_ok: + update_params(new_path, file) + params_ok = True + else: + file.unlink() + elif file.name == Z_FN: + file.rename(new_path / file.name) + else: + file.rename(new_path / (file.stem + f"_{num}" + file.suffix)) + pulse.rmdir() + + +def update_params(new_path: Path, file: Path): + params = load_toml(file) + if (p := params.get("prev_data_dir")) is not None: + p = Path(p) + params["prev_data_dir"] = str(p.parent / update_path(p.name)) + params["output_path"] = str(new_path) + save_toml(new_path / PARAM_FN, params) + file.unlink() + + +def update_path(p: str) -> str: + return re.sub(r"( ?num [0-9]+)|(u_[0-9]+ )", "", p) diff --git a/src/scgenerator/utils/variationer.py b/src/scgenerator/_utils/variationer.py similarity index 87% rename from src/scgenerator/utils/variationer.py rename to src/scgenerator/_utils/variationer.py index d53777d..fdd6368 100644 --- a/src/scgenerator/utils/variationer.py +++ b/src/scgenerator/_utils/variationer.py @@ -116,6 +116,7 @@ class VariationDescriptor(utils.HashableBaseModel): index: tuple[tuple[int, ...], ...] separator: str = "fiber" _format_registry: dict[str, Callable[..., str]] = {} + __ids: dict[int, int] = {} def __str__(self) -> str: return self.formatted_descriptor(add_identifier=False) @@ -173,7 +174,19 @@ class VariationDescriptor(utils.HashableBaseModel): raw_descr=self.raw_descr[key], index=self.index[key], separator=self.separator ) - def update_config(self, cfg: dict[str, Any]): + def update_config(self, cfg: dict[str, Any]) -> dict[str, Any]: + """updates a dictionary with the value of the descriptor + + Parameters + ---------- + cfg : dict[str, Any] + dict to be updated + + Returns + ------- + dict[str, Any] + same as cfg but with key from the descriptor added/updated. + """ return cfg | {k: v for k, v in self.raw_descr[-1]} @property @@ -188,17 +201,22 @@ class VariationDescriptor(utils.HashableBaseModel): @property def branch(self) -> "BranchDescriptor": - for i in reversed(range(len(self.raw_descr))): - for j in reversed(range(len(self.raw_descr[i]))): - if self.raw_descr[i][j][0] == "num": - del self.raw_descr[i][j] - return VariationDescriptor( - raw_descr=self.raw_descr, index=self.index, separator=self.separator - ) + descr = [] + ind = [] + for i, l in enumerate(self.raw_descr): + descr.append([]) + ind.append([]) + for j, (k, v) in enumerate(l): + if k != "num": + descr[-1].append((k, v)) + ind[-1].append(self.index[i][j]) + return BranchDescriptor(raw_descr=descr, index=ind, separator=self.separator) @property def identifier(self) -> str: - return "u_" + utils.to_62(hash(str(self.flat))) + unique_id = hash(str(self.flat)) + self.__ids.setdefault(unique_id, len(self.__ids)) + return "u_" + str(self.__ids[unique_id]) class BranchDescriptor(VariationDescriptor): @@ -208,7 +226,7 @@ class BranchDescriptor(VariationDescriptor): def identifier(self) -> str: branch_id = hash(str(self.flat)) self.__ids.setdefault(branch_id, len(self.__ids)) - return str(self.__ids[branch_id]) + return "b_" + str(self.__ids[branch_id]) @validator("raw_descr") def validate_raw_descr(cls, v): diff --git a/src/scgenerator/cli/cli.py b/src/scgenerator/cli/cli.py index 2532fdf..92b81ab 100644 --- a/src/scgenerator/cli/cli.py +++ b/src/scgenerator/cli/cli.py @@ -1,14 +1,13 @@ import argparse import os import re -import subprocess -import sys from collections import ChainMap from pathlib import Path import numpy as np -from .. import const, env, scripts, utils +from .. import const, env, scripts +from .. import _utils as utils from ..logger import get_logger from ..physics.fiber import dispersion_coefficients from ..physics.simulate import SequencialSimulations, run_simulation diff --git a/src/scgenerator/const.py b/src/scgenerator/const.py index c679236..b8843a2 100644 --- a/src/scgenerator/const.py +++ b/src/scgenerator/const.py @@ -20,7 +20,8 @@ def pbar_format(worker_id: int): SPEC1_FN = "spectrum_{}.npy" -SPECN_FN = "spectra_{}.npy" +SPECN_FN1 = "spectra_{}.npy" +SPEC1_FN_N = "spectrum_{}_{}.npy" Z_FN = "z.npy" PARAM_FN = "params.toml" PARAM_SEPARATOR = " " diff --git a/src/scgenerator/env.py b/src/scgenerator/env.py index 46c0d6e..13f3545 100644 --- a/src/scgenerator/env.py +++ b/src/scgenerator/env.py @@ -48,7 +48,7 @@ def data_folder(task_id: int) -> Optional[str]: return tmp -def get(key: str) -> Any: +def get(key: str, default=None) -> Any: str_value = os.environ.get(key) if isinstance(str_value, str): try: @@ -58,7 +58,7 @@ def get(key: str) -> Any: return t(str_value) except (ValueError, KeyError): pass - return None + return default def all_environ() -> Dict[str, str]: diff --git a/src/scgenerator/math.py b/src/scgenerator/math.py index 20db3dc..505c6b7 100644 --- a/src/scgenerator/math.py +++ b/src/scgenerator/math.py @@ -3,7 +3,7 @@ from typing import Union import numpy as np from scipy.interpolate import griddata, interp1d from scipy.special import jn_zeros -from .utils.cache import np_cache +from ._utils.cache import np_cache pi = np.pi c = 299792458.0 diff --git a/src/scgenerator/physics/__init__.py b/src/scgenerator/physics/__init__.py index 82bc01e..24c3937 100644 --- a/src/scgenerator/physics/__init__.py +++ b/src/scgenerator/physics/__init__.py @@ -10,8 +10,8 @@ from scipy.optimize import minimize_scalar from .. import math from . import fiber, materials, units, pulse -from .. import utils -from ..utils import cache +from .. import _utils +from .._utils import cache T = TypeVar("T") diff --git a/src/scgenerator/physics/fiber.py b/src/scgenerator/physics/fiber.py index ccddbb1..4a27249 100644 --- a/src/scgenerator/physics/fiber.py +++ b/src/scgenerator/physics/fiber.py @@ -7,9 +7,9 @@ from scipy.interpolate import interp1d from ..logger import get_logger -from .. import utils +from .. import _utils from ..math import abs2, argclosest, power_fact, u_nm -from ..utils.cache import np_cache +from .._utils.cache import np_cache from . import materials as mat from . import units from .units import c, pi diff --git a/src/scgenerator/physics/materials.py b/src/scgenerator/physics/materials.py index 812e60e..b9572e6 100644 --- a/src/scgenerator/physics/materials.py +++ b/src/scgenerator/physics/materials.py @@ -5,7 +5,7 @@ from scipy.integrate import cumulative_trapezoid from ..logger import get_logger from . import units -from .. import utils +from .. import _utils from .units import NA, c, kB, me, e, hbar diff --git a/src/scgenerator/physics/pulse.py b/src/scgenerator/physics/pulse.py index 6a07c70..ce8ab29 100644 --- a/src/scgenerator/physics/pulse.py +++ b/src/scgenerator/physics/pulse.py @@ -23,8 +23,6 @@ from scipy.interpolate import UnivariateSpline from scipy.optimize import minimize_scalar from scipy.optimize.optimize import OptimizeResult -from scgenerator import utils - from ..defaults import default_plotting from ..logger import get_logger from ..math import * diff --git a/src/scgenerator/physics/simulate.py b/src/scgenerator/physics/simulate.py index 477d9c5..46b9897 100644 --- a/src/scgenerator/physics/simulate.py +++ b/src/scgenerator/physics/simulate.py @@ -9,9 +9,11 @@ from typing import Any, Generator, Type, Union import numpy as np from send2trash import send2trash -from .. import env, utils +from .. import env +from .. import _utils as utils +from .._utils.utils import combine_simulations from ..logger import get_logger -from ..utils.parameter import Configuration, Parameters +from .._utils.parameter import Configuration, Parameters from . import pulse from .fiber import create_non_linear_op, fast_dispersion_op @@ -718,17 +720,9 @@ def run_simulation( sim = new_simulation(config, method) sim.run() - path_trees = utils.build_path_trees(config.fiber_paths[-1]) - final_name = env.get(env.OUTPUT_PATH) - if final_name is None: - final_name = config.final_path - - utils.merge(final_name, path_trees) - try: - send2trash(config.fiber_paths) - except (PermissionError, OSError): - get_logger(__name__).error("Could not send temporary directories to trash") + for path in config.fiber_paths: + combine_simulations(path) def new_simulation( diff --git a/src/scgenerator/physics/units.py b/src/scgenerator/physics/units.py index 80089fe..f648bb9 100644 --- a/src/scgenerator/physics/units.py +++ b/src/scgenerator/physics/units.py @@ -2,7 +2,6 @@ # 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 dataclasses import dataclass from typing import Callable, TypeVar, Union import numpy as np diff --git a/src/scgenerator/plotting.py b/src/scgenerator/plotting.py index 712bf12..6163520 100644 --- a/src/scgenerator/plotting.py +++ b/src/scgenerator/plotting.py @@ -14,8 +14,8 @@ from .const import PARAM_SEPARATOR from .defaults import default_plotting as defaults from .math import abs2, span from .physics import pulse, units -from .utils.parameter import Parameters -from .utils.utils import PlotRange, sort_axis +from ._utils.parameter import Parameters +from ._utils.utils import PlotRange, sort_axis RangeType = tuple[float, float, Union[str, Callable]] NO_LIM = object() diff --git a/src/scgenerator/scripts/__init__.py b/src/scgenerator/scripts/__init__.py index 029268f..176f236 100644 --- a/src/scgenerator/scripts/__init__.py +++ b/src/scgenerator/scripts/__init__.py @@ -12,12 +12,11 @@ from ..const import PARAM_FN, PARAM_SEPARATOR from ..physics import fiber, units from ..plotting import plot_setup from ..spectra import Pulse -from ..utils import auto_crop, open_config, save_toml, translate_parameters -from ..utils.parameter import ( +from .._utils import auto_crop, open_config, save_toml, translate_parameters +from .._utils.parameter import ( Configuration, Parameters, ) -from ..utils.variationer import VariationDescriptor def fingerprint(params: Parameters): diff --git a/src/scgenerator/scripts/slurm_submit.py b/src/scgenerator/scripts/slurm_submit.py index a977f35..3fc4b6b 100644 --- a/src/scgenerator/scripts/slurm_submit.py +++ b/src/scgenerator/scripts/slurm_submit.py @@ -9,8 +9,8 @@ from typing import Tuple import numpy as np -from ..utils import Paths -from ..utils.parameter import Configuration +from .._utils import Paths +from .._utils.parameter import Configuration def primes(n): diff --git a/src/scgenerator/spectra.py b/src/scgenerator/spectra.py index fdd6b78..82a83b8 100644 --- a/src/scgenerator/spectra.py +++ b/src/scgenerator/spectra.py @@ -1,13 +1,19 @@ +from __future__ import annotations + import os from collections.abc import Sequence from pathlib import Path -from typing import Callable, Dict, Iterable, Optional, Union +from typing import Any, Callable, Dict, Iterable, Optional, Union import matplotlib.pyplot as plt import numpy as np +from pydantic import BaseModel, DirectoryPath, root_validator from . import math -from .const import SPECN_FN +from ._utils import load_spectrum +from ._utils.parameter import Parameters +from ._utils.utils import PlotRange +from .const import SPECN_FN1, PARAM_FN, SPEC1_FN_N from .logger import get_logger from .physics import pulse, units from .plotting import ( @@ -16,9 +22,87 @@ from .plotting import ( single_position_plot, transform_2D_propagation, ) -from .utils.parameter import Parameters -from .utils.utils import PlotRange -from .utils import load_spectrum + + +class SimulationSeries: + path: Path + params: Parameters + total_length: float + total_num_steps: int + previous: SimulationSeries = None + + class Config: + arbitrary_types_allowed = True + + def __init__(self, path: os.PathLike): + self.path = Path(path) + self.params = Parameters.load(self.path / PARAM_FN) + 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") + + def fiber_map(self): + lengths = self.all_params("length") + return [ + (this[0], following[1]) for this, following in zip(lengths, [(None, 0.0)] + lengths) + ] + + def all_params(self, key: str) -> list[tuple[str, Any]]: + """returns the value of a parameter for each fiber + + Parameters + ---------- + key : str + name of the parameter + + Returns + ------- + list[tuple[str, Any]] + list of (fiber_name, param_value) tuples + """ + return list(reversed(self._all_params(key, []))) + + def accumulate_params(self, key: str) -> Any: + """returns the sum of all the values a parameter takes. Useful to + get the total length of the fiber, the total number of steps, etc. + + Parameters + ---------- + key : str + name of the parameter + + Returns + ------- + Any + final sum + """ + return sum(el[1] for el in self.all_params(key)) + + 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)) + + def _all_params(self, key: str, l: list) -> list: + l.append((self.params.name, getattr(self.params, key))) + if self.previous is not None: + return self.previous._all_params(key, l) + return l + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(path={self.path}, previous={self.previous!r})" + + def __eq__(self, other: SimulationSeries) -> bool: + return ( + self.path == other.path + and self.params == other.params + and self.previous == other.previous + ) + + def __contains__(self, other: SimulationSeries) -> bool: + if other is self or other == self: + return True + if self.previous is not None: + return other in self.previous class Spectrum(np.ndarray): @@ -129,6 +213,23 @@ class Spectrum(np.ndarray): class Pulse(Sequence): + path: Path + default_ind: Optional[int] + params: Parameters + z: np.ndarray + namx: int + t: np.ndarray + w: np.ndarray + w_order: np.ndarray + + def __new__(cls, path: os.PathLike, *args, **kwargs) -> "Pulse": + try: + if load_spectrum(Path(path) / SPECN_FN1.format(0)).ndim == 2: + return super().__new__(LegacyPulse) + except FileNotFoundError: + pass + return super().__new__(cls) + def __init__(self, path: os.PathLike, default_ind: Union[int, Iterable[int]] = None): """load a data folder as a pulse @@ -144,36 +245,6 @@ class Pulse(Sequence): FileNotFoundError path does not contain proper data """ - self.logger = get_logger(__name__) - self.path = Path(path) - self.default_ind = default_ind - - if not self.path.is_dir(): - raise FileNotFoundError(f"Folder {self.path} does not exist") - - self.params = Parameters.load(self.path / "params.toml") - self.params.compute(["name", "t", "l", "w_c", "w0", "z_targets"]) - if self.params.fiber_map is None: - self.params.fiber_map = [(0.0, self.params.name)] - - try: - self.z = np.load(os.path.join(path, "z.npy")) - except FileNotFoundError: - if self.params is not None: - self.z = self.params.z_targets - else: - raise - self.nmax = len(list(self.path.glob("spectra_*.npy"))) - if self.nmax <= 0: - raise FileNotFoundError(f"No appropriate file in specified folder {self.path}") - - self.t = self.params.t - w = math.wspace(self.t) + units.m(self.params.wavelength) - self.w_order = np.argsort(w) - self.w = w - self.wl = units.m.inv(self.w) - self.params.w = self.w - self.params.z_targets = self.z def __iter__(self): """ @@ -190,73 +261,6 @@ class Pulse(Sequence): def __getitem__(self, key) -> Spectrum: return self.all_spectra(key) - def intensity(self, unit): - if unit.type in ["WL", "FREQ", "AFREQ"]: - x_axis = unit.inv(self.w) - else: - x_axis = unit.inv(self.t) - - order = np.argsort(x_axis) - func = dict( - WL=self._to_wl_int, - FREQ=self._to_freq_int, - AFREQ=self._to_afreq_int, - TIME=self._to_time_int, - )[unit.type] - - for spec in self: - yield x_axis[order], func(spec)[:, order] - - def _to_wl_int(self, spectrum): - return units.to_WL(math.abs2(spectrum), spectrum.wl) - - def _to_freq_int(self, spectrum): - return math.abs2(spectrum) - - def _to_afreq_int(self, spectrum): - return math.abs2(spectrum) - - def _to_time_int(self, spectrum): - return math.abs2(np.fft.ifft(spectrum)) - - def amplitude(self, unit): - if unit.type in ["WL", "FREQ", "AFREQ"]: - x_axis = unit.inv(self.w) - else: - x_axis = unit.inv(self.t) - - order = np.argsort(x_axis) - func = dict( - WL=self._to_wl_amp, - FREQ=self._to_freq_amp, - AFREQ=self._to_afreq_amp, - TIME=self._to_time_amp, - )[unit.type] - - for spec in self: - yield x_axis[order], func(spec)[:, order] - - def _to_wl_amp(self, spectrum): - return ( - np.sqrt( - units.to_WL( - math.abs2(spectrum), - spectrum.wl, - ) - ) - * spectrum - / np.abs(spectrum) - ) - - def _to_freq_amp(self, spectrum): - return spectrum - - def _to_afreq_amp(self, spectrum): - return spectrum - - def _to_time_amp(self, spectrum): - return np.fft.ifft(spectrum) - def all_spectra(self, ind=None) -> Spectrum: """ loads the data already simulated. @@ -305,12 +309,7 @@ class Pulse(Sequence): return np.fft.ifft(self.all_spectra(ind=ind), axis=-1) def _load1(self, i: int): - if i < 0: - i = self.nmax + i - spec = load_spectrum(self.path / SPECN_FN.format(i)) - spec = np.atleast_2d(spec) - spec = Spectrum(spec, self.params) - return spec + pass def plot_2D( self, @@ -412,3 +411,46 @@ class Pulse(Sequence): index """ return math.argclosest(self.z, z) + + +class LegacyPulse(Pulse): + def __init__(self, path: os.PathLike, default_ind: Union[int, Iterable[int]] = None): + print("old init called", path, default_ind) + self.logger = get_logger(__name__) + self.path = Path(path) + self.default_ind = default_ind + + if not self.path.is_dir(): + raise FileNotFoundError(f"Folder {self.path} does not exist") + + self.params = Parameters.load(self.path / "params.toml") + self.params.compute(["name", "t", "l", "w_c", "w0", "z_targets"]) + if self.params.fiber_map is None: + self.params.fiber_map = [(0.0, self.params.name)] + + try: + self.z = np.load(os.path.join(path, "z.npy")) + except FileNotFoundError: + if self.params is not None: + self.z = self.params.z_targets + else: + raise + self.nmax = len(list(self.path.glob("spectra_*.npy"))) + if self.nmax <= 0: + raise FileNotFoundError(f"No appropriate file in specified folder {self.path}") + + self.t = self.params.t + w = math.wspace(self.t) + units.m(self.params.wavelength) + self.w_order = np.argsort(w) + self.w = w + self.wl = units.m.inv(self.w) + self.params.w = self.w + self.params.z_targets = self.z + + def _load1(self, i: int): + if i < 0: + i = self.nmax + i + spec = load_spectrum(self.path / SPECN_FN1.format(i)) + spec = np.atleast_2d(spec) + spec = Spectrum(spec, self.params) + return spec diff --git a/testing/long_tests/test_recovery_param_seq.py b/testing/long_tests/test_recovery_param_seq.py deleted file mode 100644 index 5d73367..0000000 --- a/testing/long_tests/test_recovery_param_seq.py +++ /dev/null @@ -1,35 +0,0 @@ -import shutil -import unittest - -import toml -from scgenerator import logger -from send2trash import send2trash - -TMP = "testing/.tmp" - - -class TestRecoveryParamSequence(unittest.TestCase): - def setUp(self): - shutil.copytree("/Users/benoitsierro/sc_tests/scgenerator_full anomalous55", TMP) - self.conf = toml.load(TMP + "/initial_config.toml") - io.set_data_folder(55, TMP) - - def test_remaining_simulations_count(self): - param_seq = initialize.RecoveryParamSequence(self.conf, 55) - self.assertEqual(5, len(param_seq)) - - def test_only_one_to_complete(self): - param_seq = initialize.RecoveryParamSequence(self.conf, 55) - i = 0 - for expected, (vary_list, params) in zip([True, False, False, False, False], param_seq): - i += 1 - self.assertEqual(expected, "recovery_last_stored" in params) - - self.assertEqual(5, i) - - def tearDown(self): - send2trash(TMP) - - -if __name__ == "__main__": - unittest.main() diff --git a/testing/test_initialize.py b/testing/test_initialize.py deleted file mode 100644 index edd3d90..0000000 --- a/testing/test_initialize.py +++ /dev/null @@ -1,216 +0,0 @@ -import unittest -from copy import deepcopy - -import numpy as np -import toml -from scgenerator import defaults, utils, math -from scgenerator.errors import * -from scgenerator.physics import pulse, units -from scgenerator.utils.parameter import Config, Parameters - - -def load_conf(name): - with open("testing/configs/" + name + ".toml") as file: - conf = toml.load(file) - return conf - - -def conf_maker(folder): - def conf(name): - return load_conf(folder + "/" + name) - - return conf - - -class TestParamSequence(unittest.TestCase): - def iterconf(self, files): - conf = conf_maker("param_sequence") - for path in files: - yield init.ParamSequence(conf(path)) - - def test_no_repeat_in_sub_folder_names(self): - for param_seq in self.iterconf(["almost_equal", "equal", "no_variations"]): - l = [] - s = [] - for vary_list, _ in utils.required_simulations(param_seq.config): - self.assertNotIn(vary_list, l) - self.assertNotIn(utils.format_variable_list(vary_list), s) - l.append(vary_list) - s.append(utils.format_variable_list(vary_list)) - - def test_no_variations_yields_only_num_and_id(self): - for param_seq in self.iterconf(["no_variations"]): - for vary_list, _ in utils.required_simulations(param_seq.config): - self.assertEqual(vary_list[1][0], "num") - self.assertEqual(vary_list[0][0], "id") - self.assertEqual(2, len(vary_list)) - - -class TestInitializeMethods(unittest.TestCase): - def test_validate_types(self): - conf = lambda s: load_conf("validate_types/" + s) - - with self.assertRaisesRegex(ValueError, r"'behaviors\[3\]' must be a str in"): - init.Config(**conf("bad2")) - - with self.assertRaisesRegex(TypeError, "value must be of type "): - init.Config(**conf("bad3")) - - with self.assertRaisesRegex(TypeError, "'parallel' is not a valid variable parameter"): - init.Config(**conf("bad4")) - - with self.assertRaisesRegex( - TypeError, "'variable intensity_noise' value must be of type " - ): - init.Config(**conf("bad5")) - - with self.assertRaisesRegex(ValueError, "'repeat' must be positive"): - init.Config(**conf("bad6")) - - with self.assertRaisesRegex( - ValueError, "variable parameter 'intensity_noise' must not be empty" - ): - init.Config(**conf("bad7")) - - self.assertIsNone(init.Config(**conf("good")).hr_w) - - def test_ensure_consistency(self): - conf = lambda s: load_conf("ensure_consistency/" + s) - with self.assertRaisesRegex( - MissingParameterError, - r"1 of '\['t0', 'width'\]' is required and no defaults have been set", - ): - init.Config(**conf("bad1")) - - with self.assertRaisesRegex( - MissingParameterError, - r"1 of '\['peak_power', 'mean_power', 'energy', 'width', 't0'\]' is required when 'soliton_num' is specified and no defaults have been set", - ): - init.Config(**conf("bad2")) - - with self.assertRaisesRegex( - MissingParameterError, - r"2 of '\['dt', 't_num', 'time_window'\]' are required and no defaults have been set", - ): - init.Config(**conf("bad3")) - - with self.assertRaisesRegex( - DuplicateParameterError, - r"got multiple values for parameter 'width'", - ): - init.Config(**conf("bad4")) - - with self.assertRaisesRegex( - MissingParameterError, - r"'capillary_thickness' is a required parameter for fiber model 'hasan' and no defaults have been set", - ): - init.Config(**conf("bad5")) - - with self.assertRaisesRegex( - MissingParameterError, - r"1 of '\['capillary_spacing', 'capillary_outer_d'\]' is required for fiber model 'hasan' and no defaults have been set", - ): - init.Config(**conf("bad6")) - - self.assertLessEqual( - {"model": "pcf"}.items(), init.Config(**conf("good1")).__dict__.items() - ) - - self.assertIsNone(init.Config(**conf("good4")).gamma) - - self.assertLessEqual( - {"raman_type": "agrawal"}.items(), - init.Config(**conf("good2")).__dict__.items(), - ) - - self.assertLessEqual( - {"name": "no name"}.items(), init.Config(**conf("good3")).__dict__.items() - ) - - self.assertLessEqual( - {"capillary_nested": 0, "capillary_resonance_strengths": []}.items(), - init.Config(**conf("good4")).__dict__.items(), - ) - - self.assertLessEqual( - dict(he_mode=(1, 1)).items(), - init.Config(**conf("good5")).__dict__.items(), - ) - - self.assertLessEqual( - dict(temperature=300, pressure=1e5, gas_name="vacuum", plasma_density=0).items(), - init.Config(**conf("good5")).__dict__.items(), - ) - - def setup_conf_custom_field(self, path) -> Parameters: - - conf = load_conf(path) - conf = Parameters(**conf) - init.build_sim_grid_in_place(conf) - return conf - - def test_setup_custom_field(self): - d = np.load("testing/configs/custom_field/init_field.npz") - t = d["time"] - field = d["field"] - conf = self.setup_conf_custom_field("custom_field/no_change") - result, conf.width, conf.peak_power, conf.energy, conf.field_0 = pulse.setup_custom_field( - conf - ) - self.assertAlmostEqual(conf.field_0.real.max(), field.real.max(), 4) - self.assertTrue(result) - - conf = self.setup_conf_custom_field("custom_field/peak_power") - result, conf.width, conf.peak_power, conf.energy, conf.field_0 = pulse.setup_custom_field( - conf - ) - conf.wavelength = pulse.correct_wavelength(conf.wavelength, conf.w_c, conf.field_0) - self.assertAlmostEqual(math.abs2(conf.field_0).max(), 20000, 4) - self.assertTrue(result) - self.assertNotAlmostEqual(conf.wavelength, 1593e-9) - - conf = self.setup_conf_custom_field("custom_field/mean_power") - result, conf.width, conf.peak_power, conf.energy, conf.field_0 = pulse.setup_custom_field( - conf - ) - self.assertAlmostEqual(np.trapz(math.abs2(conf.field_0), conf.t), 0.22 / 40e6, 4) - self.assertTrue(result) - - conf = self.setup_conf_custom_field("custom_field/recover1") - result, conf.width, conf.peak_power, conf.energy, conf.field_0 = pulse.setup_custom_field( - conf - ) - self.assertAlmostEqual(math.abs2(conf.field_0 - field).sum(), 0) - self.assertTrue(result) - - conf = self.setup_conf_custom_field("custom_field/recover2") - result, conf.width, conf.peak_power, conf.energy, conf.field_0 = pulse.setup_custom_field( - conf - ) - self.assertAlmostEqual((math.abs2(conf.field_0) / 0.9 - math.abs2(field)).sum(), 0) - self.assertTrue(result) - - conf = self.setup_conf_custom_field("custom_field/wavelength_shift1") - result = Parameters(**conf) - self.assertAlmostEqual(units.m.inv(result.w)[np.argmax(math.abs2(result.spec_0))], 1050e-9) - - conf = self.setup_conf_custom_field("custom_field/wavelength_shift1") - conf.wavelength = 1593e-9 - result = Parameters(**conf) - - conf = load_conf("custom_field/wavelength_shift2") - conf = init.Config(**conf) - for target, (variable, config) in zip( - [1050e-9, 1321e-9, 1593e-9], init.ParamSequence(conf) - ): - init.build_sim_grid_in_place(conf) - self.assertAlmostEqual( - units.m.inv(config.w)[np.argmax(math.abs2(config.spec_0))], target - ) - print(config.wavelength, target) - - -if __name__ == "__main__": - conf = conf_maker("validate_types") - - unittest.main() diff --git a/testing/test_pulse.py b/testing/test_pulse.py deleted file mode 100644 index 6a632d4..0000000 --- a/testing/test_pulse.py +++ /dev/null @@ -1,41 +0,0 @@ -import unittest -from scgenerator.physics.pulse import conform_pulse_params - - -class TestPulseMethods(unittest.TestCase): - def test_conform_pulse_params(self): - self.assertNotIn(None, conform_pulse_params("gaussian", t0=5, energy=6)) - self.assertNotIn(None, conform_pulse_params("gaussian", width=5, energy=6)) - self.assertNotIn(None, conform_pulse_params("gaussian", t0=5, peak_power=6)) - self.assertNotIn(None, conform_pulse_params("gaussian", width=5, peak_power=6)) - - self.assertEqual(4, len(conform_pulse_params("gaussian", t0=5, energy=6))) - self.assertEqual(4, len(conform_pulse_params("gaussian", width=5, energy=6))) - self.assertEqual(4, len(conform_pulse_params("gaussian", t0=5, peak_power=6))) - self.assertEqual(4, len(conform_pulse_params("gaussian", width=5, peak_power=6))) - - with self.assertRaisesRegex( - TypeError, "when soliton number is desired, both gamma and beta2 must be specified" - ): - conform_pulse_params("gaussian", t0=5, energy=6, gamma=0.01) - with self.assertRaisesRegex( - TypeError, "when soliton number is desired, both gamma and beta2 must be specified" - ): - conform_pulse_params("gaussian", t0=5, energy=6, beta2=0.01) - - self.assertEqual( - 5, len(conform_pulse_params("gaussian", t0=5, energy=6, gamma=0.01, beta2=2e-6)) - ) - self.assertEqual( - 5, len(conform_pulse_params("gaussian", width=5, energy=6, gamma=0.01, beta2=2e-6)) - ) - self.assertEqual( - 5, len(conform_pulse_params("gaussian", t0=5, peak_power=6, gamma=0.01, beta2=2e-6)) - ) - self.assertEqual( - 5, len(conform_pulse_params("gaussian", width=5, peak_power=6, gamma=0.01, beta2=2e-6)) - ) - - -if __name__ == "__main__": - unittest.main() diff --git a/testing/test_utils.py b/testing/test_utils.py deleted file mode 100644 index ef277eb..0000000 --- a/testing/test_utils.py +++ /dev/null @@ -1,67 +0,0 @@ -import unittest - -import numpy as np -import toml -from scgenerator import utils - - -def load_conf(name): - with open("testing/configs/" + name + ".toml") as file: - conf = toml.load(file) - return conf - - -def conf_maker(folder, val=True): - def conf(name): - if val: - return initialize.Config(**load_conf(folder + "/" + name)) - else: - return initialize.Config(**load_conf(folder + "/" + name)) - - return conf - - -class TestUtilsMethods(unittest.TestCase): - def test_count_variations(self): - conf = conf_maker("count_variations") - - for sim, vary in [(1, 0), (1, 1), (2, 1), (2, 0), (120, 3)]: - self.assertEqual((sim, vary), utils.count_variations(conf(f"{sim}sim_{vary}vary"))) - - def test_format_value(self): - values = [ - 122e-6, - True, - ["raman", "ss"], - np.arange(5), - 1.123, - 1.1230001, - 0.002e122, - 12.3456e-9, - ] - s = [ - "0.000122", - "True", - "raman-ss", - "0-1-2-3-4", - "1.123", - "1.1230001", - "2e+119", - "1.23456e-08", - ] - - for value, target in zip(values, s): - self.assertEqual(target, utils.format_value(value)) - - def test_override_config(self): - conf = conf_maker("override", False) - old = conf("initial_config") - new = conf("fiber2") - - over = utils.override_config(vars(new), old) - self.assertNotIn("input_transmission", over.variable) - self.assertIsNone(over.input_transmission) - - -if __name__ == "__main__": - unittest.main() diff --git a/testing/test_variationer.py b/testing/test_variationer.py new file mode 100644 index 0000000..67033e2 --- /dev/null +++ b/testing/test_variationer.py @@ -0,0 +1,54 @@ +from pydantic import main +import scgenerator as sc + + +def test_descriptor(): + # Same branch + var1 = sc.VariationDescriptor( + raw_descr=[[("num", 1), ("a", False)], [("b", 0)]], index=[[1, 0], [0]] + ) + var2 = sc.VariationDescriptor( + raw_descr=[[("num", 2), ("a", False)], [("b", 0)]], index=[[1, 0], [0]] + ) + assert var1.branch.identifier == "b_0" + assert var1.identifier != var1.branch.identifier + assert var1.identifier != var2.identifier + assert var1.branch.identifier == var2.branch.identifier + + # different branch + var3 = sc.VariationDescriptor( + raw_descr=[[("num", 2), ("a", True)], [("b", 0)]], index=[[1, 0], [0]] + ) + assert var1.branch.identifier != var3.branch.identifier + assert var1.formatted_descriptor() != var2.formatted_descriptor() + assert var1.formatted_descriptor() != var3.formatted_descriptor() + + +def test_variationer(): + var = sc.Variationer( + [ + dict(a=[1, 2], num=[0, 1, 2]), + [dict(b=["000", "111"], c=["a", "-1"])], + dict(), + dict(), + [dict(aaa=[True, False], bb=[1, 3])], + ] + ) + assert var.var_num(0) == 6 + assert var.var_num(1) == 12 + assert var.var_num() == 24 + + cfg = dict(bb=None) + branches = set() + for descr in var.iterate(): + assert descr.update_config(cfg).items() >= set(descr.raw_descr[-1]) + branches.add(descr.branch.identifier) + assert len(branches) == 8 + + +def main(): + test_descriptor() + + +if __name__ == "__main__": + main() From 6fa9056c104544657675730f57adb85db4478ad5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Sierro?= Date: Wed, 29 Sep 2021 10:32:55 +0200 Subject: [PATCH 05/11] spectra working --- src/scgenerator/__init__.py | 2 +- src/scgenerator/spectra.py | 172 ++++++++++++++++++++++++++++++++---- 2 files changed, 154 insertions(+), 20 deletions(-) diff --git a/src/scgenerator/__init__.py b/src/scgenerator/__init__.py index 6cb8fcb..25a8b04 100644 --- a/src/scgenerator/__init__.py +++ b/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 diff --git a/src/scgenerator/spectra.py b/src/scgenerator/spectra.py index 82a83b8..15686f0 100644 --- a/src/scgenerator/spectra.py +++ b/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") - - def fiber_map(self): - lengths = self.all_params("length") - return [ - (this[0], following[1]) for this, following in zip(lengths, [(None, 0.0)] + lengths) + self.z_inds = np.arange(len(self.params.z_targets)) + self.z = self.params.z_targets + if self.previous is not None: + print(f"{self.params.z_targets=}") + 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) From fcaf872a26bafe80a6482a97a35040f7c0adff2f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Sierro?= Date: Mon, 4 Oct 2021 13:29:06 +0200 Subject: [PATCH 06/11] misc --- src/scgenerator/_utils/__init__.py | 359 +------------------------- src/scgenerator/_utils/legacy.py | 100 +++++++ src/scgenerator/_utils/parameter.py | 36 ++- src/scgenerator/_utils/pbar.py | 177 +++++++++++++ src/scgenerator/_utils/utils.py | 39 ++- src/scgenerator/_utils/variationer.py | 31 ++- src/scgenerator/physics/simulate.py | 14 +- src/scgenerator/scripts/__init__.py | 4 +- src/scgenerator/spectra.py | 211 +++++++-------- 9 files changed, 485 insertions(+), 486 deletions(-) create mode 100644 src/scgenerator/_utils/legacy.py create mode 100644 src/scgenerator/_utils/pbar.py diff --git a/src/scgenerator/_utils/__init__.py b/src/scgenerator/_utils/__init__.py index 03c8859..5c653a5 100644 --- a/src/scgenerator/_utils/__init__.py +++ b/src/scgenerator/_utils/__init__.py @@ -20,11 +20,13 @@ from string import printable as str_printable from functools import cache from typing import Any, Callable, Generator, Iterable, MutableMapping, Sequence, TypeVar, Union + import numpy as np import pkg_resources as pkg import toml from tqdm import tqdm +from .pbar import PBars from ..const import PARAM_FN, PARAM_SEPARATOR, SPEC1_FN, SPECN_FN1, Z_FN, __version__ from ..env import pbar_policy from ..logger import get_logger @@ -194,34 +196,6 @@ def load_config_sequence(path: os.PathLike) -> tuple[Path, list[dict[str, Any]]] return Path(final_path), configs -def save_parameters( - params: dict[str, Any], destination_dir: Path, file_name: str = PARAM_FN -) -> Path: - """saves a parameter dictionary. Note that is does remove some entries, particularly - those that take a lot of space ("t", "w", ...) - - Parameters - ---------- - params : dict[str, Any] - dictionary to save - destination_dir : Path - destination directory - - Returns - ------- - Path - path to newly created the paramter file - """ - file_path = destination_dir / file_name - os.makedirs(file_path.parent, exist_ok=True) - - # save toml of the simulation - with open(file_path, "w") as file: - toml.dump(params, file, encoder=toml.TomlNumpyEncoder()) - - return file_path - - def load_material_dico(name: str) -> dict[str, Any]: """loads a material dictionary Parameters @@ -235,173 +209,6 @@ def load_material_dico(name: str) -> dict[str, Any]: return toml.loads(Paths.gets("materials"))[name] -def update_appended_params(source: Path, destination: Path, z: Sequence): - z_num = len(z) - params = open_config(source) - params["z_num"] = z_num - params["length"] = float(z[-1] - z[0]) - for p_name in ["recovery_data_dir", "prev_data_dir", "output_path"]: - if p_name in params: - del params[p_name] - save_toml(destination, params) - - -def build_path_trees(sim_dir: Path) -> list[PathTree]: - sim_dir = sim_dir.resolve() - path_branches: list[tuple[Path, ...]] = [] - to_check = list(sim_dir.glob("*fiber*num*")) - with PBars(len(to_check), desc="Building path trees") as pbar: - for branch in map(build_path_branch, to_check): - if branch is not None: - path_branches.append(branch) - pbar.update() - path_trees = group_path_branches(path_branches) - return path_trees - - -def build_path_branch(data_dir: Path) -> tuple[Path, ...]: - if not data_dir.is_dir(): - return None - path_branch = [data_dir] - while ( - prev_sim_path := open_config(path_branch[-1] / PARAM_FN).get("prev_data_dir") - ) is not None: - p = Path(prev_sim_path).resolve() - if not p.exists(): - p = Path(*p.parts[-2:]).resolve() - path_branch.append(p) - return tuple(reversed(path_branch)) - - -def group_path_branches(path_branches: list[tuple[Path, ...]]) -> list[PathTree]: - """groups path lists - - [ - ("a/id 0 wavelength 100 num 0"," b/id 0 wavelength 100 num 0"), - ("a/id 2 wavelength 100 num 1"," b/id 2 wavelength 100 num 1"), - ("a/id 1 wavelength 200 num 0"," b/id 1 wavelength 200 num 0"), - ("a/id 3 wavelength 200 num 1"," b/id 3 wavelength 200 num 1") - ] - -> - [ - ( - ("a/id 0 wavelength 100 num 0", "a/id 2 wavelength 100 num 1"), - ("b/id 0 wavelength 100 num 0", "b/id 2 wavelength 100 num 1"), - ) - ( - ("a/id 1 wavelength 200 num 0", "a/id 3 wavelength 200 num 1"), - ("b/id 1 wavelength 200 num 0", "b/id 3 wavelength 200 num 1"), - ) - ] - - - Parameters - ---------- - path_branches : list[tuple[Path, ...]] - each element of the list is a path to a folder containing data of one simulation - - Returns - ------- - list[PathTree] - list of PathTrees to be used in merge - """ - sort_key = lambda el: el[0] - - size = len(path_branches[0]) - out_trees_map: dict[str, dict[int, dict[int, Path]]] = {} - for branch in path_branches: - b_id = branch_id(branch) - out_trees_map.setdefault(b_id, {i: {} for i in range(size)}) - for sim_part, data_dir in enumerate(branch): - num = re.search(r"(?<=num )[0-9]+", data_dir.name)[0] - out_trees_map[b_id][sim_part][int(num)] = data_dir - - return [ - tuple( - tuple(w for _, w in sorted(v.items(), key=sort_key)) - for __, v in sorted(d.items(), key=sort_key) - ) - for d in out_trees_map.values() - ] - - -def merge_path_tree( - path_tree: PathTree, destination: Path, z_callback: Callable[[int], None] = None -): - """given a path tree, copies the file into the right location - - Parameters - ---------- - path_tree : PathTree - elements of the list returned by group_path_branches - destination : Path - dir where to save the data - """ - z_arr: list[float] = [] - - destination.mkdir(exist_ok=True) - - for i, (z, merged_spectra) in enumerate(merge_spectra(path_tree)): - z_arr.append(z) - spec_out_name = SPECN_FN1.format(i) - np.save(destination / spec_out_name, merged_spectra) - if z_callback is not None: - z_callback(i) - d = np.diff(z_arr) - d[d < 0] = 0 - z_arr = np.concatenate(([z_arr[0]], np.cumsum(d))) - np.save(destination / Z_FN, z_arr) - update_appended_params(path_tree[-1][0] / PARAM_FN, destination / PARAM_FN, z_arr) - - -def merge_spectra( - path_tree: PathTree, -) -> Generator[tuple[float, np.ndarray], None, None]: - for same_sim_paths in path_tree: - z_arr = np.load(same_sim_paths[0] / Z_FN) - for i, z in enumerate(z_arr): - spectra: list[np.ndarray] = [] - for data_dir in same_sim_paths: - spec = np.load(data_dir / SPEC1_FN.format(i)) - spectra.append(spec) - yield z, np.atleast_2d(spectra) - - -def merge(destination: os.PathLike, path_trees: list[PathTree] = None): - - destination = ensure_folder(Path(destination)) - - z_num = 0 - prev_z_num = 0 - - for i, sim_dir in enumerate(sim_dirs(path_trees)): - conf = sim_dir / "initial_config.toml" - shutil.copy( - conf, - destination / f"initial_config_{i}.toml", - ) - prev_z_num = open_config(conf).get("z_num", prev_z_num) - z_num += prev_z_num - - pbars = PBars( - len(path_trees) * z_num, "Merging", 1, worker_kwargs=dict(total=z_num, desc="current pos") - ) - for path_tree in path_trees: - pbars.reset(1) - iden_items = path_tree[-1][0].name.split()[2:] - for i, p_name in list(enumerate(iden_items))[-2::-2]: - if p_name == "num": - del iden_items[i + 1] - del iden_items[i] - iden = PARAM_SEPARATOR.join(iden_items) - merge_path_tree(path_tree, destination / iden, z_callback=lambda i: pbars.update(1)) - - -def sim_dirs(path_trees: list[PathTree]) -> Generator[Path, None, None]: - for p in path_trees[0]: - yield p[0].parent - - def save_data(data: np.ndarray, data_dir: Path, file_name: str): """saves numpy array to disk @@ -462,168 +269,6 @@ def ensure_folder(path: Path, prevent_overwrite: bool = True, mkdir=True) -> Pat path = path.parent / (folder_name + f"_{i}") -class PBars: - def __init__( - self, - task: Union[int, Iterable[T_]], - desc: str, - num_sub_bars: int = 0, - head_kwargs=None, - worker_kwargs=None, - ) -> "PBars": - - self.id = random.randint(100000, 999999) - try: - self.width = os.get_terminal_size().columns - except OSError: - self.width = 80 - if isinstance(task, abc.Iterable): - self.iterator: Iterable[T_] = iter(task) - self.num_tot: int = len(task) - else: - self.num_tot: int = task - self.iterator = None - - self.policy = pbar_policy() - if head_kwargs is None: - head_kwargs = dict() - if worker_kwargs is None: - worker_kwargs = dict( - total=1, - desc="Worker {worker_id}", - bar_format="{l_bar}{bar}" "|[{elapsed}<{remaining}, " "{rate_fmt}{postfix}]", - ) - if "print" not in pbar_policy(): - head_kwargs["file"] = worker_kwargs["file"] = StringIO() - self.width = 80 - head_kwargs["desc"] = desc - self.pbars = [tqdm(total=self.num_tot, ncols=self.width, ascii=False, **head_kwargs)] - for i in range(1, num_sub_bars + 1): - kwargs = {k: v for k, v in worker_kwargs.items()} - if "desc" in kwargs: - kwargs["desc"] = kwargs["desc"].format(worker_id=i) - self.append(tqdm(position=i, ncols=self.width, ascii=False, **kwargs)) - self.print_path = Path( - f"progress {self.pbars[0].desc.replace('/', '')} {self.id}" - ).resolve() - self.close_ev = threading.Event() - if "file" in self.policy: - self.thread = threading.Thread(target=self.print_worker, daemon=True) - self.thread.start() - - def print(self): - if "file" not in self.policy: - return - s = [] - for pbar in self.pbars: - s.append(str(pbar)) - self.print_path.write_text("\n".join(s)) - - def print_worker(self): - while True: - if self.close_ev.wait(2.0): - return - self.print() - - def __iter__(self): - with self as pb: - for thing in self.iterator: - yield thing - pb.update() - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc_value, traceback): - self.close() - - def __getitem__(self, key): - return self.pbars[key] - - def update(self, i=None, value=1): - if i is None: - for pbar in self.pbars[1:]: - pbar.update(value) - elif i > 0: - self.pbars[i].update(value) - self.pbars[0].update() - - def append(self, pbar: tqdm): - self.pbars.append(pbar) - - def reset(self, i): - self.pbars[i].update(-self.pbars[i].n) - self.print() - - def close(self): - self.print() - self.close_ev.set() - if "file" in self.policy: - self.thread.join() - for pbar in self.pbars: - pbar.close() - - -class ProgressBarActor: - def __init__(self, name: str, num_workers: int, num_steps: int) -> None: - self.counters = [0 for _ in range(num_workers + 1)] - self.p_bars = PBars( - num_steps, "Simulating " + name, num_workers, head_kwargs=dict(unit="step") - ) - - def update(self, worker_id: int, rel_pos: float = None) -> None: - """update a counter - - Parameters - ---------- - worker_id : int - id of the worker. 0 is the overall progress - rel_pos : float, optional - if None, increase the counter by one, if set, will set - the counter to the specified value (instead of incrementing it), by default None - """ - if rel_pos is None: - self.counters[worker_id] += 1 - else: - self.counters[worker_id] = rel_pos - - def update_pbars(self): - for counter, pbar in zip(self.counters, self.p_bars.pbars): - pbar.update(counter - pbar.n) - - def close(self): - self.p_bars.close() - - -def progress_worker( - name: str, num_workers: int, num_steps: int, progress_queue: multiprocessing.Queue -): - """keeps track of progress on a separate thread - - Parameters - ---------- - num_steps : int - total number of steps, used for the main progress bar (position 0) - progress_queue : multiprocessing.Queue - values are either - Literal[0] : stop the worker and close the progress bars - tuple[int, float] : worker id and relative progress between 0 and 1 - """ - with PBars( - num_steps, "Simulating " + name, num_workers, head_kwargs=dict(unit="step") - ) as pbars: - while True: - raw = progress_queue.get() - if raw == 0: - return - i, rel_pos = raw - if i > 0: - pbars[i].update(rel_pos - pbars[i].n) - pbars[0].update() - elif i == 0: - pbars[0].update(rel_pos) - - def branch_id(branch: tuple[Path, ...]) -> str: return branch[-1].name.split()[1] diff --git a/src/scgenerator/_utils/legacy.py b/src/scgenerator/_utils/legacy.py new file mode 100644 index 0000000..19c72cd --- /dev/null +++ b/src/scgenerator/_utils/legacy.py @@ -0,0 +1,100 @@ +import os +import sys +from pathlib import Path +from pprint import pprint +from typing import Any, Set + +import numpy as np +import toml + +from ..const import PARAM_FN, SPEC1_FN, SPEC1_FN_N, SPECN_FN1, Z_FN +from .parameter import Parameters +from .utils import fiber_folder, update_path, save_parameters +from .variationer import VariationDescriptor, Variationer + + +def load_config(path: os.PathLike) -> dict[str, Any]: + with open(path) as file: + d = toml.load(file) + d.setdefault("variable", {}) + return d + + +def load_config_sequence(path: os.PathLike) -> tuple[list[Path], list[dict[str, Any]]]: + paths = sorted(list(Path(path).glob("initial_config*.toml"))) + return paths, [load_config(cfg) for cfg in paths] + + +def convert_sim_folder(path: os.PathLike): + path = Path(path) + config_paths, configs = load_config_sequence(path) + master_config = dict(name=path.name, Fiber=configs) + new_fiber_paths: list[Path] = [ + path / fiber_folder(i, path.name, cfg["name"]) for i, cfg in enumerate(configs) + ] + for p in new_fiber_paths: + p.mkdir(exist_ok=True) + var = Variationer(c["variable"] for c in configs) + + paths: dict[Path, VariationDescriptor] = { + path / descr.branch.formatted_descriptor(): descr for descr in var.iterate() + } + for p in paths: + if not p.is_dir(): + raise FileNotFoundError(f"missing {p} from {path}") + processed_paths: Set[Path] = set() + for old_variation_path, descriptor in paths.items(): # fiberA=0, fiber B=0 + vary_parts = old_variation_path.name.split("fiber")[1:] + identifiers = [ + "".join("fiber" + el for el in vary_parts[: i + 1]).strip() + for i in range(len(vary_parts)) + ] + cum_z_num = 0 + for i, (fiber_path, new_identifier) in enumerate(zip(new_fiber_paths, identifiers)): + config = descriptor.update_config(configs[i], i) + new_variation_path = fiber_path / new_identifier + z_num = config["z_num"] + move = new_variation_path not in processed_paths + os.makedirs(new_variation_path, exist_ok=True) + processed_paths.add(new_variation_path) + + for spec_num in range(cum_z_num, cum_z_num + z_num): + old_spec = old_variation_path / SPECN_FN1.format(spec_num) + if move: + spec_data = np.load(old_spec) + for j, spec1 in enumerate(spec_data): + if j == 0: + np.save( + new_variation_path / SPEC1_FN.format(spec_num - cum_z_num), spec1 + ) + else: + np.save( + new_variation_path / SPEC1_FN_N.format(spec_num - cum_z_num, j), + spec1, + ) + old_spec.unlink() + if move: + if i > 0: + config["prev_data_dir"] = str( + (new_fiber_paths[i - 1] / identifiers[i - 1]).resolve() + ) + params = Parameters(**config) + params.compute() + save_parameters(params.prepare_for_dump(), new_variation_path) + cum_z_num += z_num + (old_variation_path / PARAM_FN).unlink() + (old_variation_path / Z_FN).unlink() + old_variation_path.rmdir() + + for cp in config_paths: + cp.unlink() + with open(path / "initial_config.toml", "w") as f: + toml.dump(master_config, f, encoder=toml.TomlNumpyEncoder()) + + +def main(): + convert_sim_folder(sys.argv[1]) + + +if __name__ == "__main__": + main() diff --git a/src/scgenerator/_utils/parameter.py b/src/scgenerator/_utils/parameter.py index f7d7c1a..dc436b9 100644 --- a/src/scgenerator/_utils/parameter.py +++ b/src/scgenerator/_utils/parameter.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import datetime as datetime_module import enum import inspect @@ -15,15 +17,14 @@ from typing import Any, Callable, Generator, Iterable, Literal, Optional, Sequen import numpy as np from numpy.lib import isin -from .. import math +from .. import _utils as utils +from .. import env, math +from .._utils.variationer import VariationDescriptor, Variationer from ..const import PARAM_FN, PARAM_SEPARATOR, __version__ from ..errors import EvaluatorError, NoDefaultError from ..logger import get_logger from ..physics import fiber, materials, pulse, units -from .._utils.variationer import VariationDescriptor, Variationer -from .. import _utils as utils -from .. import env -from .utils import func_rewrite, _mock_function, get_arg_names +from .utils import _mock_function, fiber_folder, func_rewrite, get_arg_names T = TypeVar("T") @@ -302,7 +303,7 @@ class Parameter: self.validator(self.name, value) instance.__dict__[self.name] = value - def display(self, num: float): + def display(self, num: float) -> str: if self.display_info is None: return str(num) else: @@ -314,10 +315,22 @@ class Parameter: @dataclass -class Parameters: +class _AbstractParameters: + @classmethod + def __init_subclass__(cls): + cls.register_param_formatters() + + @classmethod + def register_param_formatters(cls): + for k, v in cls.__dict__.items(): + if isinstance(v, Parameter): + VariationDescriptor.register_formatter(k, v.display) + + +@dataclass +class Parameters(_AbstractParameters): """ - This class defines each valid parameter's name, type and valid value. Initializing - such an obj will automatically compute all possible parameters + This class defines each valid parameter's name, type and valid value. """ # root @@ -832,7 +845,7 @@ class Configuration: self.variationer.append(config.pop("variable")) self.fiber_paths.append( utils.ensure_folder( - self.fiber_path(i, config), + self.final_path / fiber_folder(i, self.name, config["name"]), mkdir=False, prevent_overwrite=not self.overwrite, ) @@ -847,9 +860,6 @@ class Configuration: ) self.parallel = self.master_config.get("parallel", Parameters.parallel.default) - def fiber_path(self, i: int, full_config: dict[str, Any]) -> Path: - return self.final_path / PARAM_SEPARATOR.join([format(i), self.name, full_config["name"]]) - def __validate_variable(self, config: dict[str, Any]): for k, v in config.get("variable", {}).items(): p = getattr(Parameters, k) diff --git a/src/scgenerator/_utils/pbar.py b/src/scgenerator/_utils/pbar.py new file mode 100644 index 0000000..8c44e40 --- /dev/null +++ b/src/scgenerator/_utils/pbar.py @@ -0,0 +1,177 @@ +import abc +import os +import random +import multiprocessing +import threading +import typing +from io import StringIO +from pathlib import Path +from typing import Iterable, Union + +from tqdm import tqdm + +from ..env import pbar_policy + +T_ = typing.TypeVar("T_") + + +class PBars: + def __init__( + self, + task: Union[int, Iterable[T_]], + desc: str, + num_sub_bars: int = 0, + head_kwargs=None, + worker_kwargs=None, + ) -> "PBars": + + self.id = random.randint(100000, 999999) + try: + self.width = os.get_terminal_size().columns + except OSError: + self.width = 80 + if isinstance(task, abc.Iterable): + self.iterator: Iterable[T_] = iter(task) + self.num_tot: int = len(task) + else: + self.num_tot: int = task + self.iterator = None + + self.policy = pbar_policy() + if head_kwargs is None: + head_kwargs = dict() + if worker_kwargs is None: + worker_kwargs = dict( + total=1, + desc="Worker {worker_id}", + bar_format="{l_bar}{bar}" "|[{elapsed}<{remaining}, " "{rate_fmt}{postfix}]", + ) + if "print" not in pbar_policy(): + head_kwargs["file"] = worker_kwargs["file"] = StringIO() + self.width = 80 + head_kwargs["desc"] = desc + self.pbars = [tqdm(total=self.num_tot, ncols=self.width, ascii=False, **head_kwargs)] + for i in range(1, num_sub_bars + 1): + kwargs = {k: v for k, v in worker_kwargs.items()} + if "desc" in kwargs: + kwargs["desc"] = kwargs["desc"].format(worker_id=i) + self.append(tqdm(position=i, ncols=self.width, ascii=False, **kwargs)) + self.print_path = Path( + f"progress {self.pbars[0].desc.replace('/', '')} {self.id}" + ).resolve() + self.close_ev = threading.Event() + if "file" in self.policy: + self.thread = threading.Thread(target=self.print_worker, daemon=True) + self.thread.start() + + def print(self): + if "file" not in self.policy: + return + s = [] + for pbar in self.pbars: + s.append(str(pbar)) + self.print_path.write_text("\n".join(s)) + + def print_worker(self): + while True: + if self.close_ev.wait(2.0): + return + self.print() + + def __iter__(self): + with self as pb: + for thing in self.iterator: + yield thing + pb.update() + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.close() + + def __getitem__(self, key): + return self.pbars[key] + + def update(self, i=None, value=1): + if i is None: + for pbar in self.pbars[1:]: + pbar.update(value) + elif i > 0: + self.pbars[i].update(value) + self.pbars[0].update() + + def append(self, pbar: tqdm): + self.pbars.append(pbar) + + def reset(self, i): + self.pbars[i].update(-self.pbars[i].n) + self.print() + + def close(self): + self.print() + self.close_ev.set() + if "file" in self.policy: + self.thread.join() + for pbar in self.pbars: + pbar.close() + + +class ProgressBarActor: + def __init__(self, name: str, num_workers: int, num_steps: int) -> None: + self.counters = [0 for _ in range(num_workers + 1)] + self.p_bars = PBars( + num_steps, "Simulating " + name, num_workers, head_kwargs=dict(unit="step") + ) + + def update(self, worker_id: int, rel_pos: float = None) -> None: + """update a counter + + Parameters + ---------- + worker_id : int + id of the worker. 0 is the overall progress + rel_pos : float, optional + if None, increase the counter by one, if set, will set + the counter to the specified value (instead of incrementing it), by default None + """ + if rel_pos is None: + self.counters[worker_id] += 1 + else: + self.counters[worker_id] = rel_pos + + def update_pbars(self): + for counter, pbar in zip(self.counters, self.p_bars.pbars): + pbar.update(counter - pbar.n) + + def close(self): + self.p_bars.close() + + +def progress_worker( + name: str, num_workers: int, num_steps: int, progress_queue: multiprocessing.Queue +): + """keeps track of progress on a separate thread + + Parameters + ---------- + num_steps : int + total number of steps, used for the main progress bar (position 0) + progress_queue : multiprocessing.Queue + values are either + Literal[0] : stop the worker and close the progress bars + tuple[int, float] : worker id and relative progress between 0 and 1 + """ + with PBars( + num_steps, "Simulating " + name, num_workers, head_kwargs=dict(unit="step") + ) as pbars: + while True: + raw = progress_queue.get() + if raw == 0: + return + i, rel_pos = raw + if i > 0: + pbars[i].update(rel_pos - pbars[i].n) + pbars[0].update() + elif i == 0: + pbars[0].update(rel_pos) diff --git a/src/scgenerator/_utils/utils.py b/src/scgenerator/_utils/utils.py index 44eadf3..452399e 100644 --- a/src/scgenerator/_utils/utils.py +++ b/src/scgenerator/_utils/utils.py @@ -5,13 +5,14 @@ from collections import defaultdict from functools import cache from pathlib import Path from string import printable as str_printable -from typing import Callable +from typing import Any, Callable import numpy as np +import toml from pydantic import BaseModel from .._utils import load_toml, save_toml -from ..const import PARAM_FN, Z_FN +from ..const import PARAM_FN, PARAM_SEPARATOR, Z_FN from ..physics.units import get_unit @@ -184,6 +185,8 @@ def combine_simulations(path: Path, dest: Path = None): file.unlink() elif file.name == Z_FN: file.rename(new_path / file.name) + elif file.name.startswith("spectr") and num == 0: + file.rename(new_path / file.name) else: file.rename(new_path / (file.stem + f"_{num}" + file.suffix)) pulse.rmdir() @@ -199,5 +202,37 @@ def update_params(new_path: Path, file: Path): file.unlink() +def save_parameters( + params: dict[str, Any], destination_dir: Path, file_name: str = PARAM_FN +) -> Path: + """saves a parameter dictionary. Note that is does remove some entries, particularly + those that take a lot of space ("t", "w", ...) + + Parameters + ---------- + params : dict[str, Any] + dictionary to save + destination_dir : Path + destination directory + + Returns + ------- + Path + path to newly created the paramter file + """ + file_path = destination_dir / file_name + os.makedirs(file_path.parent, exist_ok=True) + + # save toml of the simulation + with open(file_path, "w") as file: + toml.dump(params, file, encoder=toml.TomlNumpyEncoder()) + + return file_path + + def update_path(p: str) -> str: return re.sub(r"( ?num [0-9]+)|(u_[0-9]+ )", "", p) + + +def fiber_folder(i: int, sim_name: str, fiber_name: str) -> str: + return PARAM_SEPARATOR.join([format(i), sim_name, fiber_name]) diff --git a/src/scgenerator/_utils/variationer.py b/src/scgenerator/_utils/variationer.py index fdd6368..5fc353a 100644 --- a/src/scgenerator/_utils/variationer.py +++ b/src/scgenerator/_utils/variationer.py @@ -2,7 +2,7 @@ from math import prod import itertools from collections.abc import MutableMapping, Sequence from pathlib import Path -from typing import Any, Callable, Generator, Iterable, Union +from typing import Any, Callable, Generator, Iterable, Optional, Union import numpy as np from pydantic import validator @@ -151,6 +151,15 @@ class VariationDescriptor(utils.HashableBaseModel): @classmethod def register_formatter(cls, p_name: str, func: Callable[..., str]): + """register a function that formats a particular parameter + + Parameters + ---------- + p_name : str + name of the parameter + func : Callable[..., str] + function that takes as single argument the value of the parameter and returns a string + """ cls._format_registry[p_name] = func def format_value(self, name: str, value) -> str: @@ -174,20 +183,24 @@ class VariationDescriptor(utils.HashableBaseModel): raw_descr=self.raw_descr[key], index=self.index[key], separator=self.separator ) - def update_config(self, cfg: dict[str, Any]) -> dict[str, Any]: + def update_config(self, cfg: dict[str, Any], index=-1) -> dict[str, Any]: """updates a dictionary with the value of the descriptor Parameters ---------- cfg : dict[str, Any] dict to be updated + index : int, optional + index of the fiber from which to apply the parameters, by default -1 Returns ------- dict[str, Any] same as cfg but with key from the descriptor added/updated. """ - return cfg | {k: v for k, v in self.raw_descr[-1]} + out_cfg = cfg.copy() + out_cfg.pop("variable", None) + return out_cfg | {k: v for k, v in self.raw_descr[index]} @property def flat(self) -> list[tuple[str, Any]]: @@ -201,8 +214,8 @@ class VariationDescriptor(utils.HashableBaseModel): @property def branch(self) -> "BranchDescriptor": - descr = [] - ind = [] + descr: list[list[tuple[str, Any]]] = [] + ind: list[list[int]] = [] for i, l in enumerate(self.raw_descr): descr.append([]) ind.append([]) @@ -218,6 +231,14 @@ class VariationDescriptor(utils.HashableBaseModel): self.__ids.setdefault(unique_id, len(self.__ids)) return "u_" + str(self.__ids[unique_id]) + @property + def parent(self) -> Optional["VariationDescriptor"]: + if len(self.raw_descr) < 2: + return None + return VariationDescriptor( + raw_descr=self.raw_descr[:-1], index=self.index[:-1], separator=self.separator + ) + class BranchDescriptor(VariationDescriptor): __ids: dict[int, int] = {} diff --git a/src/scgenerator/physics/simulate.py b/src/scgenerator/physics/simulate.py index 46b9897..3cf7e73 100644 --- a/src/scgenerator/physics/simulate.py +++ b/src/scgenerator/physics/simulate.py @@ -11,11 +11,13 @@ from send2trash import send2trash from .. import env from .. import _utils as utils -from .._utils.utils import combine_simulations +from .._utils.utils import combine_simulations, save_parameters from ..logger import get_logger from .._utils.parameter import Configuration, Parameters +from .._utils.pbar import PBars, ProgressBarActor, progress_worker from . import pulse from .fiber import create_non_linear_op, fast_dispersion_op +from scgenerator._utils import pbar try: import ray @@ -334,7 +336,7 @@ class SequentialRK4IP(RK4IP): def __init__( self, params: Parameters, - pbars: utils.PBars, + pbars: PBars, save_data=False, job_identifier="", task_id=0, @@ -490,7 +492,7 @@ class Simulations: def _run_available(self): for variable, params in self.configuration: v_list_str = variable.formatted_descriptor(True) - utils.save_parameters(params.prepare_for_dump(), Path(params.output_path)) + save_parameters(params.prepare_for_dump(), Path(params.output_path)) self.new_sim(v_list_str, params) self.finish() @@ -527,7 +529,7 @@ class SequencialSimulations(Simulations, priority=0): def __init__(self, configuration: Configuration, task_id): super().__init__(configuration, task_id=task_id) - self.pbars = utils.PBars( + self.pbars = PBars( self.configuration.total_num_steps, "Simulating " + self.configuration.final_path.name, 1, @@ -571,7 +573,7 @@ class MultiProcSimulations(Simulations, priority=1): for i in range(self.sim_jobs_per_node) ] self.p_worker = multiprocessing.Process( - target=utils.progress_worker, + target=progress_worker, args=( self.configuration.final_path.name, self.sim_jobs_per_node, @@ -660,7 +662,7 @@ class RaySimulations(Simulations, priority=2): self.pool = ray.util.ActorPool(self.propagator.remote() for _ in range(self.sim_jobs_total)) self.num_submitted = 0 self.rolling_id = 0 - self.p_actor = ray.remote(utils.ProgressBarActor).remote( + self.p_actor = ray.remote(ProgressBarActor).remote( self.configuration.final_path, self.sim_jobs_total, self.configuration.total_num_steps ) self.configuration.skip_callback = lambda num: ray.get(self.p_actor.update.remote(0, num)) diff --git a/src/scgenerator/scripts/__init__.py b/src/scgenerator/scripts/__init__.py index 176f236..11c502e 100644 --- a/src/scgenerator/scripts/__init__.py +++ b/src/scgenerator/scripts/__init__.py @@ -11,7 +11,7 @@ from .. import env, math from ..const import PARAM_FN, PARAM_SEPARATOR from ..physics import fiber, units from ..plotting import plot_setup -from ..spectra import Pulse +from ..spectra import Pulse, SimulationSeries from .._utils import auto_crop, open_config, save_toml, translate_parameters from .._utils.parameter import ( Configuration, @@ -39,7 +39,7 @@ def plot_all(sim_dir: Path, limits: list[str], show=False, **opts): ] with tqdm(total=len(dir_list) * len(limits)) as bar: for p in dir_list: - pulse = Pulse(p) + pulse = SimulationSeries(p) for left, right, unit in limits: path, fig, ax = plot_setup( pulse.path.parent diff --git a/src/scgenerator/spectra.py b/src/scgenerator/spectra.py index 15686f0..2d46224 100644 --- a/src/scgenerator/spectra.py +++ b/src/scgenerator/spectra.py @@ -13,7 +13,7 @@ from . import math from ._utils import load_spectrum from ._utils.parameter import Parameters from ._utils.utils import PlotRange -from .const import SPECN_FN1, PARAM_FN, SPEC1_FN_N +from .const import SPECN_FN1, PARAM_FN, SPEC1_FN_N, SPEC1_FN from .logger import get_logger from .physics import pulse, units from .plotting import ( @@ -24,6 +24,101 @@ from .plotting import ( ) +class Spectrum(np.ndarray): + params: Parameters + + def __new__(cls, input_array, params: Parameters): + # 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.params = params + + # Finally, we must return the newly created object: + return obj + + def __array_finalize__(self, obj): + # see InfoArray.__array_finalize__ for comments + if obj is None: + return + self.params = getattr(obj, "params", None) + + def __getitem__(self, key) -> "Spectrum": + return super().__getitem__(key) + + @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) + ) + + def measure(self) -> tuple[float, float, float]: + return pulse.measure_field(self.params.t, self.time_amp) + + class SimulationSeries: path: Path params: Parameters @@ -34,11 +129,13 @@ class SimulationSeries: fiber_positions: list[tuple[str, float]] z_inds: np.ndarray - class Config: - arbitrary_types_allowed = True - def __init__(self, path: os.PathLike): - self.path = Path(path) + self.logger = get_logger() + path = Path(path) + subdirs = [el for el in path.glob("*") if (el / PARAM_FN).exists()] + while not (path / PARAM_FN).exists() and len(subdirs) == 1: + path = subdirs[0] + self.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 @@ -50,11 +147,9 @@ class SimulationSeries: self.z_inds = np.arange(len(self.params.z_targets)) self.z = self.params.z_targets if self.previous is not None: - print(f"{self.params.z_targets=}") 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]) @@ -220,7 +315,10 @@ class SimulationSeries: np.ndarray loaded spectrum file """ - return load_spectrum(self.path / SPEC1_FN_N.format(z_ind - self.z_inds[0], sim_ind)) + if sim_ind > 0: + return load_spectrum(self.path / SPEC1_FN_N.format(z_ind - self.z_inds[0], sim_ind)) + else: + return load_spectrum(self.path / SPEC1_FN.format(z_ind - self.z_inds[0])) def _all_params(self, key: str, l: list) -> list: l.append((self.params.name, getattr(self.params, key))) @@ -250,100 +348,11 @@ class SimulationSeries: if self.previous is not None: return other in self.previous - -class Spectrum(np.ndarray): - params: Parameters - - def __new__(cls, input_array, params: Parameters): - # 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.params = params - - # Finally, we must return the newly created object: - return obj - - def __array_finalize__(self, obj): - # see InfoArray.__array_finalize__ for comments - if obj is None: - return - self.params = getattr(obj, "params", None) - - def __getitem__(self, key) -> "Spectrum": - return super().__getitem__(key) - - @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) + def __getitem__(self, key) -> Spectrum: + if isinstance(key, tuple): + return self.spectra(*key) 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) - ) - - def measure(self) -> tuple[float, float, float]: - return pulse.measure_field(self.params.t, self.time_amp) + return self.spectra(key, None) class Pulse(Sequence): From fa72a6e136fb38cea03760d1c5bf43e5fc6e805c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Sierro?= Date: Tue, 5 Oct 2021 15:55:20 +0200 Subject: [PATCH 07/11] file name still wrong in converter --- src/scgenerator/__init__.py | 19 ++- src/scgenerator/_utils/legacy.py | 31 ++-- src/scgenerator/_utils/parameter.py | 33 ++++- src/scgenerator/_utils/pbar.py | 16 ++- src/scgenerator/_utils/utils.py | 23 ++- src/scgenerator/_utils/variationer.py | 60 +++++++- src/scgenerator/physics/simulate.py | 5 +- src/scgenerator/spectra.py | 195 ++++++++++++++++---------- 8 files changed, 285 insertions(+), 97 deletions(-) diff --git a/src/scgenerator/__init__.py b/src/scgenerator/__init__.py index 25a8b04..650e52e 100644 --- a/src/scgenerator/__init__.py +++ b/src/scgenerator/__init__.py @@ -2,9 +2,24 @@ from . import math 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 .plotting import ( + mean_values_plot, + plot_spectrogram, + propagation_plot, + single_position_plot, + transform_2D_propagation, + transform_1D_values, + transform_mean_values, + get_extent, +) from .spectra import Pulse, Spectrum, SimulationSeries from ._utils import Paths, open_config, parameter from ._utils.parameter import Configuration, Parameters from ._utils.utils import PlotRange -from ._utils.variationer import Variationer, VariationDescriptor, VariationSpecsError +from ._utils.legacy import convert_sim_folder +from ._utils.variationer import ( + Variationer, + VariationDescriptor, + VariationSpecsError, + DescriptorDict, +) diff --git a/src/scgenerator/_utils/legacy.py b/src/scgenerator/_utils/legacy.py index 19c72cd..407e040 100644 --- a/src/scgenerator/_utils/legacy.py +++ b/src/scgenerator/_utils/legacy.py @@ -8,8 +8,9 @@ import numpy as np import toml from ..const import PARAM_FN, SPEC1_FN, SPEC1_FN_N, SPECN_FN1, Z_FN -from .parameter import Parameters -from .utils import fiber_folder, update_path, save_parameters +from .parameter import Configuration, Parameters +from .utils import fiber_folder, save_parameters +from .pbar import PBars from .variationer import VariationDescriptor, Variationer @@ -29,21 +30,32 @@ def convert_sim_folder(path: os.PathLike): path = Path(path) config_paths, configs = load_config_sequence(path) master_config = dict(name=path.name, Fiber=configs) + with open(path / "initial_config.toml", "w") as f: + toml.dump(master_config, f, encoder=toml.TomlNumpyEncoder()) + configuration = Configuration(path / "initial_config.toml") new_fiber_paths: list[Path] = [ path / fiber_folder(i, path.name, cfg["name"]) for i, cfg in enumerate(configs) ] for p in new_fiber_paths: p.mkdir(exist_ok=True) - var = Variationer(c["variable"] for c in configs) + repeat = configs[0].get("repeat", 1) - paths: dict[Path, VariationDescriptor] = { - path / descr.branch.formatted_descriptor(): descr for descr in var.iterate() + pbar = PBars(configuration.total_num_steps, "Converting") + + old_paths: dict[Path, VariationDescriptor] = { + path / descr.branch.formatted_descriptor(): (descr, param.final_path) + for descr, param in configuration } - for p in paths: + + # create map from old to new path + + pprint(old_paths) + quit() + for p in old_paths: if not p.is_dir(): raise FileNotFoundError(f"missing {p} from {path}") processed_paths: Set[Path] = set() - for old_variation_path, descriptor in paths.items(): # fiberA=0, fiber B=0 + for old_variation_path, descriptor in old_paths.items(): # fiberA=0, fiber B=0 vary_parts = old_variation_path.name.split("fiber")[1:] identifiers = [ "".join("fiber" + el for el in vary_parts[: i + 1]).strip() @@ -72,6 +84,9 @@ def convert_sim_folder(path: os.PathLike): new_variation_path / SPEC1_FN_N.format(spec_num - cum_z_num, j), spec1, ) + pbar.update() + else: + pbar.update(value=repeat) old_spec.unlink() if move: if i > 0: @@ -88,8 +103,6 @@ def convert_sim_folder(path: os.PathLike): for cp in config_paths: cp.unlink() - with open(path / "initial_config.toml", "w") as f: - toml.dump(master_config, f, encoder=toml.TomlNumpyEncoder()) def main(): diff --git a/src/scgenerator/_utils/parameter.py b/src/scgenerator/_utils/parameter.py index dc436b9..2d871ac 100644 --- a/src/scgenerator/_utils/parameter.py +++ b/src/scgenerator/_utils/parameter.py @@ -24,7 +24,7 @@ from ..const import PARAM_FN, PARAM_SEPARATOR, __version__ from ..errors import EvaluatorError, NoDefaultError from ..logger import get_logger from ..physics import fiber, materials, pulse, units -from .utils import _mock_function, fiber_folder, func_rewrite, get_arg_names +from .utils import _mock_function, fiber_folder, func_rewrite, get_arg_names, update_path T = TypeVar("T") @@ -75,6 +75,7 @@ VALID_VARIABLE = { "interpolation_degree", "ideal_gas", "length", + "num", } MANDATORY_PARAMETERS = [ @@ -519,6 +520,12 @@ class Parameters(_AbstractParameters): return out + @property + def final_path(self) -> Path: + if self.output_path is not None: + return update_path(self.output_path) + return None + class Rule: def __init__( @@ -777,6 +784,7 @@ class Configuration: """ fiber_configs: list[dict[str, Any]] + vary_dicts: list[dict[str, list]] master_config: dict[str, Any] fiber_paths: list[Path] num_sim: int @@ -814,9 +822,11 @@ class Configuration: self, final_config_path: os.PathLike, overwrite: bool = True, + wait: bool = False, skip_callback: Callable[[int], None] = None, ): self.logger = get_logger(__name__) + self.wait = wait self.overwrite = overwrite self.final_path, self.fiber_configs = utils.load_config_sequence(final_config_path) @@ -842,7 +852,8 @@ class Configuration: config.setdefault("name", Parameters.name.default) self.z_num += config["z_num"] fiber_names.add(config["name"]) - self.variationer.append(config.pop("variable")) + vary_dict = config.pop("variable") + self.variationer.append(vary_dict) self.fiber_paths.append( utils.ensure_folder( self.final_path / fiber_folder(i, self.name, config["name"]), @@ -850,9 +861,11 @@ class Configuration: prevent_overwrite=not self.overwrite, ) ) - self.__validate_variable(config) + self.__validate_variable(vary_dict) self.num_fibers += 1 - Evaluator.evaluate_default(config, True) + Evaluator.evaluate_default( + self.__build_base_config() | config | {k: v[0] for k, v in vary_dict.items()}, True + ) self.num_sim = self.variationer.var_num() self.total_num_steps = sum( config["z_num"] * self.variationer.var_num(i) @@ -860,8 +873,13 @@ class Configuration: ) self.parallel = self.master_config.get("parallel", Parameters.parallel.default) - def __validate_variable(self, config: dict[str, Any]): - for k, v in config.get("variable", {}).items(): + def __build_base_config(self): + cfg = self.fiber_configs[0].copy() + vary = cfg.pop("variable", {}) + return cfg | {k: v[0] for k, v in vary.items()} + + def __validate_variable(self, vary_dict: dict[str, list]): + for k, v in vary_dict.items(): p = getattr(Parameters, k) validator_list(p.validator)("variable " + k, v) if k not in VALID_VARIABLE: @@ -873,7 +891,6 @@ class Configuration: for i in range(self.num_fibers): for sim_config in self.iterate_single_fiber(i): params = Parameters(**sim_config.config) - params.compute() yield sim_config.descriptor, params def iterate_single_fiber( @@ -943,6 +960,8 @@ class Configuration: config dictionary. The only key possibly modified is 'prev_data_dir', which gets set if the simulation is partially completed """ + if not self.wait: + return self.Action.RUN, sim_config.config out_status, num = self.sim_status(sim_config.output_path, sim_config.config) if out_status == self.State.COMPLETE: return self.Action.SKIP, sim_config.config diff --git a/src/scgenerator/_utils/pbar.py b/src/scgenerator/_utils/pbar.py index 8c44e40..37db473 100644 --- a/src/scgenerator/_utils/pbar.py +++ b/src/scgenerator/_utils/pbar.py @@ -1,9 +1,9 @@ -import abc +import multiprocessing import os import random -import multiprocessing import threading import typing +from collections import abc from io import StringIO from pathlib import Path from typing import Iterable, Union @@ -24,7 +24,19 @@ class PBars: head_kwargs=None, worker_kwargs=None, ) -> "PBars": + """creates a PBars obj + Parameters + ---------- + task : int | Iterable + if int : total length of the main task + if Iterable : behaves like tqdm + desc : str + description of the main task + num_sub_bars : int + number of sub-tasks + + """ self.id = random.randint(100000, 999999) try: self.width = os.get_terminal_size().columns diff --git a/src/scgenerator/_utils/utils.py b/src/scgenerator/_utils/utils.py index 452399e..d2a4fa8 100644 --- a/src/scgenerator/_utils/utils.py +++ b/src/scgenerator/_utils/utils.py @@ -5,7 +5,7 @@ from collections import defaultdict from functools import cache from pathlib import Path from string import printable as str_printable -from typing import Any, Callable +from typing import Any, Callable, Iterator, Set import numpy as np import toml @@ -236,3 +236,24 @@ def update_path(p: str) -> str: def fiber_folder(i: int, sim_name: str, fiber_name: str) -> str: return PARAM_SEPARATOR.join([format(i), sim_name, fiber_name]) + + +def iter_simulations(path: os.PathLike) -> list[Path]: + """finds simulations folders contained in a parent directory + + Parameters + ---------- + path : os.PathLike + parent path + + Yields + ------- + Path + Absolute Path to the simulation folder + """ + paths: list[Path] = [] + for pwd, _, files in os.walk(path): + if PARAM_FN in files: + paths.append(Path(pwd)) + paths.sort(key=lambda el: el.parent.name) + return [p for p in paths if p.parent.name == paths[-1].parent.name] diff --git a/src/scgenerator/_utils/variationer.py b/src/scgenerator/_utils/variationer.py index 5fc353a..8ca2e40 100644 --- a/src/scgenerator/_utils/variationer.py +++ b/src/scgenerator/_utils/variationer.py @@ -2,14 +2,17 @@ from math import prod import itertools from collections.abc import MutableMapping, Sequence from pathlib import Path -from typing import Any, Callable, Generator, Iterable, Optional, Union +from typing import Any, Callable, Generator, Generic, Iterable, Iterator, Optional, TypeVar, Union import numpy as np from pydantic import validator +from pydantic.main import BaseModel from ..const import PARAM_SEPARATOR from . import utils +T = TypeVar("T") + class VariationSpecsError(ValueError): pass @@ -111,15 +114,15 @@ class Variationer: return max(1, prod(prod(el) for el in self.all_indices[: index + 1])) -class VariationDescriptor(utils.HashableBaseModel): +class VariationDescriptor(BaseModel): raw_descr: tuple[tuple[tuple[str, Any], ...], ...] index: tuple[tuple[int, ...], ...] separator: str = "fiber" _format_registry: dict[str, Callable[..., str]] = {} __ids: dict[int, int] = {} - def __str__(self) -> str: - return self.formatted_descriptor(add_identifier=False) + class Config: + allow_mutation = False def formatted_descriptor(self, add_identifier=False) -> str: """formats a variable list into a str such that each simulation has a unique @@ -183,6 +186,24 @@ class VariationDescriptor(utils.HashableBaseModel): raw_descr=self.raw_descr[key], index=self.index[key], separator=self.separator ) + def __str__(self) -> str: + return self.formatted_descriptor(add_identifier=False) + + def __lt__(self, other: "VariationDescriptor") -> bool: + return self.raw_descr < other.raw_descr + + def __le__(self, other: "VariationDescriptor") -> bool: + return self.raw_descr <= other.raw_descr + + def __gt__(self, other: "VariationDescriptor") -> bool: + return self.raw_descr > other.raw_descr + + def __ge__(self, other: "VariationDescriptor") -> bool: + return self.raw_descr >= other.raw_descr + + def __hash__(self) -> int: + return hash(self.raw_descr) + def update_config(self, cfg: dict[str, Any], index=-1) -> dict[str, Any]: """updates a dictionary with the value of the descriptor @@ -252,3 +273,34 @@ class BranchDescriptor(VariationDescriptor): @validator("raw_descr") def validate_raw_descr(cls, v): return tuple(tuple(el for el in variable if el[0] != "num") for variable in v) + + +class DescriptorDict(Generic[T]): + def __init__(self, dico: dict[VariationDescriptor, T] = None): + self.dico: dict[tuple[tuple[tuple[str, Any], ...], ...], tuple[VariationDescriptor, T]] = {} + if dico is not None: + for k, v in dico.items(): + self[k] = v + + def __setitem__(self, key: VariationDescriptor, value: T): + if not isinstance(key, VariationDescriptor): + raise TypeError("key must be a VariationDescriptor instance") + self.dico[key.raw_descr] = (key, value) + + def __getitem__( + self, key: Union[VariationDescriptor, tuple[tuple[tuple[str, Any], ...], ...]] + ) -> T: + if isinstance(key, VariationDescriptor): + return self.dico[key.raw_descr][1] + else: + return self.dico[key][1] + + def items(self) -> Iterator[tuple[VariationDescriptor, T]]: + for k, v in self.dico.items(): + yield k, v[1] + + def keys(self) -> list[VariationDescriptor]: + return [v[0] for v in self.dico.values()] + + def values(self) -> list[T]: + return [v[1] for v in self.dico.values()] diff --git a/src/scgenerator/physics/simulate.py b/src/scgenerator/physics/simulate.py index 3cf7e73..3864ffc 100644 --- a/src/scgenerator/physics/simulate.py +++ b/src/scgenerator/physics/simulate.py @@ -491,6 +491,7 @@ class Simulations: def _run_available(self): for variable, params in self.configuration: + params.compute() v_list_str = variable.formatted_descriptor(True) save_parameters(params.prepare_for_dump(), Path(params.output_path)) @@ -718,7 +719,7 @@ def run_simulation( config_file: os.PathLike, method: Union[str, Type[Simulations]] = None, ): - config = Configuration(config_file) + config = Configuration(config_file, wait=True) sim = new_simulation(config, method) sim.run() @@ -760,6 +761,8 @@ def parallel_RK4IP( ]: logger = get_logger(__name__) params = list(Configuration(config)) + for _, param in params: + param.compute() n = len(params) z_num = params[0][1].z_num diff --git a/src/scgenerator/spectra.py b/src/scgenerator/spectra.py index 2d46224..6a9e376 100644 --- a/src/scgenerator/spectra.py +++ b/src/scgenerator/spectra.py @@ -1,6 +1,7 @@ from __future__ import annotations import os +import warnings from collections.abc import Sequence from pathlib import Path from typing import Any, Callable, Dict, Iterable, Iterator, Optional, Union @@ -12,8 +13,8 @@ from pydantic import BaseModel, DirectoryPath, root_validator from . import math from ._utils import load_spectrum from ._utils.parameter import Parameters -from ._utils.utils import PlotRange -from .const import SPECN_FN1, PARAM_FN, SPEC1_FN_N, SPEC1_FN +from ._utils.utils import PlotRange, iter_simulations +from .const import PARAM_FN, SPEC1_FN, SPEC1_FN_N, SPECN_FN1 from .logger import get_logger from .physics import pulse, units from .plotting import ( @@ -131,11 +132,10 @@ class SimulationSeries: def __init__(self, path: os.PathLike): self.logger = get_logger() - path = Path(path) - subdirs = [el for el in path.glob("*") if (el / PARAM_FN).exists()] - while not (path / PARAM_FN).exists() and len(subdirs) == 1: - path = subdirs[0] - self.path = path + for self.path in iter_simulations(path): + break + else: + raise FileNotFoundError(f"No simulation in {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 @@ -356,24 +356,22 @@ class SimulationSeries: class Pulse(Sequence): - path: Path - default_ind: Optional[int] - params: Parameters - z: np.ndarray - namx: int - t: np.ndarray - w: np.ndarray - w_order: np.ndarray + def __new__(cls, path: os.PathLike): + warnings.warn( + "You are using the legacy version of the pulse loader. " + "Please consider updating your data with scgenerator.convert_sim_folder " + "and loading data with the SimulationSeries class" + ) + if (Path(path) / SPECN_FN1.format(0)).exists(): + return LegacyPulse(path) + return SimulationSeries(path) - def __new__(cls, path: os.PathLike, *args, **kwargs) -> "Pulse": - try: - if load_spectrum(Path(path) / SPECN_FN1.format(0)).ndim == 2: - return super().__new__(LegacyPulse) - except FileNotFoundError: - pass - return super().__new__(cls) + def __getitem__(self, key) -> Spectrum: + raise NotImplementedError() - def __init__(self, path: os.PathLike, default_ind: Union[int, Iterable[int]] = None): + +class LegacyPulse(Sequence): + def __init__(self, path: os.PathLike): """load a data folder as a pulse Parameters @@ -388,6 +386,35 @@ class Pulse(Sequence): FileNotFoundError path does not contain proper data """ + self.logger = get_logger(__name__) + self.path = Path(path) + + if not self.path.is_dir(): + raise FileNotFoundError(f"Folder {self.path} does not exist") + + self.params = Parameters.load(self.path / "params.toml") + self.params.compute(["name", "t", "l", "w_c", "w0", "z_targets"]) + if self.params.fiber_map is None: + self.params.fiber_map = [(0.0, self.params.name)] + + try: + self.z = np.load(os.path.join(path, "z.npy")) + except FileNotFoundError: + if self.params is not None: + self.z = self.params.z_targets + else: + raise + self.nmax = len(list(self.path.glob("spectra_*.npy"))) + if self.nmax <= 0: + raise FileNotFoundError(f"No appropriate file in specified folder {self.path}") + + self.t = self.params.t + w = math.wspace(self.t) + units.m(self.params.wavelength) + self.w_order = np.argsort(w) + self.w = w + self.wl = units.m.inv(self.w) + self.params.w = self.w + self.params.z_targets = self.z def __iter__(self): """ @@ -404,6 +431,73 @@ class Pulse(Sequence): def __getitem__(self, key) -> Spectrum: return self.all_spectra(key) + def intensity(self, unit): + if unit.type in ["WL", "FREQ", "AFREQ"]: + x_axis = unit.inv(self.w) + else: + x_axis = unit.inv(self.t) + + order = np.argsort(x_axis) + func = dict( + WL=self._to_wl_int, + FREQ=self._to_freq_int, + AFREQ=self._to_afreq_int, + TIME=self._to_time_int, + )[unit.type] + + for spec in self: + yield x_axis[order], func(spec)[:, order] + + def _to_wl_int(self, spectrum): + return units.to_WL(math.abs2(spectrum), spectrum.wl) + + def _to_freq_int(self, spectrum): + return math.abs2(spectrum) + + def _to_afreq_int(self, spectrum): + return math.abs2(spectrum) + + def _to_time_int(self, spectrum): + return math.abs2(np.fft.ifft(spectrum)) + + def amplitude(self, unit): + if unit.type in ["WL", "FREQ", "AFREQ"]: + x_axis = unit.inv(self.w) + else: + x_axis = unit.inv(self.t) + + order = np.argsort(x_axis) + func = dict( + WL=self._to_wl_amp, + FREQ=self._to_freq_amp, + AFREQ=self._to_afreq_amp, + TIME=self._to_time_amp, + )[unit.type] + + for spec in self: + yield x_axis[order], func(spec)[:, order] + + def _to_wl_amp(self, spectrum): + return ( + np.sqrt( + units.to_WL( + math.abs2(spectrum), + spectrum.wl, + ) + ) + * spectrum + / np.abs(spectrum) + ) + + def _to_freq_amp(self, spectrum): + return spectrum + + def _to_afreq_amp(self, spectrum): + return spectrum + + def _to_time_amp(self, spectrum): + return np.fft.ifft(spectrum) + def all_spectra(self, ind=None) -> Spectrum: """ loads the data already simulated. @@ -425,10 +519,7 @@ class Pulse(Sequence): # Check if file exists and assert how many z positions there are if ind is None: - if self.default_ind is None: - ind = range(self.nmax) - else: - ind = self.default_ind + ind = range(self.nmax) if isinstance(ind, (int, np.integer)): ind = [ind] elif isinstance(ind, (float, np.floating)): @@ -452,7 +543,12 @@ class Pulse(Sequence): return np.fft.ifft(self.all_spectra(ind=ind), axis=-1) def _load1(self, i: int): - pass + if i < 0: + i = self.nmax + i + spec = load_spectrum(self.path / SPECN_FN1.format(i)) + spec = np.atleast_2d(spec) + spec = Spectrum(spec, self.params) + return spec def plot_2D( self, @@ -554,46 +650,3 @@ class Pulse(Sequence): index """ return math.argclosest(self.z, z) - - -class LegacyPulse(Pulse): - def __init__(self, path: os.PathLike, default_ind: Union[int, Iterable[int]] = None): - print("old init called", path, default_ind) - self.logger = get_logger(__name__) - self.path = Path(path) - self.default_ind = default_ind - - if not self.path.is_dir(): - raise FileNotFoundError(f"Folder {self.path} does not exist") - - self.params = Parameters.load(self.path / "params.toml") - self.params.compute(["name", "t", "l", "w_c", "w0", "z_targets"]) - if self.params.fiber_map is None: - self.params.fiber_map = [(0.0, self.params.name)] - - try: - self.z = np.load(os.path.join(path, "z.npy")) - except FileNotFoundError: - if self.params is not None: - self.z = self.params.z_targets - else: - raise - self.nmax = len(list(self.path.glob("spectra_*.npy"))) - if self.nmax <= 0: - raise FileNotFoundError(f"No appropriate file in specified folder {self.path}") - - self.t = self.params.t - w = math.wspace(self.t) + units.m(self.params.wavelength) - self.w_order = np.argsort(w) - self.w = w - self.wl = units.m.inv(self.w) - self.params.w = self.w - self.params.z_targets = self.z - - def _load1(self, i: int): - if i < 0: - i = self.nmax + i - spec = load_spectrum(self.path / SPECN_FN1.format(i)) - spec = np.atleast_2d(spec) - spec = Spectrum(spec, self.params) - return spec From fb880077ede1fb1d28d08fa69d816bc04fb63e40 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Sierro?= Date: Fri, 8 Oct 2021 16:36:53 +0200 Subject: [PATCH 08/11] working converter --- src/scgenerator/_utils/legacy.py | 111 ++++------ src/scgenerator/_utils/parameter.py | 43 ++-- src/scgenerator/_utils/variationer.py | 41 ++-- src/scgenerator/spectra.py | 306 +------------------------- 4 files changed, 106 insertions(+), 395 deletions(-) diff --git a/src/scgenerator/_utils/legacy.py b/src/scgenerator/_utils/legacy.py index 407e040..e198f55 100644 --- a/src/scgenerator/_utils/legacy.py +++ b/src/scgenerator/_utils/legacy.py @@ -27,84 +27,67 @@ def load_config_sequence(path: os.PathLike) -> tuple[list[Path], list[dict[str, def convert_sim_folder(path: os.PathLike): - path = Path(path) + path = Path(path).resolve() config_paths, configs = load_config_sequence(path) master_config = dict(name=path.name, Fiber=configs) with open(path / "initial_config.toml", "w") as f: toml.dump(master_config, f, encoder=toml.TomlNumpyEncoder()) - configuration = Configuration(path / "initial_config.toml") - new_fiber_paths: list[Path] = [ - path / fiber_folder(i, path.name, cfg["name"]) for i, cfg in enumerate(configs) - ] - for p in new_fiber_paths: - p.mkdir(exist_ok=True) - repeat = configs[0].get("repeat", 1) - + configuration = Configuration(path / "initial_config.toml", final_output_path=path) pbar = PBars(configuration.total_num_steps, "Converting") - old_paths: dict[Path, VariationDescriptor] = { - path / descr.branch.formatted_descriptor(): (descr, param.final_path) - for descr, param in configuration - } + new_paths: dict[VariationDescriptor, Parameters] = dict(configuration) + old_paths: Set[Path] = set() + old2new: list[tuple[Path, VariationDescriptor, Parameters, tuple[int, int]]] = [] + for descriptor, params in configuration.iterate_single_fiber(-1): + old_path = path / descriptor.branch.formatted_descriptor() + if not Path(old_path).is_dir(): + raise FileNotFoundError(f"missing {old_path} from {path}. Aborting.") + old_paths.add(old_path) + for d in descriptor.iter_parents(): + z_num_start = sum(c["z_num"] for c in configs[: d.num_fibers - 1]) + z_limits = (z_num_start, z_num_start + params.z_num) + old2new.append((old_path, d, new_paths[d], z_limits)) - # create map from old to new path - - pprint(old_paths) - quit() - for p in old_paths: - if not p.is_dir(): - raise FileNotFoundError(f"missing {p} from {path}") processed_paths: Set[Path] = set() - for old_variation_path, descriptor in old_paths.items(): # fiberA=0, fiber B=0 - vary_parts = old_variation_path.name.split("fiber")[1:] - identifiers = [ - "".join("fiber" + el for el in vary_parts[: i + 1]).strip() - for i in range(len(vary_parts)) - ] - cum_z_num = 0 - for i, (fiber_path, new_identifier) in enumerate(zip(new_fiber_paths, identifiers)): - config = descriptor.update_config(configs[i], i) - new_variation_path = fiber_path / new_identifier - z_num = config["z_num"] - move = new_variation_path not in processed_paths - os.makedirs(new_variation_path, exist_ok=True) - processed_paths.add(new_variation_path) + processed_specs: Set[VariationDescriptor] = set() - for spec_num in range(cum_z_num, cum_z_num + z_num): - old_spec = old_variation_path / SPECN_FN1.format(spec_num) - if move: - spec_data = np.load(old_spec) - for j, spec1 in enumerate(spec_data): - if j == 0: - np.save( - new_variation_path / SPEC1_FN.format(spec_num - cum_z_num), spec1 - ) - else: - np.save( - new_variation_path / SPEC1_FN_N.format(spec_num - cum_z_num, j), - spec1, - ) - pbar.update() - else: - pbar.update(value=repeat) - old_spec.unlink() - if move: - if i > 0: - config["prev_data_dir"] = str( - (new_fiber_paths[i - 1] / identifiers[i - 1]).resolve() - ) - params = Parameters(**config) - params.compute() - save_parameters(params.prepare_for_dump(), new_variation_path) - cum_z_num += z_num - (old_variation_path / PARAM_FN).unlink() - (old_variation_path / Z_FN).unlink() - old_variation_path.rmdir() + for old_path, descr, new_params, (start_z, end_z) in old2new: + move_specs = descr not in processed_specs + processed_specs.add(descr) + if (parent := descr.parent) is not None: + new_params.prev_data_dir = str(new_paths[parent].final_path) + save_parameters(new_params.prepare_for_dump(), new_params.final_path) + for spec_num in range(start_z, end_z): + old_spec = old_path / SPECN_FN1.format(spec_num) + if move_specs: + _mv_specs(pbar, new_params, start_z, spec_num, old_spec) + old_spec.unlink() + if old_path not in processed_paths: + (old_path / PARAM_FN).unlink() + (old_path / Z_FN).unlink() + processed_paths.add(old_path) + + for old_path in processed_paths: + old_path.rmdir() for cp in config_paths: cp.unlink() +def _mv_specs(pbar: PBars, new_params: Parameters, start_z: int, spec_num: int, old_spec: Path): + os.makedirs(new_params.final_path, exist_ok=True) + spec_data = np.load(old_spec) + for j, spec1 in enumerate(spec_data): + if j == 0: + np.save(new_params.final_path / SPEC1_FN.format(spec_num - start_z), spec1) + else: + np.save( + new_params.final_path / SPEC1_FN_N.format(spec_num - start_z, j), + spec1, + ) + pbar.update() + + def main(): convert_sim_folder(sys.argv[1]) diff --git a/src/scgenerator/_utils/parameter.py b/src/scgenerator/_utils/parameter.py index 2d871ac..0f463d5 100644 --- a/src/scgenerator/_utils/parameter.py +++ b/src/scgenerator/_utils/parameter.py @@ -12,7 +12,18 @@ from copy import copy, deepcopy from dataclasses import asdict, dataclass, fields from functools import cache, lru_cache from pathlib import Path -from typing import Any, Callable, Generator, Iterable, Literal, Optional, Sequence, TypeVar, Union +from typing import ( + Any, + Callable, + Generator, + Iterable, + Iterator, + Literal, + Optional, + Sequence, + TypeVar, + Union, +) import numpy as np from numpy.lib import isin @@ -523,7 +534,7 @@ class Parameters(_AbstractParameters): @property def final_path(self) -> Path: if self.output_path is not None: - return update_path(self.output_path) + return Path(update_path(self.output_path)) return None @@ -820,22 +831,26 @@ class Configuration: def __init__( self, - final_config_path: os.PathLike, + config_path: os.PathLike, overwrite: bool = True, wait: bool = False, skip_callback: Callable[[int], None] = None, + final_output_path: os.PathLike = None, ): self.logger = get_logger(__name__) self.wait = wait self.overwrite = overwrite - self.final_path, self.fiber_configs = utils.load_config_sequence(final_config_path) + self.final_path, self.fiber_configs = utils.load_config_sequence(config_path) + self.final_path = env.get(env.OUTPUT_PATH, self.final_path) + if final_output_path is not None: + self.final_path = final_output_path self.final_path = utils.ensure_folder( - Path(env.get(env.OUTPUT_PATH, self.final_path)), + Path(self.final_path), mkdir=False, prevent_overwrite=not self.overwrite, ) - self.master_config = self.fiber_configs[0] + self.master_config = self.fiber_configs[0].copy() self.name = self.final_path.name self.z_num = 0 self.total_num_steps = 0 @@ -874,7 +889,7 @@ class Configuration: self.parallel = self.master_config.get("parallel", Parameters.parallel.default) def __build_base_config(self): - cfg = self.fiber_configs[0].copy() + cfg = self.master_config.copy() vary = cfg.pop("variable", {}) return cfg | {k: v[0] for k, v in vary.items()} @@ -887,15 +902,11 @@ class Configuration: if len(v) == 0: raise ValueError(f"variable parameter {k!r} must not be empty") - def __iter__(self) -> Generator[tuple[VariationDescriptor, Parameters], None, None]: + def __iter__(self) -> Iterator[tuple[VariationDescriptor, Parameters]]: for i in range(self.num_fibers): - for sim_config in self.iterate_single_fiber(i): - params = Parameters(**sim_config.config) - yield sim_config.descriptor, params + yield from self.iterate_single_fiber(i) - def iterate_single_fiber( - self, index: int - ) -> Generator["Configuration.__SimConfig", None, None]: + def iterate_single_fiber(self, index: int) -> Iterator[tuple[VariationDescriptor, Parameters]]: """iterates through the parameters of only one fiber. It takes care of recovering partially completed simulations, skipping complete ones and waiting for the previous fiber to finish @@ -909,6 +920,8 @@ class Configuration: __SimConfig configuration obj """ + if index < 0: + index = self.num_fibers + index sim_dict: dict[Path, Configuration.__SimConfig] = {} for descriptor in self.variationer.iterate(index): cfg = descriptor.update_config(self.fiber_configs[index]) @@ -929,7 +942,7 @@ class Configuration: task, config_dict = self.__decide(sim_config) if task == self.Action.RUN: sim_dict.pop(data_dir) - yield sim_config + yield sim_config.descriptor, Parameters(**sim_config.config) if "recovery_last_stored" in config_dict and self.skip_callback is not None: self.skip_callback(config_dict["recovery_last_stored"]) break diff --git a/src/scgenerator/_utils/variationer.py b/src/scgenerator/_utils/variationer.py index 8ca2e40..cbd5e8c 100644 --- a/src/scgenerator/_utils/variationer.py +++ b/src/scgenerator/_utils/variationer.py @@ -121,6 +121,19 @@ class VariationDescriptor(BaseModel): _format_registry: dict[str, Callable[..., str]] = {} __ids: dict[int, int] = {} + @classmethod + def register_formatter(cls, p_name: str, func: Callable[..., str]): + """register a function that formats a particular parameter + + Parameters + ---------- + p_name : str + name of the parameter + func : Callable[..., str] + function that takes as single argument the value of the parameter and returns a string + """ + cls._format_registry[p_name] = func + class Config: allow_mutation = False @@ -152,19 +165,6 @@ class VariationDescriptor(BaseModel): self.identifier + PARAM_SEPARATOR + self.branch.identifier + PARAM_SEPARATOR + tmp_name ) - @classmethod - def register_formatter(cls, p_name: str, func: Callable[..., str]): - """register a function that formats a particular parameter - - Parameters - ---------- - p_name : str - name of the parameter - func : Callable[..., str] - function that takes as single argument the value of the parameter and returns a string - """ - cls._format_registry[p_name] = func - def format_value(self, name: str, value) -> str: if value is True or value is False: return str(value) @@ -201,9 +201,15 @@ class VariationDescriptor(BaseModel): def __ge__(self, other: "VariationDescriptor") -> bool: return self.raw_descr >= other.raw_descr + def __eq__(self, other: "VariationDescriptor") -> bool: + return self.raw_descr == other.raw_descr + def __hash__(self) -> int: return hash(self.raw_descr) + def __contains__(self, other: "VariationDescriptor") -> bool: + return all(el in self.raw_descr for el in other.raw_descr) + def update_config(self, cfg: dict[str, Any], index=-1) -> dict[str, Any]: """updates a dictionary with the value of the descriptor @@ -223,6 +229,11 @@ class VariationDescriptor(BaseModel): out_cfg.pop("variable", None) return out_cfg | {k: v for k, v in self.raw_descr[index]} + def iter_parents(self) -> Iterator["VariationDescriptor"]: + if (p := self.parent) is not None: + yield from p.iter_parents() + yield self + @property def flat(self) -> list[tuple[str, Any]]: out = [] @@ -260,6 +271,10 @@ class VariationDescriptor(BaseModel): raw_descr=self.raw_descr[:-1], index=self.index[:-1], separator=self.separator ) + @property + def num_fibers(self) -> int: + return len(self.raw_descr) + class BranchDescriptor(VariationDescriptor): __ids: dict[int, int] = {} diff --git a/src/scgenerator/spectra.py b/src/scgenerator/spectra.py index 6a9e376..2072806 100644 --- a/src/scgenerator/spectra.py +++ b/src/scgenerator/spectra.py @@ -1,20 +1,17 @@ from __future__ import annotations import os -import warnings -from collections.abc import Sequence from pathlib import Path -from typing import Any, Callable, Dict, Iterable, Iterator, Optional, Union +from typing import Any, Callable, Iterator, Optional, Union import matplotlib.pyplot as plt import numpy as np -from pydantic import BaseModel, DirectoryPath, root_validator from . import math from ._utils import load_spectrum from ._utils.parameter import Parameters from ._utils.utils import PlotRange, iter_simulations -from .const import PARAM_FN, SPEC1_FN, SPEC1_FN_N, SPECN_FN1 +from .const import PARAM_FN, SPEC1_FN, SPEC1_FN_N from .logger import get_logger from .physics import pulse, units from .plotting import ( @@ -111,7 +108,7 @@ class Spectrum(np.ndarray): 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": + 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) ) @@ -353,300 +350,3 @@ class SimulationSeries: return self.spectra(*key) else: return self.spectra(key, None) - - -class Pulse(Sequence): - def __new__(cls, path: os.PathLike): - warnings.warn( - "You are using the legacy version of the pulse loader. " - "Please consider updating your data with scgenerator.convert_sim_folder " - "and loading data with the SimulationSeries class" - ) - if (Path(path) / SPECN_FN1.format(0)).exists(): - return LegacyPulse(path) - return SimulationSeries(path) - - def __getitem__(self, key) -> Spectrum: - raise NotImplementedError() - - -class LegacyPulse(Sequence): - def __init__(self, path: os.PathLike): - """load a data folder as a pulse - - Parameters - ---------- - path : os.PathLike - path to the data (folder containing .npy files) - default_ind : int | Iterable[int], optional - default indices to be loaded, by default None - - Raises - ------ - FileNotFoundError - path does not contain proper data - """ - self.logger = get_logger(__name__) - self.path = Path(path) - - if not self.path.is_dir(): - raise FileNotFoundError(f"Folder {self.path} does not exist") - - self.params = Parameters.load(self.path / "params.toml") - self.params.compute(["name", "t", "l", "w_c", "w0", "z_targets"]) - if self.params.fiber_map is None: - self.params.fiber_map = [(0.0, self.params.name)] - - try: - self.z = np.load(os.path.join(path, "z.npy")) - except FileNotFoundError: - if self.params is not None: - self.z = self.params.z_targets - else: - raise - self.nmax = len(list(self.path.glob("spectra_*.npy"))) - if self.nmax <= 0: - raise FileNotFoundError(f"No appropriate file in specified folder {self.path}") - - self.t = self.params.t - w = math.wspace(self.t) + units.m(self.params.wavelength) - self.w_order = np.argsort(w) - self.w = w - self.wl = units.m.inv(self.w) - self.params.w = self.w - self.params.z_targets = self.z - - def __iter__(self): - """ - similar to all_spectra but works as an iterator - """ - - self.logger.debug(f"iterating through {self.path}") - for i in range(self.nmax): - yield self._load1(i) - - def __len__(self): - return self.nmax - - def __getitem__(self, key) -> Spectrum: - return self.all_spectra(key) - - def intensity(self, unit): - if unit.type in ["WL", "FREQ", "AFREQ"]: - x_axis = unit.inv(self.w) - else: - x_axis = unit.inv(self.t) - - order = np.argsort(x_axis) - func = dict( - WL=self._to_wl_int, - FREQ=self._to_freq_int, - AFREQ=self._to_afreq_int, - TIME=self._to_time_int, - )[unit.type] - - for spec in self: - yield x_axis[order], func(spec)[:, order] - - def _to_wl_int(self, spectrum): - return units.to_WL(math.abs2(spectrum), spectrum.wl) - - def _to_freq_int(self, spectrum): - return math.abs2(spectrum) - - def _to_afreq_int(self, spectrum): - return math.abs2(spectrum) - - def _to_time_int(self, spectrum): - return math.abs2(np.fft.ifft(spectrum)) - - def amplitude(self, unit): - if unit.type in ["WL", "FREQ", "AFREQ"]: - x_axis = unit.inv(self.w) - else: - x_axis = unit.inv(self.t) - - order = np.argsort(x_axis) - func = dict( - WL=self._to_wl_amp, - FREQ=self._to_freq_amp, - AFREQ=self._to_afreq_amp, - TIME=self._to_time_amp, - )[unit.type] - - for spec in self: - yield x_axis[order], func(spec)[:, order] - - def _to_wl_amp(self, spectrum): - return ( - np.sqrt( - units.to_WL( - math.abs2(spectrum), - spectrum.wl, - ) - ) - * spectrum - / np.abs(spectrum) - ) - - def _to_freq_amp(self, spectrum): - return spectrum - - def _to_afreq_amp(self, spectrum): - return spectrum - - def _to_time_amp(self, spectrum): - return np.fft.ifft(spectrum) - - def all_spectra(self, ind=None) -> Spectrum: - """ - loads the data already simulated. - defauft shape is (z_targets, n, nt) - - Parameters - ---------- - ind : int or list of int - if only certain spectra are desired - Returns - ---------- - spectra : array of shape (nz, m, nt) - array of complex spectra (pulse at nz positions consisting - of nm simulation on a nt size grid) - """ - - self.logger.debug(f"opening {self.path}") - - # Check if file exists and assert how many z positions there are - - if ind is None: - ind = range(self.nmax) - if isinstance(ind, (int, np.integer)): - ind = [ind] - elif isinstance(ind, (float, np.floating)): - ind = [self.z_ind(ind)] - elif isinstance(ind[0], (float, np.floating)): - ind = [self.z_ind(ii) for ii in ind] - - # Load the spectra - spectra = [] - for i in ind: - spectra.append(self._load1(i)) - spectra = Spectrum(spectra, self.params) - - self.logger.debug(f"all spectra from {self.path} successfully loaded") - if len(ind) == 1: - return spectra[0] - else: - return spectra - - def all_fields(self, ind=None): - return np.fft.ifft(self.all_spectra(ind=ind), axis=-1) - - def _load1(self, i: int): - if i < 0: - i = self.nmax + i - spec = load_spectrum(self.path / SPECN_FN1.format(i)) - spec = np.atleast_2d(spec) - spec = Spectrum(spec, self.params) - return spec - - def plot_2D( - self, - 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, - ): - plot_range = PlotRange(left, right, unit) - vals = self.retrieve_plot_values(plot_range, z_pos, 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, slice(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": - vals = self.all_fields(ind=z_pos) - else: - vals = self.all_spectra(ind=z_pos) - - if sim_ind is None: - return vals - elif z_pos is None: - return vals[:, sim_ind] - else: - return vals[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.all_spectra(), 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)) - - def z_ind(self, z: float) -> int: - """return the closest z index to the given target - - Parameters - ---------- - z : float - target - - Returns - ------- - int - index - """ - return math.argclosest(self.z, z) From b3dbec1553c7829af9be13240af234a757a38448 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Sierro?= Date: Mon, 11 Oct 2021 09:16:31 +0200 Subject: [PATCH 09/11] working gas simulations --- src/scgenerator/__init__.py | 2 +- src/scgenerator/_utils/parameter.py | 2 ++ src/scgenerator/physics/fiber.py | 16 ++++++++++++++++ src/scgenerator/physics/materials.py | 4 ++-- src/scgenerator/scripts/__init__.py | 2 +- 5 files changed, 22 insertions(+), 4 deletions(-) diff --git a/src/scgenerator/__init__.py b/src/scgenerator/__init__.py index 650e52e..2699f07 100644 --- a/src/scgenerator/__init__.py +++ b/src/scgenerator/__init__.py @@ -12,7 +12,7 @@ from .plotting import ( transform_mean_values, get_extent, ) -from .spectra import Pulse, Spectrum, SimulationSeries +from .spectra import Spectrum, SimulationSeries from ._utils import Paths, open_config, parameter from ._utils.parameter import Configuration, Parameters from ._utils.utils import PlotRange diff --git a/src/scgenerator/_utils/parameter.py b/src/scgenerator/_utils/parameter.py index 0f463d5..43e1f56 100644 --- a/src/scgenerator/_utils/parameter.py +++ b/src/scgenerator/_utils/parameter.py @@ -1094,6 +1094,7 @@ default_rules: list[Rule] = [ priorities=[2, 2, 2], ), Rule("hr_w", fiber.delayed_raman_w), + Rule("n_gas_2", materials.n_gas_2), Rule("n_eff", fiber.n_eff_hasan, conditions=dict(model="hasan")), Rule("n_eff", fiber.n_eff_marcatili, conditions=dict(model="marcatili")), Rule("n_eff", fiber.n_eff_marcatili_adjusted, conditions=dict(model="marcatili_adjusted")), @@ -1109,6 +1110,7 @@ default_rules: list[Rule] = [ Rule("A_eff", fiber.A_eff_from_diam), Rule("A_eff", fiber.A_eff_hasan, conditions=dict(model="hasan")), Rule("A_eff", fiber.A_eff_from_gamma, priorities=-1), + Rule("A_eff", fiber.A_eff_marcatili, priorities=-2), Rule("A_eff_arr", fiber.A_eff_from_V, ["core_radius", "V_eff_arr"]), Rule("A_eff_arr", fiber.load_custom_A_eff), Rule("A_eff_arr", fiber.constant_A_eff_arr, priorities=-1), diff --git a/src/scgenerator/physics/fiber.py b/src/scgenerator/physics/fiber.py index 4a27249..d059eda 100644 --- a/src/scgenerator/physics/fiber.py +++ b/src/scgenerator/physics/fiber.py @@ -192,6 +192,22 @@ def n_eff_marcatili_adjusted(wl_for_disp, n_gas_2, core_radius, he_mode=(1, 1), return np.sqrt(n_gas_2 - (wl_for_disp * u / (pipi * corrected_radius)) ** 2) +def A_eff_marcatili(core_radius: float) -> float: + """Effective mode-field area for fundamental mode hollow capillaries + + Parameters + ---------- + core_radius : float + radius of the core + + Returns + ------- + float + effective mode field area + """ + return 1.5 * core_radius ** 2 + + @np_cache def n_eff_hasan( wl_for_disp: np.ndarray, diff --git a/src/scgenerator/physics/materials.py b/src/scgenerator/physics/materials.py index b9572e6..1adb5c4 100644 --- a/src/scgenerator/physics/materials.py +++ b/src/scgenerator/physics/materials.py @@ -10,9 +10,9 @@ from .units import NA, c, kB, me, e, hbar def n_gas_2( - wl_for_disp: np.ndarray, gas: str, pressure: float, temperature: float, ideal_gas: bool + wl_for_disp: np.ndarray, gas_name: str, pressure: float, temperature: float, ideal_gas: bool ): - material_dico = utils.load_material_dico(gas) + material_dico = _utils.load_material_dico(gas_name) if ideal_gas: n_gas_2 = sellmeier(wl_for_disp, material_dico, pressure, temperature) + 1 diff --git a/src/scgenerator/scripts/__init__.py b/src/scgenerator/scripts/__init__.py index 11c502e..ebc10c0 100644 --- a/src/scgenerator/scripts/__init__.py +++ b/src/scgenerator/scripts/__init__.py @@ -11,7 +11,7 @@ from .. import env, math from ..const import PARAM_FN, PARAM_SEPARATOR from ..physics import fiber, units from ..plotting import plot_setup -from ..spectra import Pulse, SimulationSeries +from ..spectra import SimulationSeries from .._utils import auto_crop, open_config, save_toml, translate_parameters from .._utils.parameter import ( Configuration, From 31fb519e883c9e0c8d14dc7514bd8549480fac93 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Sierro?= Date: Mon, 11 Oct 2021 15:38:26 +0200 Subject: [PATCH 10/11] misc --- README.md | 4 ++-- src/scgenerator/_utils/parameter.py | 6 +++--- src/scgenerator/_utils/utils.py | 17 ++++++++-------- src/scgenerator/physics/fiber.py | 31 +++++++++-------------------- src/scgenerator/scripts/__init__.py | 3 ++- src/scgenerator/spectra.py | 4 ++-- 6 files changed, 27 insertions(+), 38 deletions(-) diff --git a/README.md b/README.md index d5890e5..5405150 100644 --- a/README.md +++ b/README.md @@ -141,8 +141,8 @@ hasan : capillary_num : int number of capillaries -capillary_outer_d : float, optional if g is specified - outer diameter of the capillaries +capillary_radius : float, optional if g is specified + outer radius of the capillaries capillary_thickness : float thickness of the capillary walls diff --git a/src/scgenerator/_utils/parameter.py b/src/scgenerator/_utils/parameter.py index 43e1f56..6f4ba83 100644 --- a/src/scgenerator/_utils/parameter.py +++ b/src/scgenerator/_utils/parameter.py @@ -55,7 +55,7 @@ VALID_VARIABLE = { "effective_mode_diameter", "core_radius", "capillary_num", - "capillary_outer_d", + "capillary_radius", "capillary_thickness", "capillary_spacing", "capillary_resonance_strengths", @@ -374,7 +374,7 @@ class Parameters(_AbstractParameters): ) length: float = Parameter(non_negative(float, int)) capillary_num: int = Parameter(positive(int)) - capillary_outer_d: float = Parameter(in_range_excl(0, 1e-3)) + capillary_radius: float = Parameter(in_range_excl(0, 1e-3)) capillary_thickness: float = Parameter(in_range_excl(0, 1e-3)) capillary_spacing: float = Parameter(in_range_excl(0, 1e-3)) capillary_resonance_strengths: Iterable[float] = Parameter(num_list, default=[]) @@ -1104,7 +1104,7 @@ default_rules: list[Rule] = [ ["wl_for_disp", "pitch", "pitch_ratio"], conditions=dict(model="pcf"), ), - Rule("capillary_spacing", fiber.HCARF_gap), + Rule("capillary_spacing", fiber.capillary_spacing_hasan), # Fiber nonlinearity Rule("A_eff", fiber.A_eff_from_V), Rule("A_eff", fiber.A_eff_from_diam), diff --git a/src/scgenerator/_utils/utils.py b/src/scgenerator/_utils/utils.py index d2a4fa8..2931f09 100644 --- a/src/scgenerator/_utils/utils.py +++ b/src/scgenerator/_utils/utils.py @@ -109,11 +109,12 @@ def sort_axis( def get_arg_names(func: Callable) -> list[str]: - spec = inspect.getfullargspec(func) - args = spec.args - if spec.defaults is not None and len(spec.defaults) > 0: - args = args[: -len(spec.defaults)] - return args + # spec = inspect.getfullargspec(func) + # args = spec.args + # if spec.defaults is not None and len(spec.defaults) > 0: + # args = args[: -len(spec.defaults)] + # return args + return [k for k, v in inspect.signature(func).parameters.items() if v.default is inspect._empty] def validate_arg_names(names: list[str]): @@ -238,7 +239,7 @@ def fiber_folder(i: int, sim_name: str, fiber_name: str) -> str: return PARAM_SEPARATOR.join([format(i), sim_name, fiber_name]) -def iter_simulations(path: os.PathLike) -> list[Path]: +def simulations_list(path: os.PathLike) -> list[Path]: """finds simulations folders contained in a parent directory Parameters @@ -246,9 +247,9 @@ def iter_simulations(path: os.PathLike) -> list[Path]: path : os.PathLike parent path - Yields + Returns ------- - Path + list[Path] Absolute Path to the simulation folder """ paths: list[Path] = [] diff --git a/src/scgenerator/physics/fiber.py b/src/scgenerator/physics/fiber.py index d059eda..16ba4ed 100644 --- a/src/scgenerator/physics/fiber.py +++ b/src/scgenerator/physics/fiber.py @@ -7,7 +7,7 @@ from scipy.interpolate import interp1d from ..logger import get_logger -from .. import _utils +from .. import _utils as utils from ..math import abs2, argclosest, power_fact, u_nm from .._utils.cache import np_cache from . import materials as mat @@ -48,27 +48,6 @@ def is_dynamic_dispersion(pressure=None): return out -def HCARF_gap(core_radius: float, capillary_num: int, capillary_outer_d: float): - """computes the gap length between capillaries of a hollow core anti-resonance fiber - - Parameters - ---------- - core_radius : float - radius of the core (m) (from cented to edge of a capillary) - capillary_num : int - number of capillaries - capillary_outer_d : float - diameter of the capillaries including the wall thickness(m). The core together with the microstructure has a diameter of 2R + 2d - - Returns - ------- - gap : float - """ - return (core_radius + capillary_outer_d / 2) * 2 * np.sin( - pi / capillary_num - ) - capillary_outer_d - - def gvd_from_n_eff(n_eff: np.ndarray, wl_for_disp: np.ndarray): """computes the dispersion parameter D from an effective index of refraction n_eff Since computing gradients/derivatives of discrete arrays is not well defined on the boundary, it is @@ -208,6 +187,14 @@ def A_eff_marcatili(core_radius: float) -> float: return 1.5 * core_radius ** 2 +def capillary_spacing_hasan( + capillary_num: int, capillary_radius: float, core_radius: float +) -> float: + return ( + 2 * (capillary_radius + core_radius) * np.sin(np.pi / capillary_num) - 2 * capillary_radius + ) + + @np_cache def n_eff_hasan( wl_for_disp: np.ndarray, diff --git a/src/scgenerator/scripts/__init__.py b/src/scgenerator/scripts/__init__.py index ebc10c0..fb1f473 100644 --- a/src/scgenerator/scripts/__init__.py +++ b/src/scgenerator/scripts/__init__.py @@ -17,6 +17,7 @@ from .._utils.parameter import ( Configuration, Parameters, ) +from .._utils.utils import simulations_list def fingerprint(params: Parameters): @@ -31,7 +32,7 @@ def plot_all(sim_dir: Path, limits: list[str], show=False, **opts): opts[k] = int(v) if k in {"log", "renormalize"}: opts[k] = True if v == "True" else False - dir_list = list(p for p in sim_dir.glob("*") if p.is_dir()) + dir_list = simulations_list(sim_dir) if len(dir_list) == 0: dir_list = [sim_dir] limits = [ diff --git a/src/scgenerator/spectra.py b/src/scgenerator/spectra.py index 2072806..fe9279d 100644 --- a/src/scgenerator/spectra.py +++ b/src/scgenerator/spectra.py @@ -10,7 +10,7 @@ import numpy as np from . import math from ._utils import load_spectrum from ._utils.parameter import Parameters -from ._utils.utils import PlotRange, iter_simulations +from ._utils.utils import PlotRange, simulations_list from .const import PARAM_FN, SPEC1_FN, SPEC1_FN_N from .logger import get_logger from .physics import pulse, units @@ -129,7 +129,7 @@ class SimulationSeries: def __init__(self, path: os.PathLike): self.logger = get_logger() - for self.path in iter_simulations(path): + for self.path in simulations_list(path): break else: raise FileNotFoundError(f"No simulation in {path}") From 3f2a56d7037b562d00d38715c785df9a6f8f1642 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Sierro?= Date: Tue, 12 Oct 2021 13:18:32 +0200 Subject: [PATCH 11/11] ready to merge --- src/scgenerator/__init__.py | 2 +- src/scgenerator/_utils/__init__.py | 16 +++++++++++----- src/scgenerator/_utils/legacy.py | 7 +++++-- src/scgenerator/_utils/parameter.py | 4 ++-- src/scgenerator/physics/__init__.py | 5 ++--- src/scgenerator/physics/pulse.py | 3 ++- src/scgenerator/physics/simulate.py | 11 +++++++++++ src/scgenerator/scripts/__init__.py | 6 +++--- 8 files changed, 37 insertions(+), 17 deletions(-) diff --git a/src/scgenerator/__init__.py b/src/scgenerator/__init__.py index 2699f07..c5f3af3 100644 --- a/src/scgenerator/__init__.py +++ b/src/scgenerator/__init__.py @@ -13,7 +13,7 @@ from .plotting import ( get_extent, ) from .spectra import Spectrum, SimulationSeries -from ._utils import Paths, open_config, parameter +from ._utils import Paths, _open_config, parameter, open_single_config from ._utils.parameter import Configuration, Parameters from ._utils.utils import PlotRange from ._utils.legacy import convert_sim_folder diff --git a/src/scgenerator/_utils/__init__.py b/src/scgenerator/_utils/__init__.py index 5c653a5..b540d9e 100644 --- a/src/scgenerator/_utils/__init__.py +++ b/src/scgenerator/_utils/__init__.py @@ -103,12 +103,12 @@ def conform_toml_path(path: os.PathLike) -> str: def open_single_config(path: os.PathLike) -> dict[str, Any]: - d = open_config(path) + d = _open_config(path) f = d.pop("Fiber")[0] return d | f -def open_config(path: os.PathLike): +def _open_config(path: os.PathLike): """returns a dictionary parsed from the specified toml file This also handle having a 'INCLUDE' argument that will fill otherwise unspecified keys with what's in the INCLUDE file(s)""" @@ -170,7 +170,7 @@ def load_config_sequence(path: os.PathLike) -> tuple[Path, list[dict[str, Any]]] Parameters ---------- path : os.PathLike - path to the config toml file + path to the config toml file or a directory containing config files Returns ------- @@ -180,9 +180,15 @@ def load_config_sequence(path: os.PathLike) -> tuple[Path, list[dict[str, Any]]] one config per fiber """ - loaded_config = open_config(path) + path = Path(path) + fiber_list: list[dict[str, Any]] + if path.name.lower().endswith(".toml"): + loaded_config = _open_config(path) + fiber_list = loaded_config.pop("Fiber") + else: + loaded_config = dict(name=path.name) + fiber_list = [_open_config(p) for p in sorted(path.glob("initial_config*.toml"))] - fiber_list: list[dict[str, Any]] = loaded_config.pop("Fiber") if len(fiber_list) == 0: raise ValueError(f"No fiber in config {path}") final_path = loaded_config.get("name") diff --git a/src/scgenerator/_utils/legacy.py b/src/scgenerator/_utils/legacy.py index e198f55..fbd2c62 100644 --- a/src/scgenerator/_utils/legacy.py +++ b/src/scgenerator/_utils/legacy.py @@ -1,3 +1,4 @@ +from genericpath import exists import os import sys from pathlib import Path @@ -28,11 +29,13 @@ def load_config_sequence(path: os.PathLike) -> tuple[list[Path], list[dict[str, def convert_sim_folder(path: os.PathLike): path = Path(path).resolve() + new_root = path.parent / "sc_legagy_converter" / path.name + os.makedirs(new_root, exist_ok=True) config_paths, configs = load_config_sequence(path) master_config = dict(name=path.name, Fiber=configs) - with open(path / "initial_config.toml", "w") as f: + with open(new_root / "initial_config.toml", "w") as f: toml.dump(master_config, f, encoder=toml.TomlNumpyEncoder()) - configuration = Configuration(path / "initial_config.toml", final_output_path=path) + configuration = Configuration(path, final_output_path=new_root) pbar = PBars(configuration.total_num_steps, "Converting") new_paths: dict[VariationDescriptor, Parameters] = dict(configuration) diff --git a/src/scgenerator/_utils/parameter.py b/src/scgenerator/_utils/parameter.py index 6f4ba83..9621356 100644 --- a/src/scgenerator/_utils/parameter.py +++ b/src/scgenerator/_utils/parameter.py @@ -480,7 +480,7 @@ class Parameters(_AbstractParameters): @classmethod def load(cls, path: os.PathLike) -> "Parameters": - return cls(**utils.open_config(path)) + return cls(**utils._open_config(path)) @classmethod def load_and_compute(cls, path: os.PathLike) -> "Parameters": @@ -1011,7 +1011,7 @@ class Configuration: num = utils.find_last_spectrum_num(data_dir) if config_dict is None: try: - config_dict = utils.open_config(data_dir / PARAM_FN) + config_dict = utils._open_config(data_dir / PARAM_FN) except FileNotFoundError: self.logger.warning(f"did not find {PARAM_FN!r} in {data_dir}") return self.State.ABSENT, 0 diff --git a/src/scgenerator/physics/__init__.py b/src/scgenerator/physics/__init__.py index 24c3937..5bb6018 100644 --- a/src/scgenerator/physics/__init__.py +++ b/src/scgenerator/physics/__init__.py @@ -10,8 +10,7 @@ from scipy.optimize import minimize_scalar from .. import math from . import fiber, materials, units, pulse -from .. import _utils -from .._utils import cache +from .._utils import cache, load_material_dico T = TypeVar("T") @@ -62,7 +61,7 @@ def material_dispersion( ) return disp else: - material_dico = utils.load_material_dico(material) + material_dico = load_material_dico(material) if ideal: n_gas_2 = materials.sellmeier(wavelengths, material_dico, pressure, temperature) + 1 else: diff --git a/src/scgenerator/physics/pulse.py b/src/scgenerator/physics/pulse.py index ce8ab29..13fe4cc 100644 --- a/src/scgenerator/physics/pulse.py +++ b/src/scgenerator/physics/pulse.py @@ -810,7 +810,8 @@ def find_lobe_limits(x_axis, values, debug="", already_sorted=True): ) ax.legend() fig.savefig(out_path, bbox_inches="tight") - plt.close() + if fig is not None: + plt.close(fig) else: good_roots, left_lim, right_lim = _select_roots(d_spline, d_roots, dd_roots, fwhm_pos) diff --git a/src/scgenerator/physics/simulate.py b/src/scgenerator/physics/simulate.py index 3864ffc..0d9fef1 100644 --- a/src/scgenerator/physics/simulate.py +++ b/src/scgenerator/physics/simulate.py @@ -219,6 +219,17 @@ class RK4IP: return self.stored_spectra def irun(self) -> Generator[tuple[int, int, np.ndarray], None, None]: + """run the simulation as a generator obj + + Yields + ------- + int + current simulation step + int + current number of spectra returned + np.ndarray + spectrum + """ # Print introduction self.logger.debug( diff --git a/src/scgenerator/scripts/__init__.py b/src/scgenerator/scripts/__init__.py index fb1f473..c740781 100644 --- a/src/scgenerator/scripts/__init__.py +++ b/src/scgenerator/scripts/__init__.py @@ -12,7 +12,7 @@ from ..const import PARAM_FN, PARAM_SEPARATOR from ..physics import fiber, units from ..plotting import plot_setup from ..spectra import SimulationSeries -from .._utils import auto_crop, open_config, save_toml, translate_parameters +from .._utils import auto_crop, _open_config, save_toml, translate_parameters from .._utils.parameter import ( Configuration, Parameters, @@ -258,7 +258,7 @@ def finish_plot(fig, legend_axes, all_labels, params): def plot_helper(config_path: Path) -> Iterable[tuple[dict, list[str], Parameters]]: cc = cycler(color=[f"C{i}" for i in range(10)]) * cycler(ls=["-", "--"]) - pseq = Configuration(open_config(config_path)) + pseq = Configuration(_open_config(config_path)) for style, (variables, params) in zip(cc, pseq): lbl = [pretty_format_value(name, value) for name, value in variables[1:-1]] yield style, lbl, params @@ -267,7 +267,7 @@ def plot_helper(config_path: Path) -> Iterable[tuple[dict, list[str], Parameters def convert_params(params_file: os.PathLike): p = Path(params_file) if p.name == PARAM_FN: - d = open_config(params_file) + d = _open_config(params_file) d = translate_parameters(d) save_toml(params_file, d) print(f"converted {p}")