treams_core/smatrix/chirality/
saved.rs

1//! Fixed numeric layouts of the per-mode chirality derivative state.
2//! treams-rs extension.
3
4use super::super::saved::{invalid, product, read_complex, read_f64, total};
5use super::{ChiralityDensityResidual, OrientedChiralityResidual};
6use crate::{
7    Result,
8    saved::{Reader, SavedState, Writer},
9};
10
11impl ChiralityDensityResidual {
12    /// Byte length for `n` full and normal wavenumbers and an averaging interval.
13    pub fn state_size(n: usize) -> Result<usize> {
14        if n == 0 {
15            return Err(invalid());
16        }
17        total(&[24, product(&[32, n])?])
18    }
19}
20
21impl SavedState for ChiralityDensityResidual {
22    fn save_state(&self) -> Result<Vec<u8>> {
23        let mut writer = Writer::new(Self::state_size(self.ks.len())?);
24        writer.usize(self.ks.len());
25        for &value in self.ks.iter().chain(&self.normal) {
26            writer.complex(value);
27        }
28        self.interval.iter().for_each(|&v| writer.f64(v));
29        Ok(writer.finish())
30    }
31
32    fn from_state(bytes: &[u8]) -> Result<Self> {
33        let mut reader = Reader::new(bytes);
34        let n = reader.usize()?;
35        if bytes.len() != Self::state_size(n)? {
36            return Err(invalid());
37        }
38        let ks = (0..n)
39            .map(|_| read_complex(&mut reader))
40            .collect::<Result<Vec<_>>>()?;
41        let normal = (0..n)
42            .map(|_| read_complex(&mut reader))
43            .collect::<Result<Vec<_>>>()?;
44        let interval = [read_f64(&mut reader)?, read_f64(&mut reader)?];
45        reader.finish()?;
46        Ok(Self {
47            ks,
48            normal,
49            interval,
50        })
51    }
52}
53
54impl OrientedChiralityResidual {
55    /// Byte length for `n` transverse/normal mode vectors and static polarizations.
56    pub fn state_size(n: usize) -> Result<usize> {
57        if n == 0 {
58            return Err(invalid());
59        }
60        total(&[32, product(&[33, n])?])
61    }
62}
63
64impl SavedState for OrientedChiralityResidual {
65    fn save_state(&self) -> Result<Vec<u8>> {
66        let mut writer = Writer::new(Self::state_size(self.normal.len())?);
67        writer.usize(self.normal.len());
68        writer.usize(self.axis);
69        for &value in self.transverse.iter().flatten() {
70            writer.f64(value);
71        }
72        self.normal.iter().for_each(|&v| writer.complex(v));
73        self.polarizations.iter().for_each(|&v| writer.byte(v));
74        self.interval.iter().for_each(|&v| writer.f64(v));
75        Ok(writer.finish())
76    }
77
78    fn from_state(bytes: &[u8]) -> Result<Self> {
79        let mut reader = Reader::new(bytes);
80        let n = reader.usize()?;
81        let axis = reader.usize()?;
82        if axis > 2 || bytes.len() != Self::state_size(n)? {
83            return Err(invalid());
84        }
85        let transverse = (0..n)
86            .map(|_| Ok([read_f64(&mut reader)?, read_f64(&mut reader)?]))
87            .collect::<Result<Vec<_>>>()?;
88        let normal = (0..n)
89            .map(|_| read_complex(&mut reader))
90            .collect::<Result<Vec<_>>>()?;
91        let polarizations = (0..n)
92            .map(|_| {
93                let pol = reader.byte()?;
94                if pol > 1 { Err(invalid()) } else { Ok(pol) }
95            })
96            .collect::<Result<Vec<_>>>()?;
97        let interval = [read_f64(&mut reader)?, read_f64(&mut reader)?];
98        reader.finish()?;
99        Ok(Self {
100            transverse,
101            normal,
102            polarizations,
103            axis,
104            interval,
105        })
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112    use crate::{Complex, smatrix};
113    use nalgebra::DMatrix;
114
115    fn invalid_lengths<T: SavedState>(bytes: &[u8]) {
116        assert!(T::from_state(&bytes[..bytes.len() - 1]).is_err());
117        let mut extra = bytes.to_vec();
118        extra.push(0);
119        assert!(T::from_state(&extra).is_err());
120    }
121
122    #[test]
123    fn chirality_roundtrip_preserves_both_derivatives_and_rejects_bad_shape() {
124        let ks = vec![Complex::new(1.2, 0.1), Complex::new(1.4, 0.2)];
125        let normal = vec![Complex::new(1.1, 0.12), Complex::new(1.3, 0.23)];
126        let (_, original) =
127            smatrix::chirality_density(ks.clone(), normal.clone(), [0.1, 0.5]).unwrap();
128        let mut bytes = original.save_state().unwrap();
129        assert_eq!(
130            bytes.len(),
131            ChiralityDensityResidual::state_size(2).unwrap()
132        );
133        invalid_lengths::<ChiralityDensityResidual>(&bytes);
134        let restored = ChiralityDensityResidual::from_state(&bytes).unwrap();
135        assert_eq!(
136            original.pushforward(&ks, &normal, [-0.1, 0.2]).unwrap(),
137            restored.pushforward(&ks, &normal, [-0.1, 0.2]).unwrap()
138        );
139        let cotangent = DMatrix::from_element(3, 2, Complex::new(0.2, -0.3));
140        let (a, b) = (
141            original.pullback(&cotangent).unwrap(),
142            restored.pullback(&cotangent).unwrap(),
143        );
144        assert_eq!(a.ks, b.ks);
145        assert_eq!(a.normal, b.normal);
146        assert_eq!(a.interval.map(f64::to_bits), b.interval.map(f64::to_bits));
147        bytes[..8].copy_from_slice(&u64::MAX.to_le_bytes());
148        assert!(ChiralityDensityResidual::from_state(&bytes).is_err());
149    }
150
151    #[test]
152    fn oriented_roundtrip_preserves_both_derivatives_and_rejects_bad_shape() {
153        let transverse = vec![[0.2, -0.1], [-0.3, 0.2]];
154        let normal = vec![Complex::new(1.1, 0.12), Complex::new(1.3, 0.23)];
155        let (_, original) = smatrix::oriented_chirality(
156            transverse.clone(),
157            normal.clone(),
158            vec![0, 1],
159            0,
160            [0.1, 0.5],
161        )
162        .unwrap();
163        let mut bytes = original.save_state().unwrap();
164        assert_eq!(
165            bytes.len(),
166            OrientedChiralityResidual::state_size(2).unwrap()
167        );
168        invalid_lengths::<OrientedChiralityResidual>(&bytes);
169        let restored = OrientedChiralityResidual::from_state(&bytes).unwrap();
170        assert_eq!(
171            original
172                .pushforward(&transverse, &normal, [-0.1, 0.2])
173                .unwrap(),
174            restored
175                .pushforward(&transverse, &normal, [-0.1, 0.2])
176                .unwrap()
177        );
178        let cotangent = DMatrix::from_element(3, 2, Complex::new(0.2, -0.3));
179        let (a, b) = (
180            original.pullback(&cotangent).unwrap(),
181            restored.pullback(&cotangent).unwrap(),
182        );
183        assert_eq!(a.transverse, b.transverse);
184        assert_eq!(a.normal, b.normal);
185        assert_eq!(a.interval.map(f64::to_bits), b.interval.map(f64::to_bits));
186        bytes[..8].copy_from_slice(&u64::MAX.to_le_bytes());
187        assert!(OrientedChiralityResidual::from_state(&bytes).is_err());
188    }
189}