treams_core/cw/
polar.rs

1//! Cylindrical translation coefficients in polar coordinates, of one order difference and
2//! of one mode pair, with analytic pullbacks.
3//!
4//! Upstream: `treams.special.tl_vcw`, `tl_vcw_r` and `treams.cw.translate`.
5#![allow(clippy::indexing_slicing)] // Fixed argument arrays.
6
7use crate::{
8    Complex, Error, MAX_DEGREE, Result,
9    numerics::{Jet, broadcast, finite, parallel::Parallel},
10    special::{MAX_ORDER, Radial, pol_index, radial_jet},
11};
12
13/// Cylindrical translation in `(k_rho*rho, phi, z, kz)` with a fixed order difference,
14/// the kernel of [`tl_vcw`]. The two axial labels must move together; unequal labels are
15/// uncoupled.
16///
17/// Upstream: `treams.special.tl_vcw` and `tl_vcw_r` at equal axial wavenumbers, with
18/// `order = m - mu`.
19pub fn polar_translation(order: i32, args: [Complex; 4], radial: Radial) -> Result<Complex> {
20    Ok(cylindrical_jet::<0>(order, args, radial, None)?.value)
21}
22
23/// The phase `exp(i (order phi + kz z))` of a cylindrical translation.
24fn cylindrical_phase<const N: usize>(order: i32, phi: Jet<N>, z: Jet<N>, kz: Jet<N>) -> Jet<N> {
25    (Complex::i() * (f64::from(order) * phi + kz * z)).exp()
26}
27
28/// [`polar_translation`] with the derivatives with respect to its `N` leading arguments.
29/// A forward (`N = 0`) uses `fixed_phase` when given, the phase shared by a radial sweep.
30fn cylindrical_jet<const N: usize>(
31    order: i32,
32    args: [Complex; 4],
33    radial: Radial,
34    fixed_phase: Option<Complex>,
35) -> Result<Jet<N>> {
36    cylindrical_seeded(
37        order,
38        std::array::from_fn(|i| Jet::<N>::variable(args[i], i)),
39        radial,
40        fixed_phase,
41    )
42}
43
44/// One kernel with value-only, coordinate or directional seeds.
45fn cylindrical_seeded<const N: usize>(
46    order: i32,
47    args: [Jet<N>; 4],
48    radial: Radial,
49    fixed_phase: Option<Complex>,
50) -> Result<Jet<N>> {
51    if order.unsigned_abs() > MAX_ORDER.unsigned_abs() || args.iter().any(|&v| !v.finite()) {
52        return Err(Error::InvalidInput(
53            // 256 is MAX_ORDER.
54            "require |order difference| <= 256 and finite arguments".into(),
55        ));
56    }
57    let [kr, phi, z, kz] = args;
58    let bessel = radial_jet(order, kr, radial, false)?;
59    let phase = if N == 0
60        && let Some(phase) = fixed_phase
61    {
62        Jet::constant(phase)
63    } else {
64        cylindrical_phase(order, phi, z, kz)
65    };
66    let value = bessel * phase;
67    if !value.finite() {
68        return Err(Error::NonFinite(
69            "non-finite cylindrical translation or derivative".into(),
70        ));
71    }
72    Ok(value)
73}
74
75/// The gradients of the radial argument, azimuth, axial position and common axial
76/// wavenumber.
77fn polar_translation_pullback(
78    order: i32,
79    args: [Complex; 4],
80    radial: Radial,
81    g: Complex,
82) -> Result<[Complex; 4]> {
83    if !finite(g) {
84        return Err(Error::InvalidInput("cotangent must be finite".into()));
85    }
86    if g == Complex::default() {
87        polar_translation(order, args, radial)?;
88        return Ok([Complex::default(); 4]);
89    }
90    Ok(cylindrical_jet::<4>(order, args, radial, None)?
91        .derivative
92        .map(|d| d.conj() * g))
93}
94
95/// What [`polar_translation_array`] saves for its pullback: the orders, the radial kind
96/// and the four arguments, each stored once when it is one value for all outputs.
97///
98/// The pullback keeps every output in index order and adds the cotangents of scalar
99/// arguments in that order, so its result does not depend on the thread count.
100#[derive(Debug)]
101pub struct PolarTranslationResidual {
102    orders: Vec<i32>,
103    arguments: [Vec<Complex>; 4],
104    radial: Radial,
105    size: usize,
106}
107
108impl PolarTranslationResidual {
109    fn element(&self, i: usize) -> (i32, [Complex; 4]) {
110        (
111            broadcast::element(&self.orders, i),
112            self.arguments.each_ref().map(|a| broadcast::element(a, i)),
113        )
114    }
115
116    /// Directional derivative of the four broadcast polar arguments.
117    pub fn pushforward(&self, tangents: [&[Complex]; 4]) -> Result<Vec<Complex>> {
118        broadcast::pushforward(
119            tangents,
120            self.size,
121            "translation tangents must be finite and match output",
122            Parallel::Chunked(64),
123            |i, tangent| {
124                let (order, args) = self.element(i);
125                let jets = std::array::from_fn(|a| Jet {
126                    value: args[a],
127                    derivative: [tangent[a]],
128                });
129                Ok(cylindrical_seeded(order, jets, self.radial, None)?.derivative[0])
130            },
131        )
132    }
133
134    /// The gradients of all four continuous arguments; an argument given as one value
135    /// for all outputs gets the sum of its gradients.
136    pub fn pullback(&self, cotangent: &[Complex]) -> Result<[Vec<Complex>; 4]> {
137        broadcast::pullback(
138            cotangent,
139            self.size,
140            "cotangent must be finite and match output",
141            self.arguments.each_ref().map(Vec::len),
142            Parallel::Chunked(64),
143            |i, &g| {
144                let (order, args) = self.element(i);
145                polar_translation_pullback(order, args, self.radial, g)
146            },
147        )
148    }
149}
150
151/// Evaluate a broadcast array of cylindrical polar translation coefficients; `orders`
152/// holds order differences, as in [`polar_translation`].
153///
154/// Upstream: `treams.special.tl_vcw` and `tl_vcw_r` at equal axial wavenumbers.
155pub fn polar_translation_array(
156    orders: Vec<i32>,
157    arguments: [Vec<Complex>; 4],
158    radial: Radial,
159) -> Result<(Vec<Complex>, PolarTranslationResidual)> {
160    let residual = PolarTranslationResidual::new(orders, arguments, radial)?;
161    let size = residual.size;
162    // A radial sweep shares its angular/axial phase. Keep all original inputs
163    // in the residual so the pullback still differentiates all four arguments.
164    let fixed_phase = if size > 1
165        && residual.orders.len() == 1
166        && residual.arguments[1..].iter().all(|a| a.len() == 1)
167    {
168        let (order, args) = residual.element(0);
169        let [_, phi, z, kz] = args.map(Jet::<0>::constant);
170        Some(cylindrical_phase(order, phi, z, kz).value)
171    } else {
172        None
173    };
174    let threshold = if radial == Radial::Regular { 512 } else { 64 };
175    let value = broadcast::map(size, Parallel::Chunked(threshold), |i| {
176        let (order, args) = residual.element(i);
177        Ok(cylindrical_jet::<0>(order, args, radial, fixed_phase)?.value)
178    })?;
179    Ok((value, residual))
180}
181
182impl PolarTranslationResidual {
183    /// Retain broadcast inputs and prepare static labels without evaluating translations.
184    pub fn new(orders: Vec<i32>, arguments: [Vec<Complex>; 4], radial: Radial) -> Result<Self> {
185        // 256 is MAX_ORDER.
186        let message = "require |order difference| <= 256 and equal lengths or scalar inputs";
187        let [a, b, c, d] = arguments.each_ref().map(Vec::len);
188        let size = broadcast::size(&[orders.len(), a, b, c, d], message)?;
189        if orders
190            .iter()
191            .any(|m| m.unsigned_abs() > MAX_ORDER.unsigned_abs())
192        {
193            return Err(Error::InvalidInput(message.into()));
194        }
195        Ok(Self {
196            orders,
197            arguments,
198            radial,
199            size,
200        })
201    }
202}
203
204/// Cylindrical translation coefficient of the mode pair `(qz, m)` to `(kz, mu)`.
205///
206/// `args` is `[k_rho rho, phi, z]`, the displacement in cylindrical coordinates. Equal
207/// axial wavenumbers give [`polar_translation`] of the order difference `m - mu`, and
208/// unequal ones give zero for any displacement. The wavenumbers must be finite and the
209/// orders satisfy `|mu|, |m| <= 128` ([`MAX_DEGREE`]); the orders are `i64` so that every
210/// integer label gets this check before its conversion. At equal wavenumbers the
211/// displacement must be finite too.
212///
213/// Upstream: `treams.special.tl_vcw` (singular) and `tl_vcw_r` (regular).
214/// Differences: orders above 128, non-finite axial wavenumbers, non-finite `k_rho rho`,
215/// `phi` or `z` at equal axial wavenumbers, singular waves at `k_rho rho = 0` and values
216/// that overflow give an error, where treams returns a value, infinity or NaN.
217#[inline]
218#[allow(clippy::float_cmp)] // Axial wave labels use exact equality, as in basis expansions.
219pub fn tl_vcw(
220    kz: f64,
221    mu: i64,
222    qz: f64,
223    m: i64,
224    args: [Complex; 3],
225    radial: Radial,
226) -> Result<Complex> {
227    let bound = u64::from(MAX_DEGREE.unsigned_abs());
228    if !kz.is_finite() || !qz.is_finite() || mu.unsigned_abs() > bound || m.unsigned_abs() > bound {
229        return Err(Error::InvalidInput(
230            // 128 is MAX_DEGREE.
231            "finite axial labels and |orders| <= 128 required".into(),
232        ));
233    }
234    if kz != qz {
235        return Ok(Complex::default());
236    }
237    let order = i32::try_from(m - mu)
238        .map_err(|_| Error::InvalidInput("invalid cylindrical order".into()))?;
239    polar_translation(order, [args[0], args[1], args[2], kz.into()], radial)
240}
241
242/// Cylindrical translation coefficient of one polarized mode pair.
243///
244/// The coefficient from `(qz, m, q)` to `(kz, mu, p)` at the displacement
245/// `[k_rho rho, phi, z]` is [`tl_vcw`] for equal polarizations and zero across
246/// polarizations; the polarizations must be 0 or 1. Singular waves give zero, as in
247/// treams, where `|k_rho rho|` and `|z|` are both below `1e-16`; this self term includes
248/// coinciding expansion centres.
249///
250/// Upstream: `treams.cw.translate`.
251/// Differences: at the singular self term treams returns zero for any polarization
252/// labels, while this function rejects labels other than 0 or 1 there too. Coefficients
253/// that reach [`tl_vcw`] share its differences.
254#[inline]
255pub fn translate(
256    kz: f64,
257    mu: i64,
258    p: i64,
259    qz: f64,
260    m: i64,
261    q: i64,
262    [kr, phi, z]: [Complex; 3],
263    radial: Radial,
264) -> Result<Complex> {
265    let (p, q) = (pol_index(p)?, pol_index(q)?);
266    if p != q || (radial == Radial::Singular && kr.norm() < 1e-16 && z.norm() < 1e-16) {
267        Ok(Complex::default())
268    } else {
269        tl_vcw(kz, mu, qz, m, [kr, phi, z], radial)
270    }
271}
272
273#[cfg(test)]
274mod tests {
275    //! The shared-phase radial sweep and polar coefficients against elementwise and
276    //! Cartesian evaluations. Identities of translations are in `properties/waves.rs`.
277
278    use super::{
279        Complex, Radial, polar_translation, polar_translation_array, polar_translation_pullback,
280    };
281    use crate::{
282        cw,
283        test_support::{DEFAULT_CASES, five_point, prop_assert_close, radial_kind},
284    };
285    use proptest::{prelude::*, test_runner::TestCaseError};
286
287    // The argument strategies below have seeds in proptest-regressions/cw/polar.txt;
288    // keep their shapes when changing the checks.
289
290    /// A shrunk failure of `cylindrical_cartesian_and_polar_adjoints`, written out so that
291    /// it does not depend on the strategy that the seeds in
292    /// `proptest-regressions/cw/polar.txt` replay: singular waves of order difference 9.
293    #[test]
294    fn recorded_cylindrical_polar_case() -> Result<(), TestCaseError> {
295        check_cylindrical(-4, 5, -0.222_807_017_666_592_56, Radial::Singular)
296    }
297
298    proptest! {
299        #![proptest_config(ProptestConfig::with_cases(DEFAULT_CASES))]
300        #[test]
301        fn cylindrical_radial_sweep_preserves_values_and_all_pullbacks(
302            order in -6_i32..=6, phi in -3.0_f64..3.0, axial in -0.5_f64..0.5,
303            kz in -0.7_f64..0.7, regular in any::<bool>(), count in 2_u16..160,
304        ) {
305            check_radial_sweep(order, phi, axial, kz, radial_kind(regular), count)?;
306        }
307
308        #[test]
309        fn cylindrical_cartesian_and_polar_adjoints(
310            mu in -5_i32..6, m in -5_i32..6, phi in -3.0_f64..3.0, regular in any::<bool>(),
311        ) {
312            check_cylindrical(mu, m, phi, radial_kind(regular))?;
313        }
314    }
315
316    /// A radial sweep with fixed angle, axial position and wavenumber (the shared-phase
317    /// fast path) equals elementwise evaluations, and its pullback sums the scalar
318    /// arguments' elementwise cotangents.
319    fn check_radial_sweep(
320        order: i32,
321        phi: f64,
322        axial: f64,
323        kz: f64,
324        radial: Radial,
325        count: u16,
326    ) -> Result<(), TestCaseError> {
327        let arguments = [
328            (0..count)
329                .map(|i| Complex::new(0.5 + f64::from(i) / 100.0, 0.2))
330                .collect(),
331            vec![Complex::new(phi, 0.1)],
332            vec![axial.into()],
333            vec![kz.into()],
334        ];
335        let count = usize::from(count);
336        let (values, residual) =
337            polar_translation_array(vec![order], arguments.clone(), radial).unwrap();
338        let cotangent = vec![Complex::new(0.2, -0.3); count];
339        let mut expected = [
340            vec![],
341            vec![Complex::default()],
342            vec![Complex::default()],
343            vec![Complex::default()],
344        ];
345        prop_assert_eq!(values.len(), count);
346        for (i, value) in values.iter().enumerate() {
347            let args = [
348                arguments[0][i],
349                arguments[1][0],
350                arguments[2][0],
351                arguments[3][0],
352            ];
353            let scalar = polar_translation(order, args, radial).unwrap();
354            prop_assert_close!(*value, scalar, 2e-14 * scalar.norm().max(1.0));
355            let gradient = polar_translation_pullback(order, args, radial, cotangent[i]).unwrap();
356            expected[0].push(gradient[0]);
357            for axis in 1..4 {
358                expected[axis][0] += gradient[axis];
359            }
360        }
361        let gradients = residual.pullback(&cotangent).unwrap();
362        for (actual, expected) in gradients.iter().zip(&expected) {
363            prop_assert_eq!(actual.len(), expected.len());
364            for (actual, expected) in actual.iter().zip(expected) {
365                prop_assert_close!(*actual, *expected, 2e-13 * expected.norm().max(1.0));
366            }
367        }
368        Ok(())
369    }
370
371    /// Polar cylindrical coefficients equal Cartesian translations; the azimuth, axial
372    /// position and axial wavenumber derivatives are exactly `i n T`, `i kz T` and
373    /// `i z T` up to rounding relative to `T` (and to `n T` for the azimuth), and the
374    /// radial one matches a complex central difference.
375    fn check_cylindrical(mu: i32, m: i32, phi: f64, radial: Radial) -> Result<(), TestCaseError> {
376        let args = [Complex::new(1.3, 0.1), phi.into(), 0.3.into(), 0.2.into()];
377        let k = (args[0] * args[0] + args[3] * args[3]).sqrt();
378        let mode = |m| cw::Mode { kz: 0.2, m, pol: 0 };
379        let position = [phi.cos(), phi.sin(), 0.3];
380        let expected = cw::cartesian_translation(mode(mu), mode(m), k, position, radial).unwrap();
381        let value = polar_translation(m - mu, args, radial).unwrap();
382        prop_assert_close!(value, expected.value, 1e-10 * (1.0 + value.norm()));
383        let derivative = polar_translation_pullback(m - mu, args, radial, Complex::new(1.0, 0.0))
384            .unwrap()
385            .map(|g| g.conj());
386        let tolerance = 1e-14 * value.norm();
387        let i = Complex::i();
388        let order = f64::from(m - mu);
389        prop_assert_close!(
390            derivative[1],
391            i * order * value,
392            (1.0 + order.abs()) * tolerance
393        );
394        prop_assert_close!(derivative[2], i * args[3] * value, tolerance);
395        prop_assert_close!(derivative[3], i * args[2] * value, tolerance);
396        let direction = Complex::new(0.2, 0.1);
397        let numeric = five_point(1e-4, |t| {
398            let mut shifted = args;
399            shifted[0] += t * direction;
400            polar_translation(m - mu, shifted, radial).unwrap()
401        });
402        // Compare the full complex derivative: a real projection may nearly cancel
403        // while the high-order outgoing value is large.
404        let analytic = derivative[0] * direction;
405        prop_assert_close!(analytic, numeric, 1e-8 * (1.0 + analytic.norm()));
406        Ok(())
407    }
408}
treams_core/cw/polar.rs

treams_core/cw/
polar.rs
