treams_core/smatrix/interface/
saved.rs

1//! Fixed numeric layouts for interface and propagation derivative state.
2//! treams-rs extension.
3
4use super::super::saved::{invalid, product, read_complex, read_f64, total};
5use super::{FresnelResidual, InterfaceMatrix, InterfaceResidual, PropagationResidual};
6use crate::{
7    Complex, Result,
8    saved::{Reader, SavedState, Writer},
9};
10
11impl FresnelResidual {
12    /// Byte length of one Fresnel derivative state.
13    pub fn state_size() -> Result<usize> {
14        Ok(160)
15    }
16}
17
18impl SavedState for FresnelResidual {
19    fn save_state(&self) -> Result<Vec<u8>> {
20        let mut writer = Writer::new(Self::state_size()?);
21        for &value in self
22            .ks
23            .iter()
24            .flatten()
25            .chain(self.kzs.iter().flatten())
26            .chain(&self.zs)
27        {
28            writer.complex(value);
29        }
30        Ok(writer.finish())
31    }
32
33    fn from_state(bytes: &[u8]) -> Result<Self> {
34        if bytes.len() != Self::state_size()? {
35            return Err(invalid());
36        }
37        let mut reader = Reader::new(bytes);
38        let mut ks = [[Complex::default(); 2]; 2];
39        let mut kzs = ks;
40        let mut zs = [Complex::default(); 2];
41        for value in ks
42            .iter_mut()
43            .flatten()
44            .chain(kzs.iter_mut().flatten())
45            .chain(&mut zs)
46        {
47            *value = read_complex(&mut reader)?;
48        }
49        reader.finish()?;
50        Ok(Self { ks, kzs, zs })
51    }
52}
53
54impl InterfaceResidual {
55    /// Byte length including fixed storage for the optional 4-by-4 inverse.
56    pub fn state_size() -> Result<usize> {
57        Ok(634)
58    }
59}
60
61impl SavedState for InterfaceResidual {
62    fn save_state(&self) -> Result<Vec<u8>> {
63        let mut writer = Writer::new(Self::state_size()?);
64        for &value in self.ks.iter().flatten().chain(&self.zs) {
65            writer.complex(value);
66        }
67        self.q.iter().for_each(|&v| writer.f64(v));
68        writer.usize(self.axis);
69        writer.byte(u8::from(self.fixed_q));
70        writer.byte(u8::from(self.inverse.is_some()));
71        let inverse = self.inverse.unwrap_or_else(InterfaceMatrix::zeros);
72        for &value in inverse.iter().chain(self.result.iter()) {
73            writer.complex(value);
74        }
75        Ok(writer.finish())
76    }
77
78    fn from_state(bytes: &[u8]) -> Result<Self> {
79        if bytes.len() != Self::state_size()? {
80            return Err(invalid());
81        }
82        let mut reader = Reader::new(bytes);
83        let mut ks = [[Complex::default(); 2]; 2];
84        let mut zs = [Complex::default(); 2];
85        for value in ks.iter_mut().flatten().chain(&mut zs) {
86            *value = read_complex(&mut reader)?;
87        }
88        let q = [read_f64(&mut reader)?, read_f64(&mut reader)?];
89        let axis = reader.usize()?;
90        let fixed_q = reader.byte()?;
91        let has_inverse = reader.byte()?;
92        if axis > 2 || fixed_q > 1 || has_inverse > 1 {
93            return Err(invalid());
94        }
95        let mut inverse = InterfaceMatrix::zeros();
96        let mut result = InterfaceMatrix::zeros();
97        for value in inverse.iter_mut().chain(result.iter_mut()) {
98            *value = read_complex(&mut reader)?;
99        }
100        reader.finish()?;
101        Ok(Self {
102            ks,
103            zs,
104            q,
105            axis,
106            inverse: (has_inverse != 0).then_some(inverse),
107            result,
108            fixed_q: fixed_q != 0,
109        })
110    }
111}
112
113impl PropagationResidual {
114    /// Byte length for `n` upgoing wavevectors and their two saved phases.
115    pub fn state_size(n: usize) -> Result<usize> {
116        if n == 0 {
117            return Err(invalid());
118        }
119        total(&[32, product(&[80, n])?])
120    }
121}
122
123impl SavedState for PropagationResidual {
124    fn save_state(&self) -> Result<Vec<u8>> {
125        let mut writer = Writer::new(Self::state_size(self.vectors.len())?);
126        writer.usize(self.vectors.len());
127        for &value in self.vectors.iter().flatten() {
128            writer.complex(value);
129        }
130        self.distance.iter().for_each(|&v| writer.f64(v));
131        for &value in self.phases.iter().flatten() {
132            writer.complex(value);
133        }
134        Ok(writer.finish())
135    }
136
137    fn from_state(bytes: &[u8]) -> Result<Self> {
138        let mut reader = Reader::new(bytes);
139        let n = reader.usize()?;
140        if bytes.len() != Self::state_size(n)? {
141            return Err(invalid());
142        }
143        let mut vectors = vec![[Complex::default(); 3]; n];
144        for value in vectors.iter_mut().flatten() {
145            *value = read_complex(&mut reader)?;
146        }
147        let distance = [
148            read_f64(&mut reader)?,
149            read_f64(&mut reader)?,
150            read_f64(&mut reader)?,
151        ];
152        let mut phases = vec![[Complex::default(); 2]; n];
153        for value in phases.iter_mut().flatten() {
154            *value = read_complex(&mut reader)?;
155        }
156        reader.finish()?;
157        Ok(Self {
158            vectors,
159            distance,
160            phases,
161        })
162    }
163}
164
165#[cfg(test)]
166mod tests {
167    use super::*;
168    use crate::smatrix::{self, Blocks};
169    use nalgebra::DMatrix;
170
171    fn cotangent() -> Blocks {
172        std::array::from_fn(|_| DMatrix::from_element(2, 2, Complex::new(0.3, -0.1)))
173    }
174
175    fn invalid_lengths<T: SavedState>(bytes: &[u8]) {
176        assert!(T::from_state(&bytes[..bytes.len() - 1]).is_err());
177        let mut extra = bytes.to_vec();
178        extra.push(0);
179        assert!(T::from_state(&extra).is_err());
180    }
181
182    #[test]
183    fn fresnel_roundtrip_preserves_both_derivatives() {
184        let ks = [[Complex::new(1.2, 0.1); 2], [Complex::new(1.6, 0.2); 2]];
185        let kz = ks.map(|row| row.map(|k| (k * k - 0.2).sqrt()));
186        let zs = [Complex::from(0.9), Complex::from(1.1)];
187        let (_, original) = smatrix::fresnel(ks, kz, zs).unwrap();
188        let bytes = original.save_state().unwrap();
189        assert_eq!(bytes.len(), FresnelResidual::state_size().unwrap());
190        invalid_lengths::<FresnelResidual>(&bytes);
191        let restored = FresnelResidual::from_state(&bytes).unwrap();
192        assert_eq!(
193            original.pushforward(ks, ks, zs).unwrap(),
194            restored.pushforward(ks, ks, zs).unwrap()
195        );
196        let (a, b) = (
197            original.pullback(&cotangent()).unwrap(),
198            restored.pullback(&cotangent()).unwrap(),
199        );
200        assert_eq!(a.ks, b.ks);
201        assert_eq!(a.kz, b.kz);
202        assert_eq!(a.z, b.z);
203    }
204
205    #[test]
206    fn interface_roundtrip_preserves_inverse_and_grazing_option() {
207        let ks = [[Complex::new(1.2, 0.1); 2], [Complex::new(1.6, 0.2); 2]];
208        let zs = [Complex::from(0.9), Complex::from(1.1)];
209        let (_, original) = smatrix::interface(ks, zs, [0.2, -0.1], 1, false).unwrap();
210        let bytes = original.save_state().unwrap();
211        assert_eq!(bytes.len(), InterfaceResidual::state_size().unwrap());
212        invalid_lengths::<InterfaceResidual>(&bytes);
213        let restored = InterfaceResidual::from_state(&bytes).unwrap();
214        assert_eq!(
215            original.pushforward(ks, zs, [0.1, -0.2]).unwrap(),
216            restored.pushforward(ks, zs, [0.1, -0.2]).unwrap()
217        );
218        let (a, b) = (
219            original.pullback(&cotangent()).unwrap(),
220            restored.pullback(&cotangent()).unwrap(),
221        );
222        assert_eq!(a.ks, b.ks);
223        assert_eq!(a.z, b.z);
224        assert_eq!(a.q.map(f64::to_bits), b.q.map(f64::to_bits));
225        let (_, grazing) = smatrix::interface(
226            [[Complex::from(1.2); 2]; 2],
227            [zs[0]; 2],
228            [1.2, 0.0],
229            2,
230            true,
231        )
232        .unwrap();
233        let restored = InterfaceResidual::from_state(&grazing.save_state().unwrap()).unwrap();
234        assert!(restored.inverse.is_none());
235        assert!(restored.pullback(&cotangent()).is_err());
236    }
237
238    #[test]
239    fn propagation_roundtrip_preserves_both_derivatives_and_rejects_bad_shape() {
240        let vectors = vec![
241            [
242                Complex::from(0.2),
243                Complex::from(-0.1),
244                Complex::new(1.2, 0.1)
245            ];
246            2
247        ];
248        let (_, original) = smatrix::propagation(vectors.clone(), [0.1, -0.2, 0.3]).unwrap();
249        let mut bytes = original.save_state().unwrap();
250        assert_eq!(bytes.len(), PropagationResidual::state_size(2).unwrap());
251        invalid_lengths::<PropagationResidual>(&bytes);
252        let restored = PropagationResidual::from_state(&bytes).unwrap();
253        assert_eq!(
254            original.pushforward(&vectors, [0.2, 0.1, -0.1]).unwrap(),
255            restored.pushforward(&vectors, [0.2, 0.1, -0.1]).unwrap()
256        );
257        let (a, b) = (
258            original.pullback(&cotangent()).unwrap(),
259            restored.pullback(&cotangent()).unwrap(),
260        );
261        assert_eq!(a.vectors, b.vectors);
262        assert_eq!(a.distance.map(f64::to_bits), b.distance.map(f64::to_bits));
263        bytes[..8].copy_from_slice(&u64::MAX.to_le_bytes());
264        assert!(PropagationResidual::from_state(&bytes).is_err());
265    }
266}
treams_core/smatrix/interface/saved.rs

treams_core/smatrix/interface/
saved.rs
