From 97a19d4ffb34c624ea42b522e8177ce8ddf4025a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Sierro?= Date: Wed, 23 Jun 2021 10:06:14 +0200 Subject: [PATCH] added l to parameters, some plotting stuff --- src/scgenerator/initialize.py | 4 +- src/scgenerator/physics/units.py | 240 ++++++++++++----------------- src/scgenerator/plotting.py | 14 +- src/scgenerator/spectra.py | 109 +++++++++++-- src/scgenerator/utils/parameter.py | 1 + 5 files changed, 205 insertions(+), 163 deletions(-) diff --git a/src/scgenerator/initialize.py b/src/scgenerator/initialize.py index 5bf9ff0..3426d2f 100644 --- a/src/scgenerator/initialize.py +++ b/src/scgenerator/initialize.py @@ -635,7 +635,8 @@ def build_sim_grid( t_num = len(t) z_targets = np.linspace(0, length, z_num) w_c, w0, w, w_power_fact = update_frequency_domain(t, wavelength) - return z_targets, t, time_window, t_num, dt, w_c, w0, w, w_power_fact + l = units.To.m(w) + return z_targets, t, time_window, t_num, dt, w_c, w0, w, w_power_fact, l def build_sim_grid_in_place(params: BareParams): @@ -650,6 +651,7 @@ def build_sim_grid_in_place(params: BareParams): params.w0, params.w, params.w_power_fact, + params.l, ) = build_sim_grid( params.length, params.z_num, params.wavelength, params.time_window, params.t_num, params.dt ) diff --git a/src/scgenerator/physics/units.py b/src/scgenerator/physics/units.py index b00a201..668b8a5 100644 --- a/src/scgenerator/physics/units.py +++ b/src/scgenerator/physics/units.py @@ -2,7 +2,9 @@ # 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 typing import Callable, Union +import re +from threading import settrace +from typing import Callable, TypeVar, Union from dataclasses import dataclass from ..utils.parameter import Parameter, type_checker import numpy as np @@ -21,204 +23,147 @@ prefix = dict(P=1e12, G=1e9, M=1e6, k=1e3, d=1e-1, c=1e-2, m=1e-3, u=1e-6, n=1e- """ Below are common units. You can define your own unit function -this function must have a few porperties: -inv : function - inverse of the function. example : - um(1) -> 883651567308853.2 - um.inv(883651567308853.2) -> 1.0 -label : str - label to be displayed on plot -type : ("WL", "FREQ", "AFREQ", "TIME", "OTHER") +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") -def m(l): +class From: + pass + + +class To: + pass + + +units_map = dict() + + +def unit(tpe: str, label: str, inv: Callable = None): + def unit_maker(func): + 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.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 + + return unit_maker + + +@unit("WL", r"Wavelength $\lambda$ (m)") +def m(l: _T) -> _T: return 2 * pi * c / l -m.inv = m -m.label = r"Wavelength $\lambda$ (m)" -m.type = "WL" - - -def nm(l): +@unit("WL", r"Wavelength $\lambda$ (nm)") +def nm(l: _T) -> _T: return 2 * pi * c / (l * 1e-9) -nm.inv = nm -nm.label = r"Wavelength $\lambda$ (nm)" -nm.type = "WL" - - -def um(l): +@unit("WL", r"Wavelength $\lambda$ (μm)") +def um(l: _T) -> _T: return 2 * pi * c / (l * 1e-6) -um.inv = um -um.label = r"Wavelength $\lambda$ ($\mathrm{\mu}$m)" -um.type = "WL" - - -def THz(f): +@unit("FREQ", r"Frequency $f$ (THz)", lambda w: w / (1e12 * 2 * pi)) +def THz(f: _T) -> _T: return 1e12 * 2 * pi * f -THz.inv = lambda w: w / (1e12 * 2 * pi) -THz.label = r"Frequency $f$ (THz)" -THz.type = "FREQ" - - -def PHz(f): +@unit("FREQ", r"Frequency $f$ (PHz)", lambda w: w / (1e15 * 2 * pi)) +def PHz(f: _T) -> _T: return 1e15 * 2 * pi * f -PHz.inv = lambda w: w / (1e15 * 2 * pi) -PHz.label = r"Frequency $f$ (PHz)" -PHz.type = "FREQ" - - -def rad_s(w): +@unit("AFREQ", r"Angular frequency $\omega$ ($\frac{\mathrm{rad}}{\mathrm{s}}$)") +def rad_s(w: _T) -> _T: return w -rad_s.inv = rad_s -rad_s.label = r"Angular frequency $\omega$ ($\frac{\mathrm{rad}}{\mathrm{s}}$)" -rad_s.type = "AFREQ" - - -def Prad_s(w): +@unit( + "AFREQ", r"Angular frequency $\omega$ ($\frac{\mathrm{Prad}}{\mathrm{s}}$)", lambda w: 1e-15 * w +) +def Prad_s(w: _T) -> _T: return w * 1e15 -Prad_s.inv = lambda w: 1e-15 * w -Prad_s.label = r"Angular frequency $\omega$ ($\frac{\mathrm{Prad}}{\mathrm{s}}$)" -Prad_s.type = "AFREQ" - - -def rel_time(t): +@unit("TIME", r"relative time ${\tau}/{\tau_\mathrm{0, FWHM}}$") +def rel_time(t: _T) -> _T: return t -rel_time.inv = rel_time -rel_time.label = r"relative time ${\tau}/{\tau_\mathrm{0, FWHM}}$" -rel_time.type = "TIME" - - -def rel_freq(f): +@unit("FREQ", r"relative angular freq. $(\omega - \omega_0)/\Delta\omega_0$") +def rel_freq(f: _T) -> _T: return f -rel_freq.inv = rel_freq -rel_freq.label = r"relative angular freq. $(\omega - \omega_0)/\Delta\omega_0$" -rel_freq.type = "FREQ" - - -def s(t): +@unit("TIME", r"Time $t$ (s)") +def s(t: _T) -> _T: return t -s.inv = s -s.label = r"Time $t$ (s)" -s.type = "TIME" - - -def us(t): +@unit("TIME", r"Time $t$ (us)", lambda t: t * 1e6) +def us(t: _T) -> _T: return t * 1e-6 -us.inv = lambda t: t * 1e6 -us.label = r"Time $t$ (us)" -us.type = "TIME" - - -def ns(t): +@unit("TIME", r"Time $t$ (ns)", lambda t: t * 1e9) +def ns(t: _T) -> _T: return t * 1e-9 -ns.inv = lambda t: t * 1e9 -ns.label = r"Time $t$ (ns)" -ns.type = "TIME" - - -def ps(t): +@unit("TIME", r"Time $t$ (ps)", lambda t: t * 1e12) +def ps(t: _T) -> _T: return t * 1e-12 -ps.inv = lambda t: t * 1e12 -ps.label = r"Time $t$ (ps)" -ps.type = "TIME" - - -def fs(t): +@unit("TIME", r"Time $t$ (fs)", lambda t: t * 1e15) +def fs(t: _T) -> _T: return t * 1e-15 -fs.inv = lambda t: t * 1e15 -fs.label = r"Time $t$ (fs)" -fs.type = "TIME" - - -def inv(x): +@unit("WL", "inverse") +def inv(x: _T) -> _T: return 1 / x -inv.inv = inv -inv.label = "inverse" -inv.type = "WL" - - -def bar(p): +@unit("PRESSURE", "Pressure (bar)", lambda p: 1e-5 * p) +def bar(p: _T) -> _T: return 1e5 * p -bar.inv = lambda p: 1e-5 * p -bar.label = "Pressure (bar)" -bar.type = "PRESSURE" - - -def beta2_fs_cm(b2): +@unit("OTHER", r"$\beta_2$ (fs$^2$/cm)", lambda b2: 1e28 * b2) +def beta2_fs_cm(b2: _T) -> _T: return 1e-28 * b2 -beta2_fs_cm.inv = lambda b2: 1e28 * b2 -beta2_fs_cm.label = r"$\beta_2$ (fs$^2$/cm)" -beta2_fs_cm.type = "OTHER" - - -def beta2_ps_km(b2): +@unit("OTHER", r"$\beta_2$ (ps$^2$/km)", lambda b2: 1e27 * b2) +def beta2_ps_km(b2: _T) -> _T: return 1e-27 * b2 -beta2_ps_km.inv = lambda b2: 1e27 * b2 -beta2_ps_km.label = r"$\beta_2$ (ps$^2$/km)" -beta2_ps_km.type = "OTHER" - - -def D_ps_nm_km(D): +@unit("OTHER", r"$D$ (ps/(nm km))", lambda D: 1e6 * D) +def D_ps_nm_km(D: _T) -> _T: return 1e-6 * D -D_ps_nm_km.inv = lambda D: 1e6 * D -D_ps_nm_km.label = r"$D$ (ps/(nm km))" -D_ps_nm_km.type = "OTHER" +@unit("TEMPERATURE", r"Temperature (K)") +def K(t: _T) -> _T: + return t -units_map = dict( - nm=nm, - um=um, - m=m, - THz=THz, - PHz=PHz, - rad_s=rad_s, - Prad_s=Prad_s, - rel_freq=rel_freq, - rel_time=rel_time, - s=s, - us=us, - ns=ns, - ps=ps, - fs=fs, -) +@unit("TEMPERATURE", r"Temperature (°C)", lambda t_K: t_K - 272.15) +def C(t_C: _T) -> _T: + return t_C + 272.15 def get_unit(unit: Union[str, Callable]) -> Callable[[float], float]: @@ -325,13 +270,22 @@ def sort_axis(axis, plt_range: PlotRange): return out_ax, indices, (out_ax[0], out_ax[-1]) -def to_WL(spectrum, frep, lambda_): +def to_WL(spectrum: np.ndarray, lambda_: np.ndarray) -> np.ndarray: + """rescales the spectrum because of uneven binning when going from freq to wl + + Parameters + ---------- + spectrum : np.ndarray, shape (n, ) + intensity spectrum + lambda_ : np.ndarray, shape (n, ) + wavelength in m + + Returns + ------- + np.ndarray, shape (n, ) + intensity spectrum correctly scaled """ - When a spectrogram is displayed as function of wl instead of frequency, we - need to adjust the amplitude of each bin for the integral over the whole frequency - range to match. - """ - m = 2 * pi * c / (lambda_ ** 2) * frep * spectrum + m = 2 * pi * c / (lambda_ ** 2) * spectrum return m diff --git a/src/scgenerator/plotting.py b/src/scgenerator/plotting.py index e30a782..62d915c 100644 --- a/src/scgenerator/plotting.py +++ b/src/scgenerator/plotting.py @@ -421,9 +421,7 @@ def plot_spectrogram( new_f, ind_f, _ = units.sort_axis(params.w, f_range) values = spec[ind_t][:, ind_f] if f_range[2].type == "WL": - values = np.apply_along_axis( - units.to_WL, 1, values, params.repetition_rate, units.m(f_range[2].inv(new_f)) - ) + values = np.apply_along_axis(units.to_WL, 1, values, units.m(f_range[2].inv(new_f))) values = np.apply_along_axis(make_uniform_1D, 1, values, new_f) if time_axis == 0: @@ -528,7 +526,7 @@ def plot_results_2D( # make uniform if converting to wavelength if plt_range.unit.type == "WL": if is_spectrum: - values = np.apply_along_axis(units.to_WL, 1, values, params.repetition_rate, x_axis) + values = np.apply_along_axis(units.to_WL, 1, values, x_axis) values = np.array( [make_uniform_1D(v, x_axis, n=len(x_axis), method="linear") for v in values] ) @@ -648,7 +646,7 @@ def plot_results_1D( # make uniform if converting to wavelength if plt_range.unit.type == "WL": if is_spectrum: - values = units.to_WL(values, params.repetition_rate, units.m.inv(params.w[ind])) + values = units.to_WL(values, units.m.inv(params.w[ind])) # change the resolution if isinstance(spacing, float): @@ -810,8 +808,8 @@ def plot_avg( values *= yscaling mean_values = np.mean(values, axis=0) if plt_range.unit.type == "WL" and renormalize: - values = np.apply_along_axis(units.to_WL, 1, values, params.repetition_rate, x_axis) - mean_values = units.to_WL(mean_values, params.repetition_rate, x_axis) + values = np.apply_along_axis(units.to_WL, 1, values, x_axis) + mean_values = units.to_WL(mean_values, x_axis) # change the resolution if isinstance(spacing, float): @@ -962,7 +960,7 @@ def prepare_plot_1D(values, plt_range, x_axis, yscaling=1, spacing=1, frep=80e6) values = values[:, ind] if plt_range.unit.type == "WL": - values = np.apply_along_axis(units.to_WL, -1, values, frep, x_axis) + values = np.apply_along_axis(units.to_WL, -1, values, x_axis) if isinstance(spacing, float): new_x_axis = np.linspace(*span(x_axis), int(len(x_axis) / spacing)) diff --git a/src/scgenerator/spectra.py b/src/scgenerator/spectra.py index 30ebdc2..ed85382 100644 --- a/src/scgenerator/spectra.py +++ b/src/scgenerator/spectra.py @@ -3,24 +3,30 @@ 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 import numpy as np +from tqdm.std import Bar from . import initialize, io, math -from .physics import units +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 class Spectrum(np.ndarray): - def __new__(cls, input_array, wl, frep=1): + params: BareParams + + def __new__(cls, input_array, params: BareParams): # 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.frep = frep - obj.wl = wl + obj.params = params + # Finally, we must return the newly created object: return obj @@ -28,8 +34,91 @@ class Spectrum(np.ndarray): # see InfoArray.__array_finalize__ for comments if obj is None: return - self.frep = getattr(obj, "frep", None) - self.wl = getattr(obj, "wl", None) + self.params = getattr(obj, "params", None) + + 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) -> tuple[np.ndarray, 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) + + @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) + ) class Pulse(Sequence): @@ -112,7 +201,7 @@ class Pulse(Sequence): yield x_axis[order], func(spec)[:, order] def _to_wl_int(self, spectrum): - return units.to_WL(math.abs2(spectrum), spectrum.frep, spectrum.wl) + return units.to_WL(math.abs2(spectrum), spectrum.wl) def _to_freq_int(self, spectrum): return math.abs2(spectrum) @@ -145,7 +234,6 @@ class Pulse(Sequence): np.sqrt( units.to_WL( math.abs2(spectrum), - spectrum.frep, spectrum.wl, ) ) @@ -162,7 +250,7 @@ class Pulse(Sequence): def _to_time_amp(self, spectrum): return np.fft.ifft(spectrum) - def all_spectra(self, ind=None): + def all_spectra(self, ind=None) -> Spectrum: """ loads the data already simulated. defauft shape is (z_targets, n, nt) @@ -194,7 +282,6 @@ class Pulse(Sequence): spectra = [] for i in ind: spectra.append(self._load1(i)) - spectra = np.array(spectra) self.logger.debug(f"all spectra from {self.path} successfully loaded") if len(ind) == 1: @@ -212,7 +299,7 @@ class Pulse(Sequence): return self.cache[i] spec = np.load(self.path / SPECN_FN.format(i)) spec = np.atleast_2d(spec) - spec = Spectrum(spec, self.wl, self.params.repetition_rate) + spec = Spectrum(spec, self.params) self.cache[i] = spec return spec diff --git a/src/scgenerator/utils/parameter.py b/src/scgenerator/utils/parameter.py index 206110b..1b2ac0d 100644 --- a/src/scgenerator/utils/parameter.py +++ b/src/scgenerator/utils/parameter.py @@ -384,6 +384,7 @@ class BareParams: field_0: np.ndarray = Parameter(type_checker(np.ndarray)) spec_0: np.ndarray = Parameter(type_checker(np.ndarray)) w: np.ndarray = Parameter(type_checker(np.ndarray)) + l: np.ndarray = Parameter(type_checker(np.ndarray)) w_c: np.ndarray = Parameter(type_checker(np.ndarray)) w0: float = Parameter(positive(float)) w_power_fact: np.ndarray = Parameter(validator_list(type_checker(np.ndarray)))