treams_core/smatrix/periodic/
saved.rs

1//! Fixed saved transfer-matrix and Bloch eigensystem state.
2//!
3//! treams-rs extension.
4
5use super::{BandsResidual, PeriodicResidual};
6use crate::{
7    Complex, Result,
8    linalg::{EigResidual, SolveResidual},
9    numerics,
10    saved::{Reader, SavedState, Writer},
11    smatrix::saved::{invalid, product, read_complex, read_f64, read_matrix, total, write_matrix},
12};
13
14impl PeriodicResidual {
15    /// Byte count for saved transfer-matrix state with `n` modes per direction.
16    pub fn state_size(n: usize) -> Result<usize> {
17        let columns = product(&[2, n])?;
18        total(&[
19            8,
20            product(&[48, n, n])?,
21            SolveResidual::state_size(n, columns)?,
22        ])
23    }
24}
25
26impl SavedState for PeriodicResidual {
27    fn save_state(&self) -> Result<Vec<u8>> {
28        let n = self.top.nrows();
29        let mut writer = Writer::new(Self::state_size(n)?);
30        writer.usize(n);
31        write_matrix(&mut writer, &self.top);
32        write_matrix(&mut writer, &self.reflection);
33        writer.raw(&self.solve.save_state()?);
34        Ok(writer.finish())
35    }
36
37    fn from_state(bytes: &[u8]) -> Result<Self> {
38        let mut reader = Reader::new(bytes);
39        let n = reader.usize()?;
40        if bytes.len() != Self::state_size(n)? {
41            return Err(invalid());
42        }
43        let columns = product(&[2, n])?;
44        let top = read_matrix(&mut reader, n, columns)?;
45        let reflection = read_matrix(&mut reader, n, n)?;
46        let solve = SolveResidual::from_state(reader.raw(SolveResidual::state_size(n, columns)?)?)?;
47        if solve.shape() != (n, columns) {
48            return Err(invalid());
49        }
50        reader.finish()?;
51        Ok(Self {
52            top,
53            reflection,
54            solve,
55        })
56    }
57}
58
59impl BandsResidual {
60    /// Byte count for saved Bloch bands with `n` modes per propagation direction.
61    pub fn state_size(n: usize) -> Result<usize> {
62        let modes = product(&[2, n])?;
63        total(&[
64            16,
65            product(&[16, modes])?,
66            PeriodicResidual::state_size(n)?,
67            EigResidual::state_size(modes)?,
68        ])
69    }
70}
71
72impl SavedState for BandsResidual {
73    fn save_state(&self) -> Result<Vec<u8>> {
74        let n = self.periodic.top.nrows();
75        let mut writer = Writer::new(Self::state_size(n)?);
76        writer.usize(n);
77        writer.f64(self.period);
78        for &wavenumber in &self.wavenumbers {
79            writer.complex(wavenumber);
80        }
81        writer.raw(&self.periodic.save_state()?);
82        writer.raw(&self.eigen.save_state()?);
83        Ok(writer.finish())
84    }
85
86    fn from_state(bytes: &[u8]) -> Result<Self> {
87        let mut reader = Reader::new(bytes);
88        let n = reader.usize()?;
89        if bytes.len() != Self::state_size(n)? {
90            return Err(invalid());
91        }
92        let modes = product(&[2, n])?;
93        let period = read_f64(&mut reader)?;
94        if period <= 0.0 {
95            return Err(invalid());
96        }
97        let mut wavenumbers = numerics::filled(modes, Complex::default())?;
98        for wavenumber in &mut wavenumbers {
99            *wavenumber = read_complex(&mut reader)?;
100        }
101        let periodic = PeriodicResidual::from_state(reader.raw(PeriodicResidual::state_size(n)?)?)?;
102        let eigen = EigResidual::from_state(reader.raw(EigResidual::state_size(modes)?)?)?;
103        if periodic.top.nrows() != n || eigen.values().len() != modes {
104            return Err(invalid());
105        }
106        reader.finish()?;
107        Ok(Self {
108            periodic,
109            eigen,
110            period,
111            wavenumbers,
112        })
113    }
114}