treams_core/cw/
to_sw.rs

1//! Regular cylindrical-to-spherical conversion, including displaced positions.
2//!
3//! Upstream: `treams.cw.to_sw` and the `treams.Expand` operator from a cylindrical to a
4//! spherical basis.
5#![allow(clippy::indexing_slicing)] // Validated bases, mode polarizations and matrix dimensions.
6
7use crate::{
8    Complex, Error, Result,
9    basis::{ExpansionGradient, ModeLabel, add_pair_gradient, validate_wavenumbers},
10    cw,
11    numerics::{Jet, finite, parallel::try_fold_ordered},
12    special::{Radial, polarized_angular},
13    sw::{Basis, CartesianTranslation, Mode},
14};
15use nalgebra::DMatrix;
16use rayon::prelude::*;
17use std::collections::HashMap;
18
19/// The coefficient of the spherical mode `mode` in a regular cylindrical wave of
20/// polarization `pol` and axial wavenumber `kz` at one position, as a jet in `kz` and the
21/// wavenumber `wave`: `i^(l - m) sqrt(4 pi (2l + 1) / (l (l + 1))) sqrt((l - m)! / (l + m)!)`
22/// times the polarized angular function at `cos theta = kz / k`.
23///
24/// Upstream: the coefficient of `treams.cw.to_sw` for `m` equal on both sides.
25pub(crate) fn coefficient<const N: usize>(
26    mode: Mode,
27    kz: Jet<N>,
28    wave: Jet<N>,
29    pol: u8,
30    helicity: bool,
31) -> Result<Jet<N>> {
32    let cosine = kz / wave;
33    let sine = (1.0 - cosine * cosine).sqrt();
34    if N != 0 && sine.value == Complex::default() {
35        return Err(Error::InvalidInput(
36            "cylindrical conversion cutoff requires a limiting formulation".into(),
37        ));
38    }
39    let angular = polarized_angular(mode.l, mode.m, [cosine, sine], [mode.pol, pol], helicity);
40    let (l, m) = (mode.l, mode.m);
41    let normalization =
42        (4.0 * std::f64::consts::PI * f64::from(2 * l + 1) / f64::from(l * (l + 1))).sqrt()
43            * crate::special::factorial_ratio_sqrt(l, m);
44    Ok(Complex::i().powi(l - m) * normalization * angular)
45}
46
47/// The entries of one source column. Angular coefficients depend only on the
48/// destination mode label and translations only on the destination position and order,
49/// so each is evaluated once per column and shared by the modes of every destination
50/// position, degree and polarization.
51struct Column<'a, const N: usize> {
52    from: cw::Mode,
53    k: Complex,
54    helicity: bool,
55    destination: &'a [[f64; 3]],
56    source: [f64; 3],
57    coefficients: HashMap<Mode, Jet<N>>,
58    translations: HashMap<(usize, i32), cw::CartesianTranslation>,
59}
60
61impl<'a, const N: usize> Column<'a, N> {
62    fn new(
63        destination: &'a [[f64; 3]],
64        source: [f64; 3],
65        from: cw::Mode,
66        k: Complex,
67        helicity: bool,
68    ) -> Self {
69        Self {
70            from,
71            k,
72            helicity,
73            destination,
74            source,
75            coefficients: HashMap::new(),
76            translations: HashMap::new(),
77        }
78    }
79
80    /// The entry of destination mode `to` at position `p`, with its displacement and
81    /// wavenumber derivatives.
82    #[allow(clippy::float_cmp)] // Exact on-axis displacements determine the angular selection rule.
83    fn entry(&mut self, p: usize, to: Mode) -> Result<CartesianTranslation> {
84        let from = self.from;
85        if self.helicity && to.pol != from.pol {
86            return Ok(CartesianTranslation::default());
87        }
88        let r: [f64; 3] = std::array::from_fn(|a| self.destination[p][a] - self.source[a]);
89        // On axis only equal azimuthal orders survive; first spatial derivatives
90        // additionally couple adjacent orders. Avoid evaluating all zero entries.
91        if r[0] == 0.0 && r[1] == 0.0 && (to.m - from.m).unsigned_abs() > u32::from(N != 0) {
92            return Ok(CartesianTranslation::default());
93        }
94        let coefficient = if let Some(&coefficient) = self.coefficients.get(&to) {
95            coefficient
96        } else {
97            let wave = Jet::<N>::variable(self.k, 0);
98            let coefficient =
99                coefficient(to, Jet::constant(from.kz), wave, from.pol, self.helicity)?;
100            self.coefficients.insert(to, coefficient);
101            coefficient
102        };
103        let translation = if let Some(&translation) = self.translations.get(&(p, to.m)) {
104            translation
105        } else {
106            let to = cw::Mode { m: to.m, ..from };
107            let translation = cw::cartesian_translation(to, from, self.k, r, Radial::Regular)?;
108            self.translations.insert((p, to.m), translation);
109            translation
110        };
111        Ok(CartesianTranslation {
112            value: coefficient.value * translation.value,
113            position: translation.position.map(|g| coefficient.value * g),
114            k: coefficient.derivative.first().copied().unwrap_or_default() * translation.value
115                + coefficient.value * translation.k,
116        })
117    }
118}
119
120/// Cylindrical-to-spherical coefficient of a sphere and an axis at one position, with
121/// fixed mode labels.
122///
123/// Upstream: `treams.cw.to_sw`.
124pub fn to_sw(to: Mode, from: cw::Mode, k: Complex, helicity: bool) -> Result<Complex> {
125    to.validate()?;
126    from.validate()?;
127    if !finite(k) || k == Complex::default() {
128        return Err(Error::InvalidInput(
129            "finite nonzero wavenumber required".into(),
130        ));
131    }
132    if to.m != from.m || (helicity && to.pol != from.pol) {
133        return Ok(Complex::default());
134    }
135    Ok(coefficient(
136        to,
137        Jet::<0>::constant(from.kz),
138        Jet::constant(k),
139        from.pol,
140        helicity,
141    )?
142    .value)
143}
144
145/// What [`to_sw_matrix`] saves for its pullback: the bases, the wavenumbers and the
146/// polarization convention. The pullback recomputes the derivatives of every entry.
147///
148/// The pullback adds the per-column gradients in chunks of consecutive columns fixed by
149/// the column count, and the chunk sums in chunk order
150/// (`numerics::parallel::try_fold_ordered`), so the thread count does not change them.
151#[derive(Debug)]
152pub struct ToSwResidual {
153    destination: Basis,
154    source: cw::Basis,
155    ks: [Complex; 2],
156    helicity: bool,
157}
158
159/// Expand regular cylindrical waves around arbitrary spherical expansion positions: the
160/// matrix maps cylindrical input amplitudes to regular spherical amplitudes.
161///
162/// Upstream: `treams.operators.expand` from a cylindrical to a spherical basis. Differences:
163/// treams pairs each sphere only with the axis of the same position index and ignores the
164/// positions; this matrix couples every sphere with every axis through the translation
165/// between them.
166pub fn to_sw_matrix(
167    destination: Basis,
168    source: cw::Basis,
169    ks: [Complex; 2],
170    helicity: bool,
171) -> Result<(DMatrix<Complex>, ToSwResidual)> {
172    destination.validate()?;
173    source.validate()?;
174    validate_wavenumbers(ks, helicity, false)?;
175    let mut value = crate::numerics::zeros(destination.modes.len(), source.modes.len())?;
176    crate::threads::install(|| {
177        value
178            .as_mut_slice()
179            .par_chunks_mut(destination.modes.len())
180            .enumerate()
181            .try_for_each(|(j, column)| -> Result<()> {
182                let (q, from) = source.modes[j];
183                let k = ks[usize::from(from.pol)];
184                let mut entries = Column::<0>::new(
185                    &destination.positions,
186                    source.positions[q],
187                    from,
188                    k,
189                    helicity,
190                );
191                for (value, &(p, to)) in column.iter_mut().zip(&destination.modes) {
192                    *value = entries.entry(p, to)?.value;
193                }
194                Ok(())
195            })
196    })?;
197    Ok((
198        value,
199        ToSwResidual {
200            destination,
201            source,
202            ks,
203            helicity,
204        },
205    ))
206}
207
208impl ToSwResidual {
209    /// Saved-state size from the spherical destination and cylindrical source shapes.
210    pub fn state_size(
211        destination_modes: usize,
212        destination_positions: usize,
213        source_modes: usize,
214        source_positions: usize,
215    ) -> Result<usize> {
216        let destination = crate::saved::sw_basis_size(destination_modes, destination_positions)?;
217        let source = crate::saved::cw_basis_size(source_modes, source_positions)?;
218        destination
219            .checked_add(source)
220            .and_then(|n| n.checked_add(33))
221            .ok_or_else(crate::saved::invalid)
222    }
223    /// Spherical output and cylindrical input mode counts.
224    #[must_use]
225    pub fn shape(&self) -> (usize, usize) {
226        (self.destination.modes.len(), self.source.modes.len())
227    }
228
229    /// Directional derivative of the displaced conversion, with fixed axial labels.
230    pub fn pushforward(
231        &self,
232        destination: &[[f64; 3]],
233        source: &[[f64; 3]],
234        ks: [Complex; 2],
235    ) -> Result<DMatrix<Complex>> {
236        crate::basis::validate_expansion_tangent(
237            destination,
238            source,
239            ks,
240            (
241                self.destination.positions.len(),
242                self.source.positions.len(),
243            ),
244        )?;
245        let mut value = crate::numerics::zeros(self.shape().0, self.shape().1)?;
246        crate::threads::install(|| {
247            value
248                .as_mut_slice()
249                .par_chunks_mut(self.shape().0)
250                .enumerate()
251                .try_for_each(|(j, column)| -> Result<()> {
252                    let (q, from) = self.source.modes[j];
253                    let pol = usize::from(from.pol);
254                    let mut entries = Column::<1>::new(
255                        &self.destination.positions,
256                        self.source.positions[q],
257                        from,
258                        self.ks[pol],
259                        self.helicity,
260                    );
261                    for (out, &(p, to)) in column.iter_mut().zip(&self.destination.modes) {
262                        let entry = entries.entry(p, to)?;
263                        *out = crate::basis::pair_tangent(
264                            destination,
265                            source,
266                            ks,
267                            [p, q, pol],
268                            entry.position,
269                            entry.k,
270                        );
271                    }
272                    Ok(())
273                })
274        })?;
275        Ok(value)
276    }
277
278    /// Position and complex wavenumber gradients. Axial mode labels stay fixed.
279    pub fn pullback(&self, cotangent: &DMatrix<Complex>) -> Result<ExpansionGradient> {
280        if cotangent.shape() != self.shape() || cotangent.iter().any(|&g| !finite(g)) {
281            return Err(Error::InvalidInput("invalid conversion cotangent".into()));
282        }
283        try_fold_ordered(
284            self.source.modes.iter().collect(),
285            true,
286            || {
287                ExpansionGradient::zeros(
288                    self.destination.positions.len(),
289                    self.source.positions.len(),
290                )
291            },
292            |mut result, j, &(q, from)| -> Result<_> {
293                let pol = usize::from(from.pol);
294                let (destination, source) = (&self.destination.positions, self.source.positions[q]);
295                let mut entries =
296                    Column::<1>::new(destination, source, from, self.ks[pol], self.helicity);
297                for (i, &(p, to)) in self.destination.modes.iter().enumerate() {
298                    let cot = cotangent[(i, j)];
299                    if cot == Complex::default() {
300                        continue;
301                    }
302                    let wave = entries.entry(p, to)?;
303                    add_pair_gradient(&mut result, [p, q, pol], cot, wave.position, wave.k);
304                }
305                Ok(result)
306            },
307            |mut total, partial| {
308                total.accumulate(&partial);
309                total
310            },
311        )
312    }
313}
314
315impl crate::saved::SavedState for ToSwResidual {
316    fn save_state(&self) -> Result<Vec<u8>> {
317        let mut writer = crate::saved::Writer::new(Self::state_size(
318            self.destination.modes.len(),
319            self.destination.positions.len(),
320            self.source.modes.len(),
321            self.source.positions.len(),
322        )?);
323        crate::saved::write_sw_basis(&mut writer, &self.destination);
324        crate::saved::write_cw_basis(&mut writer, &self.source);
325        for k in self.ks {
326            writer.complex(k);
327        }
328        writer.byte(u8::from(self.helicity));
329        Ok(writer.finish())
330    }
331
332    fn from_state(bytes: &[u8]) -> Result<Self> {
333        let mut shape = crate::saved::Reader::new(bytes);
334        let (dm, dp) = crate::saved::read_sw_basis_dimensions(&mut shape)?;
335        let (sm, sp) = crate::saved::read_cw_basis_dimensions(&mut shape)?;
336        if bytes.len() != Self::state_size(dm, dp, sm, sp)? {
337            return Err(crate::saved::invalid());
338        }
339        let mut reader = crate::saved::Reader::new(bytes);
340        let destination = crate::saved::read_sw_basis(&mut reader)?;
341        let source = crate::saved::read_cw_basis(&mut reader)?;
342        let ks = [reader.complex()?, reader.complex()?];
343        let helicity = match reader.byte()? {
344            0 => false,
345            1 => true,
346            _ => return Err(crate::saved::invalid()),
347        };
348        reader.finish()?;
349        validate_wavenumbers(ks, helicity, false)?;
350        Ok(Self {
351            destination,
352            source,
353            ks,
354            helicity,
355        })
356    }
357}
358
359#[cfg(test)]
360mod tests {
361    use super::to_sw_matrix;
362    use crate::{
363        Complex, cw, sw,
364        test_support::{assert_same_bits_on_pools, bits, patterned},
365    };
366
367    /// The per-column gradients add in chunks fixed by the column count: with more
368    /// columns than chunks, they repeat bit for bit on every pool size.
369    #[test]
370    fn pullback_does_not_depend_on_the_thread_count() {
371        let modes = sw::modes(2).unwrap();
372        let destination = sw::Basis {
373            modes: (0..2)
374                .flat_map(|p| modes.iter().map(move |&mode| (p, mode)))
375                .collect(),
376            positions: vec![[0.1, -0.2, 0.3], [-0.3, 0.2, -0.1]],
377        };
378        let source = cw::Basis {
379            modes: (0..2)
380                .flat_map(|p| {
381                    [0.2, -0.35].into_iter().flat_map(move |kz| {
382                        (-4..=4).flat_map(move |m| [1, 0].map(|pol| (p, cw::Mode { kz, m, pol })))
383                    })
384                })
385                .collect(),
386            positions: vec![[0.3, 0.1, 0.0], [-0.2, 0.4, 0.2]],
387        };
388        let ks = [Complex::new(1.1, 0.02), Complex::new(1.3, 0.01)];
389        let g = patterned(destination.modes.len(), source.modes.len(), 0.7);
390        assert_same_bits_on_pools(|| {
391            let (_, residual) =
392                to_sw_matrix(destination.clone(), source.clone(), ks, true).unwrap();
393            let g = residual.pullback(&g).unwrap();
394            bits(&[&g.destination, &g.source, &g.ks])
395        });
396    }
397}
treams_core/cw/to_sw.rs

treams_core/cw/
to_sw.rs
