diff --git a/.gitignore b/.gitignore index 1e04fc0..aa77d32 100644 --- a/.gitignore +++ b/.gitignore @@ -4,6 +4,7 @@ plots* /make_*.py +/debug*.py Archive *.mp4 *.png diff --git a/src/scgenerator/evaluator.py b/src/scgenerator/evaluator.py index 4d669f6..85a14b5 100644 --- a/src/scgenerator/evaluator.py +++ b/src/scgenerator/evaluator.py @@ -1,7 +1,6 @@ import itertools from collections import defaultdict from dataclasses import dataclass -from operator import itemgetter from typing import Any, Callable, Optional, Union import numpy as np @@ -259,8 +258,8 @@ class Evaluator: value = default self.logger.info(prefix + f"using default value of {value} for {target}") self.set_value(target, value, 0) - - assert target == self.__curent_lookup.pop() + last_target = self.__curent_lookup.pop() + assert target == last_target self.__failed_rules[target] = [] if value is None and error is not None: diff --git a/src/scgenerator/math.py b/src/scgenerator/math.py index b50f519..3648519 100644 --- a/src/scgenerator/math.py +++ b/src/scgenerator/math.py @@ -354,3 +354,48 @@ def _polynom_extrapolation_in_place(y: np.ndarray, left_ind: int, right_ind: int y[:left_ind] = r_left * (y[left_ind] - y[left_ind + 1]) + y[left_ind] y[right_ind:] = r_right * (y[right_ind] - y[right_ind - 1]) + y[right_ind] return y + + +def envelope_ind( + signal: np.ndarray, dmin: int = 1, dmax: int = 1, split: bool = False +) -> tuple[np.ndarray, np.ndarray]: + """ + returns the indices of the top/bottom envelope of a signal + + Parameters + ---------- + signal : np.ndarray, shape (n,) + signal array (must be sorted) + dmin, dmax : int, optional + size of chunks for lower/upper envelope + split: bool, optional + split the signal in half along its mean, might help to generate the envelope in some cases + this has the effect of forcing the envlopes to be on either side of the dc signal + by default False + + Returns + ------- + np.ndarray, shape (m,), m < n + lower envelope indices + np.ndarray, shape (l,), l < n + upper envelope indices + """ + + local_min = (np.diff(np.sign(np.diff(signal))) > 0).nonzero()[0] + 1 + local_max = (np.diff(np.sign(np.diff(signal))) < 0).nonzero()[0] + 1 + + if split: + dc_value = np.mean(signal) + local_min = local_min[signal[local_min] < dc_value] + local_max = local_max[signal[local_max] > dc_value] + + if dmin > 1: + local_min = local_min[ + [i + np.argmin(signal[local_min[i : i + dmin]]) for i in range(0, len(local_min), dmin)] + ] + if dmax > 1: + local_max = local_max[ + [i + np.argmax(signal[local_max[i : i + dmax]]) for i in range(0, len(local_max), dmax)] + ] + + return local_min, local_max diff --git a/src/scgenerator/operators.py b/src/scgenerator/operators.py index e59a3e5..3889372 100644 --- a/src/scgenerator/operators.py +++ b/src/scgenerator/operators.py @@ -630,7 +630,7 @@ class FullFieldSPM(AbstractSPM): self.factor = self.fraction * chi3 * units.epsilon0 def __call__(self, state: CurrentState) -> np.ndarray: - return self.fraction * state.field2 * state.field + return self.factor * state.field2 * state.field ################################################## diff --git a/src/scgenerator/parameter.py b/src/scgenerator/parameter.py index 5db0785..51b8ba1 100644 --- a/src/scgenerator/parameter.py +++ b/src/scgenerator/parameter.py @@ -20,7 +20,6 @@ from .evaluator import Evaluator from .logger import get_logger from .operators import ( AbstractConservedQuantity, - EnvelopeLinearOperator, LinearOperator, NonLinearOperator, ) @@ -423,12 +422,13 @@ class Parameters: def __repr_list__(self) -> Iterator[str]: yield from (f"{k}={v}" for k, v in self.dump_dict().items()) - def dump_dict(self, compute=True) -> dict[str, Any]: + def dump_dict(self, compute=True, add_metadata=True) -> dict[str, Any]: if compute: self.compute_in_place() param = Parameters.strip_params_dict(self._param_dico) - param["datetime"] = datetime_module.datetime.now() - param["version"] = __version__ + if add_metadata: + param["datetime"] = datetime_module.datetime.now() + param["version"] = __version__ return param def compute_in_place(self, *to_compute: str): @@ -440,7 +440,8 @@ class Parameters: def compute(self, key: str) -> Any: return self._evaluator.compute(key) - def pformat(self) -> str: + def pretty_str(self) -> str: + """return a pretty formatted string describing the parameters""" return "\n".join( f"{k} = {VariationDescriptor.format_value(k, v)}" for k, v in self.dump_dict().items() ) @@ -597,7 +598,7 @@ class Configuration: self.variationer.append(config.variable) self.fiber_paths.append( utils.ensure_folder( - self.final_path / fiber_folder(i, self.name, config.fixed["name"]), + self.final_path / fiber_folder(i, self.name, Path(config.fixed["name"]).name), mkdir=False, prevent_overwrite=not self.overwrite, ) diff --git a/src/scgenerator/physics/__init__.py b/src/scgenerator/physics/__init__.py index f63631d..f59df64 100644 --- a/src/scgenerator/physics/__init__.py +++ b/src/scgenerator/physics/__init__.py @@ -58,7 +58,7 @@ def material_dispersion( disp = np.zeros(len(w)) ind = w > 0 disp[ind] = material_dispersion( - units.To.m(w[ind]), material, pressure, temperature, ideal, False + units.m.inv(w[ind]), material, pressure, temperature, ideal, False ) return disp else: @@ -102,7 +102,7 @@ def find_optimal_depth( w = w_c + w0 disp = np.zeros(len(w)) ind = w > (w0 / 10) - disp[ind] = material_dispersion(units.To.m(w[ind]), material) + disp[ind] = material_dispersion(units.m.inv(w[ind]), material) propagate = lambda z: spectrum * np.exp(-0.5j * disp * w_c ** 2 * z) integrate = lambda z: math.abs2(np.fft.ifft(propagate(z))) diff --git a/src/scgenerator/physics/units.py b/src/scgenerator/physics/units.py index 998bfd8..0d78a0b 100644 --- a/src/scgenerator/physics/units.py +++ b/src/scgenerator/physics/units.py @@ -2,7 +2,11 @@ # 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 __future__ import annotations +from collections import defaultdict + from typing import Callable, TypeVar, Union +from functools import wraps from operator import itemgetter import numpy as np from numpy import pi @@ -25,17 +29,38 @@ provided you decorate it with @unit and provide at least a type and a label types are "WL", "FREQ", "AFREQ", "TIME", "PRESSURE", "TEMPERATURE", "OTHER" """ _T = TypeVar("_T") +_UT = Callable[[_T], _T] -class From: - pass +def chain(c1: _UT, c2: _UT) -> _UT: + def chained_function(x: _T) -> _T: + return c1(c2(x)) + + return chained_function + + +class UnitMap(dict): + def __setitem__(self, new_name, new_func): + super().__setitem__(new_name, new_func) + already_here = [(name, func) for name, func in self.items() if isinstance(name, str)] + for name, func in already_here: + super().__setitem__((name, new_name), chain(self[name].inv, new_func)) + super().__setitem__((new_name, name), chain(self[new_name].inv, func)) + + +units_map: dict[str, dict[Union[str, tuple[str, str]], Unit]] = defaultdict(UnitMap) class To: - pass + def __init__(self, name: str, tpe: str): + self.name = name + self.type = tpe - -units_map = dict() + def __getattr__(self, key: str): + try: + return units_map[self.type][self.name, key] + except KeyError: + raise KeyError(f"no registered unit named {key!r} of type {self.type!r}") from None def W_to_Vm(n0: float, A_eff: float) -> float: @@ -54,22 +79,36 @@ def W_to_Vm(n0: float, A_eff: float) -> float: return 1.0 / np.sqrt(A_eff * 0.5 * epsilon0 * c * n0) +class Unit: + __func: _UT + to: To + inv: _UT + name: str + label: str + + def __init__(self, func: _UT, inv: _UT, name: str, label: str, tpe: str): + self.__func = func + self.inv = inv + self.to = To(name, tpe) + self.name = name + self.label = label + self.type = tpe + self.__name__ = name + + def __call__(self, x: _T) -> _T: + """call the original unit function""" + return self.__func(x) + + def unit(tpe: str, label: str, inv: Callable = None): - def unit_maker(func): + def unit_maker(func) -> Unit: nonlocal inv name = func.__name__ if inv is None: inv = func - setattr(From, name, func.__call__) - setattr(To, name, inv.__call__) - func.type = tpe - func.label = label - func.name = name - func.inv = inv - if name in units_map: - raise NameError(f"Two unit functions with the same name {name!r} were defined") - units_map[name] = func - return func + unit_obj = wraps(func)(Unit(func, inv, name, label, tpe)) + units_map[tpe][name] = unit_obj + return unit_obj return unit_maker @@ -89,6 +128,11 @@ def um(l: _T) -> _T: return 2 * pi * c / (l * 1e-6) +@unit("FREQ", r"Frequency $f$ (Hz)", lambda w: w / (2 * pi)) +def Hz(f: _T) -> _T: + return 2 * pi * f + + @unit("FREQ", r"Frequency $f$ (THz)", lambda w: w / (1e12 * 2 * pi)) def THz(f: _T) -> _T: return 1e12 * 2 * pi * f @@ -306,7 +350,7 @@ class PlotRange(tuple): def __iter__(self): yield self.left yield self.right - yield self.unit.__name__ + yield self.unit.name def __repr__(self): return "PlotRange" + super().__repr__() diff --git a/src/scgenerator/plotting.py b/src/scgenerator/plotting.py index 08fdb48..9c00139 100644 --- a/src/scgenerator/plotting.py +++ b/src/scgenerator/plotting.py @@ -419,6 +419,7 @@ def transform_2D_propagation( params: Parameters, log: Union[int, float, bool, str] = "1D", skip: int = 1, + y_axis=None, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """transforms raws values into plottable values @@ -455,8 +456,8 @@ def transform_2D_propagation( is_complex, x_axis, plt_range = prep_plot_axis(values, plt_range, params) if is_complex: values = abs2(values) - - y_axis = params.z_targets + if y_axis is None: + y_axis = params.z_targets x_axis, values = uniform_axis(x_axis, values, plt_range) y_axis, values.T[:] = uniform_axis(y_axis, values.T, None) @@ -1102,7 +1103,7 @@ def partial_plot(root: os.PathLike): ) t, z, values = transform_2D_propagation( - np.fft.ifft(raw_values), + params.ifft(raw_values), PlotRange(-10, 10, "ps"), params, log=False, diff --git a/src/scgenerator/scripts/__init__.py b/src/scgenerator/scripts/__init__.py index eadb2f1..e11938b 100644 --- a/src/scgenerator/scripts/__init__.py +++ b/src/scgenerator/scripts/__init__.py @@ -127,7 +127,7 @@ def plot_init( lbl = plot_1_dispersion(lim_disp, tl, tr, style, lbl, params, loss_ax) lbl = plot_1_init_spec_field(lim_field, lim_spec, bl, br, style, lbl, params) all_labels.append(lbl) - print(params.pformat()) + print(params.pretty_str()) finish_plot(fig, tr, all_labels, params) @@ -142,8 +142,8 @@ def plot_1_init_spec_field( ): field = math.abs2(params.field_0) spec = math.abs2(params.spec_0) - t = units.To.fs(params.t) - wl = units.To.nm(params.w) + t = units.fs.inv(params.t) + wl = units.nm.inv(params.w) lbl += f" max at {wl[spec.argmax()]:.1f} nm" @@ -219,7 +219,7 @@ def plot_1_dispersion( right.set_ylabel(units.D_ps_nm_km.label) # plot beta2 - left.plot(units.To.nm(params.w[m]), units.beta2_fs_cm.inv(beta_arr[m]), label=" ", **style) + left.plot(units.nm.inv(params.w[m]), units.beta2_fs_cm.inv(beta_arr[m]), label=" ", **style) left.set_ylabel(units.beta2_fs_cm.label) left.set_xlabel(units.nm.label) diff --git a/src/scgenerator/utils.py b/src/scgenerator/utils.py index 8f782fd..a2f8e34 100644 --- a/src/scgenerator/utils.py +++ b/src/scgenerator/utils.py @@ -311,8 +311,18 @@ def ensure_folder(path: Path, prevent_overwrite: bool = True, mkdir=True) -> Pat path = path.parent / (folder_name + f"_{i}") -def branch_id(branch: tuple[Path, ...]) -> str: - return branch[-1].name.split()[1] +def branch_id(branch: Path) -> tuple[int, int]: + sim_match = branch.parent.name.split()[0] + if sim_match.isdigit(): + s_int = int(sim_match) + else: + s_int = 0 + branch_match = re.search(r"(?<=b_)[0-9]+", branch.name) + if branch_match is None: + b_int = 0 + else: + b_int = int(branch_match[0]) + return s_int, b_int def find_last_spectrum_num(data_dir: Path): @@ -519,5 +529,5 @@ def simulations_list(path: os.PathLike) -> 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) + paths.sort(key=branch_id) return [p for p in paths if p.parent.name == paths[-1].parent.name] diff --git a/src/scgenerator/variationer.py b/src/scgenerator/variationer.py index 3dee911..bc9b26a 100644 --- a/src/scgenerator/variationer.py +++ b/src/scgenerator/variationer.py @@ -149,6 +149,8 @@ class VariationDescriptor(BaseModel): p = Path(value) if p.exists(): return p.stem + elif callable(value): + return getattr(value, "__name__", repr(value)) return str(value) class Config: diff --git a/testing/configs/Chang2011Fig2.toml b/testing/configs/Chang2011Fig2.toml index 957977e..e54062a 100644 --- a/testing/configs/Chang2011Fig2.toml +++ b/testing/configs/Chang2011Fig2.toml @@ -8,12 +8,12 @@ width = 30e-15 core_radius = 10e-6 model = "marcatili" gas_name = "argon" -pressure = 1e5 +pressure = 3.2e5 length = 0.1 -interpolation_range = [100e-9, 3000e-9] +interpolation_range = [120e-9, 3000e-9] full_field = true -dt = 0.1e-15 +dt = 0.05e-15 t_num = 32768 z_num = 128 -step_size = 2e-6 +step_size = 10e-6 diff --git a/testing/test_full_field.py b/testing/test_full_field.py index 9e80029..1a338a8 100644 --- a/testing/test_full_field.py +++ b/testing/test_full_field.py @@ -1,36 +1,61 @@ import warnings - -import matplotlib.pyplot as plt import numpy as np +import rediscache import scgenerator as sc from customfunc.app import PlotApp -from scgenerator.physics.simulate import RK4IP -from customfunc import pprint +from scipy.interpolate import interp1d +from tqdm import tqdm -warnings.filterwarnings("error") +# warnings.filterwarnings("error") + + +@rediscache.rcache +def get_specs(params: dict): + p = sc.Parameters(**params) + sim = sc.RK4IP(p) + return [s[-1] for s in tqdm(sim.irun(), total=p.z_num)], p.dump_dict() def main(): + params = sc.Parameters.load("testing/configs/Chang2011Fig2.toml") - x = params.l * 1e9 - o = np.argsort(x) - x = x[o] + specs, params = get_specs(params.dump_dict(add_metadata=False)) + params = sc.Parameters(**params) + rs = sc.PlotRange(100, 1500, "nm") + rt = sc.PlotRange(-500, 500, "fs") + x, o, ext = rs.sort_axis(params.w) + vmin = -50 + with PlotApp(i=(int, 0, params.z_num - 1)) as app: + spec_ax = app[0] + spec_ax.set_xlabel(rs.unit.label) + field_ax = app[1] + field_ax.set_xlabel(rt.unit.label) + x: float = 4.5 + y = sc.units.m.to.nm(x) - plt.plot(x, sc.abs2(params.spec_0[o])) - state = sc.operators.CurrentState( - params.length, 0, params.step_size, 1.0, params.ifft, params.spec_0 - ) - # expD = np.exp(state.h / 2 * params.linear_operator(state)) - # plt.plot(x, expD.imag[o], x, expD.real[o]) - plt.plot(x, sc.abs2(params.nonlinear_operator(state))[o]) - plt.yscale("log") - plt.xlim(100, 2000) - plt.show() + @app.update + def draw(i): + spec, *fields = compute(i) + spec_ax.set_line_data("spec", *spec, label=f"z = {params.z_targets[i]*1e2:.0f}cm") + for label, x, y in fields: + field_ax.set_line_data(label, x, y) - # for *_, spec in RK4IP(params).irun(): - # plt.plot(w[2:-2], sc.abs2(spec[ind])) - # plt.show() - # plt.close() + print(params) + + @app.cache + def compute(i): + xt, field = sc.transform_1D_values(params.ifft(specs[i]), rt, params) + x, spec = sc.transform_1D_values(sc.abs2(specs[i]), rs, params, log=True) + # spec = np.where(spec > vmin, spec, vmin) + field2 = sc.abs2(field) + bot, top = sc.math.envelope_ind(field2) + return (x, spec), ("field^2", xt, field2), ("envelope", xt[top], field2[top]) + + # bot, top = sc.math.envelope_ind(field) + # bot = interp1d(xt[bot], field[bot], "cubic", bounds_error=False, fill_value=0)(xt) + # top = interp1d(xt[top], field[top], "cubic", bounds_error=False, fill_value=0)(xt) + + # return ((x, spec), ("upper", xt, top), ("field", xt, field), ("lower", xt, bot)) if __name__ == "__main__":