From 400ae2fd48b10467e663740a581ae6aac398d391 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Sierro?= Date: Wed, 24 Jan 2024 15:24:46 +0100 Subject: [PATCH] working on old/new cq solver --- .gitignore | 1 + examples/compute_coherence.py | 55 +++++++++++++ examples/noisy.toml | 26 +++++++ src/scgenerator/const.py | 1 - src/scgenerator/evaluator.py | 2 +- src/scgenerator/parameter.py | 19 ++++- src/scgenerator/physics/fiber.py | 2 +- src/scgenerator/plotting.py | 16 ++-- src/scgenerator/solver.py | 128 ++++++++++++++++++++++++++++++- src/scgenerator/spectra.py | 41 +++++++++- 10 files changed, 276 insertions(+), 15 deletions(-) create mode 100644 examples/compute_coherence.py create mode 100644 examples/noisy.toml diff --git a/.gitignore b/.gitignore index b67d33c..db69881 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ +build/ .DS_store .idea .conda-env diff --git a/examples/compute_coherence.py b/examples/compute_coherence.py new file mode 100644 index 0000000..3eaef86 --- /dev/null +++ b/examples/compute_coherence.py @@ -0,0 +1,55 @@ +import multiprocessing +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np + +import scgenerator as sc + +PARAMS = Path("./examples/noisy.toml") +SEED = 2564 + + +def path(i) -> Path: + return Path(f"build/{Path(__file__).stem}_{i}.zip") + + +def propagate(i): + np.random.seed(SEED + i) + params = sc.Parameters.load("./examples/noisy.toml") + + sc.propagation(path(i), params.compile()).simulate() + + +def propagate_all(n): + to_do = [i for i in range(n) if not path(i).exists()] + with multiprocessing.Pool(4) as pool: + pool.map(propagate, to_do) + + spec, props = sc.propagation_series([path(i) for i in range(n)]) + + return spec, props + + +def main(): + n = 1 + spec, props = propagate_all(n) + print(spec.shape) + + wl, ind, _ = sc.PlotRange( + spec.wl_disp[spec.wl_disp > 0].min() * 1e9, + props.parameters.wavelength_window[1] * 1e9, + "nm", + ).sort_axis(spec.wl_disp) + + for i in range(spec.shape[1]): + fig, (left, right) = plt.subplots(1, 2) + for s in spec[:, i].time_int: + left.plot(spec.t, s) + for s in spec[:, i].wl_int: + right.plot(wl, s[ind]) + plt.show() + + +if __name__ == "__main__": + main() diff --git a/examples/noisy.toml b/examples/noisy.toml new file mode 100644 index 0000000..2638459 --- /dev/null +++ b/examples/noisy.toml @@ -0,0 +1,26 @@ +name = "Sierro2020 Noisy" +width = 50e-15 +shape = "gaussian" +peak_power = 100e3 +wavelength = 1050e-9 + +gamma = 0.018 +beta2_coefficients = [ + 1.001190e-26, + -2.131124e-41, + 3.286793e-55, + -1.290523e-69, + 1.047255e-84, + 1.696410e-98, + -9.261236e-113, + 2.149311e-127, + -2.028394e-142, +] +length = 0.1 +raman_type = "measured" +input_transmission = 1.0 + +wavelength_window = [550e-9, 2000e-9] +t_num = 8192 +quantum_noise = true +z_num = 21 diff --git a/src/scgenerator/const.py b/src/scgenerator/const.py index d513b2c..c85f84a 100644 --- a/src/scgenerator/const.py +++ b/src/scgenerator/const.py @@ -40,7 +40,6 @@ MANDATORY_PARAMETERS = { "spec_0", "c_to_a_factor", "field_0", - "mean_power", "input_transmission", "z_targets", "length", diff --git a/src/scgenerator/evaluator.py b/src/scgenerator/evaluator.py index 9fedaf4..88aa2e0 100644 --- a/src/scgenerator/evaluator.py +++ b/src/scgenerator/evaluator.py @@ -559,7 +559,7 @@ envelope_rules = default_rules + [ Rule("c_to_a_factor", pulse.c_to_a_factor), # Dispersion Rule("beta2_coefficients", fiber.auto_dispersion_coefficients), - Rule("beta2_coefficients", fiber.handle_dispersion_paramter), + Rule("beta2_coefficients", fiber.handle_dispersion_parameter), Rule("beta2_arr", fiber.dispersion_from_coefficients), Rule("beta2", lambda beta2_coefficients: beta2_coefficients[0]), Rule("beta2", lambda beta2_arr, w0_ind: beta2_arr[w0_ind]), diff --git a/src/scgenerator/parameter.py b/src/scgenerator/parameter.py index f4ba129..fcd5614 100644 --- a/src/scgenerator/parameter.py +++ b/src/scgenerator/parameter.py @@ -3,6 +3,7 @@ from __future__ import annotations import datetime as datetime_module import json import os +import tomllib import warnings from copy import copy, deepcopy from dataclasses import dataclass, field, fields @@ -429,6 +430,17 @@ class Parameters: datetime: datetime_module.datetime = Parameter(type_checker(datetime_module.datetime)) version: str = Parameter(string) + @classmethod + def from_toml(cls, s: str) -> Parameters: + decoded = tomllib.loads(s) + extras = set(decoded) - cls._p_names + if len(extras) > 0: + warnings.warn(f"extra keys (ignored) in parameter json: {extras!r}") + for e in extras: + decoded.pop(e) + + return cls(**decoded) + @classmethod def from_json(cls, s: str) -> Parameters: decoded = json.loads(s, object_hook=custom_decode_hook) @@ -442,7 +454,12 @@ class Parameters: @classmethod def load(cls, path: os.PathLike) -> Parameters: - return cls.from_json(Path(path).read_text()) + path = Path(path) + if path.suffix == ".toml": + decode = cls.from_toml + else: + decode = cls.from_json + return decode(path.read_text()) def __repr__(self) -> str: return "Parameter(" + ", ".join(self.__repr_list__()) + ")" diff --git a/src/scgenerator/physics/fiber.py b/src/scgenerator/physics/fiber.py index 6aa3ba6..0e774ca 100644 --- a/src/scgenerator/physics/fiber.py +++ b/src/scgenerator/physics/fiber.py @@ -63,7 +63,7 @@ def dispersion_slope_to_beta3( ) -def handle_dispersion_paramter( +def handle_dispersion_parameter( wavelength: float, dispersion_parameter: float, dispersion_slope: float | None = None ) -> list[float]: coeffs = [dispersion_parameter_to_beta2(dispersion_parameter, wavelength)] diff --git a/src/scgenerator/plotting.py b/src/scgenerator/plotting.py index dae2b22..0207c49 100644 --- a/src/scgenerator/plotting.py +++ b/src/scgenerator/plotting.py @@ -230,7 +230,7 @@ def create_zoom_axis( return inset -def corner_annotation(text, ax, position="tl", rel_x_offset=0.05, rel_y_offset=0.05, **text_kwargs): +def corner_annotation(text, ax, position="tl", pts_x=8, pts_y=8, **text_kwargs): """puts an annotatin in a corner of an ax Parameters ---------- @@ -250,24 +250,26 @@ def corner_annotation(text, ax, position="tl", rel_x_offset=0.05, rel_y_offset=0 # yoff = length(ylim) * rel_y_offset if position[0] == "t": - y = 1 - rel_y_offset + y = 1 + pts_y = -pts_y va = "top" else: - y = 0 + rel_y_offset + y = 0 va = "bottom" if position[1] == "l": - x = 0 + rel_x_offset + x = 0 ha = "left" else: - x = 1 - rel_x_offset + x = 1 + pts_x = -pts_x ha = "right" ax.annotate( text, (x, y), - (x, y), + (pts_x, pts_y), xycoords="axes fraction", - textcoords="axes fraction", + textcoords="offset points", verticalalignment=va, horizontalalignment=ha, **text_kwargs, diff --git a/src/scgenerator/solver.py b/src/scgenerator/solver.py index 7f95859..c7e0475 100644 --- a/src/scgenerator/solver.py +++ b/src/scgenerator/solver.py @@ -7,7 +7,7 @@ import warnings import zipfile from collections import defaultdict from pathlib import Path -from typing import Any, Iterator, Sequence +from typing import Any, Callable, Iterator, Sequence import numba import numpy as np @@ -16,6 +16,7 @@ from scgenerator.io import TimedMessage from scgenerator.logger import get_logger from scgenerator.math import abs2 from scgenerator.operators import SpecOperator, VariableQuantity +from scgenerator.physics.pulse import photon_number logger = get_logger(__name__) @@ -120,6 +121,130 @@ def pi_step_factor(error: float, last_error: float, order: int, eps: float = 0.8 return 1 + np.arctan(fac - 1) +def solve_rk4ip_cq( + spec: np.ndarray, + linear: VariableQuantity, + nonlinear: SpecOperator, + z_max: float, + atol: float, + rtol: float, + safety: float, + cq: Callable[[np.ndarray], float], + h_const: float | None = None, + targets: Sequence[float] | None = None, +) -> Iterator[tuple[np.ndarray, dict[str, Any]]]: + """ + Solve the GNLSE using a Runge-Kutta of order 4 in the interaction picture. Error estimation is + done via photon number estimation. + + Parameters + ---------- + spec : np.ndarray + initial spectrum + linear : Operator + linear operator + nonlinear : Operator + nonlinear operator + z_max : float + stop propagation when z >= z_max (the last step is not guaranteed to be exactly on z_max) + atol : float + absolute tolerance + rtol : float + relative tolerance + safety : float + safety factor when computing new step size + h_const : float | None, optional + constant step size to use, by default None (automatic step size based on atol and rtol) + + Yields + ------ + np.ndarray + last computed spectrum + dict[str, Any] + stats about the last step, including `z` + """ + if h_const is not None: + h = h_const + const_step_size = True + else: + h = 0.000664237859 # from Luna + const_step_size = False + z = 0 + stats = {} + rejected = [] + if targets is not None: + if len(targets) <= 1: + return + targets = list(sorted(set(targets))) + z = targets[0] + if not const_step_size: + h = min(h, (targets[1] - targets[0]) / 2) + targets.pop(0) + + step_ind = 0 + last_cq = cq(abs2(spec)) + error = 0 + last_error = 0 + store_next = False + + def stats(): + return dict(z=z, rejected=rejected.copy(), error=error, h=h) + + yield spec, stats() | dict(h=0) + + while True: + expD = np.exp(h * 0.5 * linear(z)) + + A_I = expD * spec + k1 = expD * (h * nonlinear(spec, z)) + k2 = h * nonlinear(A_I + 0.5 * k1, z + 0.5 * h) + k3 = h * nonlinear(A_I + 0.5 * k2, z + 0.5 * h) + k4 = h * nonlinear(expD * (A_I + k3), z + h) + + new_spec = expD * (A_I + k1 / 6 + k2 / 3 + k3 / 3) + k4 / 6 + + new_cq = cq(abs2(new_spec)) + error = abs(last_cq - new_cq) / last_cq + print(h, error) + + if error == 0: # solution is exact if no nonlinerity is included + next_h_factor = 1.5 + elif 0 < error <= rtol: + next_h_factor = safety * pi_step_factor(error, last_error, 4, 0.8) + else: + next_h_factor = max(0.1, safety * error ** (-0.25)) + + if const_step_size or error <= 1: + spec = new_spec + last_cq = new_cq + z += h + + step_ind += 1 + last_error = error + + if targets is None or store_next: + if targets is not None: + targets.pop(0) + yield spec, stats() + + rejected.clear() + if z >= z_max: + return + if const_step_size: + continue + else: + rejected.append((h, error)) + logger.info(f"{z = :.3f} rejected step {step_ind} with {h = :.2g}, {error = :.2g}") + + h = h * next_h_factor + + if targets is not None and z + h > targets[0]: + h = targets[0] - z + store_next = True + else: + store_next = False + + def solve43( spec: np.ndarray, linear: VariableQuantity, @@ -259,6 +384,7 @@ def integrate( safety: float = 0.9, targets: Sequence[float] | None = None, ) -> SimulationResult: + """legacy function""" spec0 = initial_spectrum.copy() all_spectra = [] stats = defaultdict(list) diff --git a/src/scgenerator/spectra.py b/src/scgenerator/spectra.py index 9e8624d..e129613 100644 --- a/src/scgenerator/spectra.py +++ b/src/scgenerator/spectra.py @@ -3,12 +3,13 @@ from __future__ import annotations import os import warnings from dataclasses import dataclass -from functools import cached_property +from functools import cached_property, partial from pathlib import Path -from typing import Callable, Generic, Iterable, Iterator, Sequence, TypeVar, overload +from typing import Callable, Generic, Iterator, Sequence, TypeVar, overload import numpy as np import scipy.signal as ss +from tqdm import tqdm from scgenerator import math from scgenerator.io import ( @@ -22,6 +23,7 @@ from scgenerator.io import ( from scgenerator.logger import get_logger from scgenerator.parameter import Parameters from scgenerator.physics import pulse, units +from scgenerator.solver import solve43, solve_rk4ip_cq PARAMS_FN = "params.json" ParamsOrNone = TypeVar("ParamsOrNone", Parameters, None) @@ -208,7 +210,7 @@ class Spectrum(np.ndarray): returns the coherence of the spectrum, computed by collapsing axis `axis` and aligned on the wavelength grid """ - return pulse.g12(self, axis)[..., self.wl_order] + return pulse.g12(self, axis)[..., self.l_order] def spectrogram( self, @@ -366,6 +368,39 @@ class Propagation(Generic[ParamsOrNone]): def load_all(self) -> Spectrum: return self._load_slice(slice(None)) + def simulate(self) -> Propagation: + """ + Run the simulations as specified by the Parameters + + Returns + ------- + Propagation + returns itself, so that you can use this method as a 'builder' method: + `popagation = sc.propagation(parameters).simulate()` + """ + params = self.parameters + with warnings.catch_warnings(), tqdm(total=params.z_num) as pbar: + warnings.filterwarnings("error") + for spec, _ in solve43( + params.spec_0, + params.linear_operator, + params.nonlinear_operator, + params.length, + params.tolerated_error, + params.tolerated_error, + 0.95, + targets=params.z_targets, + # cq=partial( + # pulse.photon_number, + # w=params.w, + # dw=params.w[1] - params.w[0], + # gamma=params.gamma, + # ), + ): + self.append(spec) + pbar.update() + return self + @dataclass class PropagationCollection: