treams_core/cluster/iterative/
saved.rs

1//! Exact-size callback state for a matrix-free sphere illumination.
2//!
3//! Mie interfaces and the converged solution are copied; restoring rebuilds only
4//! the angular translation plan, without evaluating Mie coefficients or GMRES.
5//! treams-rs extension.
6
7use std::sync::Arc;
8
9use nalgebra::DMatrix;
10
11use super::{IterativeResidual, IterativeSolution, IterativeSphereCluster};
12use crate::{
13    Complex, MAX_DEGREE, Result,
14    coeffs::MieResidual,
15    linalg::{Convergence, GmresOptions},
16    saved::{Reader, SavedState, Writer, invalid},
17    sw::{self, TranslationPlan},
18};
19
20impl IterativeResidual {
21    /// Callback-state bytes for an illumination with these static dimensions.
22    /// The layout keeps geometry, Mie residuals, solver options, incident and
23    /// scattered columns, and the primal convergence certificates.
24    pub fn state_size(lmax: u32, particles: usize, columns: usize) -> Result<usize> {
25        if !(1..=MAX_DEGREE.unsigned_abs()).contains(&lmax) || particles == 0 || columns == 0 {
26            return Err(invalid());
27        }
28        let lmax = usize::try_from(lmax).map_err(|_| invalid())?;
29        let mie = lmax
30            .checked_mul(MieResidual::state_size(1)?)
31            .ok_or_else(invalid)?;
32        let rows = lmax
33            .checked_add(2)
34            .and_then(|n| n.checked_mul(lmax))
35            .and_then(|n| n.checked_mul(2))
36            .and_then(|n| n.checked_mul(particles))
37            .ok_or_else(invalid)?;
38        let coefficients = rows
39            .checked_mul(columns)
40            .and_then(|n| n.checked_mul(32))
41            .ok_or_else(invalid)?;
42        // Three dimensions, k0 and four GMRES options occupy 64 bytes. Geometry
43        // needs four real scalars per particle; every report needs three scalars.
44        particles
45            .checked_mul(mie.checked_add(32).ok_or_else(invalid)?)
46            .and_then(|n| n.checked_add(coefficients))
47            .and_then(|n| columns.checked_mul(24).and_then(|c| n.checked_add(c)))
48            .and_then(|n| n.checked_add(64))
49            .ok_or_else(invalid)
50    }
51}
52
53impl SavedState for IterativeResidual {
54    fn save_state(&self) -> Result<Vec<u8>> {
55        let operator = &self.operator;
56        let lmax = operator.degrees.first().map_or(0, Vec::len);
57        let particles = operator.radii.len();
58        let columns = self.incident.ncols();
59        let size = Self::state_size(
60            u32::try_from(lmax).map_err(|_| invalid())?,
61            particles,
62            columns,
63        )?;
64        let mut writer = Writer::new(size);
65        writer.usize(lmax);
66        writer.usize(particles);
67        writer.usize(columns);
68        writer.f64(operator.k0);
69        writer.f64(self.options.rtol);
70        writer.f64(self.options.atol);
71        writer.usize(self.options.restart);
72        writer.usize(self.options.max_iterations);
73        for &radius in &operator.radii {
74            writer.f64(radius);
75        }
76        for &position in operator.positions.iter().flatten() {
77            writer.f64(position);
78        }
79        for residual in operator.degrees.iter().flatten() {
80            residual.write_state(&mut writer);
81        }
82        for &value in self.incident.iter().chain(self.solution.value.iter()) {
83            writer.complex(value);
84        }
85        for report in &self.solution.convergence {
86            writer.usize(report.iterations);
87            writer.f64(report.residual_norm);
88            writer.f64(report.rhs_norm);
89        }
90        Ok(writer.finish())
91    }
92
93    fn from_state(bytes: &[u8]) -> Result<Self> {
94        let mut reader = Reader::new(bytes);
95        let lmax = u32::try_from(reader.usize()?).map_err(|_| invalid())?;
96        let particles = reader.usize()?;
97        let columns = reader.usize()?;
98        if bytes.len() != Self::state_size(lmax, particles, columns)? {
99            return Err(invalid());
100        }
101        let k0 = reader.f64()?;
102        let options = GmresOptions {
103            rtol: reader.f64()?,
104            atol: reader.f64()?,
105            restart: reader.usize()?,
106            max_iterations: reader.usize()?,
107        };
108        options.validate()?;
109        let radii = (0..particles)
110            .map(|_| reader.f64())
111            .collect::<Result<Vec<_>>>()?;
112        let positions = (0..particles)
113            .map(|_| Ok([reader.f64()?, reader.f64()?, reader.f64()?]))
114            .collect::<Result<Vec<_>>>()?;
115        let mie_size = MieResidual::state_size(1)?;
116        let degrees = (0..particles)
117            .map(|_| {
118                (0..lmax)
119                    .map(|_| MieResidual::from_state(reader.raw(mie_size)?))
120                    .collect::<Result<Vec<_>>>()
121            })
122            .collect::<Result<Vec<_>>>()?;
123        let modes = sw::modes(lmax)?;
124        // Exact state-size validation already checked these products for overflow.
125        let dimension = modes.len() * particles;
126        let incident = read_matrix(&mut reader, dimension, columns)?;
127        let value = read_matrix(&mut reader, dimension, columns)?;
128        let convergence = (0..columns)
129            .map(|_| {
130                Ok(Convergence {
131                    iterations: reader.usize()?,
132                    residual_norm: reader.f64()?,
133                    rhs_norm: reader.f64()?,
134                })
135            })
136            .collect::<Result<Vec<_>>>()?;
137        reader.finish()?;
138        let count = usize::try_from(lmax).map_err(|_| invalid())?;
139        Ok(Self {
140            operator: Arc::new(IterativeSphereCluster {
141                k0,
142                radii,
143                positions,
144                degrees,
145                mode_degrees: crate::tmatrix::block_degrees(count).collect(),
146                plan: TranslationPlan::between(&modes, &modes, true)?,
147                dimension,
148            }),
149            incident,
150            options,
151            solution: IterativeSolution { value, convergence },
152        })
153    }
154}
155
156fn read_matrix(reader: &mut Reader<'_>, rows: usize, columns: usize) -> Result<DMatrix<Complex>> {
157    let data = (0..rows * columns)
158        .map(|_| reader.complex())
159        .collect::<Result<Vec<_>>>()?;
160    Ok(DMatrix::from_vec(rows, columns, data))
161}
162
163#[cfg(test)]
164mod tests {
165    use super::*;
166    use crate::test_support::patterned;
167
168    fn residual() -> IterativeResidual {
169        Arc::new(
170            IterativeSphereCluster::new(
171                2,
172                1.3,
173                &[0.2, 0.24],
174                &[Complex::new(2.3, 0.04), Complex::new(3.1, 0.02)],
175                &[[0.1, 0.0, -0.1], [1.2, 0.2, 0.4]],
176            )
177            .unwrap(),
178        )
179        .record(
180            patterned(32, 3, 0.3),
181            GmresOptions {
182                rtol: 2e-12,
183                ..GmresOptions::default()
184            },
185        )
186        .unwrap()
187    }
188
189    #[test]
190    fn saved_state_preserves_solution_certificates_and_reusable_derivatives() {
191        let residual = residual();
192        let state = residual.save_state().unwrap();
193        assert_eq!(state.len(), IterativeResidual::state_size(2, 2, 3).unwrap());
194        let restored = IterativeResidual::from_state(&state).unwrap();
195        assert_eq!(restored.solution.value, residual.solution.value);
196        for (actual, expected) in restored
197            .solution
198            .convergence
199            .iter()
200            .zip(&residual.solution.convergence)
201        {
202            assert_eq!(actual.iterations, expected.iterations);
203            assert_eq!(
204                actual.residual_norm.to_bits(),
205                expected.residual_norm.to_bits()
206            );
207            assert_eq!(actual.rhs_norm.to_bits(), expected.rhs_norm.to_bits());
208        }
209        let dr = [0.03, -0.02];
210        let de = [Complex::new(0.2, -0.03), Complex::new(-0.1, 0.04)];
211        let dp = [[0.03, 0.02, -0.04], [-0.02, 0.01, 0.03]];
212        let db = patterned(32, 3, -0.2);
213        let tangent = residual.pushforward(0.07, &dr, &de, &dp, &db).unwrap();
214        assert_eq!(
215            restored
216                .pushforward(0.07, &dr, &de, &dp, &db)
217                .unwrap()
218                .value,
219            tangent.value
220        );
221        let cotangent = patterned(32, 3, 0.7);
222        let actual = restored.pullback(&cotangent).unwrap();
223        let expected = residual.pullback(&cotangent).unwrap();
224        assert_eq!(actual.cluster.k0.to_bits(), expected.cluster.k0.to_bits());
225        assert_eq!(actual.cluster.radii, expected.cluster.radii);
226        assert_eq!(actual.cluster.epsilon, expected.cluster.epsilon);
227        assert_eq!(actual.cluster.positions, expected.cluster.positions);
228        assert_eq!(actual.incident, expected.incident);
229        assert_eq!(
230            restored
231                .pushforward(0.07, &dr, &de, &dp, &db)
232                .unwrap()
233                .value,
234            tangent.value
235        );
236        assert_eq!(restored.save_state().unwrap(), state);
237    }
238
239    #[test]
240    fn saved_state_checks_exact_dimensions_before_allocating() {
241        let state = residual().save_state().unwrap();
242        for end in [0, 8, 16, state.len() - 1] {
243            assert!(IterativeResidual::from_state(&state[..end]).is_err());
244        }
245        let mut extra = state.clone();
246        extra.push(0);
247        assert!(IterativeResidual::from_state(&extra).is_err());
248        for index in 0..3 {
249            let mut malformed = state.clone();
250            malformed[index * 8..(index + 1) * 8].copy_from_slice(&u64::MAX.to_le_bytes());
251            assert!(IterativeResidual::from_state(&malformed).is_err());
252        }
253        let mut nested = state;
254        // Three headers and the scalar solver settings, then two particles' geometry.
255        nested[128..136].copy_from_slice(&2_u64.to_le_bytes());
256        assert!(IterativeResidual::from_state(&nested).is_err());
257        assert!(IterativeResidual::state_size(2, usize::MAX, 1).is_err());
258        assert!(IterativeResidual::state_size(2, 1, 0).is_err());
259    }
260}
treams_core/cluster/iterative/saved.rs

treams_core/cluster/iterative/
saved.rs
