Skip to main content

comnum_math/nlse/
mod.rs

1//! The symmetric split-step Fourier solvers of the nonlinear Schrödinger equation:
2//! scalar ([`SplitStep`], [`SplitStepF32`], both built from the [`Steps`] of a span)
3//! and dual-polarisation Manakov ([`ManakovSplitStep`]).
4//!
5//! The same solver propagates a field through a fibre (`comnum_channel::optical::fiber`)
6//! and back-propagates it at the receiver (`comnum_dsp::compensation::dbp`, with the
7//! coefficients negated), so it lives here, below both.
8//!
9//! # Model
10//!
11//! $$
12//! \frac{\partial A}{\partial z} = -\frac{\alpha}{2} A
13//!   - j \frac{\beta_{2}}{2} \frac{\partial^{2} A}{\partial t^{2}}
14//!   + j \gamma |A|^{2} A
15//!     $$
16//!
17//! Each step of length $h$ is a half linear step, the nonlinear phase
18//! $\gamma |A|^{2} h_{\mathrm{NL}}$ and a half linear step, with the linear operator applied
19//! in the frequency domain:
20//!
21//! $$
22//! \tilde{A}(\omega) \leftarrow \tilde{A}(\omega)\\, e^{-\alpha h / 4}\\, e^{j \beta_{2} \omega^{2} h / 4}
23//! $$
24//!
25//! The nonlinear length $h_{\mathrm{NL}}$ is the step itself for the forward model; a
26//! back-propagation with coarse steps uses the effective length of the step instead.
27//!
28//! # The Manakov equation
29//!
30//! Averaged over the fast random rotations of the birefringence, the two polarisations
31//! $A_{x}$, $A_{y}$ obey
32//!
33//! $$
34//! \frac{\partial A_{x,y}}{\partial z} = -\frac{\alpha}{2} A_{x,y} -
35//!   j \frac{\beta_{2}}{2} \frac{\partial^{2} A_{x,y}}{\partial t^{2}} +
36//!   j \frac{8}{9} \gamma \left(|A_{x}|^{2} + |A_{y}|^{2}\right) A_{x,y}.
37//! $$
38//!
39//! [`ManakovSplitStep`] is the symmetric scheme of OptiCommPy's `manakovSSF`: a half
40//! linear step, the nonlinear phase
41//!
42//! $$
43//! \phi(t) = \frac{8}{9} \gamma\, \frac{P(z, t) + P(z + h, t)}{2}\, h, \qquad
44//! P = |A_{x}|^{2} + |A_{y}|^{2},
45//! $$
46//!
47//! (the trapezoidal rule of OptiCommPy, iterated until the step converges since
48//! $P(z + h)$ depends on the step itself, [`NonlinearStep::Trapezoidal`]; by default the
49//! power at the middle of the step, [`NonlinearStep::Midpoint`], as [`SplitStep`]) and a
50//! half linear step. The step is fixed, adapted so
51//! that the largest nonlinear phase of a step is a given rotation (the nonlinear phase
52//! rotation method of Sinkin et al., which concentrates the steps where the power is
53//! high, at the beginning of the spans), chosen by the four-wave-mixing rule of Musetti
54//! et al., or distributed logarithmically ([`logarithmic_steps`], Bosco et al.: every step
55//! carries the same integrated power). The logarithmic steps keep the spurious
56//! four-wave mixing of the split step near $1/K$ of the true one; with few steps per span
57//! they also place the nonlinear phase of a back-propagation better, with many they leave
58//! more error than equal steps (the long last steps).
59//!
60//! # Examples
61//! ```
62//! use comnum_math::nlse::{Medium, SplitStep, Steps};
63//! use comnum_math::C64;
64//! // without loss and nonlinearity, forward then negated propagation is the identity
65//! let m = Medium { alpha_np_km: 0.0, beta2_s2_km: -2.17e-23, gamma_w_km: 0.0, beta3_s3_km: 0.0 };
66//! let x: Vec<C64> = (0..64).map(|n| C64::new((n % 3) as f64, (n % 2) as f64)).collect();
67//! let mut e = x.clone();
68//! SplitStep::new(&m, Steps::new(8, 10.0), 64, 64e9).run(&mut e);
69//! SplitStep::new(&m.negated(), Steps::new(8, 10.0), 64, 64e9).run(&mut e);
70//! assert!(x.iter().zip(&e).all(|(a, b)| (a - b).norm() < 1e-9));
71//! ```
72//!
73//! # References
74//!
75//! 1. G. P. Agrawal, *Nonlinear Fiber Optics*, 5th ed., Academic Press, 2013.
76//! 2. E. Ip and J. M. Kahn, "Compensation of dispersion and nonlinear impairments
77//!    using digital backpropagation," *J. Lightw. Technol.*, 26(20), 2008.
78//! 3. D. Marcuse, C. R. Menyuk and P. K. A. Wai, "Application of the Manakov-PMD
79//!    equation to studies of signal propagation in optical fibers with randomly varying
80//!    birefringence," *J. Lightw. Technol.*, 15(9), 1735–1746, 1997.
81//! 4. O. V. Sinkin, R. Holzlöhner, J. Zweck and C. R. Menyuk, "Optimization of the
82//!    split-step Fourier method in modeling optical-fiber communications systems," *J.
83//!    Lightw. Technol.*, 21(1), 61–68, 2003.
84//! 5. G. Bosco, A. Carena, V. Curri, R. Gaudino, P. Poggiolini and S. Benedetto,
85//!    "Suppression of spurious tones induced by the split-step method in fiber systems
86//!    simulation," *IEEE Photon. Technol. Lett.*, 12(5), 489–491, 2000.
87//! 6. J. Zhang, X. Li and Z. Dong, "Digital nonlinear compensation based on the modified
88//!    logarithmic step size," *J. Lightw. Technol.*, 31(22), 3546–3555, 2013 (not in the
89//!    references: its eq. (2) as Häger & Pfister describe it).
90//! 7. C. Häger and H. D. Pfister, "Physics-based deep learning for fiber-optic
91//!    communication systems," *IEEE J. Sel. Areas Commun.*, 39(1), 280–294, 2021.
92
93use crate::fft::{self, Fft};
94use crate::{C64, cis};
95use comnum_core::{Error, Result};
96use num_complex::Complex32 as C32;
97
98mod manakov;
99
100pub use manakov::{Birefringence, ManakovSplitStep, PropagationOptions};
101
102/// The coefficients of the scalar NLSE, per kilometre.
103#[derive(Clone, Copy, Debug, PartialEq)]
104pub struct Medium {
105    /// Power attenuation `α` (Np/km).
106    pub alpha_np_km: f64,
107    /// Group-velocity dispersion `β₂` (s²/km).
108    pub beta2_s2_km: f64,
109    /// Nonlinear coefficient `γ` (1/(W·km)).
110    pub gamma_w_km: f64,
111    /// Third-order dispersion `β₃` (s³/km), 0 for none.
112    pub beta3_s3_km: f64,
113}
114
115impl Default for Medium {
116    /// No loss, no dispersion, no nonlinearity.
117    fn default() -> Medium {
118        Medium { alpha_np_km: 0.0, beta2_s2_km: 0.0, gamma_w_km: 0.0, beta3_s3_km: 0.0 }
119    }
120}
121
122impl Medium {
123    /// The medium with `α`, `β₂` and `γ` reversed: propagating through it undoes a
124    /// propagation through `self` (digital back-propagation).
125    pub fn negated(&self) -> Medium {
126        Medium {
127            alpha_np_km: -self.alpha_np_km,
128            beta2_s2_km: -self.beta2_s2_km,
129            gamma_w_km: -self.gamma_w_km,
130            beta3_s3_km: -self.beta3_s3_km,
131        }
132    }
133
134    /// The linear response `e^{-αL/2} e^{j(β₂ω²/2 + β₃ω³/6)L}` of `length_km` of the
135    /// medium at every angular frequency (loss and dispersion of the field).
136    pub fn linear_response(&self, omega: &[f64], length_km: f64) -> Vec<C64> {
137        let mut out = vec![C64::new(0.0, 0.0); omega.len()];
138        self.linear_response_into(omega, length_km, &mut out);
139        out
140    }
141
142    /// [`Medium::linear_response`] into `out` (one value per angular frequency).
143    fn linear_response_into(&self, omega: &[f64], length_km: f64, out: &mut [C64]) {
144        let (a, b2, b3) = (self.alpha_np_km, self.beta2_s2_km, self.beta3_s3_km);
145        let g = (-(a / 2.0) * length_km).exp();
146        if b3 == 0.0 {
147            for (o, wk) in out.iter_mut().zip(omega) {
148                *o = g * cis(b2 / 2.0 * wk * wk * length_km);
149            }
150        } else {
151            for (o, wk) in out.iter_mut().zip(omega) {
152                *o = g * cis((b2 / 2.0 + b3 / 6.0 * wk) * wk * wk * length_km);
153            }
154        }
155    }
156
157    /// The half linear step `e^{-αh/4} e^{jβ₂ω²h/4}` at every angular frequency.
158    fn half_step(&self, omega: &[f64], h: f64) -> Vec<C64> {
159        self.linear_response(omega, h / 2.0)
160    }
161
162    /// [`Medium::half_step`] into `out`.
163    fn half_step_into(&self, omega: &[f64], h: f64, out: &mut [C64]) {
164        self.linear_response_into(omega, h / 2.0, out);
165    }
166}
167
168/// The steps of one span: `n_steps` symmetric steps of `step_km`, the nonlinear phase
169/// of each over `nl_length_km` (the step itself for the forward model, the effective
170/// length of a coarse step for a back-propagation). Shared by [`SplitStep`] and
171/// [`SplitStepF32`], and given by a forward model to its back-propagation
172/// (`comnum_dsp::compensation::dbp::Dbp::forward`).
173#[derive(Clone, Copy, Debug, PartialEq)]
174pub struct Steps {
175    /// Number of steps.
176    pub n_steps: usize,
177    /// Length of a step (km).
178    pub step_km: f64,
179    /// Nonlinear length of a step (km).
180    pub nl_length_km: f64,
181}
182
183impl Steps {
184    /// The steps of the forward model: `n_steps` of `step_km`, whose nonlinear length is
185    /// the step.
186    pub fn new(n_steps: usize, step_km: f64) -> Steps {
187        Steps { n_steps, step_km, nl_length_km: step_km }
188    }
189
190    /// The length covered, $n h$ (km).
191    pub fn length_km(&self) -> f64 {
192        self.n_steps as f64 * self.step_km
193    }
194}
195
196/// The [`Steps`] of one span in double precision, or steps of different lengths
197/// ([`SplitStep::with_lengths`]).
198#[derive(Clone, Debug)]
199pub struct SplitStep {
200    /// The half linear steps, one per distinct step length.
201    lins: Vec<Vec<C64>>,
202    /// Every step in order: the index of its half linear step, its length and its
203    /// nonlinear length (km).
204    schedule: Vec<(usize, f64, f64)>,
205    /// `None` for an empty block (nothing to propagate).
206    plan: Option<Fft>,
207    gamma_w_km: f64,
208    steps: Steps,
209}
210
211impl SplitStep {
212    /// The `steps` through `medium` for a block of `n` samples at `fs` (no plan for an
213    /// empty block: the steps then leave it empty).
214    pub fn new(medium: &Medium, steps: Steps, n: usize, fs: f64) -> SplitStep {
215        let omega = fft::angular_frequencies(n, fs);
216        SplitStep {
217            lins: vec![medium.half_step(&omega, steps.step_km)],
218            schedule: vec![(0, steps.step_km, steps.nl_length_km); steps.n_steps],
219            plan: (n > 0).then(|| fft::plan(n)),
220            gamma_w_km: medium.gamma_w_km,
221            steps,
222        }
223    }
224
225    /// Steps of different lengths: `steps` holds, in the order they are applied, the
226    /// length of every step and its nonlinear length (km), such as the logarithmic steps
227    /// of [`logarithmic_steps`] with their effective lengths. [`SplitStep::steps`] then
228    /// gives their number, their mean length and their mean nonlinear length.
229    pub fn with_lengths(medium: &Medium, steps: &[(f64, f64)], n: usize, fs: f64) -> SplitStep {
230        let omega = fft::angular_frequencies(n, fs);
231        let mut lengths: Vec<f64> = Vec::new();
232        let mut lins = Vec::new();
233        let mut schedule = Vec::with_capacity(steps.len());
234        for &(h, nl) in steps {
235            let i = lengths.iter().position(|&l| l == h).unwrap_or_else(|| {
236                lengths.push(h);
237                lins.push(medium.half_step(&omega, h));
238                lengths.len() - 1
239            });
240            schedule.push((i, h, nl));
241        }
242        let k = steps.len().max(1) as f64;
243        let total: f64 = steps.iter().map(|s| s.0).sum();
244        let total_nl: f64 = steps.iter().map(|s| s.1).sum();
245        SplitStep {
246            lins,
247            schedule,
248            plan: (n > 0).then(|| fft::plan(n)),
249            gamma_w_km: medium.gamma_w_km,
250            steps: Steps { n_steps: steps.len(), step_km: total / k, nl_length_km: total_nl / k },
251        }
252    }
253
254    /// The steps (for [`SplitStep::with_lengths`], their number and mean lengths).
255    pub fn steps(&self) -> Steps {
256        self.steps
257    }
258
259    /// Propagate `e` (√W, time domain) through the steps, in place.
260    ///
261    /// # Panics
262    /// If `e` has another length than the block of [`SplitStep::new`].
263    pub fn run(&self, e: &mut [C64]) {
264        let Some(p) = self.plan.as_ref() else {
265            assert!(e.is_empty(), "a field for an empty block");
266            return;
267        };
268        let g = self.gamma_w_km;
269        p.forward(e);
270        for &(i, _, h) in &self.schedule {
271            mul(e, &self.lins[i]);
272            p.inverse(e);
273            for v in e.iter_mut() {
274                *v *= cis(g * v.norm_sqr() * h);
275            }
276            p.forward(e);
277            mul(e, &self.lins[i]);
278        }
279        p.inverse(e);
280    }
281
282    /// The steps for the Manakov equation: both polarisations (√W, time domain, equal
283    /// lengths) in place, with the nonlinear phase
284    /// $\frac{8}{9} \gamma (|A_{x}|^{2} + |A_{y}|^{2}) h_{\mathrm{NL}}$ at the middle of every
285    /// step (the back-propagation of a dual-polarisation link).
286    ///
287    /// # Panics
288    /// If the polarisations have different lengths, or another length than the block of
289    /// [`SplitStep::new`].
290    pub fn run_dual(&self, x: &mut [C64], y: &mut [C64]) {
291        assert_eq!(x.len(), y.len(), "the polarisations must have the same length");
292        let Some(p) = self.plan.as_ref() else {
293            assert!(x.is_empty(), "a field for an empty block");
294            return;
295        };
296        p.forward(x);
297        p.forward(y);
298        for &(i, _, h) in &self.schedule {
299            mul(x, &self.lins[i]);
300            mul(y, &self.lins[i]);
301            p.inverse(x);
302            p.inverse(y);
303            manakov_rotation(x, y, 8.0 / 9.0 * self.gamma_w_km * h);
304            p.forward(x);
305            p.forward(y);
306            mul(x, &self.lins[i]);
307            mul(y, &self.lins[i]);
308        }
309        p.inverse(x);
310        p.inverse(y);
311    }
312
313    /// [`SplitStep::run`] with a distributed gain: before the nonlinear phase and after
314    /// it, every half step also multiplies the spectrum by `gain(z_0, z_1, g)`, the gain
315    /// of the field at every frequency (FFT order) from $z_{0}$ to $z_{1}$ (km, from the
316    /// start of the steps). The medium then carries no attenuation: `gain` holds the
317    /// loss as well (the power profile of a Raman-pumped span, `comnum_channel::optical::raman`).
318    /// As in [`SplitStep::run`], the field stays in the frequency domain from the end of
319    /// a step to the start of the next: two FFTs per step.
320    ///
321    /// # Panics
322    /// If `e` has another length than the block of [`SplitStep::new`].
323    pub fn run_with_gain(&self, e: &mut [C64], gain: FieldGain<'_>) {
324        let Some(p) = self.plan.as_ref() else {
325            assert!(e.is_empty(), "a field for an empty block");
326            return;
327        };
328        let g = self.gamma_w_km;
329        let mut buf = vec![0.0; e.len()];
330        let mut z0 = 0.0;
331        p.forward(e);
332        for (k, &(i, step, h)) in self.schedule.iter().enumerate() {
333            // the start of the step: k steps of the uniform steps (as they always were
334            // computed), the sum of the steps otherwise
335            if self.lins.len() == 1 {
336                z0 = k as f64 * step;
337            }
338            mul(e, &self.lins[i]);
339            gain(z0, z0 + step / 2.0, &mut buf);
340            scale(e, &buf);
341            p.inverse(e);
342            for v in e.iter_mut() {
343                *v *= cis(g * v.norm_sqr() * h);
344            }
345            p.forward(e);
346            mul(e, &self.lins[i]);
347            gain(z0 + step / 2.0, z0 + step, &mut buf);
348            scale(e, &buf);
349            z0 += step;
350        }
351        p.inverse(e);
352    }
353}
354
355/// The gain of the field at every frequency bin (FFT order) between two positions of a
356/// span (km): `gain(z_0, z_1, out)` writes $\sqrt{G(z_{1}, f) / G(z_{0}, f)}$ into `out`.
357pub type FieldGain<'a> = &'a dyn Fn(f64, f64, &mut [f64]);
358
359/// [`SplitStep`] in single precision: the field, the FFTs, the linear operator and the
360/// nonlinear phase are computed in `f32` (about twice as fast), with the same [`Steps`]
361/// (a back-propagation included); the plans and their work space are cached per thread
362/// as those of [`fft::plan`].
363///
364/// The rounding error grows like `ε √K` with `ε ≈ 6·10⁻⁸` and `K` the number of
365/// steps: below `10⁻⁴` relative for thousands of steps.
366#[derive(Clone)]
367pub struct SplitStepF32 {
368    lin: Vec<C32>,
369    lin_n: Vec<C32>,
370    /// `None` for an empty block.
371    plan: Option<fft::FftF32>,
372    gamma_h: f32,
373    steps: Steps,
374}
375
376impl std::fmt::Debug for SplitStepF32 {
377    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
378        f.debug_struct("SplitStepF32")
379            .field("n", &self.lin.len())
380            .field("gamma_h", &self.gamma_h)
381            .field("steps", &self.steps)
382            .finish_non_exhaustive()
383    }
384}
385
386impl SplitStepF32 {
387    /// The `steps` through `medium` for a block of `n` samples at `fs` (no plan for an
388    /// empty block).
389    pub fn new(medium: &Medium, steps: Steps, n: usize, fs: f64) -> SplitStepF32 {
390        let omega = fft::angular_frequencies(n, fs);
391        let half = medium.half_step(&omega, steps.step_km);
392        // the 1/N of the inverse transform is folded into the first half step
393        let inv_n = 1.0 / n.max(1) as f64;
394        let lin = half.iter().map(|v| C32::new(v.re as f32, v.im as f32)).collect();
395        let lin_n = half.iter().map(|v| v * inv_n).map(|v| C32::new(v.re as f32, v.im as f32)).collect();
396        SplitStepF32 {
397            lin,
398            lin_n,
399            plan: (n > 0).then(|| fft::plan_f32(n)),
400            gamma_h: (medium.gamma_w_km * steps.nl_length_km) as f32,
401            steps,
402        }
403    }
404
405    /// The steps.
406    pub fn steps(&self) -> Steps {
407        self.steps
408    }
409
410    /// Propagate `e` through the steps, in place.
411    ///
412    /// # Panics
413    /// If `e` has another length than the block of [`SplitStepF32::new`].
414    pub fn run(&self, e: &mut [C32]) {
415        let Some(p) = self.plan.as_ref() else {
416            assert!(e.is_empty(), "a field for an empty block");
417            return;
418        };
419        assert_eq!(e.len(), self.lin.len(), "FFT length mismatch");
420        let n = e.len();
421        p.forward(e);
422        for _ in 0..self.steps.n_steps {
423            e.iter_mut().zip(&self.lin_n).for_each(|(u, l)| *u *= l);
424            p.inverse_unnormalized(e);
425            for v in e.iter_mut() {
426                let (sn, cs) = (self.gamma_h * v.norm_sqr()).sin_cos();
427                *v *= C32::new(cs, sn);
428            }
429            p.forward(e);
430            e.iter_mut().zip(&self.lin).for_each(|(u, l)| *u *= l);
431        }
432        p.inverse_unnormalized(e);
433        let s = 1.0 / n as f32;
434        for v in e.iter_mut() {
435            *v *= s;
436        }
437    }
438}
439
440/// The step size of [`ManakovSplitStep`].
441#[derive(Clone, Copy, Debug, PartialEq)]
442#[non_exhaustive]
443pub enum StepRule {
444    /// Steps of a fixed length (km); the last step of a span is shortened.
445    Fixed(f64),
446    /// Steps whose largest nonlinear phase $\max_{t} \phi(t)$ is this rotation (rad).
447    NonlinearPhase(f64),
448    /// This number of steps over every propagated length (a span), distributed
449    /// logarithmically ([`logarithmic_steps`], Bosco et al. 2000, eq. (6)): every step
450    /// carries the same integrated power, short where the power is high. Without loss,
451    /// equal steps.
452    Logarithmic(usize),
453    /// The four-wave-mixing-aware rule of Musetti, Serena and Bononi (2018): the first step
454    /// of every span bounds the phase mismatch of the four-wave mixing across the signal
455    /// bandwidth $B$ (eq. (13)),
456    ///
457    /// $$
458    /// h_{1} = \frac{\Phi_{\mathrm{FWM}}}{|\beta_{2}| (2\pi B)^{2}},
459    /// $$
460    ///
461    /// and the next steps grow with the attenuation by `update`. The rule does not depend on
462    /// the power: the error it leaves on the SNR is fixed by $\Phi_{\mathrm{FWM}}$ (about
463    /// 0.01 dB for 20 rad with the constant local error, 4 rad with the nonlinear-phase
464    /// update, Sec. IV-B).
465    Fwm {
466        /// $\Phi_{\mathrm{FWM}}$ (rad).
467        phase_rad: f64,
468        /// $B$ (Hz): the bandwidth of the signal; `None` measures it on the fields at the
469        /// start of the propagation ([`signal_bandwidth`]).
470        bandwidth_hz: Option<f64>,
471        /// The growth of the steps along a span.
472        update: StepUpdate,
473        /// Also bound the largest nonlinear phase of a step to this rotation (rad), the
474        /// rule of [`StepRule::NonlinearPhase`]: the four-wave-mixing rule alone lets the
475        /// steps of a single channel exceed its dispersion length at high power.
476        max_nl_phase_rad: Option<f64>,
477    },
478}
479
480/// The growth of the steps of [`StepRule::Fwm`] along a span, with the power attenuation
481/// $\alpha$ (Musetti et al. 2018, eqs. (9)–(10)).
482#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
483#[non_exhaustive]
484pub enum StepUpdate {
485    /// The constant local error of Zhang and Hayee for the symmetric step,
486    /// $h_{k+1} = h_{k} e^{\alpha h_{k} / 3}$ (eq. (10)).
487    #[default]
488    ConstantLocalError,
489    /// The update of the nonlinear-phase criterion without the fluctuations of the peak
490    /// power, $L_{\mathrm{eff}}(h_{k+1}) = L_{\mathrm{eff}}(h_{k}) e^{\alpha h_{k}}$ (eq. (9), the
491    /// logarithmic step of Bosco et al. 2000).
492    NonlinearPhase,
493}
494
495/// The distribution of the steps of a span of a back-propagation
496/// (`comnum_dsp::compensation::dbp::Dbp::distribution`).
497#[derive(Clone, Copy, Debug, Default, PartialEq)]
498#[non_exhaustive]
499pub enum StepDistribution {
500    /// Equal steps.
501    #[default]
502    Uniform,
503    /// The logarithmic steps of [`logarithmic_steps`] (Bosco et al. 2000, eq. (6)).
504    Logarithmic,
505    /// The modified logarithmic steps of Zhang, Li and Dong (2013, eq. (2)): the steps of
506    /// [`logarithmic_steps`] for the attenuation times this adjusting factor (0.4
507    /// recommended), less concentrated at the start of the span. As described by Häger
508    /// and Pfister (2021, Sec. IV-C), whose Example 1 the tests check: two steps over a
509    /// span of 100 km at 0.2 dB/km are about 30 and 70 km (15 and 85 km with the factor 1).
510    ModifiedLogarithmic(f64),
511}
512
513impl StepDistribution {
514    /// The `n` steps of a span of `length_km` with the power attenuation `alpha_np_km`
515    /// (Np/km), in the order of the propagation.
516    ///
517    /// # Errors
518    /// As [`logarithmic_steps`].
519    pub fn steps(self, alpha_np_km: f64, length_km: f64, n: usize) -> Result<Vec<f64>> {
520        match self {
521            StepDistribution::Uniform => logarithmic_steps(0.0, length_km, n),
522            StepDistribution::Logarithmic => logarithmic_steps(alpha_np_km, length_km, n),
523            StepDistribution::ModifiedLogarithmic(factor) => {
524                if !(factor > 0.0 && factor.is_finite()) {
525                    return Err(Error::Config(format!("the adjusting factor must be positive, got {factor}")));
526                }
527                logarithmic_steps(alpha_np_km * factor, length_km, n)
528            }
529        }
530    }
531}
532
533impl std::str::FromStr for StepDistribution {
534    type Err = Error;
535
536    /// `uniform`, `logarithmic` or `modified_logarithmic` (the recommended factor 0.4;
537    /// case-insensitive).
538    fn from_str(s: &str) -> Result<StepDistribution> {
539        match s.to_ascii_lowercase().as_str() {
540            "uniform" => Ok(StepDistribution::Uniform),
541            "logarithmic" => Ok(StepDistribution::Logarithmic),
542            "modified_logarithmic" => Ok(StepDistribution::ModifiedLogarithmic(0.4)),
543            _ => Err(Error::Config(format!(
544                "the step distribution must be uniform, logarithmic or modified_logarithmic, got {s:?}"
545            ))),
546        }
547    }
548}
549
550/// The end of the step $k$ of $n$ logarithmic steps over `length_km` (km from the start):
551/// $z_{k} = -\ln(1 - k\sigma)/\alpha$, $\sigma = (1 - e^{-\alpha L})/n$, $z_{n} = L$ exactly;
552/// $kL/n$ without loss.
553pub(crate) fn logarithmic_boundary(alpha_np_km: f64, length_km: f64, n: usize, k: usize) -> f64 {
554    let a = alpha_np_km;
555    if k >= n {
556        return length_km;
557    }
558    if a * length_km < 1e-12 {
559        return length_km * k as f64 / n as f64;
560    }
561    let sigma = -(-a * length_km).exp_m1() / n as f64;
562    -(-(k as f64) * sigma).ln_1p() / a
563}
564
565/// The `n` steps (km) of the logarithmic step-size distribution of Bosco et al. (2000,
566/// eq. (6)) over a fibre of `length_km` with the power attenuation $\alpha$ =
567/// `alpha_np_km` (Np/km), in the order of the propagation:
568///
569/// $$
570/// \Delta z_{k} = -\frac{1}{\alpha} \ln \frac{1 - k\sigma}{1 - (k - 1)\sigma}, \qquad
571/// \sigma = \frac{1 - e^{-\alpha L}}{n}, \qquad k = 1 \dots n
572/// $$
573///
574/// (the paper writes the field attenuation, half of $\alpha$). Every step carries the same
575/// integrated power, $\int e^{-\alpha z} dz = \sigma/\alpha$ (the same effective length): the
576/// steps are short at the start of the span, where the power is high. The phases of the
577/// spurious four-wave mixing of the split-step method then add incoherently, so that its
578/// efficiency stays below $1/n$ of the true one instead of the peaks of equal steps
579/// (Bosco et al., eq. (5) and Fig. 3). Without loss ($\alpha L < 10^{-12}$) the steps are
580/// equal; they add up to `length_km` exactly.
581///
582/// # Errors
583/// [`Error::Config`] for no step, a length that is not positive and finite, or an
584/// attenuation that is negative or not finite.
585///
586/// # Examples
587/// ```
588/// use comnum_math::nlse::logarithmic_steps;
589/// let alpha = 0.2 / (10.0 * std::f64::consts::E.log10()); // 0.2 dB/km in Np/km
590/// let h = logarithmic_steps(alpha, 80.0, 4).unwrap();
591/// assert!(h[0] < h[1] && h[1] < h[2] && h[2] < h[3]);
592/// assert!((h.iter().sum::<f64>() - 80.0).abs() < 1e-12);
593/// // the same effective length (1 - e^{-a h}) e^{-a z} / a for every step
594/// let mut z = 0.0;
595/// let leff: Vec<f64> = h.iter().map(|&s| { let l = (1.0 - (-alpha * s).exp()) * (-alpha * z).exp() / alpha; z += s; l }).collect();
596/// assert!(leff.iter().all(|l| (l - leff[0]).abs() < 1e-9));
597/// ```
598///
599/// # References
600///
601/// 1. G. Bosco, A. Carena, V. Curri, R. Gaudino, P. Poggiolini and S. Benedetto,
602///    "Suppression of spurious tones induced by the split-step method in fiber systems
603///    simulation," *IEEE Photon. Technol. Lett.*, 12(5), 489–491, 2000, eq. (6).
604pub fn logarithmic_steps(alpha_np_km: f64, length_km: f64, n: usize) -> Result<Vec<f64>> {
605    if n == 0 {
606        return Err(Error::Config("at least one step".into()));
607    }
608    if !(length_km > 0.0 && length_km.is_finite()) {
609        return Err(Error::Config(format!("the length must be positive and finite, got {length_km} km")));
610    }
611    if !(alpha_np_km >= 0.0 && alpha_np_km.is_finite()) {
612        return Err(Error::Config(format!("the attenuation must be non-negative and finite, got {alpha_np_km} Np/km")));
613    }
614    let z: Vec<f64> = (0..=n).map(|k| logarithmic_boundary(alpha_np_km, length_km, n, k)).collect();
615    Ok(z.windows(2).map(|w| w[1] - w[0]).collect())
616}
617
618impl StepRule {
619    /// The default rule: Musetti et al. for an SNR error of about 0.01 dB in WDM
620    /// ($\Phi_{\mathrm{FWM}}$ = 20 rad, the constant local error, the bandwidth measured on
621    /// the fields), and at most 0.02 rad of nonlinear phase per step, which a single
622    /// channel at high power needs: one DP-16QAM channel of 32 GBd over 10 × 80 km at
623    /// 2 to 8 dBm gets 0.3 to 0.7 dB too much SNR with the four-wave-mixing rule alone,
624    /// the converged values with both (`validation/cards/musetti2018_fig7.yaml`).
625    pub const FWM_DEFAULT: StepRule = StepRule::Fwm {
626        phase_rad: 20.0,
627        bandwidth_hz: None,
628        update: StepUpdate::ConstantLocalError,
629        max_nl_phase_rad: Some(2e-2),
630    };
631}
632
633/// The bandwidth (Hz) of the fields `fields` sampled at `fs`: the span of the frequencies
634/// whose power spectral density, summed over the fields, exceeds $10^{-4}$ of its peak
635/// ($-40$ dB). Measured on a clean signal (before any amplifier noise, which is white).
636pub fn signal_bandwidth(fields: &[&[C64]], fs: f64) -> f64 {
637    let Some(n) = fields.first().map(|f| f.len()) else { return 0.0 };
638    if n == 0 {
639        return 0.0;
640    }
641    let mut psd = vec![0.0; n];
642    for f in fields {
643        for (p, v) in psd.iter_mut().zip(fft::fft(f)) {
644            *p += v.norm_sqr();
645        }
646    }
647    let peak = psd.iter().copied().fold(0.0, f64::max);
648    if peak == 0.0 {
649        return 0.0;
650    }
651    // the bins in increasing frequency, from -fs/2
652    let signed = |k: usize| if k < n.div_ceil(2) { k as f64 } else { k as f64 - n as f64 };
653    let above: Vec<f64> = psd.iter().enumerate().filter(|(_, p)| **p > 1e-4 * peak).map(|(k, _)| signed(k)).collect();
654    let (lo, hi) = above.iter().fold((f64::INFINITY, f64::NEG_INFINITY), |(a, b), &k| (a.min(k), b.max(k)));
655    (hi - lo + 1.0) * fs / n as f64
656}
657
658/// The nonlinear step of [`ManakovSplitStep`].
659#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
660#[non_exhaustive]
661pub enum NonlinearStep {
662    /// The phase of the power at the middle of the step, after the first half of the
663    /// dispersion, as [`SplitStep`] (eq. (3) of Sinkin et al. 2003): no iteration.
664    #[default]
665    Midpoint,
666    /// OptiCommPy's `manakovSSF`: the phase of the trapezoidal mean of the power at the
667    /// two ends of the step, iterated until the field converges. When channels walk off
668    /// within a step, the mean smooths the intensity that drives the cross-phase
669    /// modulation (see [`ManakovSplitStep`]).
670    Trapezoidal,
671}
672
673impl std::str::FromStr for NonlinearStep {
674    type Err = Error;
675
676    /// `trapezoidal` or `midpoint` (case-insensitive).
677    fn from_str(s: &str) -> Result<NonlinearStep> {
678        match s.to_ascii_lowercase().as_str() {
679            "trapezoidal" => Ok(NonlinearStep::Trapezoidal),
680            "midpoint" => Ok(NonlinearStep::Midpoint),
681            _ => Err(Error::Config(format!("the nonlinear step must be trapezoidal or midpoint, got {s:?}"))),
682        }
683    }
684}
685
686/// The sampling conditions of a split-step simulation of a signal of bandwidth $B$
687/// over a dispersive link (Musetti et al. 2018, Appendix B): the sampling rate must hold
688/// the four-wave mixing that falls on the signal, at least $\frac{3}{2} B$ (all of the
689/// first-order mixing needs $3 B$), and the block must be longer than the walk-off
690/// across $B$ over the whole link, $|\beta_{2}| \, 2\pi B \, L$, or the circular
691/// convolution wraps the fast channels onto the slow ones. [`sampling_check`] computes
692/// them; [`SamplingCheck::is_sufficient`] says whether both hold.
693#[derive(Clone, Copy, Debug, PartialEq)]
694pub struct SamplingCheck {
695    /// The bandwidth of the launched signal (Hz, [`signal_bandwidth`]).
696    pub bandwidth_hz: f64,
697    /// The sampling rate (Hz).
698    pub rate_hz: f64,
699    /// The smallest rate that holds the four-wave mixing on the signal, $\frac{3}{2} B$ (Hz).
700    pub min_rate_hz: f64,
701    /// The duration of the block (s).
702    pub window_s: f64,
703    /// The walk-off across the bandwidth over the link (s).
704    pub walkoff_s: f64,
705}
706
707impl SamplingCheck {
708    /// Whether the rate and the block are sufficient.
709    pub fn is_sufficient(&self) -> bool {
710        self.rate_hz >= self.min_rate_hz && self.window_s >= self.walkoff_s
711    }
712
713    /// Nothing when [`SamplingCheck::is_sufficient`].
714    ///
715    /// # Errors
716    /// [`Error::Config`] naming what is missing otherwise.
717    pub fn ensure(&self) -> Result<()> {
718        if self.is_sufficient() {
719            return Ok(());
720        }
721        let mut what = Vec::new();
722        if self.rate_hz < self.min_rate_hz {
723            what.push(format!(
724                "the sampling rate {:.1} GHz is below 3/2 of the signal bandwidth {:.1} GHz ({:.1} GHz): the four-wave mixing aliases onto the signal (Musetti et al. 2018, Table I)",
725                self.rate_hz / 1e9,
726                self.bandwidth_hz / 1e9,
727                self.min_rate_hz / 1e9
728            ));
729        }
730        if self.window_s < self.walkoff_s {
731            what.push(format!(
732                "the block of {:.2} ns is shorter than the walk-off {:.2} ns across the bandwidth over the link: the circular convolution wraps the channels (Musetti et al. 2018, Appendix B)",
733                self.window_s * 1e9,
734                self.walkoff_s * 1e9
735            ));
736        }
737        Err(Error::Config(format!("insufficient sampling: {} (check_sampling: false to run anyway)", what.join("; "))))
738    }
739}
740
741/// The [`SamplingCheck`] of `fields` sampled at `fs` through `length_km` of `medium`.
742pub fn sampling_check(fields: &[&[C64]], fs: f64, medium: &Medium, length_km: f64) -> SamplingCheck {
743    let b = signal_bandwidth(fields, fs);
744    let n = fields.first().map_or(0, |f| f.len()) as f64;
745    SamplingCheck {
746        bandwidth_hz: b,
747        rate_hz: fs,
748        min_rate_hz: 1.5 * b,
749        window_s: n / fs,
750        walkoff_s: medium.beta2_s2_km.abs() * 2.0 * std::f64::consts::PI * b * length_km,
751    }
752}
753
754/// The nonlinear step of the Manakov equation with a phase that does not change during
755/// the step: both polarisations rotated by the same phase,
756///
757/// $$
758/// A_{x,y}(t) \leftarrow A_{x,y}(t)\, e^{j k (|A_{x}(t)|^{2} + |A_{y}(t)|^{2})},
759/// $$
760///
761/// with $k = \frac{8}{9} \gamma h_{\mathrm{NL}}$ (negated to undo it). The step of
762/// [`SplitStep::run_dual`] and of the Manakov-PMD equation of
763/// `comnum_channel::optical::pmd::ManakovPmd`.
764///
765/// # Panics
766/// If the polarisations have different lengths.
767///
768/// # Examples
769/// ```
770/// use comnum_math::{C64, nlse::manakov_rotation};
771/// let (mut x, mut y) = ([C64::new(1.0, 0.0)], [C64::new(0.0, 1.0)]);
772/// manakov_rotation(&mut x, &mut y, std::f64::consts::FRAC_PI_4);
773/// assert!((x[0] - C64::new(0.0, 1.0)).norm() < 1e-15 && (y[0] - C64::new(-1.0, 0.0)).norm() < 1e-15);
774/// ```
775pub fn manakov_rotation(x: &mut [C64], y: &mut [C64], k: f64) {
776    assert_eq!(x.len(), y.len(), "the polarisations must have the same length");
777    for (u, v) in x.iter_mut().zip(y.iter_mut()) {
778        let e = cis(k * (u.norm_sqr() + v.norm_sqr()));
779        *u *= e;
780        *v *= e;
781    }
782}
783
784fn mul(a: &mut [C64], b: &[C64]) {
785    for (u, v) in a.iter_mut().zip(b) {
786        *u *= v;
787    }
788}
789
790fn scale(a: &mut [C64], b: &[f64]) {
791    for (u, v) in a.iter_mut().zip(b) {
792        *u *= v;
793    }
794}
795
796#[cfg(test)]
797mod tests {
798    use super::*;
799
800    #[test]
801    fn an_empty_field_propagates_to_an_empty_field_instead_of_panicking() {
802        let m = Medium { alpha_np_km: 0.046, beta2_s2_km: -2.17e-23, gamma_w_km: 1.3, beta3_s3_km: 0.0 };
803        let steps = Steps { n_steps: 3, step_km: 1.0, nl_length_km: 0.8 };
804        let s = SplitStep::new(&m, steps, 0, 64e9);
805        let (mut x, mut y): (Vec<C64>, Vec<C64>) = (Vec::new(), Vec::new());
806        s.run(&mut x);
807        s.run_dual(&mut x, &mut y);
808        s.run_with_gain(&mut x, &|_, _, _| {});
809        SplitStepF32::new(&m, steps, 0, 64e9).run(&mut []);
810        assert!(x.is_empty() && y.is_empty() && s.steps() == steps);
811        assert!((steps.length_km() - 3.0).abs() < 1e-15);
812    }
813
814    #[test]
815    fn a_uniform_gain_of_the_loss_is_the_attenuation_of_the_medium() {
816        // run_with_gain on the lossless medium with the field gain e^{-α Δz / 2} at every
817        // frequency is run on the lossy medium (to the rounding of the factorisation)
818        let m = Medium { alpha_np_km: 0.046, beta2_s2_km: -2.17e-23, gamma_w_km: 1.3, beta3_s3_km: 0.0 };
819        let lossless = Medium { alpha_np_km: 0.0, ..m };
820        let x: Vec<C64> = (0..256).map(|n| cis(0.3 * n as f64) * (1e-2 * (1 + n % 5) as f64)).collect();
821        let steps = Steps::new(40, 0.5);
822        let (mut a, mut b) = (x.clone(), x);
823        SplitStep::new(&m, steps, 256, 64e9).run(&mut a);
824        let loss = |z0: f64, z1: f64, out: &mut [f64]| out.fill((-(0.046 / 2.0) * (z1 - z0)).exp());
825        SplitStep::new(&lossless, steps, 256, 64e9).run_with_gain(&mut b, &loss);
826        let peak = a.iter().map(|v| v.norm()).fold(0.0, f64::max);
827        assert!(a.iter().zip(&b).all(|(u, v)| (u - v).norm() < 1e-13 * peak));
828    }
829
830    #[test]
831    fn lossless_split_step_conserves_energy() {
832        let m = Medium { alpha_np_km: 0.0, beta2_s2_km: -2.17e-23, gamma_w_km: 1.3, beta3_s3_km: 0.0 };
833        let mut e: Vec<C64> = (0..256).map(|n| cis(0.1 * n as f64) * (1e-2 * (n % 7) as f64)).collect();
834        let e0: f64 = e.iter().map(C64::norm_sqr).sum();
835        SplitStep::new(&m, Steps::new(80, 1.0), 256, 64e9).run(&mut e);
836        let e1: f64 = e.iter().map(C64::norm_sqr).sum();
837        assert!((e1 / e0 - 1.0).abs() < 1e-12);
838    }
839
840    /// Two polarisations of about 30 GHz of bandwidth at 256 GS/s, of mean power 10 mW.
841    fn dual_field(n: usize) -> (Vec<C64>, Vec<C64>) {
842        let w = |k: usize, f: f64, ph: f64| (2.0 * std::f64::consts::PI * f * k as f64 / n as f64 + ph).sin();
843        let x = (0..n)
844            .map(|k| C64::new(w(k, 61.0, 0.3), w(k, 23.0, 1.1)) * (0.07 * (1.0 + 0.5 * w(k, 7.0, 0.0))))
845            .collect();
846        let y = (0..n)
847            .map(|k| C64::new(w(k, 41.0, 2.0), w(k, 89.0, 0.4)) * (0.07 * (1.0 + 0.5 * w(k, 5.0, 1.0))))
848            .collect();
849        (x, y)
850    }
851
852    #[test]
853    fn manakov_with_one_polarisation_is_the_scalar_equation_with_8_9_gamma() {
854        let m = Medium { alpha_np_km: 0.046, beta2_s2_km: -2.17e-23, gamma_w_km: 1.3, beta3_s3_km: 0.0 };
855        let (x0, _) = dual_field(512);
856        let (mut x, mut y) = (x0.clone(), vec![C64::new(0.0, 0.0); 512]);
857        let s = ManakovSplitStep {
858            rule: StepRule::Fixed(0.01),
859            max_iterations: 20,
860            tolerance: 1e-12,
861            ..ManakovSplitStep::new(m)
862        };
863        s.run(&mut x, &mut y, 20.0, 128e9).unwrap();
864        let mut e = x0;
865        let scalar = Medium { gamma_w_km: 8.0 / 9.0 * 1.3, ..m };
866        SplitStep::new(&scalar, Steps::new(2000, 0.01), 512, 128e9).run(&mut e);
867        let err = x.iter().zip(&e).map(|(u, v)| (u - v).norm()).fold(0.0, f64::max);
868        let peak = e.iter().map(|v| v.norm()).fold(0.0, f64::max);
869        assert!(err < 1e-4 * peak, "{}", err / peak);
870        assert!(y.iter().all(|v| v.norm() == 0.0));
871    }
872
873    #[test]
874    fn manakov_rejects_a_step_that_does_not_advance_instead_of_looping() {
875        let m = Medium { alpha_np_km: 0.0, beta2_s2_km: -2.17e-23, gamma_w_km: 1.3, beta3_s3_km: 0.0 };
876        for rule in [
877            StepRule::Fixed(0.0),
878            StepRule::Fixed(-1.0),
879            StepRule::Fixed(f64::NAN),
880            StepRule::NonlinearPhase(0.0),
881            StepRule::NonlinearPhase(-2e-2),
882        ] {
883            let (mut x, mut y) = dual_field(64);
884            let s = ManakovSplitStep { rule, ..ManakovSplitStep::new(m) };
885            assert!(matches!(s.run(&mut x, &mut y, 1.0, 256e9), Err(Error::Config(_))), "{rule:?}");
886        }
887        let s = ManakovSplitStep::new(m);
888        let (mut x, mut y) = dual_field(64);
889        assert!(matches!(s.run(&mut x, &mut y, f64::INFINITY, 256e9), Err(Error::Config(_))));
890        assert!(matches!(s.run(&mut x, &mut y[..32], 1.0, 256e9), Err(Error::Shape(_))));
891        assert_eq!(s.run(&mut x, &mut y, 0.0, 256e9).unwrap(), 0);
892        // an infinite power: a nonlinear step of zero length
893        x[0] = C64::new(f64::INFINITY, 0.0);
894        assert!(matches!(s.run(&mut x, &mut y, 1.0, 256e9), Err(Error::Numerical(_))));
895    }
896
897    #[test]
898    fn manakov_conserves_the_total_energy_and_converges_with_the_step() {
899        let m = Medium { alpha_np_km: 0.0, beta2_s2_km: -2.17e-23, gamma_w_km: 1.3, beta3_s3_km: 0.0 };
900        let run = |rule| {
901            let (mut x, mut y) = dual_field(1024);
902            let s = ManakovSplitStep { rule, ..ManakovSplitStep::new(m) };
903            let steps = s.run(&mut x, &mut y, 80.0, 256e9).unwrap();
904            (x, y, steps)
905        };
906        let (x0, y0) = dual_field(1024);
907        let e0: f64 = x0.iter().chain(&y0).map(C64::norm_sqr).sum();
908        let (xa, ya, na) = run(StepRule::NonlinearPhase(2e-2));
909        let (xb, yb, nb) = run(StepRule::NonlinearPhase(2e-4));
910        let e1: f64 = xa.iter().chain(&ya).map(C64::norm_sqr).sum();
911        assert!((e1 / e0 - 1.0).abs() < 1e-10);
912        assert!(nb > 50 * na);
913        // second order: the error of the coarse solution is small against the fine one
914        let err: f64 = xa.iter().zip(&xb).chain(ya.iter().zip(&yb)).map(|(u, v)| (u - v).norm_sqr()).sum();
915        assert!((err / e0).sqrt() < 1e-3, "{}", (err / e0).sqrt());
916    }
917
918    #[test]
919    fn the_fwm_rule_follows_eq_13_of_musetti_and_grows_with_the_loss() {
920        // two tones 200 GHz apart: a bandwidth of 200 GHz plus one bin
921        let (n, fs) = (1024, 1.024e12);
922        let tone =
923            |f: f64| -> Vec<C64> { (0..n).map(|k| cis(2.0 * std::f64::consts::PI * f * k as f64 / fs)).collect() };
924        let x: Vec<C64> = tone(-100e9).iter().zip(tone(100e9)).map(|(a, b)| a + b).collect();
925        let y = vec![C64::new(0.0, 0.0); n];
926        assert!((signal_bandwidth(&[&x, &y], fs) - 201e9).abs() < 1.0);
927        assert_eq!(signal_bandwidth(&[&y], fs), 0.0);
928        // the first step of eq. (13): 20 rad over 350 GHz of SMF is 190 m
929        let m = Medium {
930            alpha_np_km: 0.2 / 10.0 * std::f64::consts::LN_10,
931            beta2_s2_km: -2.1683e-23,
932            gamma_w_km: 1.3,
933            beta3_s3_km: 0.0,
934        };
935        let s = ManakovSplitStep {
936            rule: StepRule::Fwm {
937                phase_rad: 20.0,
938                bandwidth_hz: Some(350e9),
939                update: StepUpdate::ConstantLocalError,
940                max_nl_phase_rad: None,
941            },
942            ..ManakovSplitStep::new(m)
943        };
944        let h1 = s.step_length(&[0.0], 0.0, 100.0, 100.0, None, 350e9).unwrap();
945        assert!((h1 - 20.0 / (2.1683e-23 * (2.0 * std::f64::consts::PI * 350e9).powi(2))).abs() < 1e-12);
946        assert!((h1 - 0.190).abs() < 0.001, "{h1}");
947        // the constant local error stretches the step by e^{α h / 3}; without loss it stays
948        let h2 = s.step_length(&[0.0], h1, 100.0, h1 + 100.0, Some(h1), 350e9).unwrap();
949        assert!((h2 / h1 - (m.alpha_np_km * h1 / 3.0).exp()).abs() < 1e-12);
950        let lossless = ManakovSplitStep { medium: Medium { alpha_np_km: 0.0, ..m }, ..s };
951        assert_eq!(lossless.step_length(&[0.0], h1, 100.0, h1 + 100.0, Some(h1), 350e9).unwrap(), h1);
952        // the nonlinear-phase update keeps L_eff(h) e^{-α z}: it grows faster
953        let nlp = ManakovSplitStep {
954            rule: StepRule::Fwm {
955                phase_rad: 20.0,
956                bandwidth_hz: Some(350e9),
957                update: StepUpdate::NonlinearPhase,
958                max_nl_phase_rad: None,
959            },
960            ..s
961        };
962        let h2n = nlp.step_length(&[0.0], h1, 100.0, h1 + 100.0, Some(h1), 350e9).unwrap();
963        let leff = |h: f64| (1.0 - (-m.alpha_np_km * h).exp()) / m.alpha_np_km;
964        assert!((leff(h2n) / leff(h1) - (m.alpha_np_km * h1).exp()).abs() < 1e-12);
965        assert!(h2n > h2);
966        // a non-positive phase or bandwidth is refused
967        for rule in [
968            StepRule::Fwm {
969                phase_rad: 0.0,
970                bandwidth_hz: None,
971                update: StepUpdate::ConstantLocalError,
972                max_nl_phase_rad: None,
973            },
974            StepRule::Fwm {
975                phase_rad: 20.0,
976                bandwidth_hz: Some(0.0),
977                update: StepUpdate::ConstantLocalError,
978                max_nl_phase_rad: None,
979            },
980            StepRule::Fwm {
981                phase_rad: 20.0,
982                bandwidth_hz: None,
983                update: StepUpdate::ConstantLocalError,
984                max_nl_phase_rad: Some(0.0),
985            },
986        ] {
987            let (mut a, mut b) = (x.clone(), y.clone());
988            assert!(ManakovSplitStep { rule, ..s }.run(&mut a, &mut b, 1.0, fs).is_err());
989        }
990    }
991
992    #[test]
993    fn the_third_order_dispersion_adds_a_cubic_phase() {
994        let m = Medium { beta2_s2_km: -2.17e-23, beta3_s3_km: 1.3e-37, ..Medium::default() };
995        let omega = fft::angular_frequencies(64, 64e9);
996        let r = m.linear_response(&omega, 50.0);
997        for (v, w) in r.iter().zip(&omega) {
998            let phase = (m.beta2_s2_km / 2.0 * w * w + m.beta3_s3_km / 6.0 * w * w * w) * 50.0;
999            assert!((v - cis(phase)).norm() < 1e-12);
1000        }
1001        // negated, the response undoes itself
1002        let back = m.negated().linear_response(&omega, 50.0);
1003        assert!(r.iter().zip(&back).all(|(a, b)| ((a * b) - C64::new(1.0, 0.0)).norm() < 1e-12));
1004    }
1005
1006    #[test]
1007    fn one_field_is_the_scalar_split_step_and_fixed_steps_are_equal() {
1008        let m = Medium { alpha_np_km: 0.046, beta2_s2_km: -2.17e-23, gamma_w_km: 1.3, beta3_s3_km: 0.0 };
1009        let (x0, _) = dual_field(1024);
1010        let fs = 256e9;
1011        // 10 km with steps of at most 3 km: four steps of 2.5 km
1012        let mut scalar = x0.clone();
1013        SplitStep::new(&m, Steps::new(4, 2.5), 1024, fs).run(&mut scalar);
1014        let engine = ManakovSplitStep { rule: StepRule::Fixed(3.0), ..ManakovSplitStep::new(m) };
1015        let mut x = x0.clone();
1016        assert_eq!(engine.run_fields(&mut [&mut x], 10.0, fs, &PropagationOptions::default()).unwrap(), 4);
1017        let e: f64 = scalar.iter().map(C64::norm_sqr).sum();
1018        let d: f64 = x.iter().zip(&scalar).map(|(u, v)| (u - v).norm_sqr()).sum();
1019        assert!((d / e).sqrt() < 1e-12, "{}", (d / e).sqrt());
1020    }
1021
1022    #[test]
1023    fn the_birefringence_of_the_sections_is_their_lumped_response() {
1024        // without nonlinearity: the product of D_k R_k at every frequency
1025        let m = Medium { alpha_np_km: 0.0, beta2_s2_km: 0.0, gamma_w_km: 0.0, beta3_s3_km: 0.0 };
1026        let (n, fs) = (256, 256e9);
1027        let (x0, y0) = dual_field(n);
1028        let rot = |t: f64, p: f64| -> [C64; 4] {
1029            let (s, c) = t.sin_cos();
1030            [C64::new(c, 0.0), -cis(-p) * s, cis(p) * s, C64::new(c, 0.0)]
1031        };
1032        let rotations = [rot(0.3, 0.1), rot(1.1, -0.7), rot(-0.4, 2.0)];
1033        let dgd = [3e-12, 1e-12, 2e-12];
1034        let b = Birefringence { section_km: 2.0, rotations: &rotations, dgd_s_per_km: &dgd };
1035        let engine = ManakovSplitStep { rule: StepRule::Fixed(0.7), ..ManakovSplitStep::new(m) };
1036        let (mut x, mut y) = (x0.clone(), y0.clone());
1037        let options = PropagationOptions { gain: None, birefringence: Some(b) };
1038        let steps = engine.run_fields(&mut [&mut x, &mut y], 6.0, fs, &options).unwrap();
1039        assert_eq!(steps, 9); // three sections of three steps
1040        // the lumped response in the frequency domain
1041        let omega = fft::angular_frequencies(n, fs);
1042        let (fx, fy) = (fft::fft(&x0), fft::fft(&y0));
1043        let mut ex = vec![C64::new(0.0, 0.0); n];
1044        let mut ey = ex.clone();
1045        for k in 0..n {
1046            let (mut a, mut c) = (fx[k], fy[k]);
1047            for (r, t) in rotations.iter().zip(&dgd) {
1048                (a, c) = (r[0] * a + r[1] * c, r[2] * a + r[3] * c);
1049                let ph = 0.5 * omega[k] * t * 2.0;
1050                (a, c) = (a * cis(-ph), c * cis(ph));
1051            }
1052            ex[k] = a;
1053            ey[k] = c;
1054        }
1055        let (ex, ey) = (fft::ifft(&ex), fft::ifft(&ey));
1056        let e: f64 = ex.iter().chain(&ey).map(C64::norm_sqr).sum();
1057        let d: f64 = x.iter().zip(&ex).chain(y.iter().zip(&ey)).map(|(u, v)| (u - v).norm_sqr()).sum();
1058        assert!((d / e).sqrt() < 1e-12, "{}", (d / e).sqrt());
1059        // invalid inputs
1060        let mut z = x0.clone();
1061        assert!(engine.run_fields(&mut [&mut z], 6.0, fs, &options).is_err());
1062        assert!(engine.run_fields(&mut [&mut x, &mut y], 5.0, fs, &options).is_err());
1063        let (mut a, mut b2, mut c) = (x0.clone(), x0.clone(), x0);
1064        assert!(engine.run_fields(&mut [&mut a, &mut b2, &mut c], 1.0, fs, &PropagationOptions::default()).is_err());
1065        assert!(engine.run_fields(&mut [], 1.0, fs, &PropagationOptions::default()).is_err());
1066    }
1067
1068    #[test]
1069    fn the_midpoint_manakov_step_is_the_scalar_split_step() {
1070        // one polarisation, gamma scaled by 9/8: the same scheme as SplitStep
1071        let m = Medium { alpha_np_km: 0.046, beta2_s2_km: -2.17e-23, gamma_w_km: 1.3, beta3_s3_km: 0.0 };
1072        let (x0, _) = dual_field(1024);
1073        let (h, n_steps, fs) = (0.5, 160, 256e9);
1074        let mut scalar = x0.clone();
1075        SplitStep::new(&m, Steps::new(n_steps, h), 1024, fs).run(&mut scalar);
1076        let solver = ManakovSplitStep {
1077            rule: StepRule::Fixed(h),
1078            nonlinear: NonlinearStep::Midpoint,
1079            ..ManakovSplitStep::new(Medium { gamma_w_km: 9.0 / 8.0 * m.gamma_w_km, ..m })
1080        };
1081        let (mut x, mut y) = (x0, vec![C64::new(0.0, 0.0); 1024]);
1082        assert_eq!(solver.run(&mut x, &mut y, n_steps as f64 * h, fs).unwrap(), n_steps);
1083        let e: f64 = scalar.iter().map(C64::norm_sqr).sum();
1084        let d: f64 = x.iter().zip(&scalar).map(|(u, v)| (u - v).norm_sqr()).sum();
1085        assert!((d / e).sqrt() < 1e-12, "{}", (d / e).sqrt());
1086        assert!(y.iter().all(|v| v.norm() == 0.0));
1087        assert_eq!("Midpoint".parse::<NonlinearStep>().unwrap(), NonlinearStep::Midpoint);
1088        assert!("euler".parse::<NonlinearStep>().is_err());
1089    }
1090
1091    #[test]
1092    fn fundamental_soliton_keeps_its_shape() {
1093        // T0 = 10 ps, β₂ = -20 ps²/km, P0 = |β₂| / (γ T0²): |A(z, t)| = |A(0, t)|
1094        let (t0, b2, gamma): (f64, f64, f64) = (10e-12, -20e-24, 1.0);
1095        let p0 = b2.abs() / (gamma * t0 * t0);
1096        let (n, fs) = (1024, 2e12);
1097        let a0: Vec<C64> = (0..n)
1098            .map(|k| {
1099                let t = (k as f64 - n as f64 / 2.0) / fs;
1100                C64::new(p0.sqrt() / (t / t0).cosh(), 0.0)
1101            })
1102            .collect();
1103        let m = Medium { alpha_np_km: 0.0, beta2_s2_km: b2, gamma_w_km: gamma, beta3_s3_km: 0.0 };
1104        let mut a = a0.clone();
1105        // two dispersion lengths (L_D = T0²/|β₂| = 5 km)
1106        SplitStep::new(&m, Steps::new(1000, 0.01), n, fs).run(&mut a);
1107        let err = a.iter().zip(&a0).map(|(u, v)| (u.norm() - v.norm()).abs()).fold(0.0, f64::max);
1108        assert!(err < 1e-3 * p0.sqrt(), "relative envelope error {}", err / p0.sqrt());
1109    }
1110
1111    /// The logarithmic steps are eq. (6) of Bosco et al. (2000) written with the field
1112    /// attenuation; one step is the span, no loss gives equal steps, and the inputs are
1113    /// checked.
1114    #[test]
1115    fn logarithmic_steps_are_eq_6_of_bosco_and_check_their_inputs() {
1116        let alpha = 0.25 / (10.0 * std::f64::consts::E.log10()); // 0.25 dB/km, Fig. 3
1117        let (l, k) = (50.0, 50usize);
1118        let h = logarithmic_steps(alpha, l, k).unwrap();
1119        // eq. (6): dz_n = -1/(2 a_f) ln((1 - n s)/(1 - (n - 1) s)), s = (1 - e^{-2 a_f L})/K
1120        let af = alpha / 2.0;
1121        let sigma = (1.0 - (-2.0 * af * l).exp()) / k as f64;
1122        for (n, hn) in h.iter().enumerate() {
1123            let n = (n + 1) as f64;
1124            let bosco = -1.0 / (2.0 * af) * ((1.0 - n * sigma) / (1.0 - (n - 1.0) * sigma)).ln();
1125            assert!((hn - bosco).abs() < 1e-10 * bosco, "step {n}: {hn} against {bosco}");
1126        }
1127        assert!((h.iter().sum::<f64>() - l).abs() < 1e-12 * l);
1128        assert_eq!(logarithmic_steps(alpha, 80.0, 1).unwrap(), vec![80.0]);
1129        assert!(logarithmic_steps(0.0, 80.0, 4).unwrap().iter().all(|&s| (s - 20.0).abs() < 1e-12));
1130        assert!(logarithmic_steps(alpha, 80.0, 0).is_err());
1131        assert!(logarithmic_steps(alpha, 0.0, 4).is_err());
1132        assert!(logarithmic_steps(-1.0, 80.0, 4).is_err());
1133        assert_eq!(StepDistribution::Uniform.steps(alpha, 80.0, 4).unwrap(), vec![20.0; 4]);
1134    }
1135
1136    /// The fast path of the Manakov solver with the logarithmic rule, on one field, is the
1137    /// scalar split step with the same steps; with equal steps `with_lengths` is `new`; the
1138    /// adaptive loop (trapezoidal phase) takes the same number of steps.
1139    #[test]
1140    fn the_logarithmic_rule_runs_the_steps_of_logarithmic_steps() {
1141        let alpha = 0.2 / (10.0 * std::f64::consts::E.log10());
1142        let m = Medium { alpha_np_km: alpha, beta2_s2_km: -21.683e-24, gamma_w_km: 1.3, beta3_s3_km: 0.0 };
1143        let (n, fs) = (512usize, 64e9);
1144        let w = |k: usize, f: f64, ph: f64| (2.0 * std::f64::consts::PI * f * k as f64 / n as f64 + ph).sin();
1145        let x0: Vec<C64> = (0..n).map(|k| C64::new(w(k, 31.0, 0.2), w(k, 17.0, 1.0)) * 0.05).collect();
1146        let mut a = x0.clone();
1147        let taken = ManakovSplitStep { rule: StepRule::Logarithmic(7), ..ManakovSplitStep::new(m) }
1148            .run_fields(&mut [&mut a], 80.0, fs, &PropagationOptions::default())
1149            .unwrap();
1150        assert_eq!(taken, 7);
1151        let steps: Vec<(f64, f64)> = logarithmic_steps(alpha, 80.0, 7).unwrap().iter().map(|&h| (h, h)).collect();
1152        let mut b = x0.clone();
1153        SplitStep::with_lengths(&m, &steps, n, fs).run(&mut b);
1154        assert!(a.iter().zip(&b).all(|(u, v)| (u - v).norm() < 1e-12), "fast path against with_lengths");
1155        let mut c = x0.clone();
1156        let taken = ManakovSplitStep {
1157            rule: StepRule::Logarithmic(7),
1158            nonlinear: NonlinearStep::Trapezoidal,
1159            ..ManakovSplitStep::new(m)
1160        }
1161        .run_fields(&mut [&mut c], 80.0, fs, &PropagationOptions::default())
1162        .unwrap();
1163        assert_eq!(taken, 7);
1164        // equal steps through with_lengths: the steps of new, bit for bit
1165        let (mut d, mut e) = (x0.clone(), x0.clone());
1166        SplitStep::new(&m, Steps::new(4, 20.0), n, fs).run(&mut d);
1167        SplitStep::with_lengths(&m, &[(20.0, 20.0); 4], n, fs).run(&mut e);
1168        assert_eq!(d, e);
1169        assert!(
1170            ManakovSplitStep { rule: StepRule::Logarithmic(0), ..ManakovSplitStep::new(m) }
1171                .run_fields(&mut [&mut d], 80.0, fs, &PropagationOptions::default())
1172                .is_err()
1173        );
1174    }
1175}
1176
1177#[cfg(test)]
1178mod literature {
1179    use super::*;
1180
1181    /// The efficiency of the fictitious four-wave mixing of a split step (Bosco et al.
1182    /// 2000, eq. (5)): $\eta'' = |\sum_{n} \gamma_{n} e^{j \phi_{n}}|^{2}/(1 - e^{-\alpha L})^{2}$ with
1183    /// $\gamma_{n} = (1 - e^{-\alpha \Delta z_{n}}) e^{-\alpha z_{n-1}}$ and the phase mismatch
1184    /// $\phi_{n} = \Delta\beta\, z_{n-1}$ accumulated by the start of the step.
1185    fn fictitious_fwm(alpha: f64, steps: &[f64], dbeta: f64) -> f64 {
1186        let total = 1.0 - (-alpha * steps.iter().sum::<f64>()).exp();
1187        let mut z = 0.0;
1188        let mut sum = C64::new(0.0, 0.0);
1189        for &h in steps {
1190            let g = (1.0 - (-alpha * h).exp()) * (-alpha * z).exp();
1191            sum += C64::from_polar(g, dbeta * z);
1192            z += h;
1193        }
1194        sum.norm_sqr() / (total * total)
1195    }
1196
1197    /// Bosco et al. (2000), Sec. III and Fig. 3 (L = 50 km, 0.25 dB/km, K = 50): with the
1198    /// logarithmic steps every $\gamma_{n}$ is the same (eq. (6) imposes it), so that beyond
1199    /// the true efficiency the fictitious one stays near the floor $1/K$; with equal steps
1200    /// it has peaks of 1, the first at $\Delta\beta \Delta z = 2\pi$ (the "unrealistic
1201    /// pumps" of Fig. 2).
1202    #[test]
1203    fn bosco_2000_fig_3_logarithmic_steps_keep_the_fictitious_fwm_near_1_over_k() {
1204        let alpha = 0.25 / (10.0 * std::f64::consts::E.log10());
1205        let (l, k) = (50.0, 50usize);
1206        let log = logarithmic_steps(alpha, l, k).unwrap();
1207        let uniform = vec![l / k as f64; k];
1208        // gamma_n constant for the logarithmic steps
1209        let mut z = 0.0;
1210        let gammas: Vec<f64> = log
1211            .iter()
1212            .map(|&h| {
1213                let g = (1.0 - (-alpha * h).exp()) * (-alpha * z).exp();
1214                z += h;
1215                g
1216            })
1217            .collect();
1218        assert!(gammas.iter().all(|g| (g - gammas[0]).abs() < 1e-12 * gammas[0]));
1219        // the phase mismatch from 2 to 40 rad/km: past the decay of the true efficiency
1220        let grid: Vec<f64> = (0..3800).map(|i| 2.0 + i as f64 * 0.01).collect();
1221        let eta_log: Vec<f64> = grid.iter().map(|&b| fictitious_fwm(alpha, &log, b)).collect();
1222        let eta_uni: Vec<f64> = grid.iter().map(|&b| fictitious_fwm(alpha, &uniform, b)).collect();
1223        let peak_uniform = fictitious_fwm(alpha, &uniform, 2.0 * std::f64::consts::PI / (l / k as f64));
1224        let mean_log = eta_log.iter().sum::<f64>() / grid.len() as f64;
1225        let max_log = eta_log.iter().copied().fold(0.0, f64::max);
1226        let max_uni = eta_uni.iter().copied().fold(0.0, f64::max);
1227        assert!((peak_uniform - 1.0).abs() < 1e-9, "the peak of equal steps {peak_uniform}");
1228        assert!(max_uni > 0.99, "equal steps {max_uni}");
1229        assert!(
1230            (mean_log * k as f64 - 1.0).abs() < 0.2,
1231            "the mean of the logarithmic steps {mean_log}, floor {}",
1232            1.0 / k as f64
1233        );
1234        assert!(max_log < 0.2, "the largest of the logarithmic steps {max_log}");
1235    }
1236
1237    /// Häger & Pfister (2021), Sec. IV-C, Example 1: the modified logarithmic heuristic
1238    /// of Zhang et al. (2013, eq. (2)) with the adjusting factor 0.4 cuts a span of 100 km
1239    /// (0.2 dB/km, Table I) into steps of about 30 and 70 km, applied from the end of the
1240    /// span by the back-propagation ($\delta_{1} = \delta_{3} \approx 70$ km, $\delta_{2} = \delta_{4} \approx 30$ km).
1241    #[test]
1242    fn hager_2021_example_1_modified_logarithmic_steps() {
1243        let alpha = 0.2 / (10.0 * std::f64::consts::E.log10());
1244        let h = StepDistribution::ModifiedLogarithmic(0.4).steps(alpha, 100.0, 2).unwrap();
1245        assert!((h[0] - 30.0).abs() < 1.0 && (h[1] - 70.0).abs() < 1.0, "{h:?}");
1246        let bosco = StepDistribution::Logarithmic.steps(alpha, 100.0, 2).unwrap();
1247        assert!((bosco[0] - 14.8).abs() < 0.1, "{bosco:?}");
1248        assert_eq!(
1249            "modified_logarithmic".parse::<StepDistribution>().unwrap(),
1250            StepDistribution::ModifiedLogarithmic(0.4)
1251        );
1252        assert!(StepDistribution::ModifiedLogarithmic(0.0).steps(alpha, 100.0, 2).is_err());
1253    }
1254}