new shot noise implementation

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()

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" }]
@@ -35,13 +35,3 @@ ignore = ["E741"]
 
 [tool.ruff.pydocstyle]
 convention = "numpy"
-
-[tool.black]
-line-length = 100
-
-[tool.isort]
-profile = "black"
-skip = ["__init__.py"]
-
-[tool.pyright]
-pythonVersion = "3.11"

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

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)

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
 

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: