treams_core/smatrix/layers/
saved.rs

1//! Fixed invocation-state layout of independent planar channels.
2//!
3//! treams-rs extension. Every nested interface and Redheffer residual stores its
4//! computed matrices and factors; restoration never reruns the layer stack.
5
6use super::{Channel, LayerStackResidual, Step};
7use crate::{
8    Result,
9    saved::{Reader, SavedState, Writer},
10    smatrix::{
11        AddResidual, InterfaceResidual,
12        saved::{invalid, product, read_complex, read_f64, read_matrix, total, write_matrix},
13    },
14};
15
16impl LayerStackResidual {
17    /// Fixed state byte count for the media and independent transverse channels.
18    pub fn state_size(media: usize, channels: usize) -> Result<usize> {
19        if media < 2 || channels == 0 {
20            return Err(invalid());
21        }
22        let layers = media - 2;
23        let interface = InterfaceResidual::state_size()?;
24        let step = total(&[256, AddResidual::state_size(2)?, interface])?;
25        let channel = total(&[interface, product(&[layers, step])?])?;
26        total(&[
27            17,
28            product(&[media, 32])?,
29            product(&[channels, 16])?,
30            product(&[layers, 8])?,
31            product(&[channels, channel])?,
32        ])
33    }
34}
35
36impl SavedState for LayerStackResidual {
37    fn save_state(&self) -> Result<Vec<u8>> {
38        let mut writer = Writer::new(Self::state_size(self.ks.len(), self.channels.len())?);
39        writer.usize(self.ks.len());
40        writer.usize(self.channels.len());
41        writer.byte(u8::from(self.fixed_q));
42        for &value in self.ks.iter().flatten() {
43            writer.complex(value);
44        }
45        for &value in self.q.iter().flatten().chain(&self.thickness) {
46            writer.f64(value);
47        }
48        for channel in &self.channels {
49            writer.raw(&channel.initial.save_state()?);
50            for step in &channel.steps {
51                for &value in step.normal.iter().chain(&step.phase) {
52                    writer.complex(value);
53                }
54                for block in &step.below {
55                    write_matrix(&mut writer, block);
56                }
57                writer.raw(&step.boundary.save_state()?);
58                writer.raw(&step.interface.save_state()?);
59            }
60        }
61        Ok(writer.finish())
62    }
63
64    fn from_state(bytes: &[u8]) -> Result<Self> {
65        let mut reader = Reader::new(bytes);
66        let (media, count) = (reader.usize()?, reader.usize()?);
67        if bytes.len() != Self::state_size(media, count)? {
68            return Err(invalid());
69        }
70        let fixed_q = match reader.byte()? {
71            0 => false,
72            1 => true,
73            _ => return Err(invalid()),
74        };
75        let ks = (0..media)
76            .map(|_| Ok([read_complex(&mut reader)?, read_complex(&mut reader)?]))
77            .collect::<Result<Vec<_>>>()?;
78        let q = (0..count)
79            .map(|_| Ok([read_f64(&mut reader)?, read_f64(&mut reader)?]))
80            .collect::<Result<Vec<_>>>()?;
81        let thickness = (0..media - 2)
82            .map(|_| {
83                let value = read_f64(&mut reader)?;
84                if value < 0.0 {
85                    return Err(invalid());
86                }
87                Ok(value)
88            })
89            .collect::<Result<Vec<_>>>()?;
90        let interface_size = InterfaceResidual::state_size()?;
91        let boundary_size = AddResidual::state_size(2)?;
92        let channels = (0..count)
93            .map(|_| {
94                let initial = InterfaceResidual::from_state(reader.raw(interface_size)?)?;
95                let steps = (0..media - 2)
96                    .map(|_| {
97                        let normal = [read_complex(&mut reader)?, read_complex(&mut reader)?];
98                        let phase = [read_complex(&mut reader)?, read_complex(&mut reader)?];
99                        let below = [
100                            read_matrix(&mut reader, 2, 2)?,
101                            read_matrix(&mut reader, 2, 2)?,
102                            read_matrix(&mut reader, 2, 2)?,
103                        ];
104                        let boundary = AddResidual::from_state(reader.raw(boundary_size)?)?;
105                        if boundary.shape() != (2, 2) {
106                            return Err(invalid());
107                        }
108                        let interface = InterfaceResidual::from_state(reader.raw(interface_size)?)?;
109                        Ok(Step {
110                            normal,
111                            phase,
112                            below,
113                            boundary,
114                            interface,
115                        })
116                    })
117                    .collect::<Result<Vec<_>>>()?;
118                Ok(Channel { initial, steps })
119            })
120            .collect::<Result<Vec<_>>>()?;
121        reader.finish()?;
122        Ok(Self {
123            ks,
124            q,
125            thickness,
126            channels,
127            fixed_q,
128        })
129    }
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135    use crate::{Complex, smatrix::layer_stack, test_support::patterned};
136
137    #[test]
138    fn layer_state_roundtrip_reuses_both_derivative_directions() {
139        let c = Complex::new;
140        let ks = vec![
141            [c(1.0, 0.01); 2],
142            [c(1.5, 0.1), c(1.6, 0.1)],
143            [c(1.2, 0.02); 2],
144        ];
145        let zs = [c(1.0, 0.0), c(0.7, 0.0), c(0.9, 0.0)];
146        let q = vec![[0.2, 0.1], [0.1, -0.2]];
147        let (_, residual) = layer_stack(ks.clone(), &zs, q.clone(), &[0.4], 2, false).unwrap();
148        let state = residual.save_state().unwrap();
149        assert_eq!(state.len(), LayerStackResidual::state_size(3, 2).unwrap());
150        let restored = LayerStackResidual::from_state(&state).unwrap();
151        let tangent = residual.pushforward(&ks, &zs, &q, &[0.1]).unwrap();
152        assert_eq!(tangent, restored.pushforward(&ks, &zs, &q, &[0.1]).unwrap());
153        let cotangent = vec![std::array::from_fn(|_| patterned(2, 2, 0.3)); 2];
154        let expected = residual.pullback(cotangent.clone()).unwrap();
155        let actual = restored.pullback(cotangent).unwrap();
156        assert_eq!(expected.ks, actual.ks);
157        assert_eq!(expected.zs, actual.zs);
158        assert_eq!(expected.q, actual.q);
159        assert_eq!(expected.thickness, actual.thickness);
160        assert_eq!(tangent, restored.pushforward(&ks, &zs, &q, &[0.1]).unwrap());
161        assert!(LayerStackResidual::from_state(&state[..state.len() - 1]).is_err());
162        let mut bad = state;
163        bad[..8].copy_from_slice(&u64::MAX.to_le_bytes());
164        assert!(LayerStackResidual::from_state(&bad).is_err());
165    }
166}
treams_core/smatrix/layers/saved.rs

treams_core/smatrix/layers/
saved.rs
