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)