dedebenui/scgenerator
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Huge plotting revamp

Browse Source
This commit is contained in:
Benoît Sierro
2021-07-23 16:32:21 +02:00
parent 873d02c60f
commit ea943d7adf
11 changed files with 747 additions and 734 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/scgenerator/__init__.py
View File

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

68
src/scgenerator/initialize.py
View File

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

30
src/scgenerator/io.py
View File

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

25
src/scgenerator/physics/simulate.py
View File

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

7
src/scgenerator/physics/units.py
View File

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

1156
src/scgenerator/plotting.py
View File

File diff suppressed because it is too large Load Diff

9
src/scgenerator/scripts/__init__.py
View File

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

10
src/scgenerator/scripts/slurm_submit.py
View File

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

28
src/scgenerator/spectra.py
View File

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

64
src/scgenerator/utils/__init__.py
View File

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

14
testing/test_new_iterator.py Normal file
View File

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