diff --git a/src/scgenerator/math.py b/src/scgenerator/math.py index adf4bbb..d65271c 100644 --- a/src/scgenerator/math.py +++ b/src/scgenerator/math.py @@ -292,9 +292,7 @@ def build_sim_grid( time_window: float = None, t_num: int = None, dt: float = None, -) -> tuple[ - np.ndarray, np.ndarray, float, int, float, np.ndarray, float, np.ndarray, np.ndarray, np.ndarray -]: +) -> tuple[np.ndarray, np.ndarray, float, int, float, np.ndarray, float, np.ndarray, np.ndarray]: """computes a bunch of values that relate to the simulation grid Parameters @@ -332,8 +330,6 @@ def build_sim_grid( pump angular frequency w : np.ndarray, shape (t_num, ) actual angualr frequency grid in rad/s - w_power_fact : np.ndarray, shape (deg, t_num) - set of all the necessaray powers of w_c l : np.ndarray, shape (t_num) wavelengths in m """ @@ -343,9 +339,9 @@ def build_sim_grid( dt = t[1] - t[0] t_num = len(t) z_targets = np.linspace(0, length, z_num) - w_c, w0, w, w_power_fact = update_frequency_domain(t, wavelength, interpolation_degree) + w_c, w0, w = update_frequency_domain(t, wavelength, interpolation_degree) l = 2 * pi * c / w - return z_targets, t, time_window, t_num, dt, w_c, w0, w, w_power_fact, l + return z_targets, t, time_window, t_num, dt, w_c, w0, w, l def update_frequency_domain( @@ -365,10 +361,9 @@ def update_frequency_domain( Returns ------- Tuple[np.ndarray, float, np.ndarray, np.ndarray] - w_c, w0, w, w_power_fact + w_c, w0, w """ w_c = wspace(t) w0 = 2 * pi * c / wavelength w = w_c + w0 - w_power_fact = np.array([power_fact(w_c, k) for k in range(2, deg + 3)]) - return w_c, w0, w, w_power_fact + return w_c, w0, w diff --git a/src/scgenerator/parameter.py b/src/scgenerator/parameter.py index d33e0eb..c6729c6 100644 --- a/src/scgenerator/parameter.py +++ b/src/scgenerator/parameter.py @@ -434,7 +434,6 @@ class Parameters(_AbstractParameters): 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))) t: np.ndarray = Parameter(type_checker(np.ndarray)) L_D: float = Parameter(non_negative(float, int)) L_NL: float = Parameter(non_negative(float, int)) @@ -1045,7 +1044,7 @@ class Configuration: default_rules: list[Rule] = [ # Grid *Rule.deduce( - ["z_targets", "t", "time_window", "t_num", "dt", "w_c", "w0", "w", "w_power_fact", "l"], + ["z_targets", "t", "time_window", "t_num", "dt", "w_c", "w0", "w", "l"], math.build_sim_grid, ["time_window", "t_num", "dt"], 2, @@ -1067,7 +1066,7 @@ default_rules: list[Rule] = [ Rule("pre_field_0", pulse.initial_field, priorities=1), Rule( "field_0", - pulse.add_shot_noise, + pulse.finalize_pulse, [ "pre_field_0", "quantum_noise", @@ -1076,6 +1075,7 @@ default_rules: list[Rule] = [ "time_window", "dt", "additional_noise_factor", + "input_transmission", ], ), Rule("peak_power", pulse.E0_to_P0, ["energy", "t0", "shape"]), diff --git a/src/scgenerator/physics/fiber.py b/src/scgenerator/physics/fiber.py index 0ad35e2..9b33cae 100644 --- a/src/scgenerator/physics/fiber.py +++ b/src/scgenerator/physics/fiber.py @@ -986,10 +986,10 @@ def delayed_raman_t(t: np.ndarray, raman_type: str) -> np.ndarray: return hr_arr -def delayed_raman_w(t: np.ndarray, dt: float, raman_type: str) -> np.ndarray: +def delayed_raman_w(t: np.ndarray, raman_type: str) -> np.ndarray: """returns the delayed raman response function as function of w see delayed_raman_t for detailes""" - return fft(delayed_raman_t(t, raman_type)) * dt + return fft(delayed_raman_t(t, raman_type)) * (t[1] - t[0]) def create_non_linear_op(behaviors, w_c, w0, gamma, raman_type="stolen", f_r=0, hr_w=None): @@ -1058,7 +1058,7 @@ def create_non_linear_op(behaviors, w_c, w0, gamma, raman_type="stolen", f_r=0, return N_func -def fast_dispersion_op(w_c, beta_arr, power_fact_arr, where=slice(None), alpha=None): +def fast_dispersion_op(w_c, beta_arr, power_fact_arr, where=slice(None)): """ dispersive operator @@ -1083,10 +1083,7 @@ def fast_dispersion_op(w_c, beta_arr, power_fact_arr, where=slice(None), alpha=N out = np.zeros_like(dispersion) out[where] = dispersion[where] - if alpha is None: - return -1j * out - else: - return -1j * out - alpha / 2 + return -1j * out def _fast_disp_loop(dispersion: np.ndarray, beta_arr, power_fact_arr): diff --git a/src/scgenerator/physics/properties.py b/src/scgenerator/physics/properties.py new file mode 100644 index 0000000..442b072 --- /dev/null +++ b/src/scgenerator/physics/properties.py @@ -0,0 +1,357 @@ +""" +This file includes Dispersion, NonLinear and Loss classes to be used in the solver +Nothing except the solver should depend on this file +""" +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass, field + +import numpy as np +from scipy.interpolate import interp1d + +from . import fiber +from .. import math + + +class SpectrumDescriptor: + name: str + value: np.ndarray + + def __set__(self, instance, value): + instance.field = np.fft.ifft(value) + self.value = value + + def __get__(self, instance, owner): + return self.value + + def __delete__(self, instance): + raise AttributeError("Cannot delete Spectrum field") + + def __set_name__(self, owner, name): + self.name = name + + +@dataclass +class CurrentState: + length: float + z: float + h: float + spectrum: np.ndarray = SpectrumDescriptor() + field: np.ndarray = field(init=False) + + @property + def z_ratio(self) -> float: + return self.z / self.length + + +class NoOp: + def __init__(self, w: np.ndarray): + self.zero_arr = np.zeros_like(w) + + +################################################## +################### DISPERSION ################### +################################################## + + +class AbstractDispersion(ABC): + @abstractmethod + def __call__(self, state: CurrentState) -> np.ndarray: + """returns the dispersion in the frequency domain + + Parameters + ---------- + state : CurrentState + current state of the simulation + + Returns + ------- + np.ndarray + dispersive component + """ + + +class ConstantPolyDispersion(AbstractDispersion): + """ + dispersion approximated by fitting a polynom on the dispersion and + evaluating on the envelope + """ + + coefs: np.ndarray + w_c: np.ndarray + + def __init__( + self, + wl_for_disp: np.ndarray, + beta2_arr: np.ndarray, + w0: float, + w_c: np.ndarray, + interpolation_range: tuple[float, float] = None, + interpolation_degree: int = 8, + ): + self.coefs = fiber.dispersion_coefficients( + wl_for_disp, beta2_arr, w0, interpolation_range, interpolation_degree + ) + self.w_c = w_c + self.w_power_fact = np.array( + [math.power_fact(w_c, k) for k in range(2, interpolation_degree + 3)] + ) + + def __call__(self, state: CurrentState) -> np.ndarray: + return fiber.fast_dispersion_op(self.w_c, self.coefs, self.w_power_fact) + + +################################################## +##################### LINEAR ##################### +################################################## + + +class LinearOperator: + def __init__(self, disp: AbstractDispersion, loss: AbstractLoss): + self.disp = disp + self.loss = loss + + def __call__(self, state: CurrentState) -> np.ndarray: + """returns the linear operator to be multiplied by the spectrum in the frequency domain + + Parameters + ---------- + state : CurrentState + current state of the simulation + + Returns + ------- + np.ndarray + linear component + """ + return self.disp(state) - self.loss(state) / 2 + + +################################################## +################### NON LINEAR ################### +################################################## + +# Raman + + +class AbstractRaman(ABC): + f_r: float = 0.0 + + @abstractmethod + def __call__(self, state: CurrentState) -> np.ndarray: + """returns the raman component + + Parameters + ---------- + state : CurrentState + current state of the simulation + + Returns + ------- + np.ndarray + raman component + """ + + +class NoRaman(NoOp, AbstractRaman): + def __call__(self, state: CurrentState) -> np.ndarray: + return self.zero_arr + + +class Raman(AbstractRaman): + def __init__(self, raman_type: str, t: np.ndarray): + self.hr_w = fiber.delayed_raman_w(t, raman_type) + self.f_r = 0.245 if raman_type == "agrawal" else 0.18 + + def __call__(self, state: CurrentState) -> np.ndarray: + return self.f_r * np.fft.ifft(self.hr_w * np.fft.fft(math.abs2(state.field))) + + +# SPM + + +class AbstractSPM(ABC): + fraction: float = 1.0 + + @abstractmethod + def __call__(self, state: CurrentState) -> np.ndarray: + """returns the SPM component + + Parameters + ---------- + state : CurrentState + current state of the simulation + + Returns + ------- + np.ndarray + SPM component + """ + + +class NoSPM(NoOp, AbstractSPM): + def __call__(self, state: CurrentState) -> np.ndarray: + return self.zero_arr + + +class SPM(AbstractSPM): + def __init__(self, raman_op: AbstractRaman): + self.fraction = 1 - raman_op.f_r + + def __call__(self, state: CurrentState) -> np.ndarray: + return self.fraction * math.abs2(state.field) + + +# Selt Steepening + + +class AbstractSelfSteepening(ABC): + @abstractmethod + def __call__(self, state: CurrentState) -> np.ndarray: + """returns the self-steepening component + + Parameters + ---------- + state : CurrentState + current state of the simulation + + Returns + ------- + np.ndarray + self-steepening component + """ + + +class NoSelfSteepening(NoOp, AbstractSelfSteepening): + def __call__(self, state: CurrentState) -> np.ndarray: + return self.zero_arr + + +class SelfSteepening(AbstractSelfSteepening): + def __init__(self, w_c: np.ndarray, w0: float): + self.arr = w_c / w0 + + def __call__(self, state: CurrentState) -> np.ndarray: + return self.arr + + +# Gamma operator + + +class AbstractGamma(ABC): + @abstractmethod + def __call__(self, state: CurrentState) -> np.ndarray: + """returns the gamma component + + Parameters + ---------- + state : CurrentState + current state of the simulation + + Returns + ------- + np.ndarray + gamma component + """ + + +class NoGamma(AbstractSPM): + def __init__(self, w: np.ndarray) -> None: + self.ones_arr = np.ones_like(w) + + def __call__(self, state: CurrentState) -> np.ndarray: + return self.ones_arr + + +class ConstantGamma(AbstractSelfSteepening): + def __init__(self, gamma_arr: np.ndarray): + self.arr = gamma_arr + + def __call__(self, state: CurrentState) -> np.ndarray: + return self.arr + + +# Nonlinear combination + + +class AbstractNonLinearOperator(ABC): + @abstractmethod + def __call__(self, state: CurrentState) -> np.ndarray: + """returns the nonlinear operator applied on the spectrum in the frequency domain + + Parameters + ---------- + state : CurrentState + current state of the simulation + + Returns + ------- + np.ndarray + nonlinear component + """ + + +class EnvelopeNonLinearOperator(AbstractNonLinearOperator): + def __init__( + self, + gamma_op: AbstractGamma, + ss_op: AbstractSelfSteepening, + spm_op: AbstractSPM, + raman_op: AbstractRaman, + ): + self.gamma_op = gamma_op + + +################################################## +###################### LOSS ###################### +################################################## + + +class AbstractLoss(ABC): + @abstractmethod + def __call__(self, state: CurrentState) -> np.ndarray: + """returns the loss in the frequency domain + + Parameters + ---------- + state : CurrentState + current state of the simulation + + Returns + ------- + np.ndarray + loss in 1/m + """ + + +class ConstantLoss(AbstractLoss): + alpha_arr: np.ndarray + + def __init__(self, alpha: float, w: np.ndarray): + self.alpha_arr = alpha * np.ones_like(w) + + def __call__(self, state: CurrentState) -> np.ndarray: + return self.alpha_arr + + +class CapillaryLoss(ConstantLoss): + def __init__( + self, + l: np.ndarray, + core_radius: float, + interpolation_range: tuple[float, float], + he_mode: tuple[int, int], + ): + mask = (l < interpolation_range[1]) & (l > 0) + alpha = fiber.capillary_loss(l[mask], he_mode, core_radius) + self.alpha_arr = np.zeros_like(l) + self.alpha_arr[mask] = alpha + + +class CustomConstantLoss(ConstantLoss): + def __init__(self, l: np.ndarray, loss_file: str): + loss_data = np.load(loss_file) + wl = loss_data["wavelength"] + loss = loss_data["loss"] + self.alpha_arr = interp1d(wl, loss, fill_value=0, bounds_error=False)(l) diff --git a/src/scgenerator/physics/pulse.py b/src/scgenerator/physics/pulse.py index 53cd159..cf0815a 100644 --- a/src/scgenerator/physics/pulse.py +++ b/src/scgenerator/physics/pulse.py @@ -447,7 +447,7 @@ def shot_noise(w_c, w0, T, dt, additional_noise_factor=1.0): return out * additional_noise_factor -def add_shot_noise( +def finalize_pulse( field_0: np.ndarray, quantum_noise: bool, w_c: bool, @@ -455,10 +455,11 @@ def add_shot_noise( time_window: float, dt: float, additional_noise_factor: float, + input_transmission: float, ) -> np.ndarray: if quantum_noise: field_0 = field_0 + shot_noise(w_c, w0, time_window, dt, additional_noise_factor) - return field_0 + return np.sqrt(input_transmission) * field_0 def mean_phase(spectra): diff --git a/src/scgenerator/physics/simulate.py b/src/scgenerator/physics/simulate.py index 9da68bf..0f6094f 100644 --- a/src/scgenerator/physics/simulate.py +++ b/src/scgenerator/physics/simulate.py @@ -15,6 +15,7 @@ from ..parameter import Configuration, Parameters from ..pbar import PBars, ProgressBarActor, progress_worker from . import pulse from .fiber import create_non_linear_op, fast_dispersion_op +from .properties import CurrentState try: import ray @@ -22,18 +23,6 @@ except ModuleNotFoundError: ray = None -@dataclass -class CurrentState: - length: float - spectrum: np.ndarray - z: float - h: float - - @property - def z_ratio(self) -> float: - return self.z / self.length - - class RK4IP: def __init__( self, @@ -58,7 +47,7 @@ class RK4IP: yield from self.irun() def __len__(self) -> int: - return self.len + return self.params.z_num def set( self, @@ -79,62 +68,61 @@ class RK4IP: self.logger = get_logger(self.params.output_path) self.resuming = False - self.w_c = params.w_c - self.w = params.w - self.dw = self.w[1] - self.w[0] - self.w0 = params.w0 - self.w_power_fact = params.w_power_fact - self.alpha = params.alpha_arr - self.spec_0 = np.sqrt(params.input_transmission) * params.spec_0 - self.z_targets = params.z_targets - self.len = len(params.z_targets) - self.z_final = params.length + self.dw = self.params.w[1] - self.params.w[0] + self.z_targets = self.params.z_targets self.beta2_coefficients = ( params.beta_func if params.beta_func is not None else params.beta2_coefficients ) self.gamma = params.gamma_func if params.gamma_func is not None else params.gamma_arr - self.C_to_A_factor = (params.A_eff_arr / params.A_eff_arr[0]) ** (1 / 4) - self.behaviors = params.behaviors - self.raman_type = params.raman_type - self.hr_w = params.hr_w - self.adapt_step_size = params.adapt_step_size - self.error_ok = params.tolerated_error if self.adapt_step_size else params.step_size - self.dynamic_dispersion = params.dynamic_dispersion - self.starting_num = params.recovery_last_stored + self.C_to_A_factor = (self.params.A_eff_arr / self.params.A_eff_arr[0]) ** (1 / 4) + self.error_ok = ( + params.tolerated_error if self.params.adapt_step_size else self.params.step_size + ) self._setup_functions() self._setup_sim_parameters() def _setup_functions(self): self.N_func = create_non_linear_op( - self.behaviors, self.w_c, self.w0, self.gamma, self.raman_type, hr_w=self.hr_w + self.params.behaviors, + self.params.w_c, + self.params.w0, + self.gamma, + self.params.raman_type, + hr_w=self.params.hr_w, ) - if self.dynamic_dispersion: + if self.params.dynamic_dispersion: self.disp = lambda r: fast_dispersion_op( - self.w_c, self.beta2_coefficients(r), self.w_power_fact, alpha=self.alpha + self.params.w_c, + self.beta2_coefficients(r), + self.params.w_power_fact, + alpha=self.params.alpha_arr, ) else: self.disp = lambda r: fast_dispersion_op( - self.w_c, self.beta2_coefficients, self.w_power_fact, alpha=self.alpha + self.params.w_c, + self.beta2_coefficients, + self.params.w_power_fact, + alpha=self.params.alpha_arr, ) # Set up which quantity is conserved for adaptive step size - if self.adapt_step_size: - if "raman" in self.behaviors and self.alpha is not None: + if self.params.adapt_step_size: + if "raman" in self.params.behaviors and self.params.alpha_arr is not None: self.logger.debug("Conserved quantity : photon number with loss") self.conserved_quantity_func = lambda spectrum, h: pulse.photon_number_with_loss( - spectrum, self.w, self.dw, self.gamma, self.alpha, h + spectrum, self.params.w, self.dw, self.gamma, self.params.alpha_arr, h ) - elif "raman" in self.behaviors: + elif "raman" in self.params.behaviors: self.logger.debug("Conserved quantity : photon number without loss") self.conserved_quantity_func = lambda spectrum, h: pulse.photon_number( - spectrum, self.w, self.dw, self.gamma + spectrum, self.params.w, self.dw, self.gamma ) - elif self.alpha is not None: + elif self.params.alpha_arr is not None: self.logger.debug("Conserved quantity : energy with loss") self.conserved_quantity_func = lambda spectrum, h: pulse.pulse_energy_with_loss( - self.C_to_A_factor * spectrum, self.dw, self.alpha, h + self.C_to_A_factor * spectrum, self.dw, self.params.alpha_arr, h ) else: self.logger.debug("Conserved quantity : energy without loss") @@ -147,26 +135,29 @@ class RK4IP: def _setup_sim_parameters(self): # making sure to keep only the z that we want - self.z_stored = list(self.z_targets.copy()[0 : self.starting_num + 1]) - self.z_targets = list(self.z_targets.copy()[self.starting_num :]) + self.z_stored = list(self.z_targets.copy()[0 : self.params.recovery_last_stored + 1]) + self.z_targets = list(self.z_targets.copy()[self.params.recovery_last_stored :]) self.z_targets.sort() self.store_num = len(self.z_targets) # Initial setup of simulation parameters - self.d_w = self.w_c[1] - self.w_c[0] # resolution of the frequency grid self.z = self.z_targets.pop(0) # Setup initial values for every physical quantity that we want to track - self.current_spectrum = self.spec_0.copy() / self.C_to_A_factor - self.stored_spectra = self.starting_num * [None] + [self.current_spectrum.copy()] + self.state = CurrentState( + length=self.params.length, spectrum=self.params.spec_0.copy() / self.C_to_A_factor + ) + self.stored_spectra = self.params.recovery_last_stored * [None] + [ + self.state.spectrum.copy() + ] self.cons_qty = [ - self.conserved_quantity_func(self.current_spectrum, 0), + self.conserved_quantity_func(self.state.spectrum, 0), 0, ] self.size_fac = 2 ** (1 / 5) # Initial step size - if self.adapt_step_size: + if self.params.adapt_step_size: self.initial_h = (self.z_targets[0] - self.z) / 2 else: self.initial_h = self.error_ok @@ -179,7 +170,7 @@ class RK4IP: num : int index of the z postition """ - self._save_data(self.C_to_A_factor * self.current_spectrum, f"spectrum_{num}") + self._save_data(self.C_to_A_factor * self.state.spectrum, f"spectrum_{num}") self._save_data(self.cons_qty, f"cons_qty") self.step_saved() @@ -191,7 +182,7 @@ class RK4IP: np.ndarray spectrum """ - return self.C_to_A_factor * self.current_spectrum + return self.C_to_A_factor * self.state.spectrum def _save_data(self, data: np.ndarray, name: str): """calls the appropriate method to save data @@ -249,9 +240,9 @@ class RK4IP: yield step, len(self.stored_spectra) - 1, self.get_current_spectrum() - while self.z < self.z_final: - h_taken, h_next_step, self.current_spectrum = self.take_step( - step, h_next_step, self.current_spectrum.copy() + while self.z < self.params.length: + h_taken, h_next_step, self.state.spectrum = self.take_step( + step, h_next_step, self.state.spectrum.copy() ) self.z += h_taken @@ -262,7 +253,7 @@ class RK4IP: if store: self.logger.debug("{} steps, z = {:.4f}, h = {:.5g}".format(step, self.z, h_taken)) - self.stored_spectra.append(self.current_spectrum) + self.stored_spectra.append(self.state.spectrum) yield step, len(self.stored_spectra) - 1, self.get_current_spectrum() @@ -270,7 +261,7 @@ class RK4IP: del self.z_targets[0] # reset the constant step size after a spectrum is stored - if not self.adapt_step_size: + if not self.params.adapt_step_size: h_next_step = self.error_ok if len(self.z_targets) == 0: @@ -310,7 +301,7 @@ class RK4IP: keep = False while not keep: h = h_next_step - z_ratio = self.z / self.z_final + z_ratio = self.z / self.params.length expD = np.exp(h / 2 * self.disp(z_ratio)) @@ -321,7 +312,7 @@ class RK4IP: k4 = h * self.N_func(expD * (A_I + k3), z_ratio) new_spectrum = expD * (A_I + k1 / 6 + k2 / 3 + k3 / 3) + k4 / 6 - if self.adapt_step_size: + if self.params.adapt_step_size: self.cons_qty[step] = self.conserved_quantity_func(new_spectrum, h) curr_p_change = np.abs(self.cons_qty[step - 1] - self.cons_qty[step]) cons_qty_change_ok = self.error_ok * self.cons_qty[step - 1] @@ -364,7 +355,7 @@ class SequentialRK4IP(RK4IP): ) def step_saved(self): - self.pbars.update(1, self.z / self.z_final - self.pbars[1].n) + self.pbars.update(1, self.z / self.params.length - self.pbars[1].n) class MutliProcRK4IP(RK4IP): @@ -385,7 +376,7 @@ class MutliProcRK4IP(RK4IP): ) def step_saved(self): - self.p_queue.put((self.worker_id, self.z / self.z_final)) + self.p_queue.put((self.worker_id, self.z / self.params.length)) class RayRK4IP(RK4IP): @@ -414,7 +405,7 @@ class RayRK4IP(RK4IP): self.run() def step_saved(self): - self.p_actor.update.remote(self.worker_id, self.z / self.z_final) + self.p_actor.update.remote(self.worker_id, self.z / self.params.length) self.p_actor.update.remote(0) @@ -500,7 +491,7 @@ class Simulations: self.ensure_finised_and_complete() def _run_available(self): - for variable, params in self.configuration: + for _, params in self.configuration: params.compute() utils.save_parameters(params.prepare_for_dump(), params.output_path)