treams_core/cluster/
particles.rs

1//! Clusters of arbitrary spherical or cylindrical particle T-matrices.
2//!
3//! Upstream: `treams.TMatrix.cluster(tmats, positions).interaction.solve()` and its
4//! `TMatrixC` twin. The gradients are a treams-rs extension.
5#![allow(clippy::indexing_slicing)] // Validated particle blocks, positions and Cartesian axes.
6
7use nalgebra::DMatrix;
8
9use super::{InteractionGradient, InteractionResidual, interaction_blocks};
10use crate::sw;
11use crate::{
12    Complex, Error, Result,
13    basis::validate_wavenumbers,
14    saved::{Reader, SavedState, Writer},
15};
16
17#[derive(Debug)]
18enum ParticleExpansion {
19    Spherical(sw::ExpansionResidual),
20    Cylindrical(crate::cw::ExpansionResidual),
21}
22
23/// What [`particle_cluster`] and [`cylindrical_particle_cluster`] save for their
24/// pullback: the interaction solve, with one local block per particle, and the
25/// expansion residual of the coupling.
26#[derive(Debug)]
27pub struct ParticleClusterResidual {
28    expansion: ParticleExpansion,
29    interaction: InteractionResidual,
30}
31
32fn invalid_state() -> Error {
33    Error::InvalidInput("invalid native particle-cluster state".into())
34}
35
36/// The two sections have independent fixed layouts. Framing lets each component
37/// validate its complete byte count before restoring numerical arrays.
38fn framed_state_size(expansion: usize, interaction: usize) -> Result<usize> {
39    17_usize
40        .checked_add(expansion)
41        .and_then(|size| size.checked_add(interaction))
42        .ok_or_else(invalid_state)
43}
44
45/// Gradients of the inputs of [`particle_cluster`] and [`cylindrical_particle_cluster`],
46/// in the order of their arguments: the local matrices, the positions of the basis and
47/// the wavenumbers.
48#[derive(Debug)]
49pub struct ParticleClusterGradient {
50    /// One matrix gradient per particle, in input order.
51    pub local: Vec<DMatrix<Complex>>,
52    /// Gradients of the particle positions.
53    pub positions: Vec<[f64; 3]>,
54    /// Gradients of the embedding wavenumbers, negative helicity first.
55    pub ks: [Complex; 2],
56}
57
58/// Couple the spherical T-matrices of different particles in a common embedding
59/// medium and solve for the interacting T-matrix.
60///
61/// Basis modes must be grouped in particle order, matching the local blocks.
62/// The caller ensures the particles' enclosing spheres do not overlap.
63///
64/// Upstream: `treams.TMatrix.cluster(tmats, positions).interaction.solve()`.
65pub fn particle_cluster(
66    local: Vec<DMatrix<Complex>>,
67    basis: sw::Basis,
68    ks: [Complex; 2],
69    helicity: bool,
70) -> Result<ParticleClusterResidual> {
71    basis.validate()?;
72    validate_particle_blocks(
73        &local,
74        &basis.positions,
75        basis.modes.iter().map(|&(p, _)| p),
76    )?;
77    let (coupling, expansion) = sw::expansion(
78        basis.clone(),
79        basis,
80        ks,
81        helicity,
82        crate::special::Radial::Singular,
83    )?;
84    couple(local, coupling, ParticleExpansion::Spherical(expansion))
85}
86
87/// Solve the interaction of the local matrices through the singular expansion
88/// `coupling`. The solve owns the only coupling matrix; the expansion pullback needs
89/// just the geometry.
90fn couple(
91    local: Vec<DMatrix<Complex>>,
92    coupling: DMatrix<Complex>,
93    expansion: ParticleExpansion,
94) -> Result<ParticleClusterResidual> {
95    Ok(ParticleClusterResidual {
96        expansion,
97        interaction: interaction_blocks(local, coupling)?,
98    })
99}
100
101fn validate_particle_blocks(
102    local: &[DMatrix<Complex>],
103    positions: &[[f64; 3]],
104    mut pidxs: impl ExactSizeIterator<Item = usize>,
105) -> Result<()> {
106    if local.len() != positions.len()
107        || local.iter().map(DMatrix::nrows).sum::<usize>() != pidxs.len()
108    {
109        return Err(Error::InvalidInput(
110            "one local matrix and mode block required per position".into(),
111        ));
112    }
113    for (i, block) in local.iter().enumerate() {
114        if pidxs.by_ref().take(block.nrows()).any(|p| p != i)
115            || positions[..i].contains(&positions[i])
116        {
117            return Err(Error::InvalidInput(
118                "particle modes must be grouped at distinct positions".into(),
119            ));
120        }
121    }
122    Ok(())
123}
124
125/// Couple the cylindrical T-matrices of different particles with fixed axial mode
126/// labels and solve for the interacting T-matrix.
127///
128/// Local modes are grouped in particle order. Enclosing cylinders must not overlap.
129///
130/// Upstream: `treams.TMatrixC.cluster(tmats, positions).interaction.solve()`.
131#[allow(clippy::float_cmp)] // Exact common transverse positions are singular for cylinders.
132pub fn cylindrical_particle_cluster(
133    local: Vec<DMatrix<Complex>>,
134    basis: crate::cw::Basis,
135    ks: [Complex; 2],
136    helicity: bool,
137) -> Result<ParticleClusterResidual> {
138    basis.validate()?;
139    validate_particle_blocks(
140        &local,
141        &basis.positions,
142        basis.modes.iter().map(|&(p, _)| p),
143    )?;
144    // cw::expansion takes no polarization convention, so the parity check happens here.
145    validate_wavenumbers(ks, helicity, false)?;
146    for (i, position) in basis.positions.iter().enumerate() {
147        if basis.positions[..i].iter().any(|p| p[..2] == position[..2]) {
148            return Err(Error::InvalidInput(
149                "cylinders require distinct transverse positions".into(),
150            ));
151        }
152    }
153    let (coupling, expansion) =
154        crate::cw::expansion(basis.clone(), basis, ks, crate::special::Radial::Singular)?;
155    couple(local, coupling, ParticleExpansion::Cylindrical(expansion))
156}
157
158impl ParticleClusterResidual {
159    /// Byte count for a spherical cluster's state, from its local block sizes.
160    pub fn spherical_state_size(local_sizes: &[usize]) -> Result<usize> {
161        let (expansion, interaction) = Self::section_sizes(local_sizes, false)?;
162        framed_state_size(expansion, interaction)
163    }
164
165    /// Byte count for a cylindrical cluster's state, from its local block sizes.
166    pub fn cylindrical_state_size(local_sizes: &[usize]) -> Result<usize> {
167        let (expansion, interaction) = Self::section_sizes(local_sizes, true)?;
168        framed_state_size(expansion, interaction)
169    }
170
171    fn section_sizes(local_sizes: &[usize], cylindrical: bool) -> Result<(usize, usize)> {
172        let interaction = InteractionResidual::state_size(local_sizes)?;
173        let modes = local_sizes.iter().try_fold(0_usize, |sum, &size| {
174            sum.checked_add(size).ok_or_else(invalid_state)
175        })?;
176        let positions = local_sizes.len();
177        let expansion = if cylindrical {
178            crate::cw::ExpansionResidual::state_size(modes, positions, modes, positions)?
179        } else {
180            sw::ExpansionResidual::state_size(modes, positions, modes, positions)?
181        };
182        Ok((expansion, interaction))
183    }
184
185    /// Shapes of the recorded particle T-matrices, in particle order.
186    #[must_use = "iterate over the recorded particle shapes"]
187    pub fn local_shapes(&self) -> impl ExactSizeIterator<Item = (usize, usize)> {
188        self.interaction.local_shapes()
189    }
190
191    /// Interacting response in the particles' local multipole coordinates.
192    #[must_use]
193    pub const fn value(&self) -> &DMatrix<Complex> {
194        self.interaction.value()
195    }
196
197    /// The shape of the interacting response.
198    #[must_use]
199    pub fn shape(&self) -> (usize, usize) {
200        self.interaction.shape()
201    }
202
203    /// Directional derivative of the local matrices, particle positions and
204    /// embedding wavenumbers. Both sides of each translation share the same
205    /// position tangent; the interaction reuses its forward LU.
206    pub fn pushforward(
207        &self,
208        local: &[DMatrix<Complex>],
209        positions: &[[f64; 3]],
210        ks: [Complex; 2],
211    ) -> Result<DMatrix<Complex>> {
212        let coupling = match &self.expansion {
213            ParticleExpansion::Spherical(expansion) => {
214                expansion.pushforward(positions, positions, ks)?
215            }
216            ParticleExpansion::Cylindrical(expansion) => {
217                expansion.pushforward(positions, positions, ks)?
218            }
219        };
220        self.interaction.pushforward_blocks(local, &coupling)
221    }
222
223    /// Gradients of the local matrices, the positions and the wavenumbers from
224    /// `cotangent`, the gradient of a real loss with respect to the interacting
225    /// T-matrix.
226    ///
227    /// The position and wavenumber gradients come from the pullback of
228    /// [`sw::ExpansionResidual`] or [`cw::ExpansionResidual`](crate::cw::ExpansionResidual)
229    /// and add in their order; the LU solve and the matrix products run in faer with
230    /// the worker count of [`linalg`](crate::linalg).
231    pub fn pullback(&self, cotangent: &DMatrix<Complex>) -> Result<ParticleClusterGradient> {
232        let InteractionGradient { local, coupling } =
233            self.interaction.pullback_blocks(cotangent)?;
234        let geometry = match &self.expansion {
235            ParticleExpansion::Spherical(expansion) => expansion.pullback(&coupling)?,
236            ParticleExpansion::Cylindrical(expansion) => expansion.pullback(&coupling)?,
237        };
238        let positions = geometry
239            .destination
240            .into_iter()
241            .zip(geometry.source)
242            .map(|(a, b)| std::array::from_fn(|i| a[i] + b[i]))
243            .collect();
244        Ok(ParticleClusterGradient {
245            local,
246            positions,
247            ks: geometry.ks,
248        })
249    }
250}
251
252impl SavedState for ParticleClusterResidual {
253    fn save_state(&self) -> Result<Vec<u8>> {
254        let cylindrical = matches!(&self.expansion, ParticleExpansion::Cylindrical(_));
255        let local_sizes = self
256            .local_shapes()
257            .map(|(size, _)| size)
258            .collect::<Vec<_>>();
259        let (expansion, interaction) = Self::section_sizes(&local_sizes, cylindrical)?;
260        let mut writer = Writer::new(framed_state_size(expansion, interaction)?);
261        writer.byte(u8::from(cylindrical));
262        writer.usize(expansion);
263        writer.usize(interaction);
264        match &self.expansion {
265            ParticleExpansion::Spherical(expansion) => expansion.write_state(&mut writer),
266            ParticleExpansion::Cylindrical(expansion) => expansion.write_state(&mut writer),
267        }
268        self.interaction.write_state(&mut writer);
269        Ok(writer.finish())
270    }
271
272    fn from_state(bytes: &[u8]) -> Result<Self> {
273        let mut reader = Reader::new(bytes);
274        let kind = reader.byte()?;
275        if kind > 1 {
276            return Err(invalid_state());
277        }
278        let expansion_size = reader.usize()?;
279        let interaction_size = reader.usize()?;
280        if bytes.len() != framed_state_size(expansion_size, interaction_size)? {
281            return Err(invalid_state());
282        }
283        let expansion_bytes = reader.raw(expansion_size)?;
284        let interaction_bytes = reader.raw(interaction_size)?;
285        reader.finish()?;
286        let (expansion, shape) = if kind == 0 {
287            let residual = sw::ExpansionResidual::from_state(expansion_bytes)?;
288            let shape = residual.shape();
289            (ParticleExpansion::Spherical(residual), shape)
290        } else {
291            let residual = crate::cw::ExpansionResidual::from_state(expansion_bytes)?;
292            let shape = residual.shape();
293            (ParticleExpansion::Cylindrical(residual), shape)
294        };
295        let interaction = InteractionResidual::from_state(interaction_bytes)?;
296        if shape != interaction.shape() {
297            return Err(invalid_state());
298        }
299        Ok(Self {
300            expansion,
301            interaction,
302        })
303    }
304}
305
306#[cfg(test)]
307mod tests {
308    use super::*;
309    use crate::test_support::patterned;
310
311    fn fixture(cylindrical: bool) -> ParticleClusterResidual {
312        let positions = vec![[0.1, -0.2, 0.3], [1.7, 0.4, -0.5]];
313        let ks = [Complex::new(1.1, 0.05), Complex::new(1.4, 0.08)];
314        let local = vec![
315            patterned(1, 1, 0.2) * Complex::from(0.02),
316            patterned(2, 2, 0.3) * Complex::from(0.03),
317        ];
318        let labels = [(0, 0), (1, 0), (1, 1)];
319        if cylindrical {
320            let modes = labels
321                .map(|(p, pol)| {
322                    (
323                        p,
324                        crate::cw::Mode {
325                            kz: 0.2,
326                            m: i32::from(pol),
327                            pol,
328                        },
329                    )
330                })
331                .to_vec();
332            cylindrical_particle_cluster(local, crate::cw::Basis { modes, positions }, ks, true)
333                .unwrap()
334        } else {
335            let modes = labels
336                .map(|(p, pol)| {
337                    (
338                        p,
339                        sw::Mode {
340                            l: 1,
341                            m: i32::from(pol),
342                            pol,
343                        },
344                    )
345                })
346                .to_vec();
347            particle_cluster(local, sw::Basis { modes, positions }, ks, true).unwrap()
348        }
349    }
350
351    #[test]
352    fn saved_state_preserves_particle_blocks_values_and_derivatives() {
353        let local = [patterned(1, 1, 0.4), patterned(2, 2, 0.5)];
354        let positions = [[0.1, 0.2, -0.3], [-0.2, 0.1, 0.4]];
355        let ks = [Complex::new(0.2, -0.1), Complex::new(-0.15, 0.05)];
356        let weight = patterned(3, 3, 0.6);
357        for cylindrical in [false, true] {
358            let original = fixture(cylindrical);
359            let bytes = original.save_state().unwrap();
360            let expected_size = if cylindrical {
361                ParticleClusterResidual::cylindrical_state_size(&[1, 2])
362            } else {
363                ParticleClusterResidual::spherical_state_size(&[1, 2])
364            }
365            .unwrap();
366            assert_eq!(bytes.len(), expected_size);
367            let restored = ParticleClusterResidual::from_state(&bytes).unwrap();
368            assert_eq!(
369                restored.local_shapes().collect::<Vec<_>>(),
370                [(1, 1), (2, 2)]
371            );
372            assert_eq!(restored.value(), original.value());
373            let tangent = original.pushforward(&local, &positions, ks).unwrap();
374            assert_eq!(
375                restored.pushforward(&local, &positions, ks).unwrap(),
376                tangent
377            );
378            let actual = restored.pullback(&weight).unwrap();
379            let expected = original.pullback(&weight).unwrap();
380            assert_eq!(actual.local, expected.local);
381            assert_eq!(actual.positions, expected.positions);
382            assert_eq!(actual.ks, expected.ks);
383            assert_eq!(
384                restored.pushforward(&local, &positions, ks).unwrap(),
385                tangent
386            );
387            assert_eq!(restored.save_state().unwrap(), bytes);
388        }
389    }
390
391    #[test]
392    fn saved_state_rejects_inconsistent_section_sizes() {
393        let bytes = fixture(false).save_state().unwrap();
394        for length in [0, 1, 16, bytes.len() - 1] {
395            assert!(ParticleClusterResidual::from_state(&bytes[..length]).is_err());
396        }
397        let mut extra = bytes.clone();
398        extra.push(0);
399        assert!(ParticleClusterResidual::from_state(&extra).is_err());
400        let mut overflow = bytes;
401        overflow[1..9].copy_from_slice(&u64::MAX.to_le_bytes());
402        assert!(ParticleClusterResidual::from_state(&overflow).is_err());
403    }
404}
treams_core/cluster/particles.rs

treams_core/cluster/
particles.rs
