treams_core/cluster/
iterative.rs

1//! Matrix-free sphere clusters: GMRES solves for requested incident fields that never
2//! form the coupling matrix. treams-rs extension.
3#![allow(clippy::indexing_slicing)] // Mode, particle and Krylov indices built and checked here.
4
5mod saved;
6
7use std::sync::Arc;
8
9use faer::{Accum, MatMut, MatRef, Par, linalg::matmul::matmul};
10use nalgebra::DMatrix;
11use rayon::prelude::*;
12
13use crate::{
14    Complex, Error, Result,
15    cluster::SphereClusterGradient,
16    coeffs::{Material, Matrix2, MieResidual, mie, to_mode_order},
17    linalg::{Convergence, GmresOptions, gmres_batch, view},
18    numerics::{finite, parallel::try_fold_ordered},
19    special::Radial,
20    sw::{self, TranslationPlan},
21};
22
23/// Scattered local multipoles for only the requested incident columns.
24#[derive(Debug)]
25pub struct IterativeSolution {
26    /// Column-major matrix with one scattered-multipole column per illumination.
27    pub value: DMatrix<Complex>,
28    /// Convergence certificate for each column.
29    pub convergence: Vec<Convergence>,
30}
31
32/// The matrix-free counterpart of [`sphere_cluster`](super::sphere_cluster): the same
33/// spheres, solved with GMRES for the requested incident columns only.
34///
35/// It keeps the geometry and the local Mie coefficients, never a particle-pair matrix.
36///
37/// treams-rs extension; treams has no matrix-free cluster solver.
38#[derive(Debug)]
39pub struct IterativeSphereCluster {
40    k0: f64,
41    radii: Vec<f64>,
42    positions: Vec<[f64; 3]>,
43    degrees: Vec<Vec<MieResidual>>,
44    mode_degrees: Vec<usize>,
45    plan: TranslationPlan,
46    dimension: usize,
47}
48
49impl IterativeSphereCluster {
50    /// Homogeneous nonmagnetic spheres in vacuum, with the arguments of
51    /// [`sphere_cluster`](super::sphere_cluster). Each particle's modes follow
52    /// [`sw::modes`] in helicity convention.
53    pub fn new(
54        lmax: u32,
55        k0: f64,
56        radii: &[f64],
57        epsilon: &[Complex],
58        positions: &[[f64; 3]],
59    ) -> Result<Self> {
60        crate::cluster::validate_spheres(k0, radii, epsilon, positions)?;
61        let modes = sw::modes(lmax)?;
62        let dimension = modes
63            .len()
64            .checked_mul(radii.len())
65            .ok_or_else(|| Error::InvalidInput("cluster is too large".into()))?;
66        let degrees = radii
67            .iter()
68            .zip(epsilon)
69            .map(|(&radius, &epsilon)| {
70                let materials = crate::cluster::vacuum_sphere(epsilon);
71                (1..=lmax)
72                    .map(|l| mie(l, &[k0 * radius], &materials))
73                    .collect::<Result<Vec<_>>>()
74            })
75            .collect::<Result<Vec<_>>>()?;
76        let count = degrees.first().map_or(0, Vec::len);
77        Ok(Self {
78            k0,
79            radii: radii.to_vec(),
80            positions: positions.to_vec(),
81            degrees,
82            mode_degrees: crate::tmatrix::block_degrees(count).collect(),
83            plan: TranslationPlan::between(&modes, &modes, true)?,
84            dimension,
85        })
86    }
87
88    /// Number of local multipoles across all particles.
89    #[must_use]
90    pub const fn dimension(&self) -> usize {
91        self.dimension
92    }
93
94    /// `T input`, or `Tᴴ input` if `adjoint`, from the 2 x 2 Mie block of each degree
95    /// and order, for one or more column-major vectors.
96    ///
97    /// [`sphere`](crate::tmatrix::sphere) places the same blocks in a dense T-matrix
98    /// for [`sphere_cluster`](super::sphere_cluster); this applies them in place.
99    fn local(&self, input: &[Complex], adjoint: bool) -> Vec<Complex> {
100        self.apply_local(input, adjoint, |particle, degree| {
101            to_mode_order(self.degrees[particle][degree].value())
102        })
103    }
104
105    /// Apply value or tangent Mie blocks without forming a local T-matrix.
106    fn apply_local(
107        &self,
108        input: &[Complex],
109        adjoint: bool,
110        matrix_at: impl Fn(usize, usize) -> Matrix2,
111    ) -> Vec<Complex> {
112        let mut output = vec![Complex::default(); input.len()];
113        let modes = self.modes_per_particle();
114        for (output, input) in output
115            .chunks_exact_mut(self.dimension)
116            .zip(input.chunks_exact(self.dimension))
117        {
118            for particle in 0..self.degrees.len() {
119                for (block, &degree) in self.mode_degrees.iter().enumerate() {
120                    let matrix = matrix_at(particle, degree);
121                    let offset = particle * modes + block * 2;
122                    for i in 0..2 {
123                        for j in 0..2 {
124                            let t = if adjoint {
125                                matrix[(j, i)].conj()
126                            } else {
127                                matrix[(i, j)]
128                            };
129                            output[offset + i] += t * input[offset + j];
130                        }
131                    }
132                }
133            }
134        }
135        output
136    }
137
138    /// Multipoles per particle.
139    const fn modes_per_particle(&self) -> usize {
140        self.mode_degrees.len() * 2
141    }
142
143    /// The particle-pair coupling `C`, or `Cᴴ`, applied to `columns` column-major
144    /// vectors. Each translation block is evaluated once for all columns.
145    fn coupling(&self, input: &[Complex], columns: usize, adjoint: bool) -> Result<Vec<Complex>> {
146        self.coupling_action(input, columns, |i, j, source, destination| {
147            let (to, from) = if adjoint { (j, i) } else { (i, j) };
148            let displacement =
149                std::array::from_fn(|axis| self.positions[to][axis] - self.positions[from][axis]);
150            self.plan.apply(
151                Complex::new(self.k0, 0.0),
152                displacement,
153                Radial::Singular,
154                source,
155                destination,
156                adjoint,
157            )
158        })
159    }
160
161    /// Directional coupling action, with the same streaming storage as `C` itself.
162    fn coupling_tangent(
163        &self,
164        input: &[Complex],
165        columns: usize,
166        k0: f64,
167        positions: &[[f64; 3]],
168    ) -> Result<Vec<Complex>> {
169        self.coupling_action(input, columns, |i, j, source, destination| {
170            let displacement =
171                std::array::from_fn(|axis| self.positions[i][axis] - self.positions[j][axis]);
172            let tangent = std::array::from_fn(|axis| positions[i][axis] - positions[j][axis]);
173            self.plan.apply_pushforward(
174                Complex::new(self.k0, 0.0),
175                displacement,
176                Radial::Singular,
177                tangent,
178                Complex::new(k0, 0.0),
179                source,
180                destination,
181            )
182        })
183    }
184
185    /// Apply either the coupling or its parameter tangent, sharing each pair's
186    /// harmonic table across the incident columns without a dense pair block.
187    fn coupling_action(
188        &self,
189        input: &[Complex],
190        columns: usize,
191        apply: impl Fn(usize, usize, MatRef<'_, Complex>, MatMut<'_, Complex>) -> Result<()> + Sync,
192    ) -> Result<Vec<Complex>> {
193        let modes = self.modes_per_particle();
194        let particles = self.positions.len();
195        let input = MatRef::from_column_major_slice(input, self.dimension, columns);
196        let rows: Vec<Vec<Complex>> = crate::threads::install(|| {
197            (0..particles)
198                .into_par_iter()
199                .map(|i| -> Result<Vec<Complex>> {
200                    // Rows of particle `i`, column-major over the input columns.
201                    let mut output = vec![Complex::default(); modes * columns];
202                    for j in (0..particles).filter(|&j| j != i) {
203                        apply(
204                            i,
205                            j,
206                            input.subrows(j * modes, modes),
207                            MatMut::from_column_major_slice_mut(&mut output, modes, columns),
208                        )?;
209                    }
210                    Ok(output)
211                })
212                .collect::<Result<_>>()
213        })?;
214        let mut output = vec![Complex::default(); self.dimension * columns];
215        for (i, rows) in rows.iter().enumerate() {
216            for (column, values) in rows.chunks_exact(modes).enumerate() {
217                output[column * self.dimension + i * modes..][..modes].copy_from_slice(values);
218            }
219        }
220        Ok(output)
221    }
222
223    /// `(I - T C) input`, or `(I - T C)ᴴ input = (I - Cᴴ Tᴴ) input` if `adjoint`: the
224    /// operator that GMRES inverts.
225    fn apply(&self, input: &[Complex], columns: usize, adjoint: bool) -> Result<Vec<Complex>> {
226        let product = if adjoint {
227            self.coupling(&self.local(input, true), columns, true)?
228        } else {
229            self.local(&self.coupling(input, columns, false)?, false)
230        };
231        Ok(input.iter().zip(product).map(|(x, y)| x - y).collect())
232    }
233
234    /// Solve `(I - T C) X = rhs`, or `(I - T C)ᴴ X = rhs` if `adjoint`, one column of
235    /// `rhs` independently, sharing pair translations across bounded column batches.
236    fn solve_rhs(
237        &self,
238        rhs: &DMatrix<Complex>,
239        options: GmresOptions,
240        adjoint: bool,
241    ) -> Result<IterativeSolution> {
242        let mut value = DMatrix::zeros(self.dimension, rhs.ncols());
243        let mut convergence = Vec::with_capacity(rhs.ncols());
244        // Retain at most eight independent Krylov bases regardless of the number of
245        // illuminations. Each costly pair translation serves all active columns.
246        const BATCH_COLUMNS: usize = 8;
247        for (input, output) in rhs.as_slice().chunks(self.dimension * BATCH_COLUMNS).zip(
248            value
249                .as_mut_slice()
250                .chunks_mut(self.dimension * BATCH_COLUMNS),
251        ) {
252            let columns = input.len() / self.dimension;
253            let (answer, reports) = gmres_batch(input, columns, options, |x, columns| {
254                self.apply(x, columns, adjoint)
255            })?;
256            output.copy_from_slice(&answer);
257            convergence.extend(reports);
258        }
259        Ok(IterativeSolution { value, convergence })
260    }
261
262    /// Solve `(I - T C) X = T B` for only the columns `B` supplied in `incident`.
263    /// An unconverged solve returns an error containing the measured residual.
264    pub fn solve(
265        &self,
266        incident: &DMatrix<Complex>,
267        options: GmresOptions,
268    ) -> Result<IterativeSolution> {
269        options.validate()?;
270        if incident.nrows() != self.dimension
271            || incident.ncols() == 0
272            || incident.iter().any(|&z| !finite(z))
273        {
274            return Err(Error::InvalidInput(
275                "incident coefficients must be finite, with one row per local multipole and at least one column".into(),
276            ));
277        }
278        let rhs = DMatrix::from_vec(
279            self.dimension,
280            incident.ncols(),
281            self.local(incident.as_slice(), false),
282        );
283        self.solve_rhs(&rhs, options, false)
284    }
285
286    /// Solve like [`solve`](Self::solve) and keep what the pullback needs. The
287    /// residual shares this cluster through the [`Arc`]; the Krylov vectors are dropped.
288    pub fn record(
289        self: &Arc<Self>,
290        incident: DMatrix<Complex>,
291        options: GmresOptions,
292    ) -> Result<IterativeResidual> {
293        let solution = self.solve(&incident, options)?;
294        Ok(IterativeResidual {
295            operator: Arc::clone(self),
296            incident,
297            options,
298            solution,
299        })
300    }
301}
302
303/// What [`IterativeSphereCluster::record`] saves for its pullback: the shared
304/// cluster, the incident fields, the GMRES options and the solution.
305#[derive(Debug)]
306pub struct IterativeResidual {
307    operator: Arc<IterativeSphereCluster>,
308    incident: DMatrix<Complex>,
309    options: GmresOptions,
310    solution: IterativeSolution,
311}
312
313/// Gradients of an [`IterativeResidual`], without a dense coupling matrix or its
314/// gradient.
315///
316/// The fields follow the inputs: `cluster` holds the gradients of the arguments of
317/// [`IterativeSphereCluster::new`] in their order (`k0`, `radii`, `epsilon`,
318/// `positions`), then come those of the incident fields of
319/// [`record`](IterativeSphereCluster::record) and the convergence of the adjoint
320/// solves.
321#[derive(Debug)]
322pub struct IterativeGradient {
323    /// Gradients of the vacuum wavenumber, radii, permittivities and positions.
324    pub cluster: SphereClusterGradient,
325    /// Gradient of the incident fields.
326    pub incident: DMatrix<Complex>,
327    /// Independent convergence certificates of the adjoint solves.
328    pub convergence: Vec<Convergence>,
329}
330
331impl IterativeResidual {
332    /// The forward solution and its convergence certificates; the pullback reads the
333    /// solution.
334    #[must_use]
335    pub const fn solution(&self) -> &IterativeSolution {
336        &self.solution
337    }
338
339    /// The shape of the solution: local multipoles and incident columns.
340    #[must_use]
341    pub fn shape(&self) -> (usize, usize) {
342        self.solution.value.shape()
343    }
344
345    /// Number of particles, for binding-level tangent validation.
346    #[must_use]
347    pub fn particles(&self) -> usize {
348        self.operator.radii.len()
349    }
350
351    /// The exciting field `B + C X`, shared by both derivative directions. Reuse
352    /// the coupling action's output allocation without copying saved incident data.
353    fn response(&self) -> Result<DMatrix<Complex>> {
354        let columns = self.incident.ncols();
355        let mut response =
356            self.operator
357                .coupling(self.solution.value.as_slice(), columns, false)?;
358        for (entry, incident) in response.iter_mut().zip(self.incident.iter()) {
359            *entry += incident;
360        }
361        Ok(DMatrix::from_vec(
362            self.operator.dimension,
363            columns,
364            response,
365        ))
366    }
367
368    /// Differentiate the converged equation, retaining the matrix-free operator:
369    /// `(I - T C) dX = dT (B + C X) + T (dB + dC X)`.
370    ///
371    /// Tangent and adjoint solves use relative-only convergence: the primal `rtol`,
372    /// or the default `1e-10` if it was zero. The primal `atol` applies only to the
373    /// primal solve, so small directions retain their relative accuracy. Restart
374    /// and iteration limits are unchanged, and each solve returns an independent
375    /// residual certificate. No Krylov history is retained or differentiated.
376    pub fn pushforward(
377        &self,
378        k0: f64,
379        radii: &[f64],
380        epsilon: &[Complex],
381        positions: &[[f64; 3]],
382        incident: &DMatrix<Complex>,
383    ) -> Result<IterativeSolution> {
384        if !k0.is_finite()
385            || radii.len() != self.particles()
386            || epsilon.len() != self.particles()
387            || positions.len() != self.particles()
388            || radii
389                .iter()
390                .chain(positions.iter().flatten())
391                .any(|x| !x.is_finite())
392            || epsilon.iter().any(|&z| !finite(z))
393            || incident.shape() != self.shape()
394            || incident.iter().any(|&z| !finite(z))
395        {
396            return Err(Error::InvalidInput("invalid illumination tangents".into()));
397        }
398        let operator = &self.operator;
399        let columns = incident.ncols();
400        let value = &self.solution.value;
401        let response = self.response()?;
402        // Evaluate each directional Mie block once, then reuse it over all orders
403        // and incident columns, just as the operator reuses the value blocks.
404        let zero = Material {
405            epsilon: Complex::default(),
406            mu: Complex::default(),
407            kappa: Complex::default(),
408        };
409        let blocks = operator
410            .degrees
411            .iter()
412            .enumerate()
413            .map(|(particle, degrees)| {
414                let size = k0 * operator.radii[particle] + operator.k0 * radii[particle];
415                let materials = [
416                    Material {
417                        epsilon: epsilon[particle],
418                        ..zero
419                    },
420                    zero,
421                ];
422                degrees
423                    .iter()
424                    .map(|residual| {
425                        residual
426                            .pushforward(&[size], &materials)
427                            .map(|matrix| to_mode_order(&matrix))
428                    })
429                    .collect::<Result<Vec<_>>>()
430            })
431            .collect::<Result<Vec<_>>>()?;
432        let local_tangent = operator.apply_local(response.as_slice(), false, |particle, degree| {
433            blocks[particle][degree]
434        });
435        let mut response_tangent =
436            operator.coupling_tangent(value.as_slice(), columns, k0, positions)?;
437        for (entry, tangent) in response_tangent.iter_mut().zip(incident.iter()) {
438            *entry += tangent;
439        }
440        let mut rhs = DMatrix::from_vec(
441            operator.dimension,
442            columns,
443            operator.local(&response_tangent, false),
444        );
445        for (entry, tangent) in rhs.iter_mut().zip(local_tangent) {
446            *entry += tangent;
447        }
448        operator.solve_rhs(&rhs, self.options.for_derivative(), false)
449    }
450
451    /// Gradients of the cluster and the incident fields from `cotangent`, the gradient
452    /// of a real loss with respect to the solution.
453    ///
454    /// One adjoint GMRES solve per column gives `Y = (I - T C)⁻ᴴ G`; each particle pair
455    /// then contributes to the position and `k0` gradients. Those contributions add in
456    /// chunks of particles fixed by the particle count, and the chunk sums in chunk
457    /// order; the radius and permittivity gradients add in particle order. The thread
458    /// count changes none of the gradients.
459    /// Convergence follows the same relative-only policy as [`Self::pushforward`].
460    pub fn pullback(&self, cotangent: &DMatrix<Complex>) -> Result<IterativeGradient> {
461        if cotangent.shape() != self.shape() || cotangent.iter().any(|&z| !finite(z)) {
462            return Err(Error::InvalidInput("invalid illumination cotangent".into()));
463        }
464        let operator = &self.operator;
465        let adjoint = operator.solve_rhs(cotangent, self.options.for_derivative(), true)?;
466        let columns = self.incident.ncols();
467        let value = &self.solution.value;
468        // The local matrices multiply the response B + C X; the coupling sees T^H Y.
469        let response = self.response()?;
470        let incident = DMatrix::from_vec(
471            operator.dimension,
472            columns,
473            operator.local(adjoint.value.as_slice(), true),
474        );
475        let (positions, k0) = operator.pair_gradients(&incident, value)?;
476        let mut cluster = SphereClusterGradient {
477            k0,
478            radii: vec![0.0; operator.positions.len()],
479            epsilon: vec![Complex::default(); operator.positions.len()],
480            positions,
481        };
482        operator.local_gradients(&adjoint.value, &response, &mut cluster)?;
483        Ok(IterativeGradient {
484            cluster,
485            incident,
486            convergence: adjoint.convergence,
487        })
488    }
489}
490
491impl IterativeSphereCluster {
492    /// Position and wavenumber gradients through the coupling `C`, whose cotangent is
493    /// `(Tᴴ Y) Xᴴ` for `local_adjoint = Tᴴ Y` and the solution `value = X`, summed one
494    /// particle-pair block at a time.
495    ///
496    /// Chunks of consecutive particles `i`, fixed by the particle count, add their
497    /// pairs `(i, j)` in order into one partial sum each, and the partial sums add in
498    /// chunk order (`numerics::parallel::try_fold_ordered`), so the thread count does
499    /// not change the result.
500    fn pair_gradients(
501        &self,
502        local_adjoint: &DMatrix<Complex>,
503        value: &DMatrix<Complex>,
504    ) -> Result<(Vec<[f64; 3]>, f64)> {
505        let modes = self.modes_per_particle();
506        let particles = self.positions.len();
507        try_fold_ordered(
508            (0..particles).collect(),
509            true,
510            || (vec![[0.0; 3]; particles], 0.0),
511            |(mut positions, mut k0): (Vec<[f64; 3]>, f64), _, i: usize| -> Result<_> {
512                let mut cotangent = vec![Complex::default(); modes * modes];
513                for j in (0..particles).filter(|&j| j != i) {
514                    // The rows of particle i against the columns of particle j,
515                    // sequentially: the particles already run in parallel.
516                    matmul(
517                        MatMut::from_column_major_slice_mut(&mut cotangent, modes, modes),
518                        Accum::Replace,
519                        view(local_adjoint).subrows(i * modes, modes),
520                        view(value).subrows(j * modes, modes).adjoint(),
521                        Complex::new(1.0, 0.0),
522                        Par::Seq,
523                    );
524                    let displacement = std::array::from_fn(|axis| {
525                        self.positions[i][axis] - self.positions[j][axis]
526                    });
527                    let (position, frequency) = self
528                        .plan
529                        .pullback(
530                            Complex::new(self.k0, 0.0),
531                            displacement,
532                            Radial::Singular,
533                            &cotangent,
534                        )
535                        .map(|(position, k)| (position, k.re))?;
536                    k0 += frequency;
537                    for (axis, derivative) in position.into_iter().enumerate() {
538                        positions[i][axis] += derivative;
539                        positions[j][axis] -= derivative;
540                    }
541                }
542                Ok((positions, k0))
543            },
544            |(mut positions, k0), (partial, partial_k0)| {
545                for (a, b) in positions.iter_mut().flatten().zip(partial.iter().flatten()) {
546                    *a += b;
547                }
548                (positions, k0 + partial_k0)
549            },
550        )
551    }
552
553    /// Radius, permittivity and wavenumber gradients through the Mie coefficients,
554    /// whose local cotangent is `Y (B + C X)ᴴ` restricted to each 2 x 2 block.
555    ///
556    /// [`SphereResidual::pullback`](crate::tmatrix::SphereResidual::pullback) does the
557    /// same for the dense path of [`sphere_cluster`](super::sphere_cluster), with two
558    /// differences in order. It sums the Mie gradients of all degrees and then applies
559    /// the chain rule of the size parameter `k0 r` once; here each degree adds its own
560    /// `k0 g` and `r g`. It also swaps each 2 x 2 block to helicity order before
561    /// summing, where this sums in mode order and swaps once, which gives the same
562    /// values. The first difference changes the last bits, so the two paths agree to
563    /// rounding, and each keeps its order so that its results stay reproducible.
564    fn local_gradients(
565        &self,
566        adjoint: &DMatrix<Complex>,
567        response: &DMatrix<Complex>,
568        cluster: &mut SphereClusterGradient,
569    ) -> Result<()> {
570        let modes = self.modes_per_particle();
571        for (particle, degrees) in self.degrees.iter().enumerate() {
572            let mut degree_cotangents = vec![Matrix2::zeros(); degrees.len()];
573            for (block, &degree) in self.mode_degrees.iter().enumerate() {
574                let offset = particle * modes + block * 2;
575                for i in 0..2 {
576                    for j in 0..2 {
577                        for p in 0..adjoint.ncols() {
578                            degree_cotangents[degree][(i, j)] +=
579                                adjoint[(offset + i, p)] * response[(offset + j, p)].conj();
580                        }
581                    }
582                }
583            }
584            for (residual, cotangent) in degrees.iter().zip(degree_cotangents) {
585                let gradient = residual.pullback(&to_mode_order(&cotangent))?;
586                cluster.radii[particle] += self.k0 * gradient.sizes[0];
587                cluster.k0 += self.radii[particle] * gradient.sizes[0];
588                cluster.epsilon[particle] += gradient.epsilon[0];
589            }
590        }
591        Ok(())
592    }
593}
594
595#[cfg(test)]
596mod tests {
597    use std::{f64::consts::PI, sync::Arc};
598
599    use nalgebra::DMatrix;
600    use proptest::{prelude::*, test_runner::TestCaseError};
601
602    use super::IterativeSphereCluster;
603    use crate::{
604        Complex,
605        coeffs::Material,
606        linalg::GmresOptions,
607        sw,
608        test_support::{
609            EXPENSIVE_CASES, assert_same_bits_on_pools, bits, dot, patterned, prop_assert_close,
610            re_dot,
611        },
612    };
613
614    fn options() -> GmresOptions {
615        GmresOptions {
616            rtol: 2e-12,
617            restart: 8,
618            max_iterations: 100,
619            ..GmresOptions::default()
620        }
621    }
622
623    /// Sphere clusters in vacuum.
624    #[derive(Clone, Debug)]
625    struct Cluster {
626        lmax: u32,
627        radii: Vec<f64>,
628        epsilon: Vec<Complex>,
629        positions: Vec<[f64; 3]>,
630    }
631
632    /// One to three lossy or lossless spheres, each placed a random gap beyond touching
633    /// its predecessor, along the polar axis or in a random direction.
634    fn cluster() -> impl Strategy<Value = Cluster> {
635        let direction = prop_oneof![
636            Just([0.0, 0.0, 1.0]),
637            Just([0.0, 0.0, -1.0]),
638            (0.0..PI, -PI..PI).prop_map(|(theta, phi): (f64, f64)| {
639                [
640                    theta.sin() * phi.cos(),
641                    theta.sin() * phi.sin(),
642                    theta.cos(),
643                ]
644            }),
645        ];
646        (1_usize..=3, 1_u32..=2)
647            .prop_flat_map(move |(count, lmax)| {
648                (
649                    Just(lmax),
650                    prop::collection::vec(0.1_f64..0.35, count),
651                    prop::collection::vec(
652                        (1.2_f64..4.0, prop_oneof![Just(0.0), 0.0_f64..0.05]),
653                        count,
654                    ),
655                    prop::array::uniform3(-1.0_f64..1.0),
656                    prop::collection::vec((0.05_f64..1.0, direction.clone()), count - 1),
657                )
658            })
659            .prop_map(|(lmax, radii, epsilon, origin, steps)| {
660                let mut positions = vec![origin];
661                for (i, (gap, direction)) in steps.into_iter().enumerate() {
662                    let distance = radii[i] + radii[i + 1] + gap;
663                    let previous = positions[i];
664                    positions.push(std::array::from_fn(|a| {
665                        previous[a] + distance * direction[a]
666                    }));
667                }
668                Cluster {
669                    lmax,
670                    radii,
671                    epsilon: epsilon
672                        .into_iter()
673                        .map(|(re, im)| Complex::new(re, im))
674                        .collect(),
675                    positions,
676                }
677            })
678    }
679
680    impl Cluster {
681        fn operator(&self, k0: f64) -> IterativeSphereCluster {
682            IterativeSphereCluster::new(self.lmax, k0, &self.radii, &self.epsilon, &self.positions)
683                .unwrap()
684        }
685
686        /// The signed permutation `S` mapping `(particle, l, m, pol)` to
687        /// `(particle, l, -m, pol)` with sign `(-1)^m`.
688        fn reciprocity(&self) -> DMatrix<Complex> {
689            let modes = sw::modes(self.lmax).unwrap();
690            let modes_per_particle = modes.len();
691            let dimension = modes_per_particle * self.radii.len();
692            let mut s = DMatrix::zeros(dimension, dimension);
693            for particle in 0..self.radii.len() {
694                let offset = particle * modes_per_particle;
695                for (from, mode) in modes.iter().enumerate() {
696                    let to = modes
697                        .iter()
698                        .position(|other| {
699                            other.l == mode.l && other.m == -mode.m && other.pol == mode.pol
700                        })
701                        .unwrap();
702                    let sign = if mode.m % 2 == 0 { 1.0 } else { -1.0 };
703                    s[(offset + to, offset + from)] = Complex::from(sign);
704                }
705            }
706            s
707        }
708    }
709
710    proptest! {
711        // Five GMRES cluster solves, an adjoint solve and a dense cluster solve per case.
712        #![proptest_config(ProptestConfig::with_cases(EXPENSIVE_CASES))]
713
714        #[test]
715        fn requested_columns_match_dense_and_obey_symmetries(
716            cluster in cluster(),
717            seed in -2.0_f64..2.0,
718            scale in 0.7_f64..1.4,
719            offset in prop::array::uniform3(-1.0_f64..1.0),
720        ) {
721            check_requested_columns(&cluster, seed, scale, offset)?;
722        }
723    }
724
725    /// Matrix-free GMRES solves equal the dense cluster T-matrix applied to the
726    /// requested columns and certify their residuals; for a single sphere both equal
727    /// its Mie T-matrix and its product with the columns. The dense T-matrix obeys
728    /// Lorentz reciprocity `T = S Tᵀ S`, and so do the requested columns:
729    /// `uᵀ S X_v = vᵀ S X_u`. Solves are invariant under translation, joint scaling
730    /// and relabelling of the particles, and the pullback equals the dense one, with
731    /// the incident cotangent `Tᴴ G`, the loss pairing, zero total translation
732    /// gradient and the Euler identity of joint length and wavenumber scaling.
733    fn check_requested_columns(
734        cluster: &Cluster,
735        seed: f64,
736        scale: f64,
737        offset: [f64; 3],
738    ) -> Result<(), TestCaseError> {
739        let Cluster {
740            lmax,
741            radii,
742            epsilon,
743            positions,
744        } = cluster;
745        for i in 0..positions.len() {
746            for j in 0..i {
747                let d = (0..3).fold(0.0_f64, |d, a| d.hypot(positions[i][a] - positions[j][a]));
748                prop_assume!(d > radii[i] + radii[j] + 0.01);
749            }
750        }
751        let k0 = 1.2;
752        let operator = cluster.operator(k0);
753        let incident = patterned(operator.dimension(), 2, seed);
754        let result = operator.solve(&incident, options()).unwrap();
755        let dense = crate::cluster::sphere_cluster(*lmax, k0, radii, epsilon, positions).unwrap();
756        let t = dense.value().clone();
757        let expected = &t * &incident;
758        prop_assert_close!(&result.value, &expected, 2e-10 * expected.norm());
759        for report in &result.convergence {
760            prop_assert!(report.residual_norm <= options().rtol * report.rhs_norm);
761        }
762        if let [radius] = radii.as_slice() {
763            let material = Material {
764                epsilon: epsilon[0],
765                ..Material::default()
766            };
767            let (sphere, _) =
768                crate::tmatrix::sphere(*lmax, k0, &[*radius], &[material, Material::default()])
769                    .unwrap();
770            prop_assert_close!(&sphere, &t, 1e-13);
771            prop_assert_close!(&result.value, &(&sphere * &incident), 1e-13);
772        }
773
774        let s = cluster.reciprocity();
775        prop_assert_close!(&(&s * t.transpose() * &s), &t, 1e-12 * t.norm());
776        let (u, v) = (incident.column(0), incident.column(1));
777        let forward = (u.transpose() * &s * result.value.column(1))[0];
778        let backward = (v.transpose() * &s * result.value.column(0))[0];
779        prop_assert_close!(forward, backward, 1e-10 * u.norm() * v.norm() * t.norm());
780
781        let translated = Cluster {
782            positions: positions
783                .iter()
784                .map(|p| std::array::from_fn(|a| p[a] + offset[a]))
785                .collect(),
786            ..cluster.clone()
787        };
788        let other = translated.operator(k0).solve(&incident, options()).unwrap();
789        prop_assert_close!(&other.value, &result.value, 2e-12);
790        let scaled = Cluster {
791            radii: radii.iter().map(|r| r * scale).collect(),
792            positions: positions.iter().map(|p| p.map(|x| x * scale)).collect(),
793            ..cluster.clone()
794        };
795        let other = scaled
796            .operator(k0 / scale)
797            .solve(&incident, options())
798            .unwrap();
799        prop_assert_close!(&other.value, &result.value, 2e-11);
800        let reversed = Cluster {
801            radii: radii.iter().rev().copied().collect(),
802            epsilon: epsilon.iter().rev().copied().collect(),
803            positions: positions.iter().rev().copied().collect(),
804            ..cluster.clone()
805        };
806        let modes = operator.modes_per_particle();
807        let count = radii.len();
808        let reverse = |x: &DMatrix<Complex>| {
809            DMatrix::from_fn(x.nrows(), x.ncols(), |i, j| {
810                x[((count - 1 - i / modes) * modes + i % modes, j)]
811            })
812        };
813        let other = reversed
814            .operator(k0)
815            .solve(&reverse(&incident), options())
816            .unwrap();
817        prop_assert_close!(&other.value, &reverse(&result.value), 1e-12);
818
819        let g = patterned(operator.dimension(), 2, seed + 0.3);
820        let grad = Arc::new(operator)
821            .record(incident.clone(), options())
822            .unwrap()
823            .pullback(&g)
824            .unwrap();
825        prop_assert_close!(&grad.incident, &(t.adjoint() * &g), 3e-10 * g.norm());
826        let dense_grad = dense.pullback(&(&g * incident.adjoint())).unwrap();
827        let cluster_grad = &grad.cluster;
828        prop_assert_close!(&cluster_grad.radii, &dense_grad.radii, 3e-10);
829        prop_assert_close!(&cluster_grad.epsilon, &dense_grad.epsilon, 3e-10);
830        prop_assert_close!(&cluster_grad.positions, &dense_grad.positions, 3e-10);
831        prop_assert_close!(cluster_grad.k0, dense_grad.k0, 3e-10);
832        for axis in 0..3 {
833            let total: f64 = cluster_grad.positions.iter().map(|p| p[axis]).sum();
834            prop_assert_close!(total, 0.0, 1e-12, "axis {}", axis);
835        }
836        let loss = re_dot(&g, &result.value);
837        prop_assert_close!(loss, re_dot(&grad.incident, &incident), 3e-10);
838        let scale_gradient = dot(radii, &cluster_grad.radii)
839            + dot(
840                positions.iter().flatten(),
841                cluster_grad.positions.iter().flatten(),
842            )
843            - k0 * cluster_grad.k0;
844        prop_assert_close!(scale_gradient, 0.0, 3e-10);
845        Ok(())
846    }
847
848    /// Near-touching dielectric spheres require several restart cycles. Eleven
849    /// columns exercise a full batch and its tail, with zero and widely differing
850    /// amplitudes in both the forward and adjoint systems.
851    #[test]
852    fn restarted_batches_match_dense_for_strongly_coupled_spheres() -> Result<(), TestCaseError> {
853        let cluster = Cluster {
854            lmax: 3,
855            radii: vec![0.63, 0.60, 0.59],
856            epsilon: vec![
857                Complex::new(5.0, 0.0),
858                Complex::new(6.0, 0.05),
859                Complex::new(4.0, 0.02),
860            ],
861            positions: vec![[0.0; 3], [1.3, 0.0, 0.0], [0.65, 1.15, 0.1]],
862        };
863        let k0 = 1.4;
864        let options = GmresOptions {
865            restart: 3,
866            max_iterations: 300,
867            ..options()
868        };
869        let operator = Arc::new(cluster.operator(k0));
870        let scales = [1e-8, 0.0, 1.0, 1e8, 0.1, 7.0, 0.0, 1e-4, 1e4, 0.0, 2.0];
871        let mut incident = patterned(operator.dimension(), scales.len(), 0.4);
872        let mut g = patterned(operator.dimension(), scales.len(), 1.1);
873        for (column, scale) in scales.into_iter().enumerate() {
874            incident.column_mut(column).scale_mut(scale);
875            // Reciprocal amplitudes keep each nonzero loss contribution comparable.
876            let cotangent_scale = if [4, 9].contains(&column) {
877                0.0
878            } else if scale == 0.0 {
879                0.01
880            } else {
881                0.01 / scale
882            };
883            g.column_mut(column).scale_mut(cotangent_scale);
884        }
885        let record = operator.record(incident.clone(), options).unwrap();
886        let result = record.solution();
887        let dense = crate::cluster::sphere_cluster(
888            cluster.lmax,
889            k0,
890            &cluster.radii,
891            &cluster.epsilon,
892            &cluster.positions,
893        )
894        .unwrap();
895        let expected = dense.value() * &incident;
896        for (column, report) in result.convergence.iter().enumerate() {
897            let actual = result.value.column(column);
898            let expected = expected.column(column);
899            prop_assert_close!(
900                actual.as_slice(),
901                expected.as_slice(),
902                2e-10 * expected.norm() + f64::MIN_POSITIVE,
903                "forward column {}",
904                column
905            );
906            prop_assert!(report.residual_norm <= options.rtol * report.rhs_norm);
907            if scales[column] == 0.0 {
908                prop_assert_eq!(report.iterations, 0);
909                prop_assert!(actual.iter().all(|&z| z == Complex::default()));
910            } else {
911                prop_assert!(report.iterations > options.restart);
912            }
913        }
914        let loss = re_dot(&g, &result.value);
915        let grad = record.pullback(&g).unwrap();
916        let expected_incident = dense.value().adjoint() * &g;
917        for (column, report) in grad.convergence.iter().enumerate() {
918            let actual = grad.incident.column(column);
919            let expected = expected_incident.column(column);
920            prop_assert_close!(
921                actual.as_slice(),
922                expected.as_slice(),
923                3e-10 * g.column(column).norm() + f64::MIN_POSITIVE,
924                "adjoint column {}",
925                column
926            );
927            prop_assert!(report.residual_norm <= options.rtol * report.rhs_norm);
928            if [4, 9].contains(&column) {
929                prop_assert_eq!(report.iterations, 0);
930                prop_assert!(actual.iter().all(|&z| z == Complex::default()));
931            } else {
932                prop_assert!(report.iterations > options.restart);
933            }
934        }
935        let dense_grad = dense.pullback(&(&g * incident.adjoint())).unwrap();
936        let cluster_grad = &grad.cluster;
937        prop_assert_close!(&cluster_grad.radii, &dense_grad.radii, 3e-10);
938        prop_assert_close!(&cluster_grad.epsilon, &dense_grad.epsilon, 3e-10);
939        prop_assert_close!(&cluster_grad.positions, &dense_grad.positions, 3e-10);
940        prop_assert_close!(cluster_grad.k0, dense_grad.k0, 3e-10);
941        prop_assert_close!(loss, re_dot(&grad.incident, &incident), 3e-10);
942        for axis in 0..3 {
943            let total: f64 = cluster_grad.positions.iter().map(|p| p[axis]).sum();
944            prop_assert_close!(total, 0.0, 1e-12, "axis {}", axis);
945        }
946        let scale_gradient = dot(&cluster.radii, &cluster_grad.radii)
947            + dot(
948                cluster.positions.iter().flatten(),
949                cluster_grad.positions.iter().flatten(),
950            )
951            - k0 * cluster_grad.k0;
952        prop_assert_close!(scale_gradient, 0.0, 3e-10);
953        Ok(())
954    }
955
956    #[test]
957    fn domain_errors_are_explicit() {
958        let eps = [Complex::new(2.0, 0.0)];
959        assert!(IterativeSphereCluster::new(1, 0.0, &[0.2], &eps, &[[0.0; 3]]).is_err());
960        assert!(IterativeSphereCluster::new(1, 1.0, &[0.2], &eps, &[]).is_err());
961        assert!(
962            IterativeSphereCluster::new(
963                1,
964                1.0,
965                &[0.2, 0.3],
966                &[eps[0]; 2],
967                &[[0.0; 3], [0.4, 0.0, 0.0]]
968            )
969            .is_err()
970        );
971        let op = IterativeSphereCluster::new(1, 1.0, &[0.2], &eps, &[[0.0; 3]]).unwrap();
972        let input = DMatrix::zeros(6, 1);
973        assert!(
974            op.solve(
975                &input,
976                GmresOptions {
977                    rtol: -1.0,
978                    ..options()
979                }
980            )
981            .is_err()
982        );
983        assert!(op.solve(&DMatrix::zeros(5, 1), options()).is_err());
984        assert!(op.solve(&DMatrix::zeros(6, 0), options()).is_err());
985    }
986
987    /// The pair gradients add in chunks fixed by the particle count: with more
988    /// particles than chunks, the solve and every gradient repeat bit for bit on every
989    /// pool size.
990    #[test]
991    fn pullback_does_not_depend_on_the_thread_count() {
992        // A 6 x 4 x 3 grid of weakly scattering spheres of slightly different radii.
993        let positions: Vec<[f64; 3]> = (0..72_u32)
994            .map(|i| [i % 6, i / 6 % 4, i / 24].map(|n| 1.1 * f64::from(n)))
995            .collect();
996        let radii: Vec<f64> = (0..72_u32)
997            .map(|i| 0.2 + 0.05 * f64::from(i).sin())
998            .collect();
999        let epsilon = vec![Complex::new(2.25, 0.01); positions.len()];
1000        let cluster =
1001            Arc::new(IterativeSphereCluster::new(1, 1.0, &radii, &epsilon, &positions).unwrap());
1002        for columns in [1, 9] {
1003            let incident = patterned(cluster.dimension(), columns, 0.4);
1004            let g = patterned(cluster.dimension(), columns, 1.1);
1005            assert_same_bits_on_pools(|| {
1006                let residual = cluster.record(incident.clone(), options()).unwrap();
1007                let solution = residual.solution().value.clone();
1008                let gradient = residual.pullback(&g).unwrap();
1009                let cluster = &gradient.cluster;
1010                bits(&[
1011                    &solution,
1012                    &cluster.k0,
1013                    &cluster.radii,
1014                    &cluster.epsilon,
1015                    &cluster.positions,
1016                    &gradient.incident,
1017                ])
1018            });
1019        }
1020    }
1021}