diff --git a/src/scgenerator/_utils/__init__.py b/src/scgenerator/_utils/__init__.py index b540d9e..a2dcd30 100644 --- a/src/scgenerator/_utils/__init__.py +++ b/src/scgenerator/_utils/__init__.py @@ -91,8 +91,8 @@ def load_previous_spectrum(prev_data_dir: str) -> np.ndarray: @cache -def load_spectrum(folder: os.PathLike) -> np.ndarray: - return np.load(folder) +def load_spectrum(file: os.PathLike) -> np.ndarray: + return np.load(file) def conform_toml_path(path: os.PathLike) -> str: diff --git a/src/scgenerator/_utils/legacy.py b/src/scgenerator/_utils/legacy.py index fbd2c62..35cf4fd 100644 --- a/src/scgenerator/_utils/legacy.py +++ b/src/scgenerator/_utils/legacy.py @@ -2,7 +2,6 @@ from genericpath import exists import os import sys from pathlib import Path -from pprint import pprint from typing import Any, Set import numpy as np @@ -10,7 +9,7 @@ import toml from ..const import PARAM_FN, SPEC1_FN, SPEC1_FN_N, SPECN_FN1, Z_FN from .parameter import Configuration, Parameters -from .utils import fiber_folder, save_parameters +from .utils import save_parameters from .pbar import PBars from .variationer import VariationDescriptor, Variationer @@ -24,14 +23,22 @@ def load_config(path: os.PathLike) -> dict[str, Any]: 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] + confs = [load_config(cfg) for cfg in paths] + repeat = None + for c in confs: + nums = c["variable"].pop("num", None) + if nums is not None: + repeat = len(nums) + if repeat is not None: + confs[0]["repeat"] = repeat + return paths, confs 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) + _, configs = load_config_sequence(path) master_config = dict(name=path.name, Fiber=configs) with open(new_root / "initial_config.toml", "w") as f: toml.dump(master_config, f, encoder=toml.TomlNumpyEncoder()) @@ -43,6 +50,8 @@ def convert_sim_folder(path: os.PathLike): 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 old_path in old_paths: + continue if not Path(old_path).is_dir(): raise FileNotFoundError(f"missing {old_path} from {path}. Aborting.") old_paths.add(old_path) @@ -51,7 +60,6 @@ def convert_sim_folder(path: os.PathLike): z_limits = (z_num_start, z_num_start + params.z_num) old2new.append((old_path, d, new_paths[d], z_limits)) - processed_paths: Set[Path] = set() processed_specs: Set[VariationDescriptor] = set() for old_path, descr, new_params, (start_z, end_z) in old2new: @@ -64,17 +72,7 @@ def convert_sim_folder(path: os.PathLike): 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() + pbar.close() def _mv_specs(pbar: PBars, new_params: Parameters, start_z: int, spec_num: int, old_spec: Path): @@ -82,12 +80,10 @@ def _mv_specs(pbar: PBars, new_params: Parameters, start_z: int, spec_num: int, 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) + new_path = new_params.final_path / SPEC1_FN.format(spec_num - start_z) else: - np.save( - new_params.final_path / SPEC1_FN_N.format(spec_num - start_z, j), - spec1, - ) + new_path = new_params.final_path / SPEC1_FN_N.format(spec_num - start_z, j) + np.save(new_path, spec1) pbar.update() diff --git a/src/scgenerator/_utils/parameter.py b/src/scgenerator/_utils/parameter.py index 039cedd..ffc3215 100644 --- a/src/scgenerator/_utils/parameter.py +++ b/src/scgenerator/_utils/parameter.py @@ -747,7 +747,7 @@ class Evaluator: else: default = self.get_default(target) if default is None: - error = NoDefaultError(prefix + f"No default provided for {target}") + error = NoDefaultError(prefix + f"No default provided for {target}. Current lookup cycle : {self.__curent_lookup!r}") else: value = default self.set_value(target, value, 0) diff --git a/src/scgenerator/cli/cli.py b/src/scgenerator/cli/cli.py index 92b81ab..4d6ed39 100644 --- a/src/scgenerator/cli/cli.py +++ b/src/scgenerator/cli/cli.py @@ -4,13 +4,15 @@ import re from collections import ChainMap from pathlib import Path +import matplotlib.pyplot as plt import numpy as np -from .. import const, env, scripts from .. import _utils as utils +from .. import const, env, scripts from ..logger import get_logger from ..physics.fiber import dispersion_coefficients from ..physics.simulate import SequencialSimulations, run_simulation +from ..plotting import partial_plot try: import ray @@ -124,6 +126,13 @@ def create_parser(): convert_parser.add_argument("config", help="path to config/parameter file") convert_parser.set_defaults(func=translate_parameters) + preview_parser = subparsers.add_parser("preview", help="preview a currently running simulation") + plc_hld = "XX" + preview_parser.add_argument( + "path", help=f"path to the directory containing {const.SPEC1_FN.format(plc_hld)!r}" + ) + preview_parser.set_defaults(func=preview) + return parser @@ -165,6 +174,12 @@ def merge(args): utils.merge(output, path_trees) +def preview(args): + path = args.path + partial_plot(path) + plt.show() + + def prep_ray(): logger = get_logger(__name__) if ray: diff --git a/src/scgenerator/math.py b/src/scgenerator/math.py index 505c6b7..4b6f069 100644 --- a/src/scgenerator/math.py +++ b/src/scgenerator/math.py @@ -219,7 +219,7 @@ def make_uniform_1D(values, x_axis, n=1024, method="linear"): def all_zeros(x: np.ndarray, y: np.ndarray) -> np.ndarray: """find all the x values such that y(x)=0 with linear interpolation""" pos = np.argwhere(y[1:] * y[:-1] < 0)[:, 0] - m = (y[pos] - y[pos - 1]) / (x[pos] - x[pos - 1]) + m = (y[pos + 1] - y[pos]) / (x[pos + 1] - x[pos]) return -y[pos] / m + x[pos] diff --git a/src/scgenerator/plotting.py b/src/scgenerator/plotting.py index 6163520..9ad8ea7 100644 --- a/src/scgenerator/plotting.py +++ b/src/scgenerator/plotting.py @@ -1,4 +1,5 @@ import os +import re from pathlib import Path from typing import Any, Callable, Literal, Optional, Union @@ -10,12 +11,13 @@ from scipy.interpolate import UnivariateSpline from scipy.interpolate.interpolate import interp1d from . import math -from .const import PARAM_SEPARATOR +from ._utils import load_spectrum +from ._utils.parameter import Parameters +from ._utils.utils import PlotRange, sort_axis +from .const import PARAM_SEPARATOR, SPEC1_FN 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 RangeType = tuple[float, float, Union[str, Callable]] NO_LIM = object() @@ -1061,3 +1063,18 @@ def measure_and_annotate_fwhm( va="center", ) return right - left + + +def partial_plot(root: os.PathLike): + fig, (left, right) = plt.subplots(1, 2, figsize=(12, 8)) + path = Path(root) + spec_list = sorted( + path.glob(SPEC1_FN.format("*")), key=lambda el: int(re.search("[0-9]+", el.name)[0]) + ) + print(spec_list) + raw_values = np.array([load_spectrum(s) for s in spec_list]) + specs = units.to_log2D(math.abs2(np.fft.fftshift(raw_values))) + fields = math.abs2(np.fft.ifft(raw_values)) + left.imshow(specs, origin="lower", aspect="auto", vmin=-60, interpolation="nearest") + right.imshow(fields, origin="lower", aspect="auto", interpolation="nearest") + return left, right diff --git a/src/scgenerator/scripts/__init__.py b/src/scgenerator/scripts/__init__.py index c740781..1d5fc8f 100644 --- a/src/scgenerator/scripts/__init__.py +++ b/src/scgenerator/scripts/__init__.py @@ -131,13 +131,21 @@ def plot_init( finish_plot(fig, tr, all_labels, params) -def plot_1_init_spec_field(lim_t, lim_l, left, right, style, lbl, params): +def plot_1_init_spec_field( + lim_t: Optional[tuple[float, float]], + lim_l: Optional[tuple[float, float]], + left: plt.Axes, + right: plt.Axes, + style: dict[str, Any], + lbl: str, + params: Parameters, +): field = math.abs2(params.field_0) spec = math.abs2(params.spec_0) t = units.To.fs(params.t) wl = units.To.nm(params.w) - lbl.append(f"max at {wl[spec.argmax()]:.1f} nm") + lbl += f" max at {wl[spec.argmax()]:.1f} nm" mt = np.ones_like(t, dtype=bool) if lim_t is not None: @@ -173,9 +181,9 @@ def plot_1_dispersion( zdw = math.all_zeros(wl, beta_arr) if len(zdw) > 0: zdw = zdw[np.argmin(abs(zdw - params.wavelength))] - lbl.append(f"ZDW at {zdw:.1f}nm") + lbl += f" ZDW at {zdw*1e9:.1f}nm" else: - lbl.append("") + lbl += "" m = np.ones_like(wl, dtype=bool) if lim is None: @@ -226,18 +234,13 @@ def plot_1_dispersion( return lbl -def finish_plot(fig, legend_axes, all_labels, params): +def finish_plot(fig: plt.Figure, legend_axes: plt.Axes, all_labels: list[str], params: Parameters): fig.suptitle(params.name) plt.tight_layout() handles, _ = legend_axes.get_legend_handles_labels() - lbl_lengths = [[len(l) for l in lbl] for lbl in all_labels] - lengths = np.max(lbl_lengths, axis=0) - labels = [ - " ".join(format(l, f">{int(s)}s") for s, l in zip(lengths, lab)) for lab in all_labels - ] - legend = legend_axes.legend(handles, labels, prop=dict(size=8, family="monospace")) + legend = legend_axes.legend(handles, all_labels, prop=dict(size=8, family="monospace")) out_path = env.output_path() @@ -258,10 +261,9 @@ 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)) - 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 + for style, (descriptor, params) in zip(cc, Configuration(config_path)): + params.compute() + yield style, descriptor.branch.formatted_descriptor(), params def convert_params(params_file: os.PathLike): diff --git a/src/scgenerator/spectra.py b/src/scgenerator/spectra.py index fe9279d..eede238 100644 --- a/src/scgenerator/spectra.py +++ b/src/scgenerator/spectra.py @@ -191,25 +191,27 @@ class SimulationSeries: 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}" - ) + return self.spectra(self.z_ind(z_descr), sim_ind) else: z_ind = z_descr - - if z_ind < self.z_inds[0]: + if 0 <= z_ind < self.z_inds[0]: return self.previous.spectra(z_ind, sim_ind) + elif z_ind < 0: + z_ind = self.total_num_steps + z_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 z_ind(self, pos: float) -> int: + if self.z[0] <= pos <= self.z[-1]: + return self.z_inds[np.argmin(np.abs(self.z - pos))] + elif 0 <= pos < self.z[0]: + return self.previous.z_ind(pos) + else: + raise ValueError(f"cannot match z={pos} with max length of {self.total_length}") + def fields( self, z_descr: Union[float, int, None] = None, sim_ind: Optional[int] = 0 ) -> Spectrum: diff --git a/testing/test_all_zeros.py b/testing/test_all_zeros.py new file mode 100644 index 0000000..b39b86e --- /dev/null +++ b/testing/test_all_zeros.py @@ -0,0 +1,16 @@ +from scgenerator.math import all_zeros +import matplotlib.pyplot as plt +import numpy as np + + +def main(): + x = np.linspace(-10, 10, 30) + y = np.sin(x) + z = all_zeros(x, y) + plt.plot(x, y) + plt.plot(z, z * 0, ls="", marker="o") + plt.show() + + +if __name__ == "__main__": + main()