From 95f0d2fc3412e34c2adc90e192749bd371dc8b39 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Sierro?= <dedebenui@protonmail.com>
Date: Thu, 14 Mar 2024 15:05:15 +0100
Subject: [PATCH] added convenience operator builders

---
 src/scgenerator/operators.py | 34 ++++++++++++++++++++++++++++++++++
 1 file changed, 34 insertions(+)

diff --git a/src/scgenerator/operators.py b/src/scgenerator/operators.py
index bf2c4ae..cc410e7 100644
--- a/src/scgenerator/operators.py
+++ b/src/scgenerator/operators.py
@@ -9,6 +9,8 @@ from typing import Callable
 
 import numpy as np
 
+import scipy.fft as sfft
+
 from scgenerator import math
 from scgenerator.logger import get_logger
 from scgenerator.physics import fiber, materials, plasma, pulse, units
@@ -503,3 +505,35 @@ def full_field_nonlinear_operator(
         return 1j * fullfield_nl_prefactor(z) * fft(total_nonlinear)
 
     return operate
+
+
+##################################################
+################## CONVENIENCE ###################
+##################################################
+
+
+def no_linear() -> VariableQuantity:
+    return constant_quantity(0)
+
+
+def build_envelope_nonlinear(
+    w: np.ndarray, gamma: float, self_steepening: bool = True, raman: str | None = "measured"
+) -> tuple[VariableQuantity, SpecOperator]:
+    w0 = w[0]
+    t = math.iwspace(w)
+    w_c = w - w0
+    if self_steepening:
+        ss_op = constant_quantity(w_c / w0)
+    else:
+        ss_op = constant_quantity(0.0)
+
+    if gamma != 0:
+        raman_frac = fiber.raman_fraction(raman) if raman else 0.0
+        spm_op = envelope_spm(raman_frac)
+        hr_w = fiber.delayed_raman_w(t, raman)
+        raman_op = envelope_raman(hr_w, raman_frac, sfft.fft, sfft.ifft)
+        gamma_op = constant_quantity(np.ones(len(w)) * gamma)
+    else:
+        spm_op = raman_op = no_op_freq(len(w))
+
+    return envelope_nonlinear_operator(gamma_op, ss_op, spm_op, raman_op, sfft.fft, sfft.ifft)