molrs_compute/density/
correlation_function.rs

1//! Generic distance-binned correlation `⟨A_i · B_j⟩(r)`.
2//!
3//! Mirrors `freud.density.CorrelationFunction`
4//! ([source](https://github.com/glotzerlab/freud/blob/main/freud/density/CorrelationFunction.cc)).
5//!
6//! Given per-particle real scalar fields `A_i` and `B_j` and a neighbor
7//! list, builds the histogram `⟨A_i · B_j⟩(r)` averaged over all pairs
8//! that fall into each shell `[r, r + dr)`. Empty bins return 0.
9//!
10//! freud also supports complex-valued fields with the convention
11//! `A_i · conj(B_j)`; the real version is implemented here. The complex
12//! variant requires only a handful of extra lines and will follow when the
13//! first downstream consumer (e.g. `LocalDescriptors` in Phase 6) needs it.
14
15use molrs::spatial::neighbors::NeighborList;
16use molrs::store::frame_access::FrameAccess;
17use molrs::types::F;
18use ndarray::Array1;
19
20use crate::error::ComputeError;
21use crate::result::ComputeResult;
22use crate::traits::Compute;
23
24/// Per-frame correlation result.
25#[derive(Debug, Clone, Default)]
26pub struct CorrelationFunctionResult {
27    /// Bin edges (length `n_bins + 1`).
28    pub bin_edges: Array1<F>,
29    /// Bin centres (length `n_bins`).
30    pub bin_centers: Array1<F>,
31    /// Running pair count per bin (length `n_bins`).
32    pub bin_counts: Array1<u64>,
33    /// `⟨A · B⟩(r)` per bin (length `n_bins`); zero for empty bins.
34    pub correlation: Array1<F>,
35}
36
37impl ComputeResult for CorrelationFunctionResult {}
38
39/// Correlation-function calculator. Stateless container of bin parameters.
40#[derive(Debug, Clone)]
41pub struct CorrelationFunction {
42    n_bins: usize,
43    r_min: F,
44    r_max: F,
45    bin_width: F,
46    bin_edges: Array1<F>,
47    bin_centers: Array1<F>,
48}
49
50impl CorrelationFunction {
51    pub fn new(n_bins: usize, r_max: F, r_min: F) -> Result<Self, ComputeError> {
52        if n_bins == 0 {
53            return Err(ComputeError::OutOfRange {
54                field: "CorrelationFunction::n_bins",
55                value: "0".into(),
56            });
57        }
58        if r_min.is_nan() || r_min < 0.0 {
59            return Err(ComputeError::OutOfRange {
60                field: "CorrelationFunction::r_min",
61                value: r_min.to_string(),
62            });
63        }
64        if r_max.is_nan() || r_max <= r_min {
65            return Err(ComputeError::OutOfRange {
66                field: "CorrelationFunction::r_max",
67                value: format!("r_max={r_max}, r_min={r_min}"),
68            });
69        }
70        let bin_width = (r_max - r_min) / n_bins as F;
71        let bin_edges = Array1::from_iter((0..=n_bins).map(|i| r_min + i as F * bin_width));
72        let bin_centers =
73            Array1::from_iter((0..n_bins).map(|i| r_min + (i as F + 0.5) * bin_width));
74        Ok(Self {
75            n_bins,
76            r_min,
77            r_max,
78            bin_width,
79            bin_edges,
80            bin_centers,
81        })
82    }
83
84    pub fn n_bins(&self) -> usize {
85        self.n_bins
86    }
87    pub fn bin_edges(&self) -> &Array1<F> {
88        &self.bin_edges
89    }
90    pub fn bin_centers(&self) -> &Array1<F> {
91        &self.bin_centers
92    }
93
94    fn one_frame(
95        &self,
96        nlist: &NeighborList,
97        values_a: &[F],
98        values_b: &[F],
99    ) -> Result<CorrelationFunctionResult, ComputeError> {
100        let r_min_sq = self.r_min * self.r_min;
101        let r_max_sq = self.r_max * self.r_max;
102        let i_idx = nlist.query_point_indices();
103        let j_idx = nlist.point_indices();
104        let dist_sq = nlist.dist_sq();
105        let n_pairs = nlist.n_pairs();
106        let symmetric = matches!(
107            nlist.mode(),
108            molrs::spatial::neighbors::QueryMode::SelfQuery
109        );
110
111        let mut sum = Array1::<F>::zeros(self.n_bins);
112        let mut counts = Array1::<u64>::zeros(self.n_bins);
113
114        for k in 0..n_pairs {
115            let d2 = dist_sq[k];
116            if d2 < r_min_sq || d2 >= r_max_sq || d2 == 0.0 {
117                continue;
118            }
119            let d = d2.sqrt();
120            let b = ((d - self.r_min) / self.bin_width) as usize;
121            if b >= self.n_bins {
122                continue;
123            }
124            let i = i_idx[k] as usize;
125            let j = j_idx[k] as usize;
126            if i >= values_a.len() || j >= values_b.len() {
127                return Err(ComputeError::DimensionMismatch {
128                    expected: i.max(j) + 1,
129                    got: values_a.len().min(values_b.len()),
130                    what: "CorrelationFunction values length",
131                });
132            }
133            sum[b] += values_a[i] * values_b[j];
134            counts[b] += 1;
135            if symmetric {
136                // For a self-query list, the symmetric pair (j → i) is
137                // implicit; honour it for ⟨A · B⟩.
138                sum[b] += values_a[j] * values_b[i];
139                counts[b] += 1;
140            }
141        }
142
143        let mut correlation = Array1::<F>::zeros(self.n_bins);
144        for b in 0..self.n_bins {
145            if counts[b] > 0 {
146                correlation[b] = sum[b] / counts[b] as F;
147            }
148        }
149        Ok(CorrelationFunctionResult {
150            bin_edges: self.bin_edges.clone(),
151            bin_centers: self.bin_centers.clone(),
152            bin_counts: counts,
153            correlation,
154        })
155    }
156}
157
158/// `Args` for [`CorrelationFunction`]: parallel per-frame triplets of
159/// `(neighbor list, A values, B values)`. All three slices must have the
160/// same length as `frames`.
161pub struct CorrelationArgs<'a> {
162    pub nlists: &'a [NeighborList],
163    pub values_a: &'a [Vec<F>],
164    pub values_b: &'a [Vec<F>],
165}
166
167impl Compute for CorrelationFunction {
168    type Args<'a> = CorrelationArgs<'a>;
169    type Output = Vec<CorrelationFunctionResult>;
170
171    fn compute<'a, FA: FrameAccess + Sync + 'a>(
172        &self,
173        frames: &[&'a FA],
174        args: CorrelationArgs<'a>,
175    ) -> Result<Vec<CorrelationFunctionResult>, ComputeError> {
176        if frames.is_empty() {
177            return Err(ComputeError::EmptyInput);
178        }
179        let nf = frames.len();
180        if args.nlists.len() != nf || args.values_a.len() != nf || args.values_b.len() != nf {
181            return Err(ComputeError::DimensionMismatch {
182                expected: nf,
183                got: args
184                    .nlists
185                    .len()
186                    .min(args.values_a.len())
187                    .min(args.values_b.len()),
188                what: "CorrelationFunction frame-aligned inputs",
189            });
190        }
191        let mut out = Vec::with_capacity(nf);
192        for k in 0..nf {
193            out.push(self.one_frame(&args.nlists[k], &args.values_a[k], &args.values_b[k])?);
194        }
195        Ok(out)
196    }
197}
198
199#[cfg(test)]
200mod tests {
201    use super::*;
202    use molrs::Frame;
203    use molrs::spatial::neighbors::{LinkCell, NbListAlgo};
204    use molrs::spatial::region::simbox::SimBox;
205    use molrs::store::block::Block;
206    use ndarray::{Array1 as A1, array};
207
208    fn frame_with(positions: &[[F; 3]], box_len: F) -> Frame {
209        let x = A1::from_iter(positions.iter().map(|p| p[0]));
210        let y = A1::from_iter(positions.iter().map(|p| p[1]));
211        let z = A1::from_iter(positions.iter().map(|p| p[2]));
212        let mut block = Block::new();
213        block.insert("x", x.into_dyn()).unwrap();
214        block.insert("y", y.into_dyn()).unwrap();
215        block.insert("z", z.into_dyn()).unwrap();
216        let mut frame = Frame::new();
217        frame.insert("atoms", block);
218        frame.simbox =
219            Some(SimBox::cube(box_len, array![0.0 as F, 0.0 as F, 0.0 as F], [false; 3]).unwrap());
220        frame
221    }
222
223    fn build_nlist(frame: &Frame, cutoff: F) -> NeighborList {
224        let xp = frame
225            .get("atoms")
226            .unwrap()
227            .get("x")
228            .and_then(<F as molrs::store::block::BlockDtype>::from_column)
229            .unwrap()
230            .as_slice()
231            .unwrap()
232            .to_vec();
233        let yp = frame
234            .get("atoms")
235            .unwrap()
236            .get("y")
237            .and_then(<F as molrs::store::block::BlockDtype>::from_column)
238            .unwrap()
239            .as_slice()
240            .unwrap()
241            .to_vec();
242        let zp = frame
243            .get("atoms")
244            .unwrap()
245            .get("z")
246            .and_then(<F as molrs::store::block::BlockDtype>::from_column)
247            .unwrap()
248            .as_slice()
249            .unwrap()
250            .to_vec();
251        let n = xp.len();
252        let mut pos = ndarray::Array2::<F>::zeros((n, 3));
253        for i in 0..n {
254            pos[[i, 0]] = xp[i];
255            pos[[i, 1]] = yp[i];
256            pos[[i, 2]] = zp[i];
257        }
258        let simbox = frame.simbox.as_ref().unwrap();
259        let mut lc = LinkCell::new().cutoff(cutoff);
260        lc.build(pos.view(), simbox);
261        lc.query().clone()
262    }
263
264    #[test]
265    fn constant_values_yield_constant_correlation() {
266        // A = B = 1.0 for all particles → ⟨A·B⟩(r) = 1 in every populated bin.
267        let positions = [
268            [1.0_f64, 1.0, 1.0],
269            [2.0, 1.0, 1.0],
270            [3.0, 1.0, 1.0],
271            [4.0, 1.0, 1.0],
272        ];
273        let frame = frame_with(&positions, 20.0);
274        let nl = build_nlist(&frame, 5.0);
275        let vals = vec![1.0_f64; positions.len()];
276        let cf = CorrelationFunction::new(10, 5.0, 0.0).unwrap();
277        let r = &cf
278            .compute(
279                &[&frame],
280                CorrelationArgs {
281                    nlists: &[nl],
282                    values_a: std::slice::from_ref(&vals),
283                    values_b: std::slice::from_ref(&vals),
284                },
285            )
286            .unwrap()[0];
287        for b in 0..10 {
288            if r.bin_counts[b] > 0 {
289                assert!(
290                    (r.correlation[b] - 1.0).abs() < 1e-12,
291                    "bin {b}: got {}",
292                    r.correlation[b]
293                );
294            }
295        }
296    }
297
298    #[test]
299    fn alternating_values_average_to_minus_one() {
300        // A = B = [+1, -1, +1, -1]. Each NN pair is opposite-sign so
301        // A_i · B_j = -1; bin 0 (the only populated one for unit spacing)
302        // should average to -1.
303        let positions = [
304            [0.0_f64, 0.0, 0.0],
305            [1.0, 0.0, 0.0],
306            [2.0, 0.0, 0.0],
307            [3.0, 0.0, 0.0],
308        ];
309        let frame = frame_with(&positions, 20.0);
310        let nl = build_nlist(&frame, 1.2);
311        let vals = vec![1.0_f64, -1.0, 1.0, -1.0];
312        let cf = CorrelationFunction::new(5, 1.5, 0.0).unwrap();
313        let r = &cf
314            .compute(
315                &[&frame],
316                CorrelationArgs {
317                    nlists: &[nl],
318                    values_a: std::slice::from_ref(&vals),
319                    values_b: std::slice::from_ref(&vals),
320                },
321            )
322            .unwrap()[0];
323        // The first non-empty bin contains the unit-spacing NN pairs.
324        let first = r.bin_counts.iter().position(|&c| c > 0).unwrap();
325        assert!((r.correlation[first] + 1.0).abs() < 1e-12);
326    }
327
328    #[test]
329    fn empty_bin_returns_zero() {
330        let frame = frame_with(&[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], 20.0);
331        let nl = build_nlist(&frame, 1.5);
332        let vals = vec![1.0_f64, 1.0];
333        let cf = CorrelationFunction::new(10, 5.0, 0.0).unwrap();
334        let r = &cf
335            .compute(
336                &[&frame],
337                CorrelationArgs {
338                    nlists: &[nl],
339                    values_a: std::slice::from_ref(&vals),
340                    values_b: std::slice::from_ref(&vals),
341                },
342            )
343            .unwrap()[0];
344        // Bins beyond 1.5 should all be empty.
345        for b in 4..10 {
346            assert_eq!(r.bin_counts[b], 0);
347            assert_eq!(r.correlation[b], 0.0);
348        }
349    }
350
351    #[test]
352    fn invalid_params_error() {
353        assert!(CorrelationFunction::new(0, 1.0, 0.0).is_err());
354        assert!(CorrelationFunction::new(10, 1.0, 1.0).is_err());
355        assert!(CorrelationFunction::new(10, 0.5, 1.0).is_err());
356    }
357}
molrs_compute/density/correlation_function.rs

molrs_compute/density/
correlation_function.rs
