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

simulation runs but something's wrong

Browse Source
This commit is contained in:
Benoît Sierro
2021-10-20 11:09:24 +02:00
parent de12b0d5c1
commit d898fc9c8b
6 changed files with 213 additions and 199 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
src/scgenerator/const.py
View File

@@ -1,10 +1,9 @@
__version__ = "0.2.3rules"
from typing import Any
def pbar_format(worker_id: int):
def pbar_format(worker_id: int) -> dict[str, Any]:
if worker_id == 0:
return dict(
position=0,
@@ -80,8 +79,6 @@ MANDATORY_PARAMETERS = [
"w_c",
"w",
"w0",
"w_power_fact",
"alpha",
"spec_0",
"field_0",
"mean_power",
@@ -91,8 +88,6 @@ MANDATORY_PARAMETERS = [
"beta2_coefficients",
"gamma_arr",
"behaviors",
"raman_type",
"hr_w",
"adapt_step_size",
"tolerated_error",
"dynamic_dispersion",
@@ -100,5 +95,5 @@ MANDATORY_PARAMETERS = [
"output_path",
"repeat",
"linear_operator",
"nonlinear_op",
"nonlinear_operator",
]

11
src/scgenerator/evaluator.py
View File

@@ -7,7 +7,7 @@ import numpy as np
from . import math, operators, utils
from .const import MANDATORY_PARAMETERS
from .errors import *
from .errors import EvaluatorError, NoDefaultError
from .physics import fiber, materials, pulse, units
from .utils import _mock_function, func_rewrite, get_arg_names, get_logger
@@ -360,12 +360,13 @@ default_rules: list[Rule] = [
fiber.V_eff_step_index,
["l", "core_radius", "numerical_aperture", "interpolation_range"],
),
Rule("gamma", lambda gamma_arr: gamma_arr[0]),
Rule("gamma", lambda gamma_arr: gamma_arr[0], proprities=-1),
Rule("gamma_arr", fiber.gamma_parameter, ["n2", "w0", "A_eff_arr"]),
Rule("n2", materials.gas_n2),
Rule("n2", lambda: 2.2e-20, priorities=-1),
# Operators
Rule("gamma_op", operators.ConstantGamma),
Rule("gamma_op", operators.ConstantGamma, priorities=1),
Rule("gamma_op", operators.ConstantScalarGamma),
Rule("gamma_op", operators.NoGamma, priorities=-1),
Rule("ss_op", operators.SelfSteepening),
Rule("ss_op", operators.NoSelfSteepening, priorities=-1),
@@ -378,9 +379,7 @@ default_rules: list[Rule] = [
Rule("loss_op", operators.CapillaryLoss, priorities=2),
Rule("loss_op", operators.ConstantLoss, priorities=1),
Rule("loss_op", operators.NoLoss, priorities=-1),
Rule("disp_op", operators.ConstantPolyDispersion),
Rule("dispersion_op", operators.ConstantPolyDispersion),
Rule("linear_operator", operators.LinearOperator),
Rule("conserved_quantity", operators.ConservedQuantity),
# gas
Rule("n_gas_2", materials.n_gas_2),
]

111
src/scgenerator/operators.py
View File

@@ -5,16 +5,13 @@ Nothing except the solver should depend on this file
from __future__ import annotations
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from functools import wraps
from os import stat
from typing import Callable
import dataclasses
import numpy as np
from scipy.interpolate import interp1d
from . import math
from .logger import get_logger
from .physics import fiber, pulse
@@ -23,7 +20,9 @@ class SpectrumDescriptor:
value: np.ndarray
def __set__(self, instance, value):
instance.spec2 = math.abs2(value)
instance.field = np.fft.ifft(value)
instance.field2 = math.abs2(instance.field)
self.value = value
def __get__(self, instance, owner):
@@ -36,27 +35,39 @@ class SpectrumDescriptor:
self.name = name
@dataclass
@dataclasses.dataclass
class CurrentState:
length: float
z: float
h: float
C_to_A_factor: np.ndarray
spectrum: np.ndarray = SpectrumDescriptor()
field: np.ndarray = field(init=False)
spec2: np.ndarray = dataclasses.field(init=False)
field: np.ndarray = dataclasses.field(init=False)
field2: np.ndarray = dataclasses.field(init=False)
@property
def z_ratio(self) -> float:
return self.z / self.length
def replace(self, new_spectrum) -> CurrentState:
return CurrentState(self.length, self.z, self.h, self.C_to_A_factor, new_spectrum)
class Operator(ABC):
def __repr__(self) -> str:
return (
self.__class__.__name__
+ "("
+ ", ".join(k + "=" + repr(v) for k, v in self.__dict__.items())
+ ")"
)
value_pair_list = list(self.__dict__.items())
if len(value_pair_list) > 1:
value_pair_str_list = [k + "=" + self.__value_repr(k, v) for k, v in value_pair_list]
else:
value_pair_str_list = [self.__value_repr(value_pair_list[0][0], value_pair_list[0][1])]
return self.__class__.__name__ + "(" + ", ".join(value_pair_str_list) + ")"
def __value_repr(self, k: str, v) -> str:
if k.endswith("_const"):
return repr(v[0])
return repr(v)
@abstractmethod
def __call__(self, state: CurrentState) -> np.ndarray:
@@ -65,7 +76,7 @@ class Operator(ABC):
class NoOp:
def __init__(self, w: np.ndarray):
self.zero_arr = np.zeros_like(w)
self.arr_const = np.zeros_like(w)
##################################################
@@ -125,10 +136,10 @@ class ConstantPolyDispersion(AbstractDispersion):
##################################################
class LinearOperator:
def __init__(self, disp_op: AbstractDispersion, loss_op: AbstractLoss):
self.disp = disp_op
self.loss = loss_op
class LinearOperator(Operator):
def __init__(self, dispersion_op: AbstractDispersion, loss_op: AbstractLoss):
self.dispersion_op = dispersion_op
self.loss_op = loss_op
def __call__(self, state: CurrentState) -> np.ndarray:
"""returns the linear operator to be multiplied by the spectrum in the frequency domain
@@ -143,7 +154,7 @@ class LinearOperator:
np.ndarray
linear component
"""
return self.disp(state) - self.loss(state) / 2
return self.dispersion_op(state) - self.loss_op(state) / 2
##################################################
@@ -174,7 +185,7 @@ class AbstractRaman(Operator):
class NoRaman(NoOp, AbstractRaman):
def __call__(self, state: CurrentState) -> np.ndarray:
return self.zero_arr
return self.arr_const
class Raman(AbstractRaman):
@@ -183,7 +194,7 @@ class Raman(AbstractRaman):
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)))
return self.f_r * np.fft.ifft(self.hr_w * np.fft.fft(state.field2))
# SPM
@@ -210,7 +221,7 @@ class AbstractSPM(Operator):
class NoSPM(NoOp, AbstractSPM):
def __call__(self, state: CurrentState) -> np.ndarray:
return self.zero_arr
return self.arr_const
class SPM(AbstractSPM):
@@ -218,10 +229,10 @@ class SPM(AbstractSPM):
self.fraction = 1 - raman_op.f_r
def __call__(self, state: CurrentState) -> np.ndarray:
return self.fraction * math.abs2(state.field)
return self.fraction * state.field2
# Selt Steepening
# Self-Steepening
class AbstractSelfSteepening(Operator):
@@ -243,7 +254,7 @@ class AbstractSelfSteepening(Operator):
class NoSelfSteepening(NoOp, AbstractSelfSteepening):
def __call__(self, state: CurrentState) -> np.ndarray:
return self.zero_arr
return self.arr_const
class SelfSteepening(AbstractSelfSteepening):
@@ -274,15 +285,20 @@ class AbstractGamma(Operator):
"""
class NoGamma(AbstractSPM):
def __init__(self, w: np.ndarray) -> None:
self.ones_arr = np.ones_like(w)
class ConstantScalarGamma(AbstractGamma):
def __init__(self, gamma: np.ndarray, w: np.ndarray):
self.arr_const = gamma * np.ones_like(w)
def __call__(self, state: CurrentState) -> np.ndarray:
return self.ones_arr
return self.arr_const
class ConstantGamma(AbstractSelfSteepening):
class NoGamma(ConstantScalarGamma):
def __init__(self, w: np.ndarray) -> None:
super().__init__(0, w)
class ConstantGamma(AbstractGamma):
def __init__(self, gamma_arr: np.ndarray):
self.arr = gamma_arr
@@ -355,13 +371,13 @@ class AbstractLoss(Operator):
class ConstantLoss(AbstractLoss):
alpha_arr: np.ndarray
arr_const: np.ndarray
def __init__(self, alpha: float, w: np.ndarray):
self.alpha_arr = alpha * np.ones_like(w)
self.arr_const = alpha * np.ones_like(w)
def __call__(self, state: CurrentState = None) -> np.ndarray:
return self.alpha_arr
return self.arr_const
class NoLoss(ConstantLoss):
@@ -379,8 +395,11 @@ class CapillaryLoss(ConstantLoss):
):
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
self.arr = np.zeros_like(l)
self.arr[mask] = alpha
def __call__(self, state: CurrentState) -> np.ndarray:
return self.arr
class CustomConstantLoss(ConstantLoss):
@@ -388,7 +407,10 @@ class CustomConstantLoss(ConstantLoss):
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)
self.arr = interp1d(wl, loss, fill_value=0, bounds_error=False)(l)
def __call__(self, state: CurrentState) -> np.ndarray:
return self.arr
##################################################
@@ -397,9 +419,10 @@ class CustomConstantLoss(ConstantLoss):
class ConservedQuantity(Operator):
def __new__(
raman_op: AbstractGamma, gamma_op: AbstractGamma, loss_op: AbstractLoss, w: np.ndarray
):
@classmethod
def create(
cls, raman_op: AbstractGamma, gamma_op: AbstractGamma, loss_op: AbstractLoss, w: np.ndarray
) -> ConservedQuantity:
logger = get_logger(__name__)
raman = not isinstance(raman_op, NoRaman)
loss = not isinstance(raman_op, NoLoss)
@@ -435,7 +458,7 @@ class PhotonNumberLoss(ConservedQuantity):
def __call__(self, state: CurrentState) -> float:
return pulse.photon_number_with_loss(
state.spectrum, self.w, self.dw, self.gamma_op(state), self.loss_op(state), state.h
state.spec2, self.w, self.dw, self.gamma_op(state), self.loss_op(state), state.h
)
@@ -446,7 +469,7 @@ class PhotonNumberNoLoss(ConservedQuantity):
self.gamma_op = gamma_op
def __call__(self, state: CurrentState) -> float:
return pulse.photon_number(state.spectrum, self.w, self.dw, self.gamma_op(state))
return pulse.photon_number(state.spec2, self.w, self.dw, self.gamma_op(state))
class EnergyLoss(ConservedQuantity):
@@ -455,7 +478,9 @@ class EnergyLoss(ConservedQuantity):
self.loss_op = loss_op
def __call__(self, state: CurrentState) -> float:
return pulse.pulse_energy_with_loss(state.spectrum, self.dw, self.loss_op(state), state.h)
return pulse.pulse_energy_with_loss(
math.abs2(state.C_to_A_factor * state.spectrum), self.dw, self.loss_op(state), state.h
)
class EnergyNoLoss(ConservedQuantity):
@@ -463,4 +488,4 @@ class EnergyNoLoss(ConservedQuantity):
self.dw = w[1] - w[0]
def __call__(self, state: CurrentState) -> float:
return pulse.pulse_energy(state.spectrum, self.dw)
return pulse.pulse_energy(math.abs2(state.C_to_A_factor * state.spectrum), self.dw)

16
src/scgenerator/parameter.py
View File

@@ -30,7 +30,9 @@ def type_checker(*types):
def _type_checker_wrapper(validator, n=None):
if isinstance(validator, str) and n is not None:
_name = validator
validator = lambda *args: None
def validator(*args):
pass
def _type_checker_wrapped(name, n):
if not isinstance(n, types):
@@ -73,7 +75,7 @@ def in_range_incl(_min, _max):
def boolean(name, n):
if not n is True and not n is False:
if n is not True and n is not False:
raise ValueError(f"{name!r} must be True or False")
@@ -118,7 +120,7 @@ def literal(*l):
@type_checker(str)
def _string(name, s):
if not s in l:
if s not in l:
raise ValueError(f"{name!r} must be a str in {l}")
return _string
@@ -407,7 +409,7 @@ class Parameters:
dico : dict
dictionary
"""
forbiden_keys = [
forbiden_keys = {
"w_c",
"w_power_fact",
"field_0",
@@ -420,7 +422,9 @@ class Parameters:
"alpha",
"gamma_arr",
"A_eff_arr",
]
"nonlinear_op",
"linear_op",
}
types = (np.ndarray, float, int, str, list, tuple, dict)
out = {}
for key, value in dico.items():
@@ -695,7 +699,7 @@ class Configuration:
cfg | dict(variable=self.variationer.all_dicts[i])
for i, cfg in enumerate(self.fiber_configs)
]
utils.save_toml(self.final_path / f"initial_config.toml", dict(name=self.name, Fiber=cfgs))
utils.save_toml(self.final_path / "initial_config.toml", dict(name=self.name, Fiber=cfgs))
@property
def first(self) -> Parameters:

81
src/scgenerator/physics/pulse.py
View File

@@ -11,7 +11,7 @@ n is the number of spectra at the same z position and nt is the size of the time
import itertools
import os
from dataclasses import astuple, dataclass, fields
from dataclasses import astuple, dataclass
from pathlib import Path
from typing import Literal, Tuple, TypeVar
@@ -19,13 +19,12 @@ import matplotlib.pyplot as plt
import numpy as np
from numpy import pi
from numpy.fft import fft, fftshift, ifft
from scipy.interpolate import UnivariateSpline
from scipy.interpolate import UnivariateSpline, interp1d
from scipy.optimize import minimize_scalar
from scipy.optimize.optimize import OptimizeResult
from ..defaults import default_plotting
from ..logger import get_logger
from ..math import *
from .. import math
from . import units
c = 299792458.0
@@ -156,9 +155,9 @@ def modify_field_ratio(
"""
ratio = 1
if target_energy is not None:
ratio *= np.sqrt(target_energy / np.trapz(abs2(field), t))
ratio *= np.sqrt(target_energy / np.trapz(math.abs2(field), t))
elif target_power is not None:
ratio *= np.sqrt(target_power / abs2(field).max())
ratio *= np.sqrt(target_power / math.abs2(field).max())
if intensity_noise is not None:
d_int, _ = technical_noise(intensity_noise, noise_correlation)
@@ -327,10 +326,10 @@ def load_and_adjust_field_file(
) -> np.ndarray:
field_0 = load_field_file(field_file, t)
if energy is not None:
curr_energy = np.trapz(abs2(field_0), t)
curr_energy = np.trapz(math.abs2(field_0), t)
field_0 = field_0 * np.sqrt(energy / curr_energy)
elif peak_power is not None:
ratio = np.sqrt(peak_power / abs2(field_0).max())
ratio = np.sqrt(peak_power / math.abs2(field_0).max())
field_0 = field_0 * ratio
else:
raise ValueError(f"Not enough parameters specified to load {field_file} correctly")
@@ -356,7 +355,7 @@ def correct_wavelength(init_wavelength: float, w_c: np.ndarray, field_0: np.ndar
finds a new wavelength parameter such that the maximum of the spectrum corresponding
to field_0 is located at init_wavelength
"""
delta_w = w_c[np.argmax(abs2(np.fft.fft(field_0)))]
delta_w = w_c[np.argmax(math.abs2(np.fft.fft(field_0)))]
return units.m.inv(units.m(init_wavelength) - delta_w)
@@ -382,21 +381,20 @@ def gaussian_pulse(t, t0, P0, offset=0):
return np.sqrt(P0) * np.exp(-(((t - offset) / t0) ** 2))
def photon_number(spectrum, w, dw, gamma) -> float:
return np.sum(1 / gamma * abs2(spectrum) / w * dw)
def photon_number(spec2, w, dw, gamma) -> float:
return np.sum(1 / gamma * spec2 / w * dw)
def photon_number_with_loss(spectrum, w, dw, gamma, alpha, h) -> float:
spec2 = abs2(spectrum)
def photon_number_with_loss(spec2, w, dw, gamma, alpha, h) -> float:
return np.sum(1 / gamma * spec2 / w * dw) - h * np.sum(alpha / gamma * spec2 / w * dw)
def pulse_energy(spectrum, dw) -> float:
return np.sum(abs2(spectrum) * dw)
def pulse_energy(spec2, dw) -> float:
return np.sum(spec2 * dw)
def pulse_energy_with_loss(spectrum, dw, alpha, h) -> float:
spec2 = abs2(spectrum)
def pulse_energy_with_loss(spec2, dw, alpha, h) -> float:
spec2 = spec2
return np.sum(spec2 * dw) - h * np.sum(alpha * spec2 * dw)
@@ -539,7 +537,7 @@ def ideal_compressed_pulse(spectra):
compressed : 1D array
time envelope of the compressed field
"""
return abs2(fftshift(ifft(np.sqrt(np.mean(abs2(spectra), axis=0)))))
return math.abs2(fftshift(ifft(np.sqrt(np.mean(math.abs2(spectra), axis=0)))))
def spectrogram(time, values, t_res=256, t_win=24e-12, gate_width=200e-15, shift=False):
@@ -571,7 +569,7 @@ def spectrogram(time, values, t_res=256, t_win=24e-12, gate_width=200e-15, shift
spec = np.zeros((t_res, len(time)))
for i, delay in enumerate(delays):
masked = values * np.exp(-(((time - delay) / gate_width) ** 2))
spec[i] = abs2(fft(masked))
spec[i] = math.abs2(fft(masked))
if shift:
spec[i] = fftshift(spec[i])
return spec, delays
@@ -592,7 +590,7 @@ def g12(values):
# Create all the possible pairs of values
n = len(values)
field_pairs = itertools.combinations(values, 2)
mean_spec = np.mean(abs2(values), axis=0)
mean_spec = np.mean(math.abs2(values), axis=0)
mask = mean_spec > 1e-15 * mean_spec.max()
corr = np.zeros_like(values[0])
for pair in field_pairs:
@@ -618,7 +616,7 @@ def avg_g12(values):
if len(values.shape) > 2:
pass
avg_values = np.mean(abs2(values), axis=0)
avg_values = np.mean(math.abs2(values), axis=0)
coherence = g12(values)
return np.sum(coherence * avg_values) / np.sum(avg_values)
@@ -765,7 +763,6 @@ def find_lobe_limits(x_axis, values, debug="", already_sorted=True):
spline_4 : scipy.interpolate.UnivariateSpline
order 4 spline that interpolate values around the peak
"""
logger = get_logger(__name__)
if not already_sorted:
x_axis, values = x_axis.copy(), values.copy()
@@ -914,7 +911,7 @@ def _detailed_find_lobe_limits(
else:
iterations += 1
mam = (spline_4(peak_pos), argclosest(x_axis, peak_pos))
mam = (spline_4(peak_pos), math.argclosest(x_axis, peak_pos))
small_spline, spline_4, spline_5, d_spline, d_roots, dd_roots, l_ind, r_ind = setup_splines(
x_axis, values, mam
@@ -941,7 +938,7 @@ def _detailed_find_lobe_limits(
color = default_plotting["color_cycle"]
if debug != "":
newx = np.linspace(*span(x_axis[l_ind : r_ind + 1]), 1000)
newx = np.linspace(*math.span(x_axis[l_ind : r_ind + 1]), 1000)
ax.plot(x_axis[l_ind - 5 : r_ind + 6], values[l_ind - 5 : r_ind + 6], c=color[0])
ax.plot(newx, spline_5(newx), c=color[1])
ax.scatter(fwhm_pos, spline_4(fwhm_pos), marker="+", label="all fwhm", c=color[2])
@@ -1000,25 +997,27 @@ def measure_properties(spectra, t, compress=True, return_limits=False, debug="")
list of tuples of the form ([left_lobe_lim, right_lobe_lim, lobe_pos], [left_hm, right_hm])
"""
if compress:
fields = abs2(compress_pulse(spectra))
fields = math.abs2(compress_pulse(spectra))
else:
fields = abs2(ifft(spectra))
fields = math.abs2(ifft(spectra))
field = np.mean(fields, axis=0)
ideal_field = abs2(fftshift(ifft(np.sqrt(np.mean(abs2(spectra), axis=0)))))
ideal_field = math.abs2(fftshift(ifft(np.sqrt(np.mean(math.abs2(spectra), axis=0)))))
# Isolate whole central lobe of bof mean and ideal field
lobe_lim, fwhm_lim, _, big_spline = find_lobe_limits(t, field, debug)
lobe_lim_i, _, _, big_spline_i = find_lobe_limits(t, ideal_field, debug)
# Compute quality factor
energy_fraction = (big_spline.integral(*span(lobe_lim[:2]))) / np.trapz(field, x=t)
energy_fraction_i = (big_spline_i.integral(*span(lobe_lim_i[:2]))) / np.trapz(ideal_field, x=t)
energy_fraction = (big_spline.integral(*math.span(lobe_lim[:2]))) / np.trapz(field, x=t)
energy_fraction_i = (big_spline_i.integral(*math.span(lobe_lim_i[:2]))) / np.trapz(
ideal_field, x=t
)
qf = energy_fraction / energy_fraction_i
# Compute mean coherence
mean_g12 = avg_g12(spectra)
fwhm_abs = length(fwhm_lim)
fwhm_abs = math.length(fwhm_lim)
# To compute amplitude and fwhm fluctuations, we need to measure every single peak
P0 = []
@@ -1030,7 +1029,7 @@ def measure_properties(spectra, t, compress=True, return_limits=False, debug="")
lobe_lim, fwhm_lim, _, big_spline = find_lobe_limits(t, f, debug)
all_lims.append((lobe_lim, fwhm_lim))
P0.append(big_spline(lobe_lim[2]))
fwhm.append(length(fwhm_lim))
fwhm.append(math.length(fwhm_lim))
t_offset.append(lobe_lim[2])
energies.append(np.trapz(fields, t))
fwhm_var = np.std(fwhm) / np.mean(fwhm)
@@ -1063,7 +1062,7 @@ def rin_curve(spectra: np.ndarray) -> np.ndarray:
RIN curve
"""
if np.iscomplexobj(spectra):
A2 = abs2(spectra)
A2 = math.abs2(spectra)
else:
A2 = spectra
return np.std(A2, axis=-2) / np.mean(A2, axis=-2)
@@ -1072,11 +1071,11 @@ def rin_curve(spectra: np.ndarray) -> np.ndarray:
def measure_field(t: np.ndarray, field: np.ndarray) -> Tuple[float, float, float]:
"""returns fwhm, peak_power, energy"""
if np.iscomplexobj(field):
intensity = abs2(field)
intensity = math.abs2(field)
else:
intensity = field
_, fwhm_lim, _, _ = find_lobe_limits(t, intensity)
fwhm = length(fwhm_lim)
fwhm = math.length(fwhm_lim)
peak_power = intensity.max()
energy = np.trapz(intensity, t)
return fwhm, peak_power, energy
@@ -1101,10 +1100,12 @@ def remove_2nd_order_dispersion(
np.ndarray, shape (n, )
spectrum with 2nd order dispersion removed
"""
propagate = lambda z: spectrum * np.exp(-0.5j * beta2 * w_c ** 2 * z)
def propagate(z):
return spectrum * np.exp(-0.5j * beta2 * w_c ** 2 * z)
def score(z):
return -np.max(abs2(np.fft.ifft(propagate(z))))
return -np.max(math.abs2(np.fft.ifft(propagate(z))))
opti = minimize_scalar(score, bracket=(max_z, 0))
return propagate(opti.x), opti
@@ -1127,8 +1128,12 @@ def remove_2nd_order_dispersion2(
np.ndarray, shape (n, )
spectrum with 2nd order dispersion removed
"""
propagate = lambda gdd: spectrum * np.exp(-0.5j * w_c ** 2 * 1e-30 * gdd)
integrate = lambda gdd: abs2(np.fft.ifft(propagate(gdd)))
def propagate(gdd):
return spectrum * np.exp(-0.5j * w_c ** 2 * 1e-30 * gdd)
def integrate(gdd):
return math.abs2(np.fft.ifft(propagate(gdd)))
def score(gdd):
return -np.sum(integrate(gdd) ** 6)

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

@@ -1,11 +1,10 @@
import multiprocessing
import multiprocessing.connection
import os
import random
from datetime import datetime
from pathlib import Path
from typing import Any, Generator, Type, Union
from typing import Any, Generator, Type, Union, Optional
from logging import Logger
import numpy as np
from .. import utils
@@ -21,24 +20,39 @@ except ModuleNotFoundError:
class RK4IP:
params: Parameters
save_data: bool
data_dir: Optional[Path]
logger: Logger
dw: float
z_targets: list[float]
z_store: list[float]
z: float
store_num: int
error_ok: float
size_fac: float
cons_qty: list[float]
state: CurrentState
conserved_quantity_func: ConservedQuantity
stored_spectra: list[np.ndarray]
def __init__(
self,
params: Parameters,
save_data=False,
task_id=0,
):
"""A 1D solver using 4th order Runge-Kutta in the interaction picture
Parameters
----------
Parameters
params : Parameters
parameters of the simulation
save_data : bool, optional
save calculated spectra to disk, by default False
task_id : int, optional
unique identifier of the session, by default 0
"""
self.set(params, save_data, task_id)
self.set(params, save_data)
def __iter__(self) -> Generator[tuple[int, int, np.ndarray], None, None]:
yield from self.irun()
@@ -50,10 +64,8 @@ class RK4IP:
self,
params: Parameters,
save_data=False,
task_id=0,
):
self.params = params
self.id = task_id
self.save_data = save_data
if self.save_data:
@@ -62,31 +74,28 @@ class RK4IP:
else:
self.data_dir = None
self.logger = get_logger(self.params.output_path)
self.resuming = False
self.logger = get_logger(self.params.output_path.name)
self.dw = self.params.w[1] - self.params.w[0]
self.z_targets = self.params.z_targets
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.set_cons_qty()
self._setup_sim_parameters()
def _setup_functions(self):
def set_cons_qty(self):
# Set up which quantity is conserved for adaptive step size
if self.params.adapt_step_size:
self.conserved_quantity_func = ConservedQuantity(
self.conserved_quantity_func = ConservedQuantity.create(
self.params.nonlinear_operator.raman_op,
self.params.nonlinear_operator.gamma_op,
self.params.linear_operator.loss_op,
self.params.w,
)
else:
self.logger.debug(f"Using constant step size of {1e6*self.error_ok:.3f}")
self.logger.debug(f"Using constant step size of {1e6*self.error_ok:.3f}μm")
self.conserved_quantity_func = NoConservedQuantity()
def _setup_sim_parameters(self):
@@ -96,17 +105,20 @@ class RK4IP:
self.z_targets.sort()
self.store_num = len(self.z_targets)
# Setup initial values for every physical quantity that we want to track
C_to_A_factor = (self.params.A_eff_arr / self.params.A_eff_arr[0]) ** (1 / 4)
z = self.z_targets.pop(0)
# Initial step size
if self.params.adapt_step_size:
initial_h = (self.z_targets[0] - self.z) / 2
initial_h = (self.z_targets[0] - z) / 2
else:
initial_h = self.error_ok
# Setup initial values for every physical quantity that we want to track
self.state = CurrentState(
length=self.params.length,
z=self.z_targets.pop(0),
z=z,
h=initial_h,
spectrum=self.params.spec_0.copy() / self.C_to_A_factor,
C_to_A_factor=C_to_A_factor,
spectrum=self.params.spec_0.copy() / C_to_A_factor,
)
self.stored_spectra = self.params.recovery_last_stored * [None] + [
self.state.spectrum.copy()
@@ -117,7 +129,6 @@ class RK4IP:
]
self.size_fac = 2 ** (1 / 5)
def _save_current_spectrum(self, num: int):
"""saves the spectrum and the corresponding cons_qty array
@@ -126,11 +137,11 @@ class RK4IP:
num : int
index of the z postition
"""
self._save_data(self.C_to_A_factor * self.state.spectrum, f"spectrum_{num}")
self._save_data(self.cons_qty, f"cons_qty")
self._save_data(self.get_current_spectrum(), f"spectrum_{num}")
self._save_data(self.cons_qty, "cons_qty")
self.step_saved()
def get_current_spectrum(self) -> tuple[int, np.ndarray]:
def get_current_spectrum(self) -> np.ndarray:
"""returns the current spectrum
Returns
@@ -138,7 +149,7 @@ class RK4IP:
np.ndarray
spectrum
"""
return self.C_to_A_factor * self.state.spectrum
return self.state.C_to_A_factor * self.state.spectrum
def _save_data(self, data: np.ndarray, name: str):
"""calls the appropriate method to save data
@@ -190,35 +201,32 @@ class RK4IP:
# Start of the integration
step = 1
h_taken = self.initial_h
h_next_step = self.initial_h
store = False # store a spectrum
yield step, len(self.stored_spectra) - 1, self.get_current_spectrum()
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()
)
while self.state.z < self.params.length:
h_taken = self.take_step(step)
self.z += h_taken
step += 1
self.cons_qty.append(0)
# Whether the current spectrum has to be stored depends on previous step
if store:
self.logger.debug("{} steps, z = {:.4f}, h = {:.5g}".format(step, self.z, h_taken))
self.logger.debug(
"{} steps, z = {:.4f}, h = {:.5g}".format(step, self.state.z, h_taken)
)
self.stored_spectra.append(self.state.spectrum)
yield step, len(self.stored_spectra) - 1, self.get_current_spectrum()
self.z_stored.append(self.z)
self.z_stored.append(self.state.z)
del self.z_targets[0]
# reset the constant step size after a spectrum is stored
if not self.params.adapt_step_size:
h_next_step = self.error_ok
self.state.h = self.error_ok
if len(self.z_targets) == 0:
break
@@ -226,50 +234,39 @@ class RK4IP:
# if the next step goes over a position at which we want to store
# a spectrum, we shorten the step to reach this position exactly
if self.z + h_next_step >= self.z_targets[0]:
if self.state.z + self.state.h >= self.z_targets[0]:
store = True
h_next_step = self.z_targets[0] - self.z
self.state.h = self.z_targets[0] - self.state.z
def take_step(
self, step: int, h_next_step: float, current_spectrum: np.ndarray
) -> tuple[float, float, np.ndarray]:
def take_step(self, step: int) -> tuple[float, float, np.ndarray]:
"""computes a new spectrum, whilst adjusting step size if required, until the error estimation
validates the new spectrum
validates the new spectrum. Saves the result in the internal state attribute
Parameters
----------
step : int
index of the current
h_next_step : float
candidate step size
current_spectrum : np.ndarray
spectrum of the last step taken
Returns
-------
h : float
step sized used
h_next_step : float
candidate next step size
new_spectrum : np.ndarray
new spectrum
"""
keep = False
while not keep:
h = h_next_step
z_ratio = self.z / self.params.length
h = self.state.h
expD = np.exp(h / 2 * self.disp(z_ratio))
expD = np.exp(h / 2 * self.params.linear_operator(self.state))
A_I = expD * current_spectrum
k1 = expD * (h * self.N_func(current_spectrum, z_ratio))
k2 = h * self.N_func(A_I + k1 / 2, z_ratio)
k3 = h * self.N_func(A_I + k2 / 2, z_ratio)
k4 = h * self.N_func(expD * (A_I + k3), z_ratio)
new_spectrum = expD * (A_I + k1 / 6 + k2 / 3 + k3 / 3) + k4 / 6
A_I = expD * self.state.spectrum
k1 = expD * (h * self.params.nonlinear_operator(self.state))
k2 = h * self.params.nonlinear_operator(self.state.replace(A_I + k1 / 2))
k3 = h * self.params.nonlinear_operator(self.state.replace(A_I + k2 / 2))
k4 = h * self.params.nonlinear_operator(self.state.replace(expD * (A_I + k3)))
new_state = self.state.replace(expD * (A_I + k1 / 6 + k2 / 3 + k3 / 3) + k4 / 6)
if self.params.adapt_step_size:
self.cons_qty[step] = self.conserved_quantity_func(new_spectrum, h)
self.cons_qty[step] = self.conserved_quantity_func(new_state)
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]
@@ -289,7 +286,10 @@ class RK4IP:
h_next_step = h
else:
keep = True
return h, h_next_step, new_spectrum
self.state = new_state
self.state.h = h_next_step
self.state.z += h
return h
def step_saved(self):
pass
@@ -301,17 +301,15 @@ class SequentialRK4IP(RK4IP):
params: Parameters,
pbars: PBars,
save_data=False,
task_id=0,
):
self.pbars = pbars
super().__init__(
params,
save_data=save_data,
task_id=task_id,
)
def step_saved(self):
self.pbars.update(1, self.z / self.params.length - self.pbars[1].n)
self.pbars.update(1, self.state.z / self.params.length - self.pbars[1].n)
class MutliProcRK4IP(RK4IP):
@@ -321,18 +319,16 @@ class MutliProcRK4IP(RK4IP):
p_queue: multiprocessing.Queue,
worker_id: int,
save_data=False,
task_id=0,
):
self.worker_id = worker_id
self.p_queue = p_queue
super().__init__(
params,
save_data=save_data,
task_id=task_id,
)
def step_saved(self):
self.p_queue.put((self.worker_id, self.z / self.params.length))
self.p_queue.put((self.worker_id, self.state.z / self.params.length))
class RayRK4IP(RK4IP):
@@ -345,23 +341,21 @@ class RayRK4IP(RK4IP):
p_actor,
worker_id: int,
save_data=False,
task_id=0,
):
self.worker_id = worker_id
self.p_actor = p_actor
super().set(
params,
save_data=save_data,
task_id=task_id,
)
def set_and_run(self, v):
params, p_actor, worker_id, save_data, task_id = v
self.set(params, p_actor, worker_id, save_data, task_id)
params, p_actor, worker_id, save_data = v
self.set(params, p_actor, worker_id, save_data)
self.run()
def step_saved(self):
self.p_actor.update.remote(self.worker_id, self.z / self.params.length)
self.p_actor.update.remote(self.worker_id, self.state.z / self.params.length)
self.p_actor.update.remote(0)
@@ -391,7 +385,7 @@ class Simulations:
@classmethod
def new(
cls, configuration: Configuration, task_id, method: Union[str, Type["Simulations"]] = None
cls, configuration: Configuration, method: Union[str, Type["Simulations"]] = None
) -> "Simulations":
"""Prefered method to create a new simulations object
@@ -403,27 +397,24 @@ class Simulations:
if method is not None:
if isinstance(method, str):
method = Simulations.simulation_methods_dict[method]
return method(configuration, task_id)
return method(configuration)
elif configuration.num_sim > 1 and configuration.parallel:
return Simulations.get_best_method()(configuration, task_id)
return Simulations.get_best_method()(configuration)
else:
return SequencialSimulations(configuration, task_id)
return SequencialSimulations(configuration)
def __init__(self, configuration: Configuration, task_id=0):
def __init__(self, configuration: Configuration):
"""
Parameters
----------
configuration : scgenerator.Configuration obj
parameter sequence
task_id : int, optional
a unique id that identifies the simulation, by default 0
data_folder : str, optional
path to the folder where data is saved, by default "scgenerator/"
"""
if not self.is_available():
raise RuntimeError(f"{self.__class__} is currently not available")
self.logger = get_logger(__name__)
self.id = int(task_id)
self.configuration = configuration
@@ -470,7 +461,7 @@ class Simulations:
def ensure_finised_and_complete(self):
while not self.finished_and_complete():
self.logger.warning(f"Something wrong happened, running again to finish simulation")
self.logger.warning("Something wrong happened, running again to finish simulation")
self._run_available()
def stop(self):
@@ -482,8 +473,8 @@ class SequencialSimulations(Simulations, priority=0):
def is_available(cls):
return True
def __init__(self, configuration: Configuration, task_id):
super().__init__(configuration, task_id=task_id)
def __init__(self, configuration: Configuration):
super().__init__(configuration)
self.pbars = PBars(
self.configuration.total_num_steps,
"Simulating " + self.configuration.final_path.name,
@@ -493,7 +484,7 @@ class SequencialSimulations(Simulations, priority=0):
def new_sim(self, params: Parameters):
self.logger.info(f"{self.configuration.final_path} : launching simulation")
SequentialRK4IP(params, self.pbars, save_data=True, task_id=self.id).run()
SequentialRK4IP(params, self.pbars, save_data=True).run()
def stop(self):
pass
@@ -507,8 +498,8 @@ class MultiProcSimulations(Simulations, priority=1):
def is_available(cls):
return True
def __init__(self, configuration: Configuration, task_id):
super().__init__(configuration, task_id=task_id)
def __init__(self, configuration: Configuration):
super().__init__(configuration)
if configuration.worker_num is not None:
self.sim_jobs_per_node = configuration.worker_num
else:
@@ -519,7 +510,7 @@ class MultiProcSimulations(Simulations, priority=1):
self.workers = [
multiprocessing.Process(
target=MultiProcSimulations.worker,
args=(self.id, i + 1, self.queue, self.progress_queue),
args=(i + 1, self.queue, self.progress_queue),
)
for i in range(self.sim_jobs_per_node)
]
@@ -556,7 +547,6 @@ class MultiProcSimulations(Simulations, priority=1):
@staticmethod
def worker(
task_id,
worker_id: int,
queue: multiprocessing.JoinableQueue,
p_queue: multiprocessing.Queue,
@@ -572,7 +562,6 @@ class MultiProcSimulations(Simulations, priority=1):
p_queue,
worker_id,
save_data=True,
task_id=task_id,
).run()
queue.task_done()
@@ -590,9 +579,8 @@ class RaySimulations(Simulations, priority=2):
def __init__(
self,
configuration: Configuration,
task_id=0,
):
super().__init__(configuration, task_id)
super().__init__(configuration)
nodes = ray.nodes()
self.logger.info(
@@ -629,7 +617,6 @@ class RaySimulations(Simulations, priority=2):
self.p_actor,
self.rolling_id + 1,
True,
self.id,
),
)
self.num_submitted += 1
@@ -679,9 +666,8 @@ def new_simulation(
method: Union[str, Type[Simulations]] = None,
) -> Simulations:
logger = get_logger(__name__)
task_id = random.randint(1e9, 1e12)
logger.info(f"running {configuration.final_path}")
return Simulations.new(configuration, task_id, method)
return Simulations.new(configuration, method)
def __parallel_RK4IP_worker(