From a9f77cea34222cf97490ea8e28a0834b99aadfb5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Sierro?= Date: Fri, 8 Mar 2024 11:15:10 +0100 Subject: [PATCH] new shot noise implementation --- examples/shot_noise_versions.py | 41 +++++++++++++++ pyproject.toml | 2 +- src/scgenerator/math.py | 1 - src/scgenerator/parameter.py | 23 ++++----- src/scgenerator/physics/pulse.py | 89 ++++++++++++++++++++++++++------ src/scgenerator/spectra.py | 29 +++++++---- 6 files changed, 146 insertions(+), 39 deletions(-) create mode 100644 examples/shot_noise_versions.py diff --git a/examples/shot_noise_versions.py b/examples/shot_noise_versions.py new file mode 100644 index 0000000..dec146e --- /dev/null +++ b/examples/shot_noise_versions.py @@ -0,0 +1,41 @@ +import itertools +import scgenerator as sc +import matplotlib.pyplot as plt +import numpy as np + +plt.rcParams["font.family"] = "serif" + + +def main(): + t = np.linspace(-100, 100, 512) * 1e-15 + w = sc.wspace(t) + sc.units.nm_rads(1056) + + fig, (tops, bots) = plt.subplots(2, 4, constrained_layout=True, figsize=(11, 5)) + + spec_kw = dict(cmap="Blues_r", marker=".") + time_kw = dict(cmap="Greens_r", marker=".") + + for (const, phase), top, bot in zip( + itertools.product((True, False), (True, False)), tops, bots + ): + sn = sc.pulse.shot_noise(w, const, phase) + sn_t = np.fft.ifft(sn) + top.scatter(sn.real, sn.imag, c=np.abs(w), **spec_kw) + bot.scatter(sn_t.real, sn_t.imag, c=t[::-1], **time_kw) + top.set_aspect(1, adjustable="datalim") + bot.set_aspect(1, adjustable="datalim") + top.axis("off") + bot.axis("off") + top.set_title( + f"variance~$\\omega{'_0' if const else ''}$\n{'phase only' if phase else 'Re + Im'}" + ) + + sc.plotting.corner_annotation("↑Im\n→Re", tops[0]) + sc.plotting.corner_annotation(r"$\tilde{A}(\omega)$", tops[0], "bl") + sc.plotting.corner_annotation("$A(t)$", bots[0], "bl") + fig.suptitle("Different shot noise implementations, shading according to $|\\omega|$, $t$") + plt.show() + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index 8fbfa8b..25c6500 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "scgenerator" -version = "0.4.7" +version = "0.4.8" description = "Simulate nonlinear pulse propagation in optical fibers" readme = "README.md" authors = [{ name = "Benoit Sierro", email = "benoit.sierro@iap.unibe.ch" }] diff --git a/src/scgenerator/math.py b/src/scgenerator/math.py index cdd354b..1301f00 100644 --- a/src/scgenerator/math.py +++ b/src/scgenerator/math.py @@ -3,7 +3,6 @@ collection of purely mathematical function """ import math -import platform import warnings from dataclasses import dataclass from functools import cache diff --git a/src/scgenerator/parameter.py b/src/scgenerator/parameter.py index 4649b12..0013d8c 100644 --- a/src/scgenerator/parameter.py +++ b/src/scgenerator/parameter.py @@ -6,7 +6,7 @@ import os import tomllib import warnings from copy import copy, deepcopy -from dataclasses import dataclass, field, fields +from dataclasses import dataclass, field from functools import lru_cache, wraps from math import isnan from pathlib import Path @@ -19,6 +19,7 @@ from scgenerator.evaluator import Evaluator, EvaluatorError from scgenerator.io import CustomEncoder, DataFile, custom_decode_hook from scgenerator.operators import Qualifier, SpecOperator, VariableQuantity from scgenerator.physics.units import unit_formatter +from scgenerator.physics.pulse import ShotNoiseParameter T = TypeVar("T") DISPLAY_FUNCTIONS: dict[str, Callable[[float], str]] = {} @@ -59,7 +60,7 @@ def type_checker(*types): if isinstance(validator, str) and n is not None: _name = validator - def validator(*args): + def validator(*_): pass @wraps(validator) @@ -214,14 +215,10 @@ def func_validator(name, n): raise TypeError(f"{name!r} must be callable") -# classes - - class Parameter: def __init__( self, - validator: Callable[[str, Any], None], - converter: Callable = None, + validator: Callable[[str, Any], Any | None], default=None, unit: str | None = None, can_pickle: bool = True, @@ -234,9 +231,8 @@ class Parameter: validator : Callable[[str, Any], None] signature : validator(name, value) must raise a ValueError when value doesn't fit the criteria checked by - validator. name is passed to validator to be included in the error message - converter : Callable, optional - converts a valid value (for example, str.lower), by default None + validator. name is passed to validator to be included in the error message. + a validator can return a value, in which case it acts as a converter. default : callable, optional factory function for a default value (for example, list), by default None display_info : tuple[str, int], optional @@ -244,7 +240,6 @@ class Parameter: example : ("W", 3) will mean the value 1.120045e6 is displayed as '1.12MW' """ self._validator = validator - self.converter = converter self.default = default self.unit = unit self.can_pickle = can_pickle @@ -260,7 +255,7 @@ class Parameter: if self.unit is not None: DISPLAY_FUNCTIONS[self.name] = unit_formatter(self.unit, DECIMALS_DISPLAY) - def __get__(self, instance: Parameters, owner): + def __get__(self, instance: Parameters, _): if instance is None: return self return instance._param_dico.get(self.name) @@ -371,7 +366,9 @@ class Parameters: delay: float = Parameter(type_checker(*number), unit="s") # Behaviors to include - quantum_noise: bool = Parameter(boolean, default=False) + quantum_noise: ShotNoiseParameter = Parameter( + ShotNoiseParameter.validate, default=ShotNoiseParameter() + ) self_steepening: bool = Parameter(boolean, default=True) ideal_gas: bool = Parameter(boolean, default=False) photoionization: bool = Parameter(boolean, default=False) diff --git a/src/scgenerator/physics/pulse.py b/src/scgenerator/physics/pulse.py index ba41134..721a2ce 100644 --- a/src/scgenerator/physics/pulse.py +++ b/src/scgenerator/physics/pulse.py @@ -9,10 +9,13 @@ of shape `([what, ever,] n, nt)` (could be just `(n, nt)` for 2D sc-ordered arra n is the number of spectra at the same z position and nt is the size of the time/frequency grid """ +from __future__ import annotations + import os from dataclasses import astuple, dataclass from pathlib import Path -from typing import Literal, Tuple, TypeVar +from typing import Literal, Tuple, TypeVar, Any, Sequence +from operator import itemgetter import matplotlib.pyplot as plt import numba @@ -96,6 +99,48 @@ class PulseProperties: return [cls(*a) for a in arr] +class ShotNoiseParameter(tuple): + """ + convenience class to handle shot noise parameters when specifying it in the main Parameters + class. Refer to the `shot_noise` function for information of parameters. + """ + + constant_amplitude: bool = property(itemgetter(0)) + phase_only: bool = property(itemgetter(1)) + __slots__ = [] + + @classmethod + def validate(cls, _: str, params: Any): + if isinstance(params, Sequence): + params = tuple(params) + if params in {(), (False,), False}: + return cls() + elif params in {(True,), True}: + return cls(False, True) + else: + return cls(*params) + + def __new__(cls, p1: tuple | bool | None = None, p2: bool | None = None): + if isinstance(p1, cls): + return p1 + elif isinstance(p1, Sequence): + return cls(*p1) + elif p1 is None and p2 is None: + return tuple.__new__(cls, ()) + elif isinstance(p1, bool) and isinstance(p2, bool): + return tuple.__new__(cls, (p1, p2)) + else: + raise TypeError( + f"{cls.__name__} can either be a 0-tuple or a 2-tuple of bool, got {p1=}, {p2=}" + ) + + def __bool__(self) -> bool: + return len(self) > 0 + + +no_shot_noise = ShotNoiseParameter() + + def initial_full_field( t: np.ndarray, shape: str, @@ -562,41 +607,55 @@ def technical_noise(rms_noise: float, noise_correlation: float = -0.4): return psy, 1 + noise_correlation * (psy - 1) -def shot_noise(w: np.ndarray, dt: float): +def shot_noise( + w: np.ndarray, constant_amplitude: bool = False, phase_only: bool = True +) -> np.ndarray: """ Parameters ---------- w : array, shape (n,) angular frequencies - dt : float - resolution of time grid + constant_amplitude: bool, optional + if True, the shot noise spectrum variance is a proportional to the pump angular frequency + over the entire frequency grid. Otherwise (default), it is proportional to each frequency + bin. + phase_only: bool, optional + if True (default), only the phase of the complex shot noise amplitude is random. If False, + the real parts and imaginary parts are generated independently (this doesn't change the + variance). Returns ------- np.ndarray, shape (n,) - noise field to be added on top of initial field in time domain + noise spectrum, scaled such that |`ifft(shot_noise(w))`|^2 represents instantaneous power + in W. """ + n = len(w) dw = w[1] - w[0] - - rand_phase = np.random.rand(len(w)) * 2 * pi - shot_noise_phase = np.exp(-1j * rand_phase) - - shot_noise_amplitude = np.sqrt(0.5 * units.hbar * np.abs(w) / dw) - - return ifft(shot_noise_amplitude * shot_noise_phase / dt * np.sqrt(2 * pi)) + dt = 2 * pi / (dw * n) + if constant_amplitude: + std = np.sqrt(0.5 * units.hbar * np.abs(w[0]) / dw) + else: + std = np.sqrt(0.5 * units.hbar * np.abs(w) / dw) + fac = std / dt * np.sqrt(2 * pi) + if phase_only: + return fac * np.exp(2j * pi * np.random.rand(n)) + else: + real = np.random.randn(n) * np.sqrt(0.5) + imag = np.random.randn(n) * np.sqrt(0.5) * 1j + return fac * (real + imag) def finalize_pulse( pre_field_0: np.ndarray, - quantum_noise: bool, + quantum_noise: ShotNoiseParameter, w: np.ndarray, - dt: float, input_transmission: float, ) -> np.ndarray: pre_field_0 *= np.sqrt(input_transmission) if quantum_noise: - pre_field_0 = pre_field_0 + shot_noise(w, dt) + pre_field_0 = pre_field_0 + np.fft.ifft(shot_noise(w, *quantum_noise)) return pre_field_0 diff --git a/src/scgenerator/spectra.py b/src/scgenerator/spectra.py index 85c56f5..45fafda 100644 --- a/src/scgenerator/spectra.py +++ b/src/scgenerator/spectra.py @@ -202,19 +202,30 @@ class Spectrum(np.ndarray): def energy(self) -> np.ndarray: return np.trapz(self.time_int, x=self.t, axis=-1) - def mask_wl(self, pos: float, width: float, preserve_sn: bool = False) -> Spectrum: + def mask_wl( + self, pos: float, width: float, shot_noise: pulse.ShotNoiseParameter | None = None + ) -> Spectrum: """ Filters the spectrum with a bandpass centered at `pos` of FWHM `width`. - The FWHM is taken on the intensity profile, not on the complex amplitude + The FWHM is taken on the intensity profile (as would be in an experiment), not on the + complex amplitude + + Parameters + ---------- + pos : float + peak of the gaussian bandpass filter in m + width : float + fwhm of the bandpass filter, in m + sn_args : (bool, bool) | None, optional + if specified, shot noise is compensated in the attenuation process and this tuple + corresponds to (constant_amplitude, phase_only) in `pulse.shot_noise`. By default, shot + noise is not compensated. """ band = np.exp(-(((self.l - pos) / (pulse.fwhm_to_T0_fac["gaussian"] * width)) ** 2)) - if preserve_sn: - dt = self.t[1] - self.t[0] - sn = np.fft.fft( - np.reshape( - [pulse.shot_noise(self.w, dt) for _ in range(np.prod(self.shape[:-1]))], - self.shape, - ) + if shot_noise: + sn = np.reshape( + [pulse.shot_noise(self.w, *shot_noise) for _ in range(np.prod(self.shape[:-1]))], + self.shape, ) return self * band + np.sqrt(1 - band**2) * sn else: