treams_core/special/
legendre.rs

1//! Associated Legendre functions of integer degree and the vector-wave angular functions
2//! `pi` and `tau`, with analytic argument derivatives. Real-degree Legendre values come
3//! from the Ferrers functions in `ferrers`.
4//!
5//! Upstream: `treams.special.lpmv`, `pi_fun` and `tau_fun` (`_waves.pyx`).
6
7use std::f64::consts::PI;
8
9use super::helicity_sign;
10use crate::{
11    Complex, Error, MAX_DEGREE, Result,
12    numerics::{
13        Jet,
14        broadcast::{self, element},
15        finite,
16        parallel::Parallel,
17    },
18};
19
20/// Elementwise angular evaluations and pullbacks run in parallel from this many elements.
21const PARALLEL: Parallel = Parallel::AtLeast(1024);
22
23/// Associated Legendre functions `F_l^m` after factoring out `sin(theta)^|m|`, visited
24/// for every `l` in `|m|..=max` in order from one run of the degree recurrence. Keeping
25/// the factor explicit avoids 0/0 and cancellation in angular functions near the poles.
26pub(crate) fn legendre_factors<const N: usize>(
27    max: i32,
28    m: i32,
29    z: Jet<N>,
30    mut visit: impl FnMut(Jet<N>),
31) {
32    let order = m.abs();
33    if max < order {
34        return;
35    }
36    let mut p = Jet::constant(1.0);
37    if m >= 0 {
38        for k in 1..=m {
39            p *= -f64::from(2 * k - 1);
40        }
41    } else {
42        for k in 1..=-m {
43            p *= 1.0 / f64::from(2 * k);
44        }
45    }
46    visit(p);
47    if max == order {
48        return;
49    }
50    let mut prev = p;
51    p *= (f64::from(2 * order + 1) / f64::from(order - m + 1)) * z;
52    visit(p);
53    for k in (order + 2)..=max {
54        let next =
55            (f64::from(2 * k - 1) * z * p - f64::from(k + m - 1) * prev) * (1.0 / f64::from(k - m));
56        prev = p;
57        p = next;
58        visit(p);
59    }
60}
61
62/// `F_l^m` and `F_(l-1)^m` of [`legendre_factors`], with `F_(l-1)^m = 0` for `l = |m|`
63/// and both zero for `|m| > l`.
64fn legendre_factor_pair<const N: usize>(l: i32, m: i32, z: Jet<N>) -> [Jet<N>; 2] {
65    let mut pair = [Jet::default(); 2];
66    legendre_factors(l, m, z, |p| pair = [p, pair[0]]);
67    pair
68}
69
70/// Associated Legendre polynomial `F_l^m` after factoring out `sin(theta)^|m|`.
71pub(crate) fn legendre_factor<const N: usize>(l: i32, m: i32, z: Jet<N>) -> Jet<N> {
72    let [value, _] = legendre_factor_pair(l, m, z);
73    value
74}
75
76/// `tau_l^m / sin(theta)^(|m|-1)` for `m != 0` from one Legendre recurrence, by DLMF
77/// 14.10.5: `sin(theta) tau = l cos(theta) P_l^m - (l + m) P_(l-1)^m`.
78fn tau_factor<const N: usize>(
79    l: i32,
80    m: i32,
81    [upper, lower]: [Jet<N>; 2],
82    cosine: Jet<N>,
83) -> Jet<N> {
84    f64::from(l) * cosine * upper - f64::from(l + m) * lower
85}
86
87/// The angular functions `[pi, tau]` from the polar cosine and sine, with the sine
88/// factor kept explicit at the poles.
89fn angular_jets<const N: usize>(l: i32, m: i32, cosine: Jet<N>, sine: Jet<N>) -> [Jet<N>; 2] {
90    if m == 0 {
91        // tau_l^0 = P_l^1.
92        return [Jet::default(), sine * legendre_factor(l, 1, cosine)];
93    }
94    let pair = legendre_factor_pair(l, m, cosine);
95    let power = sine.powi(m.abs() - 1);
96    [
97        f64::from(m) * power * pair[0],
98        power * tau_factor(l, m, pair, cosine),
99    ]
100}
101
102/// Angular factor coupling the vector spherical wave `(l, m)` of polarization
103/// `mode_pol` to a plane wave of polarization `pol` whose direction has the polar
104/// cosine and sine given: `tau -+ pi` for negative (0) or positive (1) helicity; in
105/// parity polarization `tau` for the mode's own polarization and `pi` for the other.
106/// The polarization module states this rule for the vector waves themselves.
107pub(crate) fn polarized_angular<const N: usize>(
108    l: i32,
109    m: i32,
110    [cosine, sine]: [Jet<N>; 2],
111    [mode_pol, pol]: [u8; 2],
112    helicity: bool,
113) -> Jet<N> {
114    let [pi, tau] = angular_jets(l, m, cosine, sine);
115    if helicity {
116        tau + helicity_sign(pol) * pi
117    } else if mode_pol == pol {
118        tau
119    } else {
120        pi
121    }
122}
123
124/// Associated Legendre polynomial or one of the two vector-wave angular functions.
125///
126/// Upstream: `treams.special.lpmv`, `pi_fun` and `tau_fun`.
127#[derive(Clone, Copy, Debug)]
128pub enum Angular {
129    /// Associated Legendre P, including its Condon-Shortley phase.
130    Legendre,
131    /// m P / sqrt(1-z^2), with analytic polar limits.
132    Pi,
133    /// Polar-angle derivative of P(cos(theta)).
134    Tau,
135}
136
137/// `sin(theta)^power = (1 - z^2)^(power / 2)` at the cosine `z`.
138fn sine_power<const N: usize>(z: Jet<N>, power: i32) -> Jet<N> {
139    let w = 1.0 - z * z;
140    if power % 2 == 0 {
141        w.powi(power / 2)
142    } else if power > 1 && w.value == Complex::default() {
143        // Odd powers >=3 have zero value and first derivative at the poles.
144        Jet::default()
145    } else {
146        w.powi(power / 2) * w.sqrt()
147    }
148}
149
150/// `F_l^m(z) sin(theta)^power`: the Legendre function for `power = |m|`, and `pi / m`
151/// for `power = |m| - 1`.
152fn angular_legendre<const N: usize>(l: i32, m: i32, z: Jet<N>, power: i32) -> Jet<N> {
153    if m.abs() > l {
154        return Jet::default();
155    }
156    legendre_factor(l, m, z) * sine_power(z, power)
157}
158
159/// The angular function `kind` of degree `l` and order `m` at `z`, with its derivative
160/// for `N = 1`, after checking the labels and the argument.
161#[allow(clippy::cast_possible_truncation)] // Integer labels are checked and bounded before conversion.
162fn angular_jet<const N: usize>(l: f64, m: f64, z: Complex, kind: Angular) -> Result<Jet<N>> {
163    if !finite(z)
164        || !l.is_finite()
165        || !m.is_finite()
166        || m.fract() != 0.0
167        || !(0.0..=f64::from(MAX_DEGREE)).contains(&l)
168    {
169        return Err(Error::InvalidInput(
170            // 128 is MAX_DEGREE.
171            "angular functions require 0 <= degree <= 128, integer order and finite argument"
172                .into(),
173        ));
174    }
175    if l.fract() != 0.0 && (!matches!(kind, Angular::Legendre) || z.im != 0.0) {
176        return Err(Error::InvalidInput(
177            "noninteger degrees require real Legendre arguments".into(),
178        ));
179    }
180    if m.abs() > l {
181        return Ok(Jet::default());
182    }
183    if l.fract() != 0.0 {
184        let (value, derivative) = if N == 0 {
185            super::ferrers::ferrers_real_degree::<false>(l, m as i32, z.re)?
186        } else {
187            super::ferrers::ferrers_real_degree::<true>(l, m as i32, z.re)?
188        };
189        return Ok(Jet::<N>::variable(z, 0).chain(value.into(), derivative.into()));
190    }
191    let (l, m) = (l as i32, m as i32);
192    let z = Jet::<N>::variable(z, 0);
193    let value = match kind {
194        Angular::Legendre => angular_legendre(l, m, z, m.abs()),
195        Angular::Pi if m == 0 => Jet::default(),
196        Angular::Pi => f64::from(m) * angular_legendre(l, m, z, m.abs() - 1),
197        Angular::Tau if m == 0 => angular_legendre(l, 1, z, 1),
198        Angular::Tau => {
199            tau_factor(l, m, legendre_factor_pair(l, m, z), z) * sine_power(z, m.abs() - 1)
200        }
201    };
202    if !value.finite() {
203        return Err(Error::SpecialFunction(
204            if N == 0 {
205                "non-finite angular result"
206            } else {
207                "angular argument derivative is undefined or non-finite"
208            }
209            .into(),
210        ));
211    }
212    Ok(value)
213}
214
215/// Angular function at a cosine argument; real Legendre values also accept real degrees.
216///
217/// Upstream: `treams.special.lpmv` (which takes `(order, degree, x)`), `pi_fun` and
218/// `tau_fun`. Differences: degrees outside `[0, 128]`, non-integer orders, and
219/// non-integer degrees with `pi`, `tau` or a complex argument give an error.
220pub fn angular_value(l: f64, m: f64, z: Complex, kind: Angular) -> Result<Complex> {
221    Ok(angular_jet::<0>(l, m, z, kind)?.value)
222}
223
224/// Associated Legendre function `P_degree^order(x)` of a real argument, as a real value.
225///
226/// Outside `[-1, 1]` the odd orders with `|order| <= degree` have imaginary values: this
227/// function rejects them, and [`angular_value`] returns them for a complex argument.
228/// Orders with `|order| > degree` give zero.
229///
230/// Upstream: `treams.special.lpmv` with real arguments, which takes `(order, degree, x)`.
231/// Differences: for an integer degree, `x` outside `[-1, 1]` and `0 < |order| <= degree`,
232/// treams returns NaN, while this function returns the real value for even orders and an
233/// error for odd orders. Degrees outside `[0, 128]`, non-integer orders, non-finite `x`
234/// and non-integer degrees with `x` outside `(-1, 1]` give an error, where treams returns
235/// a value or NaN.
236#[inline]
237pub fn lpmv_real(degree: f64, order: f64, x: f64) -> Result<f64> {
238    let value = angular_value(degree, order, x.into(), Angular::Legendre)?;
239    if value.im == 0.0 {
240        Ok(value.re)
241    } else {
242        Err(Error::InvalidInput(
243            "real lpmv requires -1 <= x <= 1 for odd orders |order| <= degree; \
244             pass complex z for the continuation"
245                .into(),
246        ))
247    }
248}
249
250/// What [`angular_array`] saves for its pullback: the degrees, orders, arguments and
251/// function. The pullback recomputes the local derivatives.
252#[derive(Debug)]
253pub struct AngularResidual {
254    degrees: Vec<f64>,
255    orders: Vec<f64>,
256    arguments: Vec<Complex>,
257    kind: Angular,
258    size: usize,
259}
260
261/// Broadcast angular values and a complex-argument residual; labels stay fixed.
262///
263/// Upstream: the functions of [`angular_value`] over arrays.
264pub fn angular_array(
265    degrees: Vec<f64>,
266    orders: Vec<f64>,
267    arguments: Vec<Complex>,
268    kind: Angular,
269) -> Result<(Vec<Complex>, AngularResidual)> {
270    let residual = AngularResidual::new(degrees, orders, arguments, kind)?;
271    let values = broadcast::map(residual.size, PARALLEL, |i| {
272        angular_value(
273            element(&residual.degrees, i),
274            element(&residual.orders, i),
275            element(&residual.arguments, i),
276            kind,
277        )
278    })?;
279    Ok((values, residual))
280}
281
282impl AngularResidual {
283    /// Keep the broadcast inputs needed by derivatives without evaluating values.
284    pub fn new(
285        degrees: Vec<f64>,
286        orders: Vec<f64>,
287        arguments: Vec<Complex>,
288        kind: Angular,
289    ) -> Result<Self> {
290        let size = broadcast::size(
291            &[degrees.len(), orders.len(), arguments.len()],
292            "angular arrays must have equal lengths or scalar inputs",
293        )?;
294        Ok(Self {
295            degrees,
296            orders,
297            arguments,
298            kind,
299            size,
300        })
301    }
302
303    fn argument_derivative(&self, i: usize) -> Result<Complex> {
304        let jet = angular_jet::<1>(
305            element(&self.degrees, i),
306            element(&self.orders, i),
307            element(&self.arguments, i),
308            self.kind,
309        )?;
310        Ok(jet.derivative.first().copied().unwrap_or_default())
311    }
312
313    /// Directional derivative of every value with respect to its complex argument.
314    pub fn pushforward(&self, tangents: [&[Complex]; 1]) -> Result<Vec<Complex>> {
315        broadcast::pushforward(
316            tangents,
317            self.size,
318            "angular tangent must be finite and broadcast to output",
319            PARALLEL,
320            |i, [tangent]| {
321                if tangent == Complex::default() {
322                    return Ok(tangent);
323                }
324                Ok(tangent * self.argument_derivative(i)?)
325            },
326        )
327    }
328
329    /// The gradient with respect to the arguments, given the gradient `cotangent` with
330    /// respect to the values. A zero cotangent skips derivatives that are singular at
331    /// the poles.
332    pub fn pullback(&self, cotangent: &[Complex]) -> Result<Vec<Complex>> {
333        let [gradient] = broadcast::pullback(
334            cotangent,
335            self.size,
336            "angular cotangent must be finite and match output",
337            [self.arguments.len()],
338            PARALLEL,
339            |i, &g| {
340                if g == Complex::default() {
341                    return Ok([g]);
342                }
343                Ok([g * self.argument_derivative(i)?.conj()])
344            },
345        )?;
346        Ok(gradient)
347    }
348}
349
350/// `[cos(theta), sin(theta)]` with their derivatives. The sine takes the principal
351/// branch (nonnegative real part) and is exactly zero at the poles, where `theta` is a
352/// floating-point multiple of `PI`.
353#[allow(clippy::float_cmp)] // Exact coordinate pole labels.
354pub(crate) fn polar_trig<const N: usize>(theta: Jet<N>) -> [Jet<N>; 2] {
355    let cosine = theta.chain(theta.value.cos(), -theta.value.sin());
356    let mut sine = theta.chain(theta.value.sin(), theta.value.cos());
357    // Associated Legendre functions use the principal sine factor. Computing
358    // its magnitude from sin(theta) preserves small angles lost in 1-cos²(theta).
359    if sine.value.re < 0.0 {
360        sine = -sine;
361    }
362    // Multipole translations amplify sin(PI)'s rounding residue. An exact
363    // floating-point multiple of the PI constant denotes the coordinate pole.
364    if theta.value.im == 0.0 && theta.value.re % PI == 0.0 {
365        sine.value = Complex::default();
366    }
367    [cosine, sine]
368}
treams_core/special/legendre.rs

treams_core/special/
legendre.rs
