molrs_core/
grid.rs

1//! Uniform spatial grid storing multiple named scalar arrays.
2//!
3//! A [`Grid`] holds any number of named scalar fields (e.g. `electron_density`,
4//! `spin_density`) that all share the same spatial grid: same dimensions,
5//! same origin, and same cell vectors. This maps naturally onto VASP CHGCAR
6//! (2–4 arrays on one grid) and Gaussian cube files (one array per file,
7//! possibly combined by the caller).
8//!
9//! Grid positions are **not stored** — they are computed on demand from the
10//! cell and dim fields.
11
12use std::collections::HashMap;
13
14use ndarray::{Array3, ArrayD};
15
16use crate::error::MolRsError;
17use crate::types::F;
18
19/// A collection of named scalar arrays on a shared uniform spatial grid.
20///
21/// Cell vectors are stored as columns (same convention as [`SimBox`][crate::region::simbox::SimBox]).
22/// The Cartesian position of voxel `(i, j, k)` is:
23/// ```text
24/// origin + (i/nx)*cell_col0 + (j/ny)*cell_col1 + (k/nz)*cell_col2
25/// ```
26#[derive(Debug, Clone, PartialEq)]
27pub struct Grid {
28    /// Grid dimensions `[nx, ny, nz]`.
29    pub dim: [usize; 3],
30    /// Cartesian origin in Ångström.
31    pub origin: [F; 3],
32    /// Cell matrix: columns are the three lattice vectors (Å).
33    pub cell: [[F; 3]; 3],
34    /// Periodic boundary flags for each axis.
35    pub pbc: [bool; 3],
36    /// Named scalar arrays stored in row-major `(ix, iy, iz)` order.
37    arrays: HashMap<String, Vec<F>>,
38}
39
40impl Grid {
41    /// Create an empty grid with the given spatial definition.
42    pub fn new(dim: [usize; 3], origin: [F; 3], cell: [[F; 3]; 3], pbc: [bool; 3]) -> Self {
43        Self {
44            dim,
45            origin,
46            cell,
47            pbc,
48            arrays: HashMap::new(),
49        }
50    }
51
52    /// Total number of voxels: `nx * ny * nz`.
53    pub fn total(&self) -> usize {
54        self.dim[0] * self.dim[1] * self.dim[2]
55    }
56
57    /// Insert (or replace) a named array.
58    ///
59    /// `data` must have length `nx * ny * nz` in row-major `(ix, iy, iz)` order.
60    pub fn insert(&mut self, name: impl Into<String>, data: Vec<F>) -> Result<(), MolRsError> {
61        let expected = self.total();
62        let name = name.into();
63        if data.len() != expected {
64            return Err(MolRsError::validation(format!(
65                "grid array '{}' length mismatch: expected {}, got {}",
66                name,
67                expected,
68                data.len()
69            )));
70        }
71        self.arrays.insert(name, data);
72        Ok(())
73    }
74
75    /// Return a named array reshaped to `(nx, ny, nz)`, or `None` if absent.
76    pub fn get(&self, name: &str) -> Option<ArrayD<F>> {
77        self.arrays.get(name).map(|data| {
78            Array3::from_shape_vec([self.dim[0], self.dim[1], self.dim[2]], data.clone())
79                .expect("grid shape matches stored data")
80                .into_dyn()
81        })
82    }
83
84    /// Borrow the raw flat slice for a named array, or `None` if absent.
85    pub fn get_raw(&self, name: &str) -> Option<&[F]> {
86        self.arrays.get(name).map(|v| v.as_slice())
87    }
88
89    /// Whether a named array is present.
90    pub fn contains(&self, name: &str) -> bool {
91        self.arrays.contains_key(name)
92    }
93
94    /// Number of named arrays stored in this grid.
95    pub fn len(&self) -> usize {
96        self.arrays.len()
97    }
98
99    /// Returns `true` if no arrays are stored.
100    pub fn is_empty(&self) -> bool {
101        self.arrays.is_empty()
102    }
103
104    /// Iterate over `(name, flat_data)` pairs.
105    pub fn raw_arrays(&self) -> impl Iterator<Item = (&str, &[F])> {
106        self.arrays.iter().map(|(k, v)| (k.as_str(), v.as_slice()))
107    }
108
109    /// Iterate over array names.
110    pub fn keys(&self) -> impl Iterator<Item = &str> {
111        self.arrays.keys().map(|s| s.as_str())
112    }
113
114    /// Compute Cartesian position of voxel `(ix, iy, iz)`.
115    pub fn voxel_position(&self, ix: usize, iy: usize, iz: usize) -> [F; 3] {
116        let fx = ix as F / self.dim[0] as F;
117        let fy = iy as F / self.dim[1] as F;
118        let fz = iz as F / self.dim[2] as F;
119        [
120            self.origin[0] + fx * self.cell[0][0] + fy * self.cell[1][0] + fz * self.cell[2][0],
121            self.origin[1] + fx * self.cell[0][1] + fy * self.cell[1][1] + fz * self.cell[2][1],
122            self.origin[2] + fx * self.cell[0][2] + fy * self.cell[1][2] + fz * self.cell[2][2],
123        ]
124    }
125}
126
127#[cfg(test)]
128mod tests {
129    use super::*;
130
131    #[test]
132    fn insert_validates_length() {
133        let mut g = Grid::new(
134            [2, 2, 2],
135            [0.0; 3],
136            [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]],
137            [false; 3],
138        );
139        assert!(g.insert("rho", vec![0.0; 7]).is_err());
140        assert!(g.insert("rho", vec![0.0; 8]).is_ok());
141    }
142
143    #[test]
144    fn get_returns_shaped_array() {
145        let mut g = Grid::new(
146            [2, 3, 4],
147            [0.0; 3],
148            [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]],
149            [false; 3],
150        );
151        g.insert("rho", (0..24).map(|x| x as F).collect()).unwrap();
152        let arr = g.get("rho").unwrap();
153        assert_eq!(arr.shape(), &[2, 3, 4]);
154    }
155}
molrs_core/grid.rs

molrs_core/
grid.rs
