molcrafts-molpack 0.1.0

Faithful Rust port of Packmol molecular packing (split off from molrs)
Documentation
//! Reusable temporary buffers for objective/gradient and movebad paths.

use molrs::types::F;

/// Reusable mutable buffers shared across packing iterations.
pub struct WorkBuffers {
    /// Cartesian gradient accumulator used by objective gradient evaluation.
    pub gxcar: Vec<[F; 3]>,
    /// Per-rayon-worker scratch gradient buffers for the parallel pair-gradient
    /// path. Flat, thread-major: worker `t` owns the contiguous region
    /// `[t*ntotat .. (t+1)*ntotat)`. Each worker accumulates its half-stencil
    /// pair forces here race-free (its region is touched by no other
    /// concurrently-running task), then the regions are reduced into [`gxcar`].
    /// Reusing one persistent buffer keeps the hot path allocation-free; it is
    /// zeroed (in parallel) at the start of each gradient evaluation. Sized
    /// lazily to `nthreads * ntotat`.
    ///
    /// [`gxcar`]: Self::gxcar
    #[cfg(feature = "rayon")]
    pub grad_partials: Vec<[F; 3]>,
    /// Reused per-active-molecule descriptor list `(itype, icart0, ilubar,
    /// ilugan)` for the phase-structured `expand_molecules` /
    /// `project_cartesian_gradient` passes. Both rebuild it from the current
    /// `comptype` each call (cheap index arithmetic, no trig); persisting the
    /// `Vec` keeps that hot rebuild allocation-free.
    pub mol_descs: Vec<(usize, usize, usize, usize)>,
    /// Temporary radius backup used by movebad/radius scaling paths.
    pub radiuswork: Vec<F>,
    /// Per-molecule score buffer used by flashsort/movebad ranking.
    pub fmol: Vec<F>,
    /// Index permutation buffer reused by flashsort in movebad.
    pub flash_ind: Vec<usize>,
    /// Histogram bucket buffer reused by flashsort.
    pub flash_l: Vec<usize>,
    /// Last x-vector whose expanded Cartesian geometry is still resident in `PackContext`.
    pub cached_x: Vec<F>,
    /// Active-type mask associated with `cached_x`.
    pub cached_comptype: Vec<bool>,
    /// Whether the cached geometry was built in init1 mode.
    pub cached_init1: bool,
    /// Cell grid signature for the cached geometry.
    pub cached_ncells: [usize; 3],
    pub cached_cell_length: [F; 3],
    pub cached_pbc_min: [F; 3],
    pub cached_pbc_length: [F; 3],
    pub cached_pbc_periodic: [bool; 3],
    /// Whether the cached geometry metadata is valid.
    pub cached_geometry_valid: bool,
}

impl WorkBuffers {
    pub fn new(ntotat: usize) -> Self {
        Self {
            gxcar: vec![[0.0; 3]; ntotat],
            #[cfg(feature = "rayon")]
            grad_partials: Vec::new(),
            mol_descs: Vec::new(),
            radiuswork: vec![0.0; ntotat],
            fmol: Vec::new(),
            flash_ind: Vec::new(),
            flash_l: Vec::new(),
            cached_x: Vec::new(),
            cached_comptype: Vec::new(),
            cached_init1: false,
            cached_ncells: [0; 3],
            cached_cell_length: [0.0; 3],
            cached_pbc_min: [0.0; 3],
            cached_pbc_length: [0.0; 3],
            cached_pbc_periodic: [false; 3],
            cached_geometry_valid: false,
        }
    }

    pub fn ensure_atom_capacity(&mut self, ntotat: usize) {
        if self.gxcar.len() != ntotat {
            self.gxcar.resize(ntotat, [0.0; 3]);
        }
        if self.radiuswork.len() != ntotat {
            self.radiuswork.resize(ntotat, 0.0);
        }
        self.cached_geometry_valid = false;
    }

    #[allow(clippy::too_many_arguments)]
    pub fn matches_cached_geometry(
        &self,
        x: &[F],
        comptype: &[bool],
        init1: bool,
        ncells: [usize; 3],
        cell_length: [F; 3],
        pbc_min: [F; 3],
        pbc_length: [F; 3],
        pbc_periodic: [bool; 3],
    ) -> bool {
        self.cached_geometry_valid
            && self.cached_init1 == init1
            && self.cached_ncells == ncells
            && self.cached_cell_length == cell_length
            && self.cached_pbc_min == pbc_min
            && self.cached_pbc_length == pbc_length
            && self.cached_pbc_periodic == pbc_periodic
            && self.cached_x == x
            && self.cached_comptype == comptype
    }

    #[allow(clippy::too_many_arguments)]
    pub fn update_cached_geometry(
        &mut self,
        x: &[F],
        comptype: &[bool],
        init1: bool,
        ncells: [usize; 3],
        cell_length: [F; 3],
        pbc_min: [F; 3],
        pbc_length: [F; 3],
        pbc_periodic: [bool; 3],
    ) {
        self.cached_x.clear();
        self.cached_x.extend_from_slice(x);
        self.cached_comptype.clear();
        self.cached_comptype.extend_from_slice(comptype);
        self.cached_init1 = init1;
        self.cached_ncells = ncells;
        self.cached_cell_length = cell_length;
        self.cached_pbc_min = pbc_min;
        self.cached_pbc_length = pbc_length;
        self.cached_pbc_periodic = pbc_periodic;
        self.cached_geometry_valid = true;
    }
}