treams_core/cw/
expansion_lattice_state.rs

1//! Saved periodic geometry and coupling labels; restoring performs no Ewald sum.
2//! treams-rs extension.
3
4use super::{Couplings, LatticeExpansionResidual};
5use crate::{
6    Result,
7    basis::validate_wavenumbers,
8    lattice::BlochLattice,
9    numerics::finite,
10    saved::{self, Reader, SavedState, Writer},
11};
12
13impl LatticeExpansionResidual {
14    /// Fixed byte count from the destination and source basis shapes.
15    pub fn state_size(
16        destination_modes: usize,
17        destination_positions: usize,
18        source_modes: usize,
19        source_positions: usize,
20    ) -> Result<usize> {
21        saved::cw_basis_size(destination_modes, destination_positions)?
22            .checked_add(saved::cw_basis_size(source_modes, source_positions)?)
23            .and_then(|size| size.checked_add(BlochLattice::STATE_SIZE + 48))
24            .ok_or_else(saved::invalid)
25    }
26}
27
28impl SavedState for LatticeExpansionResidual {
29    fn save_state(&self) -> Result<Vec<u8>> {
30        let mut writer = Writer::new(Self::state_size(
31            self.destination.modes.len(),
32            self.destination.positions.len(),
33            self.source.modes.len(),
34            self.source.positions.len(),
35        )?);
36        saved::write_cw_basis(&mut writer, &self.destination);
37        saved::write_cw_basis(&mut writer, &self.source);
38        for k in self.ks {
39            writer.complex(k);
40        }
41        writer.complex(self.eta);
42        self.lattice.write_state(&mut writer);
43        Ok(writer.finish())
44    }
45
46    fn from_state(bytes: &[u8]) -> Result<Self> {
47        let mut shape = Reader::new(bytes);
48        let (dm, dp) = saved::read_cw_basis_dimensions(&mut shape)?;
49        let (sm, sp) = saved::read_cw_basis_dimensions(&mut shape)?;
50        if bytes.len() != Self::state_size(dm, dp, sm, sp)? {
51            return Err(saved::invalid());
52        }
53        let mut reader = Reader::new(bytes);
54        let destination = saved::read_cw_basis(&mut reader)?;
55        let source = saved::read_cw_basis(&mut reader)?;
56        let ks = [reader.complex()?, reader.complex()?];
57        let eta = reader.complex()?;
58        let lattice = BlochLattice::read_state(&mut reader)?;
59        reader.finish()?;
60        validate_wavenumbers(ks, true, false)?;
61        if !finite(eta) || lattice.dimension() > 2 {
62            return Err(saved::invalid());
63        }
64        let couplings = Couplings::new(&destination, &source, ks)?;
65        Ok(Self {
66            destination,
67            source,
68            ks,
69            lattice,
70            eta,
71            couplings,
72        })
73    }
74}
treams_core/cw/expansion_lattice_state.rs

treams_core/cw/
expansion_lattice_state.rs
