molpack/context/
work_buffers.rs

1//! Reusable temporary buffers for objective/gradient and movebad paths.
2
3use molrs::types::F;
4
5/// Reusable mutable buffers shared across packing iterations.
6pub struct WorkBuffers {
7    /// Cartesian gradient accumulator used by objective gradient evaluation.
8    pub gxcar: Vec<[F; 3]>,
9    /// Per-rayon-worker scratch gradient buffers for the parallel pair-gradient
10    /// path. Flat, thread-major: worker `t` owns the contiguous region
11    /// `[t*ntotat .. (t+1)*ntotat)`. Each worker accumulates its half-stencil
12    /// pair forces here race-free (its region is touched by no other
13    /// concurrently-running task), then the regions are reduced into [`gxcar`].
14    /// Reusing one persistent buffer keeps the hot path allocation-free; it is
15    /// zeroed (in parallel) at the start of each gradient evaluation. Sized
16    /// lazily to `nthreads * ntotat`.
17    ///
18    /// [`gxcar`]: Self::gxcar
19    #[cfg(feature = "rayon")]
20    pub grad_partials: Vec<[F; 3]>,
21    /// Reused per-active-molecule descriptor list `(itype, icart0, ilubar,
22    /// ilugan)` for the phase-structured `expand_molecules` /
23    /// `project_cartesian_gradient` passes. Both rebuild it from the current
24    /// `comptype` each call (cheap index arithmetic, no trig); persisting the
25    /// `Vec` keeps that hot rebuild allocation-free.
26    pub mol_descs: Vec<(usize, usize, usize, usize)>,
27    /// Temporary radius backup used by movebad/radius scaling paths.
28    pub radiuswork: Vec<F>,
29    /// Per-molecule score buffer used by flashsort/movebad ranking.
30    pub fmol: Vec<F>,
31    /// Index permutation buffer reused by flashsort in movebad.
32    pub flash_ind: Vec<usize>,
33    /// Histogram bucket buffer reused by flashsort.
34    pub flash_l: Vec<usize>,
35    /// Last x-vector whose expanded Cartesian geometry is still resident in `PackContext`.
36    pub cached_x: Vec<F>,
37    /// Active-type mask associated with `cached_x`.
38    pub cached_comptype: Vec<bool>,
39    /// Whether the cached geometry was built in init1 mode.
40    pub cached_init1: bool,
41    /// Cell grid signature for the cached geometry.
42    pub cached_ncells: [usize; 3],
43    pub cached_cell_length: [F; 3],
44    pub cached_pbc_min: [F; 3],
45    pub cached_pbc_length: [F; 3],
46    pub cached_pbc_periodic: [bool; 3],
47    /// Whether the cached geometry metadata is valid.
48    pub cached_geometry_valid: bool,
49}
50
51impl WorkBuffers {
52    pub fn new(ntotat: usize) -> Self {
53        Self {
54            gxcar: vec![[0.0; 3]; ntotat],
55            #[cfg(feature = "rayon")]
56            grad_partials: Vec::new(),
57            mol_descs: Vec::new(),
58            radiuswork: vec![0.0; ntotat],
59            fmol: Vec::new(),
60            flash_ind: Vec::new(),
61            flash_l: Vec::new(),
62            cached_x: Vec::new(),
63            cached_comptype: Vec::new(),
64            cached_init1: false,
65            cached_ncells: [0; 3],
66            cached_cell_length: [0.0; 3],
67            cached_pbc_min: [0.0; 3],
68            cached_pbc_length: [0.0; 3],
69            cached_pbc_periodic: [false; 3],
70            cached_geometry_valid: false,
71        }
72    }
73
74    pub fn ensure_atom_capacity(&mut self, ntotat: usize) {
75        if self.gxcar.len() != ntotat {
76            self.gxcar.resize(ntotat, [0.0; 3]);
77        }
78        if self.radiuswork.len() != ntotat {
79            self.radiuswork.resize(ntotat, 0.0);
80        }
81        self.cached_geometry_valid = false;
82    }
83
84    #[allow(clippy::too_many_arguments)]
85    pub fn matches_cached_geometry(
86        &self,
87        x: &[F],
88        comptype: &[bool],
89        init1: bool,
90        ncells: [usize; 3],
91        cell_length: [F; 3],
92        pbc_min: [F; 3],
93        pbc_length: [F; 3],
94        pbc_periodic: [bool; 3],
95    ) -> bool {
96        self.cached_geometry_valid
97            && self.cached_init1 == init1
98            && self.cached_ncells == ncells
99            && self.cached_cell_length == cell_length
100            && self.cached_pbc_min == pbc_min
101            && self.cached_pbc_length == pbc_length
102            && self.cached_pbc_periodic == pbc_periodic
103            && self.cached_x == x
104            && self.cached_comptype == comptype
105    }
106
107    #[allow(clippy::too_many_arguments)]
108    pub fn update_cached_geometry(
109        &mut self,
110        x: &[F],
111        comptype: &[bool],
112        init1: bool,
113        ncells: [usize; 3],
114        cell_length: [F; 3],
115        pbc_min: [F; 3],
116        pbc_length: [F; 3],
117        pbc_periodic: [bool; 3],
118    ) {
119        self.cached_x.clear();
120        self.cached_x.extend_from_slice(x);
121        self.cached_comptype.clear();
122        self.cached_comptype.extend_from_slice(comptype);
123        self.cached_init1 = init1;
124        self.cached_ncells = ncells;
125        self.cached_cell_length = cell_length;
126        self.cached_pbc_min = pbc_min;
127        self.cached_pbc_length = pbc_length;
128        self.cached_pbc_periodic = pbc_periodic;
129        self.cached_geometry_valid = true;
130    }
131}
molpack/context/work_buffers.rs

molpack/context/
work_buffers.rs
