From 376abf628f82f116e9e83c97756fa301c4910468 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Sierro?= <dedebenui@protonmail.com>
Date: Mon, 29 Apr 2024 09:44:04 +0200
Subject: [PATCH] added chirp to Gaussian pulses

---
 src/scgenerator/parameter.py     |  5 +++--
 src/scgenerator/physics/pulse.py | 24 ++++++++++++++++--------
 tests/test_pulse.py              | 23 +++++++++++++++++++++++
 3 files changed, 42 insertions(+), 10 deletions(-)

diff --git a/src/scgenerator/parameter.py b/src/scgenerator/parameter.py
index 0013d8c..6f40121 100644
--- a/src/scgenerator/parameter.py
+++ b/src/scgenerator/parameter.py
@@ -361,9 +361,10 @@ class Parameters:
     wavelength: float = Parameter(in_range_incl(100e-9, 10000e-9), unit="m")
     intensity_noise: float = Parameter(in_range_incl(0, 1), unit="%")
     noise_correlation: float = Parameter(in_range_incl(-10, 10))
-    width: float = Parameter(in_range_excl(0, 1e-9), unit="s")
-    t0: float = Parameter(in_range_excl(0, 1e-9), unit="s")
+    width: float = Parameter(positive(*number), unit="s")
+    t0: float = Parameter(positive(*number), unit="s")
     delay: float = Parameter(type_checker(*number), unit="s")
+    chirp: float = Parameter(type_checker(*number), unit="s")
 
     # Behaviors to include
     quantum_noise: ShotNoiseParameter = Parameter(
diff --git a/src/scgenerator/physics/pulse.py b/src/scgenerator/physics/pulse.py
index fc7cee7..8096a3a 100644
--- a/src/scgenerator/physics/pulse.py
+++ b/src/scgenerator/physics/pulse.py
@@ -12,10 +12,11 @@ n is the number of spectra at the same z position and nt is the size of the time
 from __future__ import annotations
 
 import os
+import warnings
 from dataclasses import astuple, dataclass
-from pathlib import Path
-from typing import Literal, Tuple, TypeVar, Any, Sequence
 from operator import itemgetter
+from pathlib import Path
+from typing import Any, Literal, Sequence, Tuple, TypeVar
 
 import matplotlib.pyplot as plt
 import numba
@@ -185,7 +186,12 @@ def initial_full_field(
 
 
 def initial_field_envelope(
-    t: np.ndarray, shape: str, t0: float, peak_power: float, delay: float | None = None
+    t: np.ndarray,
+    shape: str,
+    t0: float,
+    peak_power: float,
+    delay: float | None = None,
+    chirp: float = 0.0,
 ) -> np.ndarray:
     """
     returns the initial field
@@ -219,9 +225,9 @@ def initial_field_envelope(
         t = t - delay
 
     if shape == "gaussian":
-        return gaussian_pulse(t, t0, peak_power)
+        return gaussian_pulse(t, t0, peak_power, chirp)
     elif shape == "sech":
-        return sech_pulse(t, t0, peak_power)
+        return sech_pulse(t, t0, peak_power, chirp)
     else:
         raise ValueError(f"shape '{shape}' not understood")
 
@@ -521,7 +527,9 @@ def P0_to_E0(peak_power: float, t0: float, shape: str):
     return peak_power * t0 * P0T0_to_E0_fac[shape]
 
 
-def sech_pulse(t: np.ndarray, t0: float, P0: float):
+def sech_pulse(t: np.ndarray, t0: float, P0: float, chirp: float = 0.0):
+    if chirp != 0.0:
+        warnings.warn("Chirp is not implemented for `sech_pulse`")
     arg = t / t0
     ind = (arg < 700) & (arg > -700)
     out = np.zeros_like(t)
@@ -529,8 +537,8 @@ def sech_pulse(t: np.ndarray, t0: float, P0: float):
     return out
 
 
-def gaussian_pulse(t: np.ndarray, t0: float, P0: float):
-    return np.sqrt(P0) * np.exp(-((t / t0) ** 2))
+def gaussian_pulse(t: np.ndarray, t0: float, P0: float, chirp: float = 0.0):
+    return np.sqrt(P0) * np.exp(-((t / t0) ** 2) * (1 - 1j * chirp))
 
 
 def pulse_envelope(
diff --git a/tests/test_pulse.py b/tests/test_pulse.py
index dcf9fa2..0580e52 100644
--- a/tests/test_pulse.py
+++ b/tests/test_pulse.py
@@ -9,3 +9,26 @@ def test_pulse_envelope():
 
     assert np.trapz(sc.pulse.pulse_envelope(t, 10, energy=3) ** 2, x=t) == pytest.approx(3)
     assert (sc.pulse.pulse_envelope(t, 10, peak_power=3) ** 2).max() == pytest.approx(3)
+
+
+def test_chirp():
+    t = np.linspace(-10, 10, 1024)
+    f = sc.pulse.gaussian_pulse(t, 1, 1, chirp=1.0)
+    c = np.unwrap(np.angle(f))
+    c -= c.min()
+    assert pytest.approx(c, rel=1e-2, abs=1e-4) == t**2
+
+    params = sc.Parameters(
+        time_window=20,
+        t_num=1024,
+        peak_power=1,
+        wavelength=800e-9,
+        shape="gaussian",
+        chirp=1,
+        t0=1,
+        quantum_noise=False,
+    )
+    t2, f2 = params.compute("t", "field_0")
+    c2 = np.unwrap(np.angle(f2))
+    c2 -= c2.min()
+    assert pytest.approx(c2, rel=1e-2, abs=1e-4) == t2**2