Skip to main content

uff_relax/
spatial.rs

1use glam::DVec3;
2use crate::cell::{UnitCell, CellType};
3
4pub struct CellList {
5    pub cells: Vec<Vec<usize>>,
6    pub dx: usize,
7    pub dy: usize,
8    pub dz: usize,
9    pub cell_size: DVec3,
10    pub min_p: DVec3,
11}
12
13impl CellList {
14    pub fn build(positions: &[DVec3], cell: &UnitCell, cutoff: f64) -> Self {
15        let (min_p, max_p) = match cell.cell_type {
16            CellType::None => {
17                let mut min = positions[0];
18                let mut max = positions[0];
19                for &p in positions {
20                    min = min.min(p);
21                    max = max.max(p);
22                }
23                (min, max)
24            }
25            CellType::Orthorhombic { size } => (DVec3::ZERO, size),
26            CellType::Triclinic { matrix } => {
27                // For triclinic, we use a simple bounding box of the cell vectors
28                let mut min = DVec3::ZERO;
29                let mut max = DVec3::ZERO;
30                let corners = [
31                    DVec3::ZERO,
32                    matrix.col(0),
33                    matrix.col(1),
34                    matrix.col(2),
35                    matrix.col(0) + matrix.col(1),
36                    matrix.col(0) + matrix.col(2),
37                    matrix.col(1) + matrix.col(2),
38                    matrix.col(0) + matrix.col(1) + matrix.col(2),
39                ];
40                for &c in &corners {
41                    min = min.min(c);
42                    max = max.max(c);
43                }
44                (min, max)
45            }
46        };
47
48        let span = max_p - min_p;
49        let dx = (span.x / cutoff).floor() as usize;
50        let dy = (span.y / cutoff).floor() as usize;
51        let dz = (span.z / cutoff).floor() as usize;
52        
53        let dx = dx.max(1);
54        let dy = dy.max(1);
55        let dz = dz.max(1);
56        
57        let cell_size = DVec3::new(span.x / dx as f64, span.y / dy as f64, span.z / dz as f64);
58        let mut cells = vec![Vec::new(); dx * dy * dz];
59
60        for (i, &p) in positions.iter().enumerate() {
61            let rel = p - min_p;
62            let ix = ((rel.x / cell_size.x) as usize).min(dx - 1);
63            let iy = ((rel.y / cell_size.y) as usize).min(dy - 1);
64            let iz = ((rel.z / cell_size.z) as usize).min(dz - 1);
65            cells[ix * dy * dz + iy * dz + iz].push(i);
66        }
67
68        Self { cells, dx, dy, dz, cell_size, min_p }
69    }
70}
71
72#[cfg(test)]
73mod tests {
74    use super::*;
75
76    #[test]
77    fn test_cell_list_build() {
78        let positions = vec![
79            DVec3::new(1.0, 1.0, 1.0),
80            DVec3::new(1.1, 1.1, 1.1),
81            DVec3::new(5.0, 5.0, 5.0),
82        ];
83        let cell = UnitCell::new_none();
84        let cutoff = 2.0;
85        let cl = CellList::build(&positions, &cell, cutoff);
86        
87        // Check that points close together are in the same or adjacent cells
88        // In this case (1,1,1) and (1.1, 1.1, 1.1) should be very close.
89        assert!(cl.cells.iter().any(|c| c.len() >= 2 || (c.len() == 1)));
90    }
91}