treams_core/tmatrix/metric/
saved.rs

1//! Fixed-shape metric gradient state, including deferred nonsmoothness errors.
2//! treams-rs extension.
3
4use super::MetricResidual;
5use crate::saved::{Reader, SavedState, Writer, invalid};
6use crate::{Complex, DerivativeError, Error, Result};
7use nalgebra::DMatrix;
8
9// Status zero means a saved gradient; the remaining tags identify the only
10// errors chirality and SvdvalsResidual::pullback can save in MetricResidual.
11// Keep the existing tags 1..=4 and payload layout when wording changes.
12const GRADIENT_ERRORS: [DerivativeError; 4] = [
13    DerivativeError::ZeroChiralityContrast,
14    DerivativeError::InvalidSingularValueCotangent,
15    DerivativeError::UnresolvedSingularValue,
16    DerivativeError::UnequalSingularValueWeights,
17];
18
19impl MetricResidual {
20    /// Number of bytes in the fixed derivative state for a square metric input.
21    pub fn state_size(dimension: usize) -> Result<usize> {
22        if dimension == 0 {
23            return Err(invalid());
24        }
25        dimension
26            .checked_mul(dimension)
27            .and_then(|count| count.checked_mul(16))
28            .and_then(|bytes| bytes.checked_add(25))
29            .ok_or_else(invalid)
30    }
31}
32
33impl SavedState for MetricResidual {
34    fn save_state(&self) -> Result<Vec<u8>> {
35        let status = match &self.gradient {
36            Ok(_) => 0,
37            Err(Error::Derivative(reason)) => GRADIENT_ERRORS
38                .iter()
39                .zip(1_u8..)
40                .find_map(|(known, status)| (known == reason).then_some(status))
41                .ok_or_else(invalid)?,
42            Err(_) => return Err(invalid()),
43        };
44        let mut writer = Writer::new(Self::state_size(self.dimension)?);
45        writer.usize(self.dimension);
46        writer.byte(status);
47        for gradient in self.ks_gradient {
48            writer.f64(gradient);
49        }
50        match &self.gradient {
51            Ok(gradient) => {
52                for &value in gradient.iter() {
53                    writer.complex(value);
54                }
55            }
56            Err(_) => {
57                // An undefined gradient has the same shape as a valid one and
58                // reserves canonical zeros instead of variable-length text.
59                for _ in 0..self.dimension * self.dimension {
60                    writer.complex(Complex::default());
61                }
62            }
63        }
64        Ok(writer.finish())
65    }
66
67    fn from_state(bytes: &[u8]) -> Result<Self> {
68        let mut reader = Reader::new(bytes);
69        let dimension = reader.usize()?;
70        if bytes.len() != Self::state_size(dimension)? {
71            return Err(invalid());
72        }
73        let status = reader.byte()?;
74        let error = if status == 0 {
75            None
76        } else {
77            Some(
78                *GRADIENT_ERRORS
79                    .get(usize::from(status) - 1)
80                    .ok_or_else(invalid)?,
81            )
82        };
83        let ks_gradient = [reader.f64()?, reader.f64()?];
84        let gradient = if let Some(reason) = error {
85            for _ in 0..dimension * dimension {
86                if reader.complex()? != Complex::default() {
87                    return Err(invalid());
88                }
89            }
90            Err(Error::Derivative(reason))
91        } else {
92            let values = (0..dimension * dimension)
93                .map(|_| reader.complex())
94                .collect::<Result<Vec<_>>>()?;
95            Ok(DMatrix::from_vec(dimension, dimension, values))
96        };
97        reader.finish()?;
98        Ok(Self {
99            gradient,
100            ks_gradient,
101            dimension,
102        })
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    #![allow(clippy::float_cmp)] // Serialization preserves numerical bits exactly.
109
110    use super::*;
111    use crate::tmatrix::metric::{Metric, metric};
112
113    #[test]
114    fn roundtrip_preserves_repeated_pushforward_and_pullback() {
115        let matrix = DMatrix::from_row_slice(
116            2,
117            2,
118            &[
119                Complex::new(-0.3, 0.03),
120                Complex::new(0.02, 0.01),
121                Complex::new(0.04, -0.02),
122                Complex::new(-0.2, -0.01),
123            ],
124        );
125        let tangent = matrix.map(|z| Complex::new(z.im, z.re));
126        for kind in [
127            Metric::CircularDichroism,
128            Metric::DualityBreaking,
129            Metric::Chirality,
130        ] {
131            let (_, residual) = metric(&matrix, &[0, 1], [1.0, 1.3], kind).unwrap();
132            let bytes = residual.save_state().unwrap();
133            assert_eq!(bytes.len(), MetricResidual::state_size(2).unwrap());
134            let restored = MetricResidual::from_state(&bytes).unwrap();
135            assert_eq!(restored.shape(), (2, 2));
136            assert_eq!(restored.save_state().unwrap(), bytes);
137            for weight in [1.2, 0.0, -0.7, 1.2] {
138                let expected = residual.pullback(weight).unwrap();
139                let actual = restored.pullback(weight).unwrap();
140                assert_eq!(actual.matrix, expected.matrix);
141                assert_eq!(actual.ks, expected.ks);
142                assert_eq!(
143                    restored.pushforward(&tangent, [0.2, -0.1]).unwrap(),
144                    residual.pushforward(&tangent, [0.2, -0.1]).unwrap()
145                );
146            }
147        }
148    }
149
150    #[test]
151    fn roundtrip_preserves_deferred_errors_and_zero_directions() {
152        for reason in GRADIENT_ERRORS {
153            let message = reason.to_string();
154            let residual = MetricResidual {
155                gradient: Err(Error::Derivative(reason)),
156                ks_gradient: [0.0; 2],
157                dimension: 2,
158            };
159            let bytes = residual.save_state().unwrap();
160            assert!(bytes[25..].iter().all(|&byte| byte == 0));
161            let restored = MetricResidual::from_state(&bytes).unwrap();
162            assert_eq!(restored.save_state().unwrap(), bytes);
163            for _ in 0..2 {
164                assert_eq!(restored.pullback(1.0).unwrap_err().to_string(), message);
165                assert_eq!(
166                    restored
167                        .pushforward(&DMatrix::identity(2, 2), [0.0; 2])
168                        .unwrap_err()
169                        .to_string(),
170                    message
171                );
172                assert_eq!(restored.pullback(0.0).unwrap().matrix, DMatrix::zeros(2, 2));
173                assert_eq!(
174                    restored
175                        .pushforward(&DMatrix::zeros(2, 2), [0.0; 2])
176                        .unwrap(),
177                    0.0
178                );
179            }
180        }
181    }
182
183    #[test]
184    fn metric_errors_survive_roundtrip() {
185        for (entries, expected_value, message) in [
186            ([0.5, 0.0, 0.0, 0.5], 0.0, "zero contrast"),
187            ([0.0, 0.0, 0.0, 0.5], 1.0, "below numerical resolution"),
188            (
189                [0.3, 0.0, 0.1, 0.2],
190                (1.0_f64 / 7.0).sqrt(),
191                "below numerical resolution",
192            ),
193        ] {
194            let matrix = DMatrix::from_row_slice(2, 2, &entries.map(|x| Complex::new(x, 0.0)));
195            let (value, residual) = metric(&matrix, &[0, 1], [1.0; 2], Metric::Chirality).unwrap();
196            assert!((value - expected_value).abs() < 1e-15);
197            let expected = residual.pullback(1.0).unwrap_err().to_string();
198            assert!(expected.contains(message), "{expected}");
199            let bytes = residual.save_state().unwrap();
200            let restored = MetricResidual::from_state(&bytes).unwrap();
201            assert_eq!(restored.pullback(1.0).unwrap_err().to_string(), expected);
202            assert_eq!(
203                restored
204                    .pushforward(&matrix, [0.0; 2])
205                    .unwrap_err()
206                    .to_string(),
207                expected
208            );
209            assert_eq!(restored.pullback(0.0).unwrap().matrix, DMatrix::zeros(2, 2));
210            assert_eq!(
211                restored
212                    .pushforward(&DMatrix::zeros(2, 2), [0.0; 2])
213                    .unwrap(),
214                0.0
215            );
216            assert_eq!(restored.save_state().unwrap(), bytes);
217        }
218    }
219
220    #[test]
221    fn rejects_invalid_shape_status_and_length_before_allocation() {
222        let (_, residual) = metric(
223            &DMatrix::identity(2, 2),
224            &[0, 1],
225            [1.0; 2],
226            Metric::Chirality,
227        )
228        .unwrap();
229        let bytes = residual.save_state().unwrap();
230        for len in 0..bytes.len() {
231            assert!(MetricResidual::from_state(&bytes[..len]).is_err());
232        }
233        let mut trailing = bytes.clone();
234        trailing.push(0);
235        assert!(MetricResidual::from_state(&trailing).is_err());
236        for dimension in [0_u64, u64::MAX, u64::MAX / 16] {
237            let mut invalid = bytes.clone();
238            invalid[..8].copy_from_slice(&dimension.to_le_bytes());
239            assert!(MetricResidual::from_state(&invalid).is_err());
240        }
241        let mut invalid = bytes.clone();
242        invalid[8] = 255;
243        assert!(MetricResidual::from_state(&invalid).is_err());
244        let mut invalid = bytes;
245        invalid[25..33].copy_from_slice(&1.0_f64.to_le_bytes());
246        assert!(MetricResidual::from_state(&invalid).is_err());
247        assert!(MetricResidual::state_size(usize::MAX).is_err());
248    }
249}
treams_core/tmatrix/metric/saved.rs

treams_core/tmatrix/metric/
saved.rs
