treams_core/pw/
permute.rs

1//! Polarization coefficients of plane waves under cyclic permutations of the Cartesian
2//! axes.
3//!
4//! Upstream: `treams.pw.permute_xyz` and its kernels `_cxyz_to_yzx_*` (one turn) and
5//! `_cxyz_to_zxy_*` (two turns, `inverse=True`).
6#![allow(clippy::indexing_slicing)] // Validated vectors, polarizations and cotangent shapes.
7
8use nalgebra::DMatrix;
9
10use super::polarization::{Direction, polarization_jet, transverse_values, wavenumbers};
11use crate::{
12    Complex, Error, Result,
13    numerics::{
14        Jet, finite,
15        parallel::{PARALLEL_ITEMS, try_fill_chunks},
16    },
17    special::{check_pol, helicity_sign},
18};
19
20/// What [`permutation`] saves for its pullback: the wavevectors, the polarizations, the
21/// number of turns and the polarization convention.
22///
23/// The pullback writes each mode's gradient on its own, so its result does not depend
24/// on the thread count.
25#[derive(Debug)]
26pub struct PermutationResidual {
27    pub(super) vectors: Vec<[Complex; 3]>,
28    pub(super) polarizations: Vec<u8>,
29    pub(super) turns: usize,
30    pub(super) helicity: bool,
31}
32
33/// The parity coefficients `[same, cross]` of the plane wave `vector` under `turns`
34/// cyclic turns of the axes, as jets in the wavevector components.
35///
36/// After one turn the permuted wavevector is `(kz, kx, ky)`, after two `(ky, kz, kx)`.
37/// `same` maps a parity polarization to itself, `cross` to the other one (upstream
38/// `polout == polin` and `polout != polin`). Zero turns give `[1, 0]`.
39///
40/// Three branches compute the pair:
41///
42/// 1. The closed form of upstream, `[-y z, -i x k] / (sqrt(x^2 + y^2) sqrt(z^2 + x^2))`
43///    for one turn and its twin for two, when the largest component lies within 1e±70
44///    and the squared modulus of the denominator within 1e±140. The first bound keeps
45///    the squares of the components finite; the second keeps the direction away from
46///    an axis and the derivative jets, which divide by the denominator squared, finite.
47/// 2. The [`Direction`] frames of the original and the permuted wavevector, which scale
48///    extreme magnitudes, when both have a nonzero transverse part.
49/// 3. On the axis of either frame, the products of the polarization vectors in the
50///    axis gauge of [`polarization_jet`].
51fn permutation_pair<const N: usize>(vector: [Complex; 3], turns: usize) -> Result<[Jet<N>; 2]> {
52    if turns == 0 {
53        wavenumbers(vector)?;
54        return Ok([Jet::constant(1.0), Jet::default()]);
55    }
56    // Both coefficients share one denominator, so the closed form needs no separate
57    // normalization of the four transverse components.
58    let scale = vector
59        .iter()
60        .map(|v| v.re.abs().max(v.im.abs()))
61        .fold(0.0, f64::max);
62    if (1e-70..=1e70).contains(&scale) && vector.iter().all(|&v| finite(v)) {
63        let [x, y, z] = std::array::from_fn(|i| Jet::<N>::variable(vector[i], i));
64        let transverse = (x * x + y * y).sqrt();
65        let destination = if turns == 1 {
66            (z * z + x * x).sqrt()
67        } else {
68            (y * y + z * z).sqrt()
69        };
70        let denominator = transverse * destination;
71        let k = (x * x + y * y + z * z).sqrt();
72        if (1e-140..=1e140).contains(&denominator.value.norm_sqr()) && k.value != Complex::default()
73        {
74            let inverse = Jet::constant(1.0) / denominator;
75            return Ok(if turns == 1 {
76                [-y * z * inverse, -Complex::i() * x * k * inverse]
77            } else {
78                [-x * z * inverse, Complex::i() * y * k * inverse]
79            });
80        }
81    }
82    let rotated = std::array::from_fn(|axis| vector[(axis + 3 - turns) % 3]);
83    let source_direction = Direction::<N>::new(vector)?;
84    let (transverse, xy) = transverse_values([rotated[0], rotated[1]])?;
85    // Cyclic rotation preserves the full norm; only the transverse frame changes.
86    let destination =
87        Direction::<N>::from_parts(rotated, source_direction.k.value, transverse, xy)?;
88    if source_direction.transverse.value != Complex::default()
89        && destination.transverse.value != Complex::default()
90    {
91        let unpermute = |mut jet: Jet<N>| {
92            jet.derivative = std::array::from_fn(|j| jet.derivative[(j + turns) % 3]);
93            jet
94        };
95        let pair = if turns == 1 {
96            (
97                -source_direction.xy[1] * unpermute(destination.xy[0]),
98                -Complex::i()
99                    * source_direction.xy[0]
100                    * (source_direction.k / unpermute(destination.transverse)),
101            )
102        } else {
103            (
104                -source_direction.xy[0] * unpermute(destination.xy[1]),
105                Complex::i()
106                    * source_direction.xy[1]
107                    * (source_direction.k / unpermute(destination.transverse)),
108            )
109        };
110        return Ok(pair.into());
111    }
112    // At an axis, evaluate the defined Cartesian polarization gauges directly.
113    let source = polarization_jet::<N>(vector, 0, false)?;
114    let mut result = [Jet::default(); 2];
115    for (p, output) in result.iter_mut().enumerate() {
116        // The algebraic dual of the parity pair is (-M, N).
117        let dual_pol = u8::from(p == 1);
118        let sign = if p == 0 { -1.0 } else { 1.0 };
119        let dual = polarization_jet::<N>(rotated, dual_pol, false)?;
120        for (axis, mut component) in dual.into_iter().enumerate() {
121            component.derivative = std::array::from_fn(|j| component.derivative[(j + turns) % 3]);
122            *output += sign * component * source[(axis + 3 - turns) % 3];
123        }
124    }
125    Ok(result)
126}
127
128fn permutation_polarization<const N: usize>(
129    pair: [Jet<N>; 2],
130    pol: u8,
131    helicity: bool,
132) -> [Jet<N>; 2] {
133    let [same, cross] = pair;
134    std::array::from_fn(|p| {
135        if helicity {
136            if p == usize::from(pol) {
137                same + helicity_sign(pol) * cross
138            } else {
139                Jet::default()
140            }
141        } else if p == usize::from(pol) {
142            same
143        } else {
144            cross
145        }
146    })
147}
148
149/// Polarization coefficients of plane modes under `turns` cyclic permutations of the
150/// axes, shape (2, modes).
151///
152/// Each column contains both output polarizations for one input mode; matching
153/// direction labels and construction of a dense basis operator are separate.
154///
155/// Upstream: `treams.pw.permute_xyz` on arrays of plane modes.
156pub fn permutation(
157    vectors: Vec<[Complex; 3]>,
158    polarizations: Vec<u8>,
159    turns: usize,
160    helicity: bool,
161) -> Result<(DMatrix<Complex>, PermutationResidual)> {
162    if vectors.is_empty()
163        || polarizations.len() != vectors.len()
164        || polarizations.iter().any(|&p| p > 1)
165    {
166        return Err(Error::InvalidInput(
167            "require matching nonempty vectors and polarizations".into(),
168        ));
169    }
170    let turns = turns % 3;
171    let mut value = DMatrix::zeros(2, vectors.len());
172    let parallel = vectors.len() >= PARALLEL_ITEMS;
173    try_fill_chunks(value.as_mut_slice(), 4, parallel, |group, output| {
174        let first = 2 * group;
175        let pair = permutation_pair::<0>(vectors[first], turns)?;
176        for (offset, column) in output.chunks_mut(2).enumerate() {
177            let i = first + offset;
178            let pair = if offset == 0 || vectors[i] == vectors[first] {
179                pair
180            } else {
181                permutation_pair::<0>(vectors[i], turns)?
182            };
183            for (out, coefficient) in
184                column
185                    .iter_mut()
186                    .zip(permutation_polarization(pair, polarizations[i], helicity))
187            {
188                *out = coefficient.value;
189            }
190        }
191        Ok(())
192    })?;
193    if value.iter().any(|&z| !finite(z)) {
194        return Err(Error::NonFinite("plane permutation overflow".into()));
195    }
196    Ok((
197        value,
198        PermutationResidual {
199            vectors,
200            polarizations,
201            turns,
202            helicity,
203        },
204    ))
205}
206
207impl PermutationResidual {
208    /// Output polarizations and input plane modes: the shape `(2, modes)` of the
209    /// coefficients.
210    #[must_use]
211    pub fn shape(&self) -> (usize, usize) {
212        (2, self.vectors.len())
213    }
214
215    /// Validate a wavevector tangent for the recorded operation.
216    pub fn validate_tangents(&self, vectors: &[[Complex; 3]]) -> Result<()> {
217        if vectors.len() != self.vectors.len() || vectors.iter().flatten().any(|&k| !finite(k)) {
218            return Err(Error::InvalidInput(
219                "invalid plane-permutation tangent".into(),
220            ));
221        }
222        Ok(())
223    }
224
225    /// Contract the polarization coefficient jets with one wavevector direction.
226    pub fn pushforward(&self, vectors: &[[Complex; 3]]) -> Result<DMatrix<Complex>> {
227        self.validate_tangents(vectors)?;
228        let mut tangent = DMatrix::zeros(2, self.vectors.len());
229        if self.turns == 0 {
230            return Ok(tangent);
231        }
232        // Adjacent modes share their polarization jet when their directions agree,
233        // while each mode still carries its independent wavevector tangent.
234        try_fill_chunks(
235            tangent.as_mut_slice(),
236            4,
237            self.vectors.len() >= PARALLEL_ITEMS,
238            |group, output| -> Result<()> {
239                let first = 2 * group;
240                let mut previous = None;
241                for (offset, column) in output.chunks_mut(2).enumerate() {
242                    let i = first + offset;
243                    if vectors[i].iter().all(|&v| v == Complex::default()) {
244                        continue;
245                    }
246                    let pair = if let Some((vector, pair)) = previous
247                        && vector == self.vectors[i]
248                    {
249                        pair
250                    } else {
251                        let pair = permutation_pair::<3>(self.vectors[i], self.turns)?;
252                        previous = Some((self.vectors[i], pair));
253                        pair
254                    };
255                    for (out, coefficient) in column.iter_mut().zip(permutation_polarization(
256                        pair,
257                        self.polarizations[i],
258                        self.helicity,
259                    )) {
260                        *out = coefficient
261                            .derivative
262                            .iter()
263                            .zip(vectors[i])
264                            .map(|(&d, v)| d * v)
265                            .sum();
266                    }
267                }
268                Ok(())
269            },
270        )?;
271        Ok(tangent)
272    }
273
274    /// Wavevector gradients for a `cotangent` of the coefficients' shape. Direction
275    /// derivatives at an axial polarization gauge are undefined.
276    pub fn pullback(&self, cotangent: &DMatrix<Complex>) -> Result<Vec<[Complex; 3]>> {
277        if cotangent.shape() != self.shape() || cotangent.iter().any(|&z| !finite(z)) {
278            return Err(Error::InvalidInput(
279                "invalid plane-permutation cotangent".into(),
280            ));
281        }
282        let modes = self.vectors.len();
283        if self.turns == 0 {
284            return Ok(vec![[Complex::default(); 3]; modes]);
285        }
286        // As in the forward pass, adjacent modes usually share one direction.
287        let mut gradient = vec![[Complex::default(); 3]; modes];
288        let parallel = modes >= PARALLEL_ITEMS;
289        try_fill_chunks(&mut gradient, 2, parallel, |group, output| {
290            let first = 2 * group;
291            let pair = permutation_pair::<3>(self.vectors[first], self.turns)?;
292            for (offset, gradient) in output.iter_mut().enumerate() {
293                let i = first + offset;
294                let pair = if offset == 0 || self.vectors[i] == self.vectors[first] {
295                    pair
296                } else {
297                    permutation_pair::<3>(self.vectors[i], self.turns)?
298                };
299                let coefficients =
300                    permutation_polarization(pair, self.polarizations[i], self.helicity);
301                *gradient = std::array::from_fn(|axis| {
302                    (0..2)
303                        .map(|p| cotangent[(p, i)] * coefficients[p].derivative[axis].conj())
304                        .sum()
305                });
306            }
307            Ok(())
308        })?;
309        Ok(gradient)
310    }
311}
312
313/// One cyclic-coordinate polarization coefficient without allocating a matrix: the
314/// coefficient of polarization `destination` in the permuted plane wave for an input of
315/// polarization `source`.
316///
317/// Upstream: `treams.pw.permute_xyz`; `inverse=False` is `turns = 1` and `inverse=True`
318/// is `turns = 2`.
319pub fn permute_xyz(
320    vector: [Complex; 3],
321    destination: u8,
322    source: u8,
323    turns: usize,
324    helicity: bool,
325) -> Result<Complex> {
326    check_pol(destination)?;
327    check_pol(source)?;
328    let value =
329        permutation_polarization(permutation_pair::<0>(vector, turns % 3)?, source, helicity)
330            [usize::from(destination)]
331        .value;
332    if !finite(value) {
333        return Err(Error::NonFinite("non-finite plane permutation".into()));
334    }
335    Ok(value)
336}
treams_core/pw/permute.rs

treams_core/pw/
permute.rs
