From 1efd75a2619ab277961782238e602abca1b3c915 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Sierro?= <benoitsierro@bluewin.ch>
Date: Wed, 29 Sep 2021 15:11:06 +0200
Subject: [PATCH] additional quantum noise factor

---
 src/scgenerator/physics/__init__.py |  8 ++++++--
 src/scgenerator/physics/pulse.py    | 18 +++++++++++++-----
 src/scgenerator/utils/parameter.py  | 11 ++++++++++-
 3 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/src/scgenerator/physics/__init__.py b/src/scgenerator/physics/__init__.py
index 82bc01e..36709a8 100644
--- a/src/scgenerator/physics/__init__.py
+++ b/src/scgenerator/physics/__init__.py
@@ -114,7 +114,9 @@ def find_optimal_depth(
     return propagate(opti.x), opti
 
 
-def propagate_field(t: np.ndarray, field: np.ndarray, z: float, material: str) -> np.ndarray:
+def propagate_field(
+    t: np.ndarray, field: np.ndarray, z: float, material: str, center_wl_nm: float = 1540.0
+) -> np.ndarray:
     """propagates a field through bulk material
 
     Parameters
@@ -127,6 +129,8 @@ def propagate_field(t: np.ndarray, field: np.ndarray, z: float, material: str) -
         distance to propagate in m
     material : str
         material name
+    center_wl_nm : float, optional
+        center wavelength of the grid in nm, by default 1540
 
     Returns
     -------
@@ -134,6 +138,6 @@ def propagate_field(t: np.ndarray, field: np.ndarray, z: float, material: str) -
         propagated field
     """
     w_c = math.wspace(t)
-    l = units.m(w_c + units.nm(1540))
+    l = units.m(w_c + units.nm(center_wl_nm))
     disp = material_dispersion(l, material)
     return np.fft.ifft(np.fft.fft(field) * np.exp(0.5j * disp * w_c ** 2 * z))
diff --git a/src/scgenerator/physics/pulse.py b/src/scgenerator/physics/pulse.py
index 6a07c70..f44191e 100644
--- a/src/scgenerator/physics/pulse.py
+++ b/src/scgenerator/physics/pulse.py
@@ -414,7 +414,7 @@ def technical_noise(rms_noise, noise_correlation=-0.4):
     return psy, 1 + noise_correlation * (psy - 1)
 
 
-def shot_noise(w_c, w0, T, dt):
+def shot_noise(w_c, w0, T, dt, additional_noise_factor=1.0):
     """
 
     Parameters
@@ -427,23 +427,31 @@ def shot_noise(w_c, w0, T, dt):
             length of the time windows
         dt : float
             resolution of time grid
+        additional_noise_factor : float
+            resulting noise spectrum is multiplied by this number
 
     Returns
-    ----------
+    -------
         out : 1D array of size len(w_c)
             noise field to be added on top of initial field in time domain
     """
     rand_phase = np.random.rand(len(w_c)) * 2 * pi
     A_oppm = np.sqrt(hbar * (np.abs(w_c + w0)) * T) * np.exp(-1j * rand_phase)
     out = ifft(A_oppm / dt * np.sqrt(2 * pi))
-    return out
+    return out * additional_noise_factor
 
 
 def add_shot_noise(
-    field_0: np.ndarray, quantum_noise: bool, w_c: bool, w0: float, time_window: float, dt: float
+    field_0: np.ndarray,
+    quantum_noise: bool,
+    w_c: bool,
+    w0: float,
+    time_window: float,
+    dt: float,
+    additional_noise_factor: float,
 ) -> np.ndarray:
     if quantum_noise:
-        field_0 = field_0 + shot_noise(w_c, w0, time_window, dt)
+        field_0 = field_0 + shot_noise(w_c, w0, time_window, dt, additional_noise_factor)
     return field_0
 
 
diff --git a/src/scgenerator/utils/parameter.py b/src/scgenerator/utils/parameter.py
index 6822aa7..b23450d 100644
--- a/src/scgenerator/utils/parameter.py
+++ b/src/scgenerator/utils/parameter.py
@@ -376,6 +376,7 @@ class Parameters:
     energy: float = Parameter(positive(float, int), display_info=(1e6, "μJ"))
     soliton_num: float = Parameter(non_negative(float, int))
     quantum_noise: bool = Parameter(boolean, default=False)
+    additional_noise_factor: float = Parameter(positive(float, int), default=1)
     shape: str = Parameter(literal("gaussian", "sech"), default="gaussian")
     wavelength: float = Parameter(in_range_incl(100e-9, 3000e-9), display_info=(1e9, "nm"))
     intensity_noise: float = Parameter(in_range_incl(0, 1), display_info=(1e2, "%"), default=0)
@@ -1383,7 +1384,15 @@ default_rules: list[Rule] = [
     Rule(
         "field_0",
         pulse.add_shot_noise,
-        ["pre_field_0", "quantum_noise", "w_c", "w0", "time_window", "dt"],
+        [
+            "pre_field_0",
+            "quantum_noise",
+            "w_c",
+            "w0",
+            "time_window",
+            "dt",
+            "additional_noise_factor",
+        ],
     ),
     Rule("peak_power", pulse.E0_to_P0, ["energy", "t0", "shape"]),
     Rule("peak_power", pulse.soliton_num_to_peak_power),