molpack/restraint/collective/
mod.rs

1//! Collective (group-level) restraints — the **geometric distribution-matching**
2//! family.
3//!
4//! Where a per-atom [`AtomRestraint`](crate::restraint::AtomRestraint) sees one atom at a
5//! time and contributes an independent external field `∑ᵢ U(xᵢ)`, a
6//! [`Restraint`] sees *every* copy of a species at once and returns a
7//! single penalty whose gradient is **coupled across the whole group**. That
8//! coupling is what lets a species *follow* a target spatial distribution: a
9//! per-atom field built from a target density is minimised by collapsing every
10//! atom onto the density's mode, whereas a distribution-distance penalty is
11//! minimised when the empirical distribution *equals* the target.
12//!
13//! # Structure of this family
14//!
15//! Every member matches a target **distribution** of a scalar reaction
16//! coordinate ξ defined by a **geometry**, via the squared 1-D Wasserstein
17//! (sorted-CDF) metric ([`engine`]). The two axes are orthogonal:
18//!
19//! - **geometry** ([`geometry`]) — maps Cartesian coordinates to ξ and scatters
20//!   `∂L/∂ξ` back onto them: `plane` (ξ = signed distance to a plane → a slab),
21//!   `point` (ξ = distance to a centre → a spherical shell), …
22//! - **distribution** — the target quantile function `q(p) = F⁻¹(p)`: Gaussian,
23//!   exponential, …
24//!
25//! Concrete types are the cross product, named `<Distribution><Geometry>`
26//! and implementing [`Restraint`] directly (no wrapper, no builder —
27//! same direction-3 convention as the per-atom restraints):
28//! [`GaussianPlane`], [`GaussianPoint`], … Adding a distribution is a new
29//! quantile function; adding a geometry is a new ξ/scatter pair; a new concrete
30//! type then composes the two through the shared [`engine`].
31//!
32//! **Gradient convention** mirrors [`AtomRestraint`](crate::restraint::AtomRestraint):
33//! `fg` accumulates `∂L/∂coords[i]` INTO `grads[i]` with `+=`. `coords` and
34//! `grads` have equal length (one entry per atom in the group, same order).
35
36use molrs::types::F;
37
38// ============================================================================
39// Trait
40// ============================================================================
41
42/// Group-level penalty over all copies of one species.
43///
44/// Unlike [`AtomRestraint`](crate::restraint::AtomRestraint), which is evaluated once
45/// per atom with only that atom's coordinate, a `Restraint` is
46/// evaluated once per group with the coordinates of *all* copies. Its gradient
47/// may therefore couple every particle to every other — exactly what a
48/// distribution-matching penalty needs.
49pub trait Restraint: Send + Sync + std::fmt::Debug {
50    /// Penalty value for the group's current configuration.
51    ///
52    /// `coords[i]` is the Cartesian position of the `i`-th atom in the group.
53    fn f(&self, coords: &[[F; 3]], scale: F, scale2: F) -> F;
54
55    /// Fused value + gradient. Accumulates `∂L/∂coords[i]` INTO `grads[i]`
56    /// with `+=`; returns the same value `f` would. `grads.len() == coords.len()`.
57    fn fg(&self, coords: &[[F; 3]], scale: F, scale2: F, grads: &mut [[F; 3]]) -> F;
58
59    /// If `false`, the scheduler serializes this restraint (Python-backed
60    /// collective restraints MUST return `false`).
61    fn is_parallel_safe(&self) -> bool {
62        true
63    }
64
65    /// Human-readable identifier.
66    fn name(&self) -> &'static str {
67        std::any::type_name::<Self>()
68    }
69}
70
71mod engine;
72mod exponential;
73mod gaussian;
74mod geometry;
75mod tabulated;
76
77pub use exponential::{ExponentialPlane, ExponentialPoint};
78pub use gaussian::{GaussianPlane, GaussianPoint};
79pub use tabulated::{TabulatedPlane, TabulatedPoint};
80
81/// Shared test helpers for the concrete `<Distribution><Geometry>` types:
82/// a dependency-free RNG and a finite-difference gradient check.
83#[cfg(test)]
84pub(super) mod testutil {
85    use super::Restraint;
86    use molrs::types::F;
87
88    /// Deterministic xorshift64* uniform in `[lo, hi)` — no external dep.
89    pub(crate) fn rng_uniform(seed: &mut u64, lo: F, hi: F) -> F {
90        let mut x = *seed;
91        x ^= x >> 12;
92        x ^= x << 25;
93        x ^= x >> 27;
94        *seed = x;
95        let u = (x.wrapping_mul(0x2545F4914F6CDD1D) >> 11) as F / (1u64 << 53) as F;
96        lo + u * (hi - lo)
97    }
98
99    /// Central finite-difference check of the analytic gradient along every axis.
100    /// Coordinates must be spread enough that a 1e-6 perturbation never crosses a
101    /// rank swap (where the sorted-CDF objective is non-smooth).
102    pub(crate) fn assert_fd_grad(r: &dyn Restraint, coords: &[[F; 3]]) {
103        let mut analytic = vec![[0.0 as F; 3]; coords.len()];
104        r.fg(coords, 1.0, 1.0, &mut analytic);
105        let eps = 1e-6;
106        for i in 0..coords.len() {
107            for k in 0..3 {
108                let mut plus = coords.to_vec();
109                let mut minus = coords.to_vec();
110                plus[i][k] += eps;
111                minus[i][k] -= eps;
112                let fd = (r.f(&plus, 1.0, 1.0) - r.f(&minus, 1.0, 1.0)) / (2.0 * eps);
113                assert!(
114                    (fd - analytic[i][k]).abs() < 1e-4,
115                    "{} atom {i} axis {k}: fd={fd}, analytic={}",
116                    r.name(),
117                    analytic[i][k]
118                );
119            }
120        }
121    }
122}
molpack/restraint/collective/mod.rs

molpack/restraint/collective/
mod.rs
