treams_core/sw/
coupling.rs

1//! Degree terms `(p, weight)` of the spherical translation coefficients, from the
2//! Wigner 3j couplings of a source and a destination mode.
3//!
4//! Upstream: `treams.special._tl_vsw_helper` and the degree sums of `tl_vsw_A` and
5//! `tl_vsw_B`.
6#![allow(clippy::indexing_slicing)] // Fixed indices of the two Wigner 3j rows.
7
8use super::Mode;
9use crate::Complex;
10use crate::special::{Wigner3jRow, helicity_sign};
11
12/// Degree-`p` coefficient of the coupling of source `(l, m)` to destination
13/// `(lambda, -mu)` with `q = p - cross`. Zero unless both Wigner 3j symbols it multiplies,
14/// `(l lambda p; m mu -m-mu)` from `rows[0] = Wigner3jRow::new(l, lambda, m, mu)` and
15/// `(l lambda q; 0 0 0)` from `rows[1] = Wigner3jRow::new(l, lambda, 0, 0)`, can be
16/// nonzero.
17///
18/// Upstream: `treams.special._tl_vsw_helper`.
19fn tl_vsw_term(
20    l: i32,
21    m: i32,
22    lambda: i32,
23    mu: i32,
24    p: i32,
25    q: i32,
26    rows: [&Wigner3jRow; 2],
27) -> Complex {
28    if p < (m + mu).abs().max((l - lambda).abs())
29        || p > l + lambda
30        || q < (l - lambda).abs()
31        || q > l + lambda
32        || (q + l + lambda) % 2 != 0
33    {
34        return Complex::default();
35    }
36    f64::from(2 * p + 1)
37        * Complex::i().powi(lambda - l + p)
38        * crate::special::factorial_ratio_sqrt(p, m + mu)
39        * rows[0].get(p)
40        * rows[1].get(q)
41}
42
43/// Which of the two coupling kinds of a [`Coupling`] a polarization pair uses.
44#[derive(Clone, Copy, Debug, PartialEq, Eq)]
45pub(crate) struct Kinds {
46    /// The `tl_vsw_A` terms.
47    same: bool,
48    /// The `tl_vsw_B` terms.
49    cross: bool,
50}
51
52impl Kinds {
53    /// Both kinds, for plans that serve every polarization pair.
54    pub(crate) const BOTH: Self = Self {
55        same: true,
56        cross: true,
57    };
58
59    /// The kinds that couple source polarization `from` to destination polarization `to`.
60    /// The parity basis couples equal polarizations through `same` and opposite ones
61    /// through `cross`; the helicity basis couples equal helicities through both.
62    const fn of(to: u8, from: u8, helicity: bool) -> Self {
63        let equal = to == from;
64        Self {
65            same: equal,
66            cross: if helicity { equal } else { !equal },
67        }
68    }
69}
70
71/// The degree terms `(p, weight)` coupling a source `(l, m)` to a destination
72/// `(lambda, mu)`, for all four polarization pairs. The parity basis couples equal
73/// polarizations through `same` and opposite ones through `cross`; the helicity basis
74/// couples equal helicities through both, with `cross` signed by `2 pol - 1`.
75///
76/// `same` holds the terms of `treams.special.tl_vsw_A`, `cross` those of `tl_vsw_B`.
77#[derive(Debug)]
78pub(crate) struct Coupling {
79    same: Vec<(i32, Complex)>,
80    cross: Vec<(i32, Complex)>,
81}
82
83impl Coupling {
84    /// Evaluates the `same` and `cross` terms selected by `kinds`; `zeros` is
85    /// `Wigner3jRow::new(l, lambda, 0, 0)`, shared by all orders.
86    pub(crate) fn new(
87        (l, m): (i32, i32),
88        (lambda, mu): (i32, i32),
89        zeros: &Wigner3jRow,
90        kinds: Kinds,
91    ) -> Self {
92        let orders = Wigner3jRow::new(l, lambda, m, -mu);
93        let rows = [&orders, zeros];
94        let sign = if m % 2 == 0 { 1.0 } else { -1.0 };
95        let pref = 0.5
96            * sign
97            * (f64::from((2 * l + 1) * (2 * lambda + 1))
98                / (f64::from(l * (l + 1)) * f64::from(lambda * (lambda + 1))))
99            .sqrt();
100        let [same, cross] = [false, true].map(|cross| {
101            if !(if cross { kinds.cross } else { kinds.same }) {
102                return Vec::new();
103            }
104            degrees(l, m, lambda, mu, cross)
105                .filter_map(|p| {
106                    let factor = if cross {
107                        (f64::from(l + lambda + 1 + p)
108                            * f64::from(l + lambda + 1 - p)
109                            * f64::from(p + lambda - l)
110                            * f64::from(p - lambda + l))
111                        .sqrt()
112                    } else {
113                        f64::from(l * (l + 1) + lambda * (lambda + 1) - p * (p + 1))
114                    };
115                    let weight = pref
116                        * factor
117                        * tl_vsw_term(l, m, lambda, -mu, p, p - i32::from(cross), rows);
118                    (weight.norm_sqr() > 0.0).then_some((p, weight))
119                })
120                .collect()
121        });
122        Self { same, cross }
123    }
124
125    /// Terms `(p, weight)` from source polarization `from` to destination polarization
126    /// `to`, with same-polarization terms first.
127    pub(crate) fn terms(
128        &self,
129        to: u8,
130        from: u8,
131        helicity: bool,
132    ) -> impl Iterator<Item = (i32, Complex)> + '_ {
133        let kinds = Kinds::of(to, from, helicity);
134        let sign = if helicity { helicity_sign(from) } else { 1.0 };
135        let same = if kinds.same { &self.same[..] } else { &[] };
136        let cross = if kinds.cross { &self.cross[..] } else { &[] };
137        same.iter()
138            .copied()
139            .chain(cross.iter().map(move |&(p, weight)| (p, sign * weight)))
140    }
141}
142
143/// Degrees `p` that `tl_vsw_term` admits for source `(l, m)` and destination
144/// `(lambda, mu)`, descending in steps of two. `formal/Formal/SelectionRules.lean` proves
145/// this range.
146pub(crate) fn degrees(
147    l: i32,
148    m: i32,
149    lambda: i32,
150    mu: i32,
151    cross: bool,
152) -> impl ExactSizeIterator<Item = i32> {
153    let start = l + lambda - i32::from(cross);
154    let end = (lambda - l)
155        .abs()
156        .saturating_add(i32::from(cross))
157        .max((m - mu).abs());
158    // `(end..=start).rev().step_by(2)`, whose length is known.
159    (0..(start - end + 2).max(0) / 2).map(move |k| start - 2 * k)
160}
161
162/// The most [`terms`] from `from` to `to`: one per admitted degree of each selected
163/// kind. Terms of zero weight are left out, so there can be fewer.
164pub(crate) fn term_bound(to: Mode, from: Mode, helicity: bool) -> usize {
165    let Kinds { same, cross } = Kinds::of(to.pol, from.pol, helicity);
166    let count = |kind| degrees(from.l, from.m, to.l, to.m, kind).len();
167    usize::from(same) * count(false) + usize::from(cross) * count(true)
168}
169
170/// Terms `(p, m - mu, weight)` of the translation from source `(l, m)` to destination
171/// `(lambda, mu)`: `sum weight z_p(k r) P_p^(m - mu) e^(i (m - mu) phi)`.
172pub(crate) fn terms(to: Mode, from: Mode, helicity: bool) -> Vec<(i32, i32, Complex)> {
173    let kinds = Kinds::of(to.pol, from.pol, helicity);
174    if !kinds.same && !kinds.cross {
175        return Vec::new();
176    }
177    let zeros = Wigner3jRow::new(from.l, to.l, 0, 0);
178    Coupling::new((from.l, from.m), (to.l, to.m), &zeros, kinds)
179        .terms(to.pol, from.pol, helicity)
180        .map(|(p, weight)| (p, from.m - to.m, weight))
181        .collect()
182}
183
184#[cfg(test)]
185mod tests {
186    //! The degree selection rules of the coupling terms against the Lean model.
187    //! Identities of the translations themselves are in `properties/waves.rs`.
188
189    use super::{Complex, Mode, Wigner3jRow, degrees, term_bound, terms, tl_vsw_term};
190    use crate::test_support::table;
191
192    /// The cases of `formal/Golden.lean`'s `degrees`, in its order: `[l, m, lambda, mu, cross]`.
193    fn lean_degree_cases() -> Vec<[i32; 5]> {
194        (1..=5)
195            .flat_map(|l| (1..=5).map(move |lambda| (l, lambda)))
196            .flat_map(|(l, lambda)| (-l..=l).map(move |m| (l, lambda, m)))
197            .flat_map(|(l, lambda, m)| (-lambda..=lambda).map(move |mu| (l, lambda, m, mu)))
198            .flat_map(|(l, lambda, m, mu)| [0, 1].map(|cross| [l, m, lambda, mu, cross]))
199            .collect()
200    }
201
202    #[test]
203    fn degrees_match_lean_model() {
204        // `just formal` keeps this file equal to `Treams.SelectionRules.termDegrees`.
205        let cases = table::<i32, i32>(include_str!(concat!(
206            env!("CARGO_MANIFEST_DIR"),
207            "/../../formal/golden/degrees.txt"
208        )));
209        let grid = lean_degree_cases();
210        assert_eq!(cases.len(), grid.len());
211        for ((key, expected), case) in cases.iter().zip(grid) {
212            assert_eq!(key, &case);
213            let [l, m, lambda, mu, cross] = case;
214            let actual: Vec<_> = degrees(l, m, lambda, mu, cross == 1).collect();
215            assert_eq!(&actual, expected, "{key:?}");
216        }
217    }
218
219    #[test]
220    fn skipped_degrees_have_zero_coefficients() {
221        // Rust counterpart of `SelectionRules.skipped_degree_vanishes`: `terms` evaluates
222        // `tl_vsw_term(l, m, lambda, -mu, p, p - cross, rows)` only at `degrees`, so every
223        // other degree must be rejected by `tl_vsw_term` itself.
224        for [l, m, lambda, mu, cross] in lean_degree_cases() {
225            let visited: Vec<_> = degrees(l, m, lambda, mu, cross == 1).collect();
226            let rows = [
227                &Wigner3jRow::new(l, lambda, m, -mu),
228                &Wigner3jRow::new(l, lambda, 0, 0),
229            ];
230            for p in (0..=l + lambda + 1).filter(|p| !visited.contains(p)) {
231                let coefficient = tl_vsw_term(l, m, lambda, -mu, p, p - cross, rows);
232                assert_eq!(
233                    coefficient,
234                    Complex::default(),
235                    "{l} {m} {lambda} {mu} {cross}: {p}"
236                );
237            }
238        }
239    }
240
241    #[test]
242    fn terms_use_admitted_degrees_and_table_indices() {
243        // Same- and cross-polarization couplings in both bases (`pol` only flips signs).
244        let polarizations = [(true, 1, 1), (false, 0, 0), (false, 0, 1)];
245        for [l, m, lambda, mu, _] in lean_degree_cases().into_iter().step_by(2) {
246            for (helicity, from_pol, to_pol) in polarizations {
247                let from = Mode {
248                    l,
249                    m,
250                    pol: from_pol,
251                };
252                let to = Mode {
253                    l: lambda,
254                    m: mu,
255                    pol: to_pol,
256                };
257                let admitted = |p, cross| degrees(l, m, lambda, mu, cross).any(|q| q == p);
258                let terms = terms(to, from, helicity);
259                assert!(terms.len() <= term_bound(to, from, helicity));
260                for (p, order, _) in terms {
261                    let case = format!("{from:?} -> {to:?}, helicity {helicity}: {p} {order}");
262                    assert_eq!(order, m - mu, "{case}");
263                    let expected = if helicity {
264                        admitted(p, false) || admitted(p, true)
265                    } else {
266                        admitted(p, from_pol != to_pol)
267                    };
268                    assert!(expected, "{case}");
269                    // `Harmonics.table_index`: `p * p + p + order` lies in the plan table.
270                    assert!(order.abs() <= p && p <= l + lambda, "{case}");
271                }
272            }
273        }
274    }
275}
treams_core/sw/coupling.rs

treams_core/sw/
coupling.rs
