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" __version__ = "0.2.3rules"
from typing import Any from typing import Any
def pbar_format(worker_id: int): def pbar_format(worker_id: int) -> dict[str, Any]:
if worker_id == 0: if worker_id == 0:
return dict( return dict(
position=0, position=0,
@@ -80,8 +79,6 @@ MANDATORY_PARAMETERS = [
"w_c", "w_c",
"w", "w",
"w0", "w0",
"w_power_fact",
"alpha",
"spec_0", "spec_0",
"field_0", "field_0",
"mean_power", "mean_power",
@@ -91,8 +88,6 @@ MANDATORY_PARAMETERS = [
"beta2_coefficients", "beta2_coefficients",
"gamma_arr", "gamma_arr",
"behaviors", "behaviors",
"raman_type",
"hr_w",
"adapt_step_size", "adapt_step_size",
"tolerated_error", "tolerated_error",
"dynamic_dispersion", "dynamic_dispersion",
@@ -100,5 +95,5 @@ MANDATORY_PARAMETERS = [
"output_path", "output_path",
"repeat", "repeat",
"linear_operator", "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 . import math, operators, utils
from .const import MANDATORY_PARAMETERS from .const import MANDATORY_PARAMETERS
from .errors import * from .errors import EvaluatorError, NoDefaultError
from .physics import fiber, materials, pulse, units from .physics import fiber, materials, pulse, units
from .utils import _mock_function, func_rewrite, get_arg_names, get_logger 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, fiber.V_eff_step_index,
["l", "core_radius", "numerical_aperture", "interpolation_range"], ["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("gamma_arr", fiber.gamma_parameter, ["n2", "w0", "A_eff_arr"]),
Rule("n2", materials.gas_n2), Rule("n2", materials.gas_n2),
Rule("n2", lambda: 2.2e-20, priorities=-1), Rule("n2", lambda: 2.2e-20, priorities=-1),
# Operators # 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("gamma_op", operators.NoGamma, priorities=-1),
Rule("ss_op", operators.SelfSteepening), Rule("ss_op", operators.SelfSteepening),
Rule("ss_op", operators.NoSelfSteepening, priorities=-1), 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.CapillaryLoss, priorities=2),
Rule("loss_op", operators.ConstantLoss, priorities=1), Rule("loss_op", operators.ConstantLoss, priorities=1),
Rule("loss_op", operators.NoLoss, 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("linear_operator", operators.LinearOperator),
Rule("conserved_quantity", operators.ConservedQuantity), 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 __future__ import annotations
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from dataclasses import dataclass, field import dataclasses
from functools import wraps
from os import stat
from typing import Callable
import numpy as np import numpy as np
from scipy.interpolate import interp1d from scipy.interpolate import interp1d
from . import math from . import math
from .logger import get_logger from .logger import get_logger
from .physics import fiber, pulse from .physics import fiber, pulse
@@ -23,7 +20,9 @@ class SpectrumDescriptor:
value: np.ndarray value: np.ndarray
def __set__(self, instance, value): def __set__(self, instance, value):
instance.spec2 = math.abs2(value)
instance.field = np.fft.ifft(value) instance.field = np.fft.ifft(value)
instance.field2 = math.abs2(instance.field)
self.value = value self.value = value
def __get__(self, instance, owner): def __get__(self, instance, owner):
@@ -36,27 +35,39 @@ class SpectrumDescriptor:
self.name = name self.name = name
@dataclass @dataclasses.dataclass
class CurrentState: class CurrentState:
length: float length: float
z: float z: float
h: float h: float
C_to_A_factor: np.ndarray
spectrum: np.ndarray = SpectrumDescriptor() 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 @property
def z_ratio(self) -> float: def z_ratio(self) -> float:
return self.z / self.length 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): class Operator(ABC):
def __repr__(self) -> str: def __repr__(self) -> str:
return ( value_pair_list = list(self.__dict__.items())
self.__class__.__name__ if len(value_pair_list) > 1:
+ "(" value_pair_str_list = [k + "=" + self.__value_repr(k, v) for k, v in value_pair_list]
+ ", ".join(k + "=" + repr(v) for k, v in self.__dict__.items()) 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 @abstractmethod
def __call__(self, state: CurrentState) -> np.ndarray: def __call__(self, state: CurrentState) -> np.ndarray:
@@ -65,7 +76,7 @@ class Operator(ABC):
class NoOp: class NoOp:
def __init__(self, w: np.ndarray): 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: class LinearOperator(Operator):
def __init__(self, disp_op: AbstractDispersion, loss_op: AbstractLoss): def __init__(self, dispersion_op: AbstractDispersion, loss_op: AbstractLoss):
self.disp = disp_op self.dispersion_op = dispersion_op
self.loss = loss_op self.loss_op = loss_op
def __call__(self, state: CurrentState) -> np.ndarray: def __call__(self, state: CurrentState) -> np.ndarray:
"""returns the linear operator to be multiplied by the spectrum in the frequency domain """returns the linear operator to be multiplied by the spectrum in the frequency domain
@@ -143,7 +154,7 @@ class LinearOperator:
np.ndarray np.ndarray
linear component 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): class NoRaman(NoOp, AbstractRaman):
def __call__(self, state: CurrentState) -> np.ndarray: def __call__(self, state: CurrentState) -> np.ndarray:
return self.zero_arr return self.arr_const
class Raman(AbstractRaman): class Raman(AbstractRaman):
@@ -183,7 +194,7 @@ class Raman(AbstractRaman):
self.f_r = 0.245 if raman_type == "agrawal" else 0.18 self.f_r = 0.245 if raman_type == "agrawal" else 0.18
def __call__(self, state: CurrentState) -> np.ndarray: 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 # SPM
@@ -210,7 +221,7 @@ class AbstractSPM(Operator):
class NoSPM(NoOp, AbstractSPM): class NoSPM(NoOp, AbstractSPM):
def __call__(self, state: CurrentState) -> np.ndarray: def __call__(self, state: CurrentState) -> np.ndarray:
return self.zero_arr return self.arr_const
class SPM(AbstractSPM): class SPM(AbstractSPM):
@@ -218,10 +229,10 @@ class SPM(AbstractSPM):
self.fraction = 1 - raman_op.f_r self.fraction = 1 - raman_op.f_r
def __call__(self, state: CurrentState) -> np.ndarray: 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): class AbstractSelfSteepening(Operator):
@@ -243,7 +254,7 @@ class AbstractSelfSteepening(Operator):
class NoSelfSteepening(NoOp, AbstractSelfSteepening): class NoSelfSteepening(NoOp, AbstractSelfSteepening):
def __call__(self, state: CurrentState) -> np.ndarray: def __call__(self, state: CurrentState) -> np.ndarray:
return self.zero_arr return self.arr_const
class SelfSteepening(AbstractSelfSteepening): class SelfSteepening(AbstractSelfSteepening):
@@ -274,15 +285,20 @@ class AbstractGamma(Operator):
""" """
class NoGamma(AbstractSPM): class ConstantScalarGamma(AbstractGamma):
def __init__(self, w: np.ndarray) -> None: def __init__(self, gamma: np.ndarray, w: np.ndarray):
self.ones_arr = np.ones_like(w) self.arr_const = gamma * np.ones_like(w)
def __call__(self, state: CurrentState) -> np.ndarray: 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): def __init__(self, gamma_arr: np.ndarray):
self.arr = gamma_arr self.arr = gamma_arr
@@ -355,13 +371,13 @@ class AbstractLoss(Operator):
class ConstantLoss(AbstractLoss): class ConstantLoss(AbstractLoss):
alpha_arr: np.ndarray arr_const: np.ndarray
def __init__(self, alpha: float, w: 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: def __call__(self, state: CurrentState = None) -> np.ndarray:
return self.alpha_arr return self.arr_const
class NoLoss(ConstantLoss): class NoLoss(ConstantLoss):
@@ -379,8 +395,11 @@ class CapillaryLoss(ConstantLoss):
): ):
mask = (l < interpolation_range[1]) & (l > 0) mask = (l < interpolation_range[1]) & (l > 0)
alpha = fiber.capillary_loss(l[mask], he_mode, core_radius) alpha = fiber.capillary_loss(l[mask], he_mode, core_radius)
self.alpha_arr = np.zeros_like(l) self.arr = np.zeros_like(l)
self.alpha_arr[mask] = alpha self.arr[mask] = alpha
def __call__(self, state: CurrentState) -> np.ndarray:
return self.arr
class CustomConstantLoss(ConstantLoss): class CustomConstantLoss(ConstantLoss):
@@ -388,7 +407,10 @@ class CustomConstantLoss(ConstantLoss):
loss_data = np.load(loss_file) loss_data = np.load(loss_file)
wl = loss_data["wavelength"] wl = loss_data["wavelength"]
loss = loss_data["loss"] 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): class ConservedQuantity(Operator):
def __new__( @classmethod
raman_op: AbstractGamma, gamma_op: AbstractGamma, loss_op: AbstractLoss, w: np.ndarray def create(
): cls, raman_op: AbstractGamma, gamma_op: AbstractGamma, loss_op: AbstractLoss, w: np.ndarray
) -> ConservedQuantity:
logger = get_logger(__name__) logger = get_logger(__name__)
raman = not isinstance(raman_op, NoRaman) raman = not isinstance(raman_op, NoRaman)
loss = not isinstance(raman_op, NoLoss) loss = not isinstance(raman_op, NoLoss)
@@ -435,7 +458,7 @@ class PhotonNumberLoss(ConservedQuantity):
def __call__(self, state: CurrentState) -> float: def __call__(self, state: CurrentState) -> float:
return pulse.photon_number_with_loss( 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 self.gamma_op = gamma_op
def __call__(self, state: CurrentState) -> float: 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): class EnergyLoss(ConservedQuantity):
@@ -455,7 +478,9 @@ class EnergyLoss(ConservedQuantity):
self.loss_op = loss_op self.loss_op = loss_op
def __call__(self, state: CurrentState) -> float: 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): class EnergyNoLoss(ConservedQuantity):
@@ -463,4 +488,4 @@ class EnergyNoLoss(ConservedQuantity):
self.dw = w[1] - w[0] self.dw = w[1] - w[0]
def __call__(self, state: CurrentState) -> float: 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): def _type_checker_wrapper(validator, n=None):
if isinstance(validator, str) and n is not None: if isinstance(validator, str) and n is not None:
_name = validator _name = validator
validator = lambda *args: None
def validator(*args):
pass
def _type_checker_wrapped(name, n): def _type_checker_wrapped(name, n):
if not isinstance(n, types): if not isinstance(n, types):
@@ -73,7 +75,7 @@ def in_range_incl(_min, _max):
def boolean(name, n): 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") raise ValueError(f"{name!r} must be True or False")
@@ -118,7 +120,7 @@ def literal(*l):
@type_checker(str) @type_checker(str)
def _string(name, s): 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}") raise ValueError(f"{name!r} must be a str in {l}")
return _string return _string
@@ -407,7 +409,7 @@ class Parameters:
dico : dict dico : dict
dictionary dictionary
""" """
forbiden_keys = [ forbiden_keys = {
"w_c", "w_c",
"w_power_fact", "w_power_fact",
"field_0", "field_0",
@@ -420,7 +422,9 @@ class Parameters:
"alpha", "alpha",
"gamma_arr", "gamma_arr",
"A_eff_arr", "A_eff_arr",
] "nonlinear_op",
"linear_op",
}
types = (np.ndarray, float, int, str, list, tuple, dict) types = (np.ndarray, float, int, str, list, tuple, dict)
out = {} out = {}
for key, value in dico.items(): for key, value in dico.items():
@@ -695,7 +699,7 @@ class Configuration:
cfg | dict(variable=self.variationer.all_dicts[i]) cfg | dict(variable=self.variationer.all_dicts[i])
for i, cfg in enumerate(self.fiber_configs) 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 @property
def first(self) -> Parameters: 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 itertools
import os import os
from dataclasses import astuple, dataclass, fields from dataclasses import astuple, dataclass
from pathlib import Path from pathlib import Path
from typing import Literal, Tuple, TypeVar from typing import Literal, Tuple, TypeVar
@@ -19,13 +19,12 @@ import matplotlib.pyplot as plt
import numpy as np import numpy as np
from numpy import pi from numpy import pi
from numpy.fft import fft, fftshift, ifft 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 import minimize_scalar
from scipy.optimize.optimize import OptimizeResult from scipy.optimize.optimize import OptimizeResult
from ..defaults import default_plotting from ..defaults import default_plotting
from ..logger import get_logger from .. import math
from ..math import *
from . import units from . import units
c = 299792458.0 c = 299792458.0
@@ -156,9 +155,9 @@ def modify_field_ratio(
""" """
ratio = 1 ratio = 1
if target_energy is not None: 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: 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: if intensity_noise is not None:
d_int, _ = technical_noise(intensity_noise, noise_correlation) d_int, _ = technical_noise(intensity_noise, noise_correlation)
@@ -327,10 +326,10 @@ def load_and_adjust_field_file(
) -> np.ndarray: ) -> np.ndarray:
field_0 = load_field_file(field_file, t) field_0 = load_field_file(field_file, t)
if energy is not None: 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) field_0 = field_0 * np.sqrt(energy / curr_energy)
elif peak_power is not None: 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 field_0 = field_0 * ratio
else: else:
raise ValueError(f"Not enough parameters specified to load {field_file} correctly") 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 finds a new wavelength parameter such that the maximum of the spectrum corresponding
to field_0 is located at init_wavelength 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) 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)) return np.sqrt(P0) * np.exp(-(((t - offset) / t0) ** 2))
def photon_number(spectrum, w, dw, gamma) -> float: def photon_number(spec2, w, dw, gamma) -> float:
return np.sum(1 / gamma * abs2(spectrum) / w * dw) return np.sum(1 / gamma * spec2 / w * dw)
def photon_number_with_loss(spectrum, w, dw, gamma, alpha, h) -> float: def photon_number_with_loss(spec2, w, dw, gamma, alpha, h) -> float:
spec2 = abs2(spectrum)
return np.sum(1 / gamma * spec2 / w * dw) - h * np.sum(alpha / gamma * spec2 / w * dw) return np.sum(1 / gamma * spec2 / w * dw) - h * np.sum(alpha / gamma * spec2 / w * dw)
def pulse_energy(spectrum, dw) -> float: def pulse_energy(spec2, dw) -> float:
return np.sum(abs2(spectrum) * dw) return np.sum(spec2 * dw)
def pulse_energy_with_loss(spectrum, dw, alpha, h) -> float: def pulse_energy_with_loss(spec2, dw, alpha, h) -> float:
spec2 = abs2(spectrum) spec2 = spec2
return np.sum(spec2 * dw) - h * np.sum(alpha * spec2 * dw) return np.sum(spec2 * dw) - h * np.sum(alpha * spec2 * dw)
@@ -539,7 +537,7 @@ def ideal_compressed_pulse(spectra):
compressed : 1D array compressed : 1D array
time envelope of the compressed field 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): 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))) spec = np.zeros((t_res, len(time)))
for i, delay in enumerate(delays): for i, delay in enumerate(delays):
masked = values * np.exp(-(((time - delay) / gate_width) ** 2)) masked = values * np.exp(-(((time - delay) / gate_width) ** 2))
spec[i] = abs2(fft(masked)) spec[i] = math.abs2(fft(masked))
if shift: if shift:
spec[i] = fftshift(spec[i]) spec[i] = fftshift(spec[i])
return spec, delays return spec, delays
@@ -592,7 +590,7 @@ def g12(values):
# Create all the possible pairs of values # Create all the possible pairs of values
n = len(values) n = len(values)
field_pairs = itertools.combinations(values, 2) 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() mask = mean_spec > 1e-15 * mean_spec.max()
corr = np.zeros_like(values[0]) corr = np.zeros_like(values[0])
for pair in field_pairs: for pair in field_pairs:
@@ -618,7 +616,7 @@ def avg_g12(values):
if len(values.shape) > 2: if len(values.shape) > 2:
pass pass
avg_values = np.mean(abs2(values), axis=0) avg_values = np.mean(math.abs2(values), axis=0)
coherence = g12(values) coherence = g12(values)
return np.sum(coherence * avg_values) / np.sum(avg_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 spline_4 : scipy.interpolate.UnivariateSpline
order 4 spline that interpolate values around the peak order 4 spline that interpolate values around the peak
""" """
logger = get_logger(__name__)
if not already_sorted: if not already_sorted:
x_axis, values = x_axis.copy(), values.copy() x_axis, values = x_axis.copy(), values.copy()
@@ -914,7 +911,7 @@ def _detailed_find_lobe_limits(
else: else:
iterations += 1 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( small_spline, spline_4, spline_5, d_spline, d_roots, dd_roots, l_ind, r_ind = setup_splines(
x_axis, values, mam x_axis, values, mam
@@ -941,7 +938,7 @@ def _detailed_find_lobe_limits(
color = default_plotting["color_cycle"] color = default_plotting["color_cycle"]
if debug != "": 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(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.plot(newx, spline_5(newx), c=color[1])
ax.scatter(fwhm_pos, spline_4(fwhm_pos), marker="+", label="all fwhm", c=color[2]) 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]) list of tuples of the form ([left_lobe_lim, right_lobe_lim, lobe_pos], [left_hm, right_hm])
""" """
if compress: if compress:
fields = abs2(compress_pulse(spectra)) fields = math.abs2(compress_pulse(spectra))
else: else:
fields = abs2(ifft(spectra)) fields = math.abs2(ifft(spectra))
field = np.mean(fields, axis=0) 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 # Isolate whole central lobe of bof mean and ideal field
lobe_lim, fwhm_lim, _, big_spline = find_lobe_limits(t, field, debug) 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) lobe_lim_i, _, _, big_spline_i = find_lobe_limits(t, ideal_field, debug)
# Compute quality factor # Compute quality factor
energy_fraction = (big_spline.integral(*span(lobe_lim[:2]))) / np.trapz(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(*span(lobe_lim_i[:2]))) / np.trapz(ideal_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 qf = energy_fraction / energy_fraction_i
# Compute mean coherence # Compute mean coherence
mean_g12 = avg_g12(spectra) 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 # To compute amplitude and fwhm fluctuations, we need to measure every single peak
P0 = [] 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) lobe_lim, fwhm_lim, _, big_spline = find_lobe_limits(t, f, debug)
all_lims.append((lobe_lim, fwhm_lim)) all_lims.append((lobe_lim, fwhm_lim))
P0.append(big_spline(lobe_lim[2])) P0.append(big_spline(lobe_lim[2]))
fwhm.append(length(fwhm_lim)) fwhm.append(math.length(fwhm_lim))
t_offset.append(lobe_lim[2]) t_offset.append(lobe_lim[2])
energies.append(np.trapz(fields, t)) energies.append(np.trapz(fields, t))
fwhm_var = np.std(fwhm) / np.mean(fwhm) fwhm_var = np.std(fwhm) / np.mean(fwhm)
@@ -1063,7 +1062,7 @@ def rin_curve(spectra: np.ndarray) -> np.ndarray:
RIN curve RIN curve
""" """
if np.iscomplexobj(spectra): if np.iscomplexobj(spectra):
A2 = abs2(spectra) A2 = math.abs2(spectra)
else: else:
A2 = spectra A2 = spectra
return np.std(A2, axis=-2) / np.mean(A2, axis=-2) 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]: def measure_field(t: np.ndarray, field: np.ndarray) -> Tuple[float, float, float]:
"""returns fwhm, peak_power, energy""" """returns fwhm, peak_power, energy"""
if np.iscomplexobj(field): if np.iscomplexobj(field):
intensity = abs2(field) intensity = math.abs2(field)
else: else:
intensity = field intensity = field
_, fwhm_lim, _, _ = find_lobe_limits(t, intensity) _, fwhm_lim, _, _ = find_lobe_limits(t, intensity)
fwhm = length(fwhm_lim) fwhm = math.length(fwhm_lim)
peak_power = intensity.max() peak_power = intensity.max()
energy = np.trapz(intensity, t) energy = np.trapz(intensity, t)
return fwhm, peak_power, energy return fwhm, peak_power, energy
@@ -1101,10 +1100,12 @@ def remove_2nd_order_dispersion(
np.ndarray, shape (n, ) np.ndarray, shape (n, )
spectrum with 2nd order dispersion removed 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): 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)) opti = minimize_scalar(score, bracket=(max_z, 0))
return propagate(opti.x), opti return propagate(opti.x), opti
@@ -1127,8 +1128,12 @@ def remove_2nd_order_dispersion2(
np.ndarray, shape (n, ) np.ndarray, shape (n, )
spectrum with 2nd order dispersion removed 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): def score(gdd):
return -np.sum(integrate(gdd) ** 6) return -np.sum(integrate(gdd) ** 6)

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

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