treams_core/pw/
saved.rs

1//! Fixed numerical layouts for plane-wave derivative state.
2//! treams-rs extension.
3
4use super::{ExpansionResidual, FieldResidual, PermutationResidual, PhasesResidual};
5use crate::{
6    Complex, Result,
7    basis::MultipoleBasis,
8    saved::{
9        Reader, SavedState, Writer, cw_basis_size, invalid, read_basis, read_complex,
10        read_f64 as read_real, sw_basis_size, write_basis,
11    },
12};
13
14fn size(header: usize, parts: &[(usize, usize)]) -> Result<usize> {
15    parts.iter().try_fold(header, |total, &(count, width)| {
16        count
17            .checked_mul(width)
18            .and_then(|part| total.checked_add(part))
19            .ok_or_else(invalid)
20    })
21}
22
23fn read_bool(reader: &mut Reader<'_>) -> Result<bool> {
24    match reader.byte()? {
25        0 => Ok(false),
26        1 => Ok(true),
27        _ => Err(invalid()),
28    }
29}
30
31fn write_points(writer: &mut Writer, points: &[[f64; 3]]) {
32    for &value in points.iter().flatten() {
33        writer.f64(value);
34    }
35}
36
37fn read_points(reader: &mut Reader<'_>, count: usize) -> Result<Vec<[f64; 3]>> {
38    (0..count)
39        .map(|_| Ok([read_real(reader)?, read_real(reader)?, read_real(reader)?]))
40        .collect()
41}
42
43fn write_vectors(writer: &mut Writer, vectors: &[[Complex; 3]]) {
44    for &value in vectors.iter().flatten() {
45        writer.complex(value);
46    }
47}
48
49fn read_vectors(reader: &mut Reader<'_>, count: usize) -> Result<Vec<[Complex; 3]>> {
50    (0..count)
51        .map(|_| {
52            Ok([
53                read_complex(reader)?,
54                read_complex(reader)?,
55                read_complex(reader)?,
56            ])
57        })
58        .collect()
59}
60
61fn write_polarizations(writer: &mut Writer, polarizations: &[u8]) {
62    for &value in polarizations {
63        writer.byte(value);
64    }
65}
66
67fn read_polarizations(reader: &mut Reader<'_>, count: usize) -> Result<Vec<u8>> {
68    (0..count)
69        .map(|_| {
70            let pol = reader.byte()?;
71            if pol <= 1 { Ok(pol) } else { Err(invalid()) }
72        })
73        .collect()
74}
75
76impl PhasesResidual {
77    /// Byte count for phases at `points` sample points and `modes` plane waves.
78    pub fn state_size(points: usize, modes: usize) -> Result<usize> {
79        size(16, &[(points, 24), (modes, 48)])
80    }
81}
82
83impl SavedState for PhasesResidual {
84    fn save_state(&self) -> Result<Vec<u8>> {
85        let mut writer = Writer::new(Self::state_size(self.points.len(), self.vectors.len())?);
86        writer.usize(self.points.len());
87        writer.usize(self.vectors.len());
88        write_points(&mut writer, &self.points);
89        write_vectors(&mut writer, &self.vectors);
90        Ok(writer.finish())
91    }
92
93    fn from_state(bytes: &[u8]) -> Result<Self> {
94        let mut reader = Reader::new(bytes);
95        let (points, modes) = (reader.usize()?, reader.usize()?);
96        if modes == 0 || bytes.len() != Self::state_size(points, modes)? {
97            return Err(invalid());
98        }
99        let points = read_points(&mut reader, points)?;
100        let vectors = read_vectors(&mut reader, modes)?;
101        reader.finish()?;
102        Ok(Self { points, vectors })
103    }
104}
105
106impl FieldResidual {
107    /// Byte count for plane fields; weighted fields retain one amplitude per mode.
108    pub fn state_size(points: usize, modes: usize, weighted: bool) -> Result<usize> {
109        size(
110            19,
111            &[(points, 24), (modes, 49 + 16 * usize::from(weighted))],
112        )
113    }
114}
115
116impl SavedState for FieldResidual {
117    fn save_state(&self) -> Result<Vec<u8>> {
118        let mut writer = Writer::new(Self::state_size(
119            self.points.len(),
120            self.vectors.len(),
121            self.coefficients.is_some(),
122        )?);
123        writer.usize(self.points.len());
124        writer.usize(self.vectors.len());
125        writer.byte(u8::from(self.coefficients.is_some()));
126        writer.byte(u8::from(self.helicity));
127        writer.byte(u8::from(self.fixed_vectors));
128        write_points(&mut writer, &self.points);
129        write_vectors(&mut writer, &self.vectors);
130        write_polarizations(&mut writer, &self.polarizations);
131        if let Some(coefficients) = &self.coefficients {
132            for &value in coefficients {
133                writer.complex(value);
134            }
135        }
136        Ok(writer.finish())
137    }
138
139    fn from_state(bytes: &[u8]) -> Result<Self> {
140        let mut reader = Reader::new(bytes);
141        let (points, modes) = (reader.usize()?, reader.usize()?);
142        let weighted = read_bool(&mut reader)?;
143        let helicity = read_bool(&mut reader)?;
144        let fixed_vectors = read_bool(&mut reader)?;
145        if modes == 0 || bytes.len() != Self::state_size(points, modes, weighted)? {
146            return Err(invalid());
147        }
148        let points = read_points(&mut reader, points)?;
149        let vectors = read_vectors(&mut reader, modes)?;
150        let polarizations = read_polarizations(&mut reader, modes)?;
151        let coefficients = weighted
152            .then(|| {
153                (0..modes)
154                    .map(|_| read_complex(&mut reader))
155                    .collect::<Result<Vec<_>>>()
156            })
157            .transpose()?;
158        reader.finish()?;
159        Ok(Self {
160            vectors,
161            polarizations,
162            points,
163            coefficients,
164            helicity,
165            fixed_vectors,
166        })
167    }
168}
169
170impl PermutationResidual {
171    /// Byte count for a permutation of `modes` plane waves.
172    pub fn state_size(modes: usize) -> Result<usize> {
173        size(17, &[(modes, 49)])
174    }
175}
176
177impl SavedState for PermutationResidual {
178    fn save_state(&self) -> Result<Vec<u8>> {
179        let mut writer = Writer::new(Self::state_size(self.vectors.len())?);
180        writer.usize(self.vectors.len());
181        writer.usize(self.turns);
182        writer.byte(u8::from(self.helicity));
183        write_vectors(&mut writer, &self.vectors);
184        write_polarizations(&mut writer, &self.polarizations);
185        Ok(writer.finish())
186    }
187
188    fn from_state(bytes: &[u8]) -> Result<Self> {
189        let mut reader = Reader::new(bytes);
190        let modes = reader.usize()?;
191        let turns = reader.usize()?;
192        let helicity = read_bool(&mut reader)?;
193        if modes == 0 || turns > 2 || bytes.len() != Self::state_size(modes)? {
194            return Err(invalid());
195        }
196        let vectors = read_vectors(&mut reader, modes)?;
197        let polarizations = read_polarizations(&mut reader, modes)?;
198        reader.finish()?;
199        Ok(Self {
200            vectors,
201            polarizations,
202            turns,
203            helicity,
204        })
205    }
206}
207
208impl ExpansionResidual {
209    /// Byte count for expansion into a spherical or cylindrical multipole basis.
210    pub fn state_size(
211        multipoles: usize,
212        positions: usize,
213        modes: usize,
214        cylindrical: bool,
215    ) -> Result<usize> {
216        let basis = if cylindrical {
217            cw_basis_size(multipoles, positions)?
218        } else {
219            sw_basis_size(multipoles, positions)?
220        };
221        size(11, &[(basis, 1), (modes, 49)])
222    }
223}
224
225impl SavedState for ExpansionResidual {
226    fn save_state(&self) -> Result<Vec<u8>> {
227        let mut writer = Writer::new(Self::state_size(
228            self.basis.len(),
229            self.basis.positions().len(),
230            self.vectors.len(),
231            matches!(self.basis, MultipoleBasis::Cylindrical(_)),
232        )?);
233        writer.usize(self.vectors.len());
234        writer.byte(u8::from(self.helicity));
235        writer.byte(u8::from(self.fixed_vectors));
236        write_basis(&mut writer, &self.basis);
237        write_vectors(&mut writer, &self.vectors);
238        write_polarizations(&mut writer, &self.polarizations);
239        Ok(writer.finish())
240    }
241
242    fn from_state(bytes: &[u8]) -> Result<Self> {
243        let mut reader = Reader::new(bytes);
244        let modes = reader.count(49)?;
245        let helicity = read_bool(&mut reader)?;
246        let fixed_vectors = read_bool(&mut reader)?;
247        let basis = read_basis(&mut reader)?;
248        let cylindrical = matches!(basis, MultipoleBasis::Cylindrical(_));
249        if modes == 0
250            || bytes.len()
251                != Self::state_size(basis.len(), basis.positions().len(), modes, cylindrical)?
252        {
253            return Err(invalid());
254        }
255        let vectors = read_vectors(&mut reader, modes)?;
256        if cylindrical && vectors.iter().any(|vector| vector[2].im != 0.0) {
257            return Err(invalid());
258        }
259        let polarizations = read_polarizations(&mut reader, modes)?;
260        reader.finish()?;
261        Ok(Self {
262            basis,
263            vectors,
264            polarizations,
265            helicity,
266            fixed_vectors,
267        })
268    }
269}
270
271#[cfg(test)]
272mod tests {
273    use super::*;
274    use crate::{
275        pw::{expansion, field, permutation, phases},
276        test_support::{bits, patterned},
277    };
278
279    fn restore<R: SavedState>(residual: &R, size: usize) -> R {
280        let bytes = residual.save_state().unwrap();
281        assert_eq!(bytes.len(), size);
282        let restored = R::from_state(&bytes).unwrap();
283        assert_eq!(restored.save_state().unwrap(), bytes);
284        assert!(R::from_state(bytes.split_last().unwrap().1).is_err());
285        restored
286    }
287
288    fn vectors() -> Vec<[Complex; 3]> {
289        vec![[
290            Complex::new(0.4, 0.0),
291            Complex::new(0.3, 0.0),
292            Complex::new(0.8, 0.0),
293        ]]
294    }
295
296    #[test]
297    fn restored_phases_preserve_both_derivatives() {
298        let (_, residual) = phases(vec![[0.2, -0.4, 0.7]], vectors()).unwrap();
299        let restored = restore(&residual, PhasesResidual::state_size(1, 1).unwrap());
300        let dpoints = [[0.2, 0.1, -0.3]];
301        let dvectors = [[Complex::new(0.1, -0.2); 3]];
302        assert_eq!(
303            residual.pushforward(&dpoints, &dvectors).unwrap(),
304            restored.pushforward(&dpoints, &dvectors).unwrap()
305        );
306        let a = residual.pullback(&patterned(1, 1, 0.8)).unwrap();
307        let b = restored.pullback(&patterned(1, 1, 0.8)).unwrap();
308        assert_eq!(
309            bits(&[&a.points, &a.vectors]),
310            bits(&[&b.points, &b.vectors])
311        );
312    }
313
314    #[test]
315    fn restored_fields_preserve_weighted_and_operator_derivatives() {
316        for weighted in [true, false] {
317            let coefficients = weighted.then(|| vec![Complex::new(0.7, -0.2)]);
318            let (_, residual) = field(
319                vectors(),
320                vec![1],
321                vec![[0.2, -0.4, 0.7]],
322                coefficients,
323                true,
324                false,
325            )
326            .unwrap();
327            let restored = restore(
328                &residual,
329                FieldResidual::state_size(1, 1, weighted).unwrap(),
330            );
331            let coefficients = if weighted {
332                vec![Complex::new(0.2, 0.1)]
333            } else {
334                vec![]
335            };
336            let dpoints = [[0.2, 0.1, -0.3]];
337            let dvectors = [[Complex::new(0.1, -0.2); 3]];
338            assert_eq!(
339                residual
340                    .pushforward(&coefficients, &dpoints, &dvectors)
341                    .unwrap(),
342                restored
343                    .pushforward(&coefficients, &dpoints, &dvectors)
344                    .unwrap()
345            );
346            let a = residual.pullback(&patterned(3, 1, 0.8)).unwrap();
347            let b = restored.pullback(&patterned(3, 1, 0.8)).unwrap();
348            assert_eq!(
349                bits(&[&a.coefficients, &a.points, &a.vectors]),
350                bits(&[&b.coefficients, &b.points, &b.vectors])
351            );
352        }
353    }
354
355    #[test]
356    fn restored_permutations_preserve_both_derivatives() {
357        let (_, residual) = permutation(vectors(), vec![1], 2, true).unwrap();
358        let restored = restore(&residual, PermutationResidual::state_size(1).unwrap());
359        let dvectors = [[Complex::new(0.1, -0.2); 3]];
360        assert_eq!(
361            residual.pushforward(&dvectors).unwrap(),
362            restored.pushforward(&dvectors).unwrap()
363        );
364        assert_eq!(
365            residual.pullback(&patterned(2, 1, 0.8)).unwrap(),
366            restored.pullback(&patterned(2, 1, 0.8)).unwrap()
367        );
368    }
369
370    #[test]
371    fn restored_expansions_preserve_both_basis_families() {
372        let spherical = MultipoleBasis::Spherical(crate::sw::Basis {
373            modes: vec![(0, crate::sw::Mode { l: 1, m: 0, pol: 1 })],
374            positions: vec![[0.2, -0.4, 0.7]],
375        });
376        let cylindrical = MultipoleBasis::Cylindrical(crate::cw::Basis {
377            modes: vec![(
378                0,
379                crate::cw::Mode {
380                    kz: 0.8,
381                    m: 1,
382                    pol: 1,
383                },
384            )],
385            positions: vec![[0.2, -0.4, 0.7]],
386        });
387        for basis in [spherical, cylindrical] {
388            let cylindrical = matches!(basis, MultipoleBasis::Cylindrical(_));
389            let (_, residual) = expansion(basis, vectors(), vec![1], true, false).unwrap();
390            let restored = restore(
391                &residual,
392                ExpansionResidual::state_size(1, 1, 1, cylindrical).unwrap(),
393            );
394            let dpoints = [[0.2, 0.1, -0.3]];
395            let dvectors = [[
396                Complex::new(0.1, -0.2),
397                Complex::new(0.2, 0.1),
398                Complex::default(),
399            ]];
400            assert_eq!(
401                residual.pushforward(&dpoints, &dvectors).unwrap(),
402                restored.pushforward(&dpoints, &dvectors).unwrap()
403            );
404            let a = residual.pullback(&patterned(1, 1, 0.8)).unwrap();
405            let b = restored.pullback(&patterned(1, 1, 0.8)).unwrap();
406            assert_eq!(
407                bits(&[&a.positions, &a.vectors]),
408                bits(&[&b.positions, &b.vectors])
409            );
410        }
411    }
412}
treams_core/pw/saved.rs

treams_core/pw/
saved.rs
