treams_core/tmatrix/cylinder/
saved.rs

1//! Fixed derivative state with one complete Mie state per output block.
2//!
3//! Mirrored blocks copy their shared solve when saved, so the byte count depends
4//! only on input shapes. Restoring this state never solves a boundary problem.
5//! treams-rs extension.
6
7use super::{Block, CylinderResidual};
8use crate::{
9    MAX_DEGREE, Result,
10    coeffs::MieCylResidual,
11    saved::{Reader, SavedState, Writer, invalid},
12};
13
14impl CylinderResidual {
15    /// Byte count of saved derivative state for these input shape parameters.
16    pub fn state_size(kz_count: usize, mmax: usize, boundaries: usize) -> Result<usize> {
17        if kz_count == 0 || mmax > MAX_DEGREE.unsigned_abs() as usize {
18            return Err(invalid());
19        }
20        let blocks = kz_count.checked_mul(2 * mmax + 1).ok_or_else(invalid)?;
21        MieCylResidual::state_size(boundaries)?
22            .checked_add(9)
23            .and_then(|per_block| per_block.checked_mul(blocks))
24            .and_then(|states| states.checked_add(24))
25            .ok_or_else(invalid)
26    }
27}
28
29impl SavedState for CylinderResidual {
30    fn save_state(&self) -> Result<Vec<u8>> {
31        let mmax = (self.blocks.len() / self.kz_count - 1) / 2;
32        let mut writer = Writer::new(Self::state_size(self.kz_count, mmax, self.boundaries)?);
33        writer.usize(self.kz_count);
34        writer.usize(self.boundaries);
35        writer.usize(self.blocks.len());
36        for block in &self.blocks {
37            writer.usize(block.kz_index);
38            writer.byte(u8::from(block.mirrored));
39            self.solves[block.solve].write_state(&mut writer);
40        }
41        Ok(writer.finish())
42    }
43
44    fn from_state(bytes: &[u8]) -> Result<Self> {
45        let mut reader = Reader::new(bytes);
46        let kz_count = reader.usize()?;
47        let boundaries = reader.usize()?;
48        let block_count = reader.usize()?;
49        if kz_count == 0 || block_count == 0 || !block_count.is_multiple_of(kz_count) {
50            return Err(invalid());
51        }
52        let orders = block_count / kz_count;
53        if orders.is_multiple_of(2)
54            || Self::state_size(kz_count, (orders - 1) / 2, boundaries)? != bytes.len()
55        {
56            return Err(invalid());
57        }
58        let mut blocks = Vec::with_capacity(block_count);
59        let mut solves = Vec::with_capacity(block_count);
60        for solve in 0..block_count {
61            let kz_index = reader.usize()?;
62            if kz_index >= kz_count {
63                return Err(invalid());
64            }
65            let mirrored = match reader.byte()? {
66                0 => false,
67                1 => true,
68                _ => return Err(invalid()),
69            };
70            let residual = MieCylResidual::read_state(&mut reader)?;
71            if residual.boundaries() != boundaries {
72                return Err(invalid());
73            }
74            blocks.push(Block {
75                kz_index,
76                solve,
77                mirrored,
78            });
79            solves.push(residual);
80        }
81        reader.finish()?;
82        Ok(Self {
83            blocks,
84            solves,
85            kz_count,
86            boundaries,
87        })
88    }
89}
90
91#[cfg(test)]
92mod tests {
93    #![allow(clippy::float_cmp)] // Serialization preserves numerical bits exactly.
94
95    use super::*;
96    use crate::{Complex, coeffs::Material, tmatrix::cylinder};
97    use nalgebra::DMatrix;
98
99    #[test]
100    fn saved_cylinder_preserves_mirrors_and_independent_kz_directions() -> Result<()> {
101        let materials = [
102            Material {
103                epsilon: Complex::new(2.2, 0.1),
104                kappa: Complex::new(0.06, 0.0),
105                ..Material::default()
106            },
107            Material::default(),
108        ];
109        let (_, original) = cylinder(&[-0.2, 0.0, 0.2], 1, 1.3, &[0.7], &materials)?;
110        let bytes = original.save_state()?;
111        assert_eq!(bytes.len(), CylinderResidual::state_size(3, 1, 1)?);
112        let restored = CylinderResidual::from_state(&bytes)?;
113        let direction = [Material {
114            epsilon: Complex::new(0.2, 0.1),
115            mu: Complex::new(-0.1, 0.0),
116            kappa: Complex::new(0.03, 0.0),
117        }; 2];
118        let kzs = [0.1, -0.2, 0.3];
119        let tangent = original.pushforward(&kzs, 0.2, &[-0.1], &direction)?;
120        let (rows, cols) = original.shape();
121        let cotangent = DMatrix::from_element(rows, cols, Complex::new(0.3, -0.2));
122        let expected = original.pullback(&cotangent)?;
123        for _ in 0..2 {
124            assert_eq!(
125                restored.pushforward(&kzs, 0.2, &[-0.1], &direction)?,
126                tangent
127            );
128            let actual = restored.pullback(&cotangent)?;
129            assert_eq!(actual.kzs, expected.kzs);
130            assert_eq!(actual.k0, expected.k0);
131            assert_eq!(actual.layers.radii, expected.layers.radii);
132            assert_eq!(actual.layers.epsilon, expected.layers.epsilon);
133            assert_eq!(actual.layers.mu, expected.layers.mu);
134            assert_eq!(actual.layers.kappa, expected.layers.kappa);
135        }
136        assert_eq!(restored.save_state()?, bytes);
137        Ok(())
138    }
139
140    #[test]
141    fn saved_cylinder_checks_lengths_and_indices() -> Result<()> {
142        let (_, residual) = cylinder(&[0.2], 0, 1.0, &[0.5], &[Material::default(); 2])?;
143        let bytes = residual.save_state()?;
144        assert!(CylinderResidual::from_state(&bytes[..bytes.len() - 1]).is_err());
145        let mut extended = bytes.clone();
146        extended.push(0);
147        assert!(CylinderResidual::from_state(&extended).is_err());
148        for offset in [0, 8, 16, 24, 33] {
149            let mut corrupt = bytes.clone();
150            corrupt[offset..offset + 8].copy_from_slice(&u64::MAX.to_le_bytes());
151            assert!(CylinderResidual::from_state(&corrupt).is_err());
152        }
153        let mut corrupt = bytes;
154        corrupt[32] = 2;
155        assert!(CylinderResidual::from_state(&corrupt).is_err());
156        assert!(CylinderResidual::state_size(usize::MAX, 1, 1).is_err());
157        assert!(CylinderResidual::state_size(1, 1, usize::MAX).is_err());
158        Ok(())
159    }
160}