treams_core/sw/
periodic_to_cw.rs

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

treams_core/sw/
periodic_to_cw.rs
