1use super::{
5 FieldGain, Medium, NonlinearStep, StepRule, StepUpdate, logarithmic_boundary, manakov_rotation, mul, scale,
6 signal_bandwidth,
7};
8use crate::fft;
9use crate::{C64, cis};
10use comnum_core::{Error, Result, ensure, same_len};
11
12#[derive(Clone, Copy, Debug, PartialEq)]
20pub struct ManakovSplitStep {
21 pub medium: Medium,
23 pub rule: StepRule,
25 pub nonlinear: NonlinearStep,
27 pub max_iterations: usize,
29 pub tolerance: f64,
31}
32
33#[derive(Clone, Copy, Debug)]
40pub struct Birefringence<'a> {
41 pub section_km: f64,
43 pub rotations: &'a [[C64; 4]],
45 pub dgd_s_per_km: &'a [f64],
47}
48
49#[derive(Clone, Copy, Default)]
52pub struct PropagationOptions<'a> {
53 pub gain: Option<FieldGain<'a>>,
56 pub birefringence: Option<Birefringence<'a>>,
58}
59
60impl std::fmt::Debug for PropagationOptions<'_> {
61 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
62 f.debug_struct("PropagationOptions")
63 .field("gain", &self.gain.map(|_| "FieldGain"))
64 .field("birefringence", &self.birefringence)
65 .finish()
66 }
67}
68
69struct Work {
71 omega: Vec<f64>,
73 p0: Vec<f64>,
75 p1: Vec<f64>,
77 phi: Vec<f64>,
79 lin: Vec<C64>,
81 half1: Vec<Vec<C64>>,
83 half2: Vec<Vec<C64>>,
84 gain: Vec<f64>,
85 h: Vec<Vec<C64>>,
88 f: Vec<Vec<C64>>,
89 c: Vec<Vec<C64>>,
90}
91
92impl ManakovSplitStep {
93 pub fn new(medium: Medium) -> ManakovSplitStep {
100 ManakovSplitStep {
101 medium,
102 rule: StepRule::FWM_DEFAULT,
103 nonlinear: NonlinearStep::Midpoint,
104 max_iterations: 10,
105 tolerance: 1e-5,
106 }
107 }
108
109 pub fn run(&self, x: &mut [C64], y: &mut [C64], length_km: f64, fs: f64) -> Result<usize> {
116 self.run_fields(&mut [x, y], length_km, fs, &PropagationOptions::default())
117 }
118
119 pub fn run_with_gain(
126 &self,
127 x: &mut [C64],
128 y: &mut [C64],
129 length_km: f64,
130 fs: f64,
131 gain: FieldGain<'_>,
132 ) -> Result<usize> {
133 self.run_fields(&mut [x, y], length_km, fs, &PropagationOptions { gain: Some(gain), birefringence: None })
134 }
135
136 pub fn run_fields(
151 &self,
152 fields: &mut [&mut [C64]],
153 length_km: f64,
154 fs: f64,
155 options: &PropagationOptions<'_>,
156 ) -> Result<usize> {
157 let nf = fields.len();
158 if nf == 0 || nf > 2 {
159 return Err(Error::Shape(format!("a propagation takes one or two fields, got {nf}")));
160 }
161 let n = fields[0].len();
162 if nf == 2 {
163 same_len(n, fields[1].len(), "the polarisations")?;
164 }
165 self.rule_parameter()?;
166 ensure(length_km.is_finite(), || format!("the length must be finite, got {length_km} km"))?;
167 if let Some(b) = &options.birefringence {
168 check_birefringence(b, nf, length_km)?;
169 }
170 if n == 0 || length_km <= 0.0 {
171 return Ok(0);
172 }
173 if let (StepRule::Fixed(_) | StepRule::Logarithmic(_), NonlinearStep::Midpoint) = (self.rule, self.nonlinear) {
174 return Ok(self.run_fixed_midpoint(fields, length_km, fs, options));
175 }
176 let p = fft::plan(n);
177 let zero = C64::new(0.0, 0.0);
178 let mut w = Work {
179 omega: fft::angular_frequencies(n, fs),
180 p0: vec![0.0; n],
181 p1: vec![0.0; n],
182 phi: vec![0.0; n],
183 lin: vec![zero; n],
184 half1: vec![vec![zero; n]; nf],
185 half2: vec![vec![zero; n]; nf],
186 gain: vec![0.0; n],
187 h: vec![vec![zero; n]; nf],
188 f: vec![vec![zero; n]; nf],
189 c: fields.iter().map(|u| u.to_vec()).collect(),
190 };
191 let k = if nf == 2 { 8.0 / 9.0 } else { 1.0 } * self.medium.gamma_w_km;
192 let bandwidth = match self.rule {
193 StepRule::Fwm { bandwidth_hz: Some(b), .. } => b,
194 StepRule::Fwm { bandwidth_hz: None, .. } => {
195 let views: Vec<&[C64]> = fields.iter().map(|u| &**u).collect();
196 signal_bandwidth(&views, fs)
197 }
198 _ => 0.0,
199 };
200 let mut z = 0.0;
201 let mut steps = 0;
202 let mut previous = None;
203 let mut section: Option<usize> = None;
205 while z < length_km {
206 let mut stretch_end = length_km;
207 if let Some(b) = &options.birefringence {
208 let current = ((z / b.section_km) * (1.0 + 1e-12)).floor() as usize;
209 let current = current.min(b.rotations.len() - 1);
210 if section != Some(current) {
211 let (a, rest) = fields.split_at_mut(1);
214 rotate(&b.rotations[current], a[0], rest[0]);
215 let (a, rest) = w.c.split_at_mut(1);
216 rotate(&b.rotations[current], &mut a[0], &mut rest[0]);
217 section = Some(current);
218 }
219 stretch_end = (((current + 1) as f64) * b.section_km).min(length_km);
220 }
221 total_power_of(fields, &mut w.p0);
222 total_power_of(&w.c, &mut w.p1);
223 trapezoid(&w.p0, &w.p1, k, &mut w.phi);
224 let h = self.step_length(&w.phi, z, stretch_end - z, length_km, previous, bandwidth)?;
225 previous = Some(h);
226 self.half_steps(&mut w, z, h, section, options);
227 for (o, u) in w.h.iter_mut().zip(fields.iter()) {
228 o.copy_from_slice(u);
229 }
230 for (u, half) in w.h.iter_mut().zip(&w.half1) {
231 p.forward(u);
232 mul(u, half);
233 p.inverse(u);
234 }
235 if self.nonlinear == NonlinearStep::Midpoint {
236 kerr_rotation(&mut w.h, k * h);
238 for (u, half) in w.h.iter_mut().zip(&w.half2) {
239 p.forward(u);
240 mul(u, half);
241 p.inverse(u);
242 }
243 for ((u, c), hh) in fields.iter_mut().zip(w.c.iter_mut()).zip(&w.h) {
244 u.copy_from_slice(hh);
245 c.copy_from_slice(hh);
246 }
247 } else {
248 for it in 0..self.max_iterations.max(1) {
249 for (o, v) in w.f.iter_mut().zip(&w.h) {
250 for ((a, b), ph) in o.iter_mut().zip(v).zip(&w.phi) {
251 *a = b * cis(ph * h);
252 }
253 }
254 for (u, half) in w.f.iter_mut().zip(&w.half2) {
255 p.forward(u);
256 mul(u, half);
257 p.inverse(u);
258 }
259 let diff: f64 =
260 w.f.iter().zip(&w.c).flat_map(|(a, b)| a.iter().zip(b)).map(|(u, v)| (u - v).norm_sqr()).sum();
261 let norm: f64 = w.c.iter().flatten().map(C64::norm_sqr).sum();
262 std::mem::swap(&mut w.c, &mut w.f);
263 if norm == 0.0 || (diff / norm).sqrt() < self.tolerance || it + 1 == self.max_iterations {
264 break;
265 }
266 total_power_of(&w.c, &mut w.p1);
267 trapezoid(&w.p0, &w.p1, k, &mut w.phi);
268 }
269 for (u, c) in fields.iter_mut().zip(&w.c) {
270 u.copy_from_slice(c);
271 }
272 }
273 z += h;
274 steps += 1;
275 }
276 Ok(steps)
277 }
278
279 fn run_fixed_midpoint(
285 &self,
286 fields: &mut [&mut [C64]],
287 length_km: f64,
288 fs: f64,
289 options: &PropagationOptions<'_>,
290 ) -> usize {
291 let (nf, n) = (fields.len(), fields[0].len());
292 let p = fft::plan(n);
293 let zero = C64::new(0.0, 0.0);
294 let mut w = Work {
295 omega: fft::angular_frequencies(n, fs),
296 p0: Vec::new(),
297 p1: Vec::new(),
298 phi: Vec::new(),
299 lin: vec![zero; n],
300 half1: vec![vec![zero; n]; nf],
301 half2: vec![vec![zero; n]; nf],
302 gain: vec![0.0; n],
303 h: fields.iter().map(|u| u.to_vec()).collect(),
304 f: Vec::new(),
305 c: Vec::new(),
306 };
307 let k = if nf == 2 { 8.0 / 9.0 } else { 1.0 } * self.medium.gamma_w_km;
308 let (n_stretches, stretch_km) = match &options.birefringence {
310 Some(b) => (b.rotations.len(), b.section_km),
311 None => (1, length_km),
312 };
313 for u in &mut w.h {
314 p.forward(u);
315 }
316 let mut steps = 0;
317 for s in 0..n_stretches {
318 let z_start = s as f64 * stretch_km;
319 let z_end = if s + 1 == n_stretches { length_km } else { (s + 1) as f64 * stretch_km };
320 if let Some(b) = &options.birefringence {
321 let (a, rest) = w.h.split_at_mut(1);
324 rotate(&b.rotations[s], &mut a[0], &mut rest[0]);
325 }
326 let section = options.birefringence.is_some().then_some(s);
327 let stretch = self.predetermined_steps(z_start, z_end, length_km);
328 let mut current = f64::NAN;
331 let mut z_log = z_start;
332 for (i, &h) in stretch.iter().enumerate() {
333 if h.to_bits() != current.to_bits() {
334 self.half_steps(&mut w, z_start, h, section, &PropagationOptions { gain: None, ..*options });
335 current = h;
336 }
337 let z = match self.rule {
339 StepRule::Logarithmic(_) => z_log,
340 _ => z_start + i as f64 * h,
341 };
342 z_log += h;
343 for (u, half) in w.h.iter_mut().zip(&w.half1) {
344 mul(u, half);
345 }
346 if let Some(g) = options.gain {
347 g(z, z + h / 2.0, &mut w.gain);
348 for u in &mut w.h {
349 scale(u, &w.gain);
350 }
351 }
352 for u in &mut w.h {
353 p.inverse(u);
354 }
355 kerr_rotation_scalar_order(&mut w.h, k, h);
356 for (u, half) in w.h.iter_mut().zip(&w.half2) {
357 p.forward(u);
358 mul(u, half);
359 }
360 if let Some(g) = options.gain {
361 g(z + h / 2.0, z + h, &mut w.gain);
362 for u in &mut w.h {
363 scale(u, &w.gain);
364 }
365 }
366 steps += 1;
367 }
368 }
369 for (u, spectrum) in fields.iter_mut().zip(w.h.iter_mut()) {
370 p.inverse(spectrum);
371 u.copy_from_slice(spectrum);
372 }
373 steps
374 }
375
376 fn half_steps(&self, w: &mut Work, z: f64, h: f64, section: Option<usize>, options: &PropagationOptions<'_>) {
379 self.medium.half_step_into(&w.omega, h, &mut w.lin);
380 let nf = w.half1.len();
381 let dgd = match (options.birefringence, section) {
382 (Some(b), Some(s)) => b.dgd_s_per_km[s],
383 _ => 0.0,
384 };
385 for i in 0..nf {
386 let sign = if i == 0 { -1.0 } else { 1.0 };
388 for ((a, b), (l, wk)) in w.half1[i].iter_mut().zip(w.half2[i].iter_mut()).zip(w.lin.iter().zip(&w.omega)) {
389 let d = if dgd == 0.0 { *l } else { l * cis(sign * 0.25 * wk * dgd * h) };
390 *a = d;
391 *b = d;
392 }
393 }
394 if let Some(g) = options.gain {
395 for (half, z0, z1) in [(&mut w.half1, z, z + h / 2.0), (&mut w.half2, z + h / 2.0, z + h)] {
396 g(z0, z1, &mut w.gain);
397 for v in half.iter_mut() {
398 for (o, b) in v.iter_mut().zip(&w.gain) {
399 *o *= b;
400 }
401 }
402 }
403 }
404 }
405
406 fn predetermined_steps(&self, z_start: f64, z_end: f64, length_km: f64) -> Vec<f64> {
411 match self.rule {
412 StepRule::Logarithmic(n) => {
413 let a = self.medium.alpha_np_km;
414 let eps = 1e-12 * length_km;
415 let mut z = vec![z_start];
416 z.extend(
417 (1..n)
418 .map(|k| logarithmic_boundary(a, length_km, n, k))
419 .filter(|&b| b > z_start + eps && b < z_end - eps),
420 );
421 z.push(z_end);
422 z.windows(2).map(|p| p[1] - p[0]).collect()
423 }
424 StepRule::Fixed(hf) => {
425 let (n, h) = equal_steps(z_end - z_start, hf);
426 vec![h; n]
427 }
428 _ => unreachable!("only the rules of predetermined steps"),
429 }
430 }
431
432 fn rule_parameter(&self) -> Result<f64> {
435 let (step, what) = match self.rule {
436 StepRule::Fixed(h) => (h, "fixed step (km)"),
437 StepRule::NonlinearPhase(r) => (r, "nonlinear rotation (rad)"),
438 StepRule::Logarithmic(n) => (n as f64, "number of logarithmic steps"),
439 StepRule::Fwm { phase_rad, bandwidth_hz, max_nl_phase_rad, .. } => {
440 if let Some(b) = bandwidth_hz {
441 ensure(b > 0.0 && b.is_finite(), || format!("the bandwidth must be positive, got {b} Hz"))?;
442 }
443 if let Some(r) = max_nl_phase_rad {
444 ensure(r > 0.0, || format!("the nonlinear rotation must be positive, got {r}"))?;
445 }
446 (phase_rad, "four-wave-mixing phase (rad)")
447 }
448 };
449 ensure(step > 0.0, || format!("the {what} must be positive, got {step}"))?;
450 Ok(step)
451 }
452
453 fn fwm_first_step(&self, phase_rad: f64, b: f64) -> f64 {
456 let w = 2.0 * std::f64::consts::PI * b;
457 let h = phase_rad / (self.medium.beta2_s2_km.abs() * w * w);
458 if h.is_finite() { h } else { f64::INFINITY }
459 }
460
461 pub(super) fn step_length(
469 &self,
470 phi: &[f64],
471 z: f64,
472 rest: f64,
473 length_km: f64,
474 previous: Option<f64>,
475 bandwidth: f64,
476 ) -> Result<f64> {
477 let h = match self.rule {
478 StepRule::Fixed(hf) => equal_steps(rest, hf).1,
479 StepRule::Logarithmic(n) => {
480 let a = self.medium.alpha_np_km;
481 let eps = 1e-12 * length_km;
482 let next = (1..=n)
483 .map(|k| logarithmic_boundary(a, length_km, n, k))
484 .find(|&b| b > z + eps)
485 .unwrap_or(length_km);
486 (next - z).min(rest)
487 }
488 StepRule::NonlinearPhase(rot) => {
489 let m = phi.iter().copied().fold(0.0, f64::max);
490 if m > 0.0 && rest >= rot / m { rot / m } else { rest }
491 }
492 StepRule::Fwm { phase_rad, update, max_nl_phase_rad, .. } => {
493 let a = self.medium.alpha_np_km;
494 let h = match previous {
495 None => self.fwm_first_step(phase_rad, bandwidth),
496 Some(hk) if a <= 0.0 => hk,
497 Some(hk) => match update {
498 StepUpdate::ConstantLocalError => hk * (a * hk / 3.0).exp(),
499 StepUpdate::NonlinearPhase => {
500 let l = (1.0 - (-a * hk).exp()) / a * (a * hk).exp();
502 if a * l < 1.0 { -(1.0 - a * l).ln() / a } else { f64::INFINITY }
503 }
504 },
505 };
506 let m = phi.iter().copied().fold(0.0, f64::max);
507 let nl = match max_nl_phase_rad {
508 Some(rot) if m > 0.0 => rot / m,
509 _ => f64::INFINITY,
510 };
511 h.min(nl).min(rest)
512 }
513 };
514 if z + h <= z || h.is_nan() {
517 return Err(Error::Numerical(format!("the split step vanishes at z = {z} km (step {h} km)")));
518 }
519 Ok(h)
520 }
521}
522
523fn equal_steps(rest: f64, hf: f64) -> (usize, f64) {
527 let ratio = rest / hf;
528 let whole = ratio.round();
529 if whole >= 1.0 && (whole - ratio).abs() <= 1e-9 * ratio {
530 (whole as usize, hf)
531 } else {
532 let n = ratio.ceil().max(1.0);
533 (n as usize, rest / n)
534 }
535}
536
537fn check_birefringence(b: &Birefringence<'_>, nf: usize, length_km: f64) -> Result<()> {
540 if nf != 2 {
541 return Err(Error::Shape("a birefringence needs two polarisations".into()));
542 }
543 let n = b.rotations.len();
544 let covers = b.section_km > 0.0 && ((n as f64) * b.section_km - length_km).abs() <= 1e-9 * length_km.max(1.0);
545 if n == 0 || b.dgd_s_per_km.len() != n || !covers {
546 return Err(Error::Shape(format!(
547 "{n} sections of {} km and {} DGDs for a length of {length_km} km",
548 b.section_km,
549 b.dgd_s_per_km.len()
550 )));
551 }
552 Ok(())
553}
554
555fn rotate(m: &[C64; 4], x: &mut [C64], y: &mut [C64]) {
557 for (u, v) in x.iter_mut().zip(y.iter_mut()) {
558 let (a, b) = (m[0] * *u + m[1] * *v, m[2] * *u + m[3] * *v);
559 *u = a;
560 *v = b;
561 }
562}
563
564fn total_power_of<T: AsRef<[C64]>>(fields: &[T], out: &mut [f64]) {
566 out.fill(0.0);
567 for f in fields {
568 for (o, v) in out.iter_mut().zip(f.as_ref()) {
569 *o += v.norm_sqr();
570 }
571 }
572}
573
574fn kerr_rotation_scalar_order(fields: &mut [Vec<C64>], k: f64, h: f64) {
578 match fields {
579 [x] => {
580 for v in x.iter_mut() {
581 *v *= cis(k * v.norm_sqr() * h);
582 }
583 }
584 [x, y] => manakov_rotation(x, y, k * h),
585 _ => {}
586 }
587}
588
589fn kerr_rotation(fields: &mut [Vec<C64>], k: f64) {
592 match fields {
593 [x] => {
594 for v in x.iter_mut() {
595 *v *= cis(k * v.norm_sqr());
596 }
597 }
598 [x, y] => manakov_rotation(x, y, k),
599 _ => {}
600 }
601}
602
603fn trapezoid(p0: &[f64], p1: &[f64], k: f64, phi: &mut [f64]) {
605 for (f, (u, v)) in phi.iter_mut().zip(p0.iter().zip(p1)) {
606 *f = k * (u + v) / 2.0;
607 }
608}