molrs_compute/diffraction/
diffraction_pattern.rs

1// 2-D grid kernel reads naturally with double index loops; FFT2 helper
2// also needs explicit row/column index ordering.
3#![allow(clippy::needless_range_loop)]
4
5//! 2-D diffraction pattern (FFT of a projected density image).
6//!
7//! Mirrors `freud.diffraction.DiffractionPattern`
8//! ([source](https://github.com/glotzerlab/freud/blob/main/freud/diffraction/DiffractionPattern.cc)).
9//!
10//! Builds a 2-D Gaussian-smeared image of the projected particle positions,
11//! takes a 2-D FFT (composed from `rustfft`'s 1-D plans), and returns the
12//! power spectrum `|F(k)|²`. Projection axis defaults to `+z`, mapping
13//! particles onto the `xy` plane.
14//!
15//! This is the first analyzer in the port to use `rustfft 6` for real 2-D
16//! Fourier work; the FFT planner is built once per frame and reused for
17//! both row and column passes.
18//!
19//! # Conventions
20//!
21//! - The output image is FFT-shifted so that `k = 0` sits at the centre
22//!   `(n_grid / 2, n_grid / 2)`, matching freud (and the usual
23//!   `numpy.fft.fftshift` convention).
24//! - Square grid `(n_grid × n_grid)`; rectangular grids are a follow-up.
25//! - Orthorhombic boxes only (matches `freud.DiffractionPattern.compute`).
26
27use molrs::spatial::region::simbox::BoxKind;
28use molrs::store::frame_access::FrameAccess;
29use molrs::types::F;
30use ndarray::Array2;
31use rustfft::FftPlanner;
32use rustfft::num_complex::Complex as RfComplex;
33
34use crate::error::ComputeError;
35use crate::result::ComputeResult;
36use crate::traits::Compute;
37use crate::util::get_positions_ref;
38
39/// Per-frame diffraction-pattern result.
40#[derive(Debug, Clone, Default)]
41pub struct DiffractionPatternResult {
42    /// Power spectrum `|F(k)|²`, shape `(n_grid, n_grid)`. FFT-shifted so
43    /// the k=0 component is at the centre.
44    pub diffraction: Array2<F>,
45    /// Real-space Gaussian-smeared image used as the FFT input
46    /// (`(n_grid, n_grid)`).
47    pub image: Array2<F>,
48}
49
50impl ComputeResult for DiffractionPatternResult {}
51
52/// `DiffractionPattern` analyzer.
53#[derive(Debug, Clone, Copy)]
54pub struct DiffractionPattern {
55    n_grid: usize,
56    sigma: F,
57    /// Projection axis: 0 = x, 1 = y, 2 = z (default).
58    axis: usize,
59}
60
61impl DiffractionPattern {
62    /// New analyzer with `n_grid × n_grid` pixels and Gaussian smearing
63    /// `sigma` (in box-length units).
64    pub fn new(n_grid: usize, sigma: F) -> Result<Self, ComputeError> {
65        if n_grid == 0 {
66            return Err(ComputeError::OutOfRange {
67                field: "DiffractionPattern::n_grid",
68                value: "0".into(),
69            });
70        }
71        if sigma.is_nan() || sigma <= 0.0 {
72            return Err(ComputeError::OutOfRange {
73                field: "DiffractionPattern::sigma",
74                value: sigma.to_string(),
75            });
76        }
77        Ok(Self {
78            n_grid,
79            sigma,
80            axis: 2,
81        })
82    }
83
84    /// Set the projection axis (0/1/2 = x/y/z).
85    pub fn with_axis(mut self, axis: usize) -> Self {
86        assert!(axis < 3);
87        self.axis = axis;
88        self
89    }
90
91    pub fn n_grid(&self) -> usize {
92        self.n_grid
93    }
94    pub fn sigma(&self) -> F {
95        self.sigma
96    }
97
98    fn one_frame<FA: FrameAccess>(
99        &self,
100        frame: &FA,
101    ) -> Result<DiffractionPatternResult, ComputeError> {
102        let simbox = frame.simbox_ref().ok_or(ComputeError::MissingSimBox)?;
103        let lens = match simbox.kind() {
104            BoxKind::Ortho { len, .. } => [len[0], len[1], len[2]],
105            BoxKind::Triclinic => {
106                return Err(ComputeError::OutOfRange {
107                    field: "DiffractionPattern::simbox",
108                    value: "triclinic boxes not supported".into(),
109                });
110            }
111        };
112        let (a0, a1) = match self.axis {
113            0 => (1, 2),
114            1 => (0, 2),
115            _ => (0, 1),
116        };
117        let l0 = lens[a0];
118        let l1 = lens[a1];
119
120        let (xs_p, ys_p, zs_p) = get_positions_ref(frame)?;
121        let coords = [xs_p.slice(), ys_p.slice(), zs_p.slice()];
122
123        let n = self.n_grid;
124        let d0 = l0 / n as F;
125        let d1 = l1 / n as F;
126        let r_max = 3.0 * self.sigma;
127        let half_k0 = (r_max / d0).ceil() as isize;
128        let half_k1 = (r_max / d1).ceil() as isize;
129        let two_sigma_sq = 2.0 * self.sigma * self.sigma;
130        let pref = (two_sigma_sq * std::f64::consts::PI).powf(-1.0);
131        let r_max_sq = r_max * r_max;
132
133        let mut image = Array2::<F>::zeros((n, n));
134        let origin = simbox.origin_view();
135        let o0 = origin[a0];
136        let o1 = origin[a1];
137
138        for k in 0..coords[0].len() {
139            let p0 = coords[a0][k];
140            let p1 = coords[a1][k];
141            let c0 = ((p0 - o0) / d0).floor() as isize;
142            let c1 = ((p1 - o1) / d1).floor() as isize;
143            for i0 in (c0 - half_k0)..=(c0 + half_k0) {
144                let v0 = o0 + (i0 as F + 0.5) * d0 - p0;
145                if v0.abs() > r_max {
146                    continue;
147                }
148                let g0 = i0.rem_euclid(n as isize) as usize;
149                for i1 in (c1 - half_k1)..=(c1 + half_k1) {
150                    let v1 = o1 + (i1 as F + 0.5) * d1 - p1;
151                    let r2 = v0 * v0 + v1 * v1;
152                    if r2 > r_max_sq {
153                        continue;
154                    }
155                    let g1 = i1.rem_euclid(n as isize) as usize;
156                    image[[g0, g1]] += pref * (-r2 / two_sigma_sq).exp();
157                }
158            }
159        }
160
161        let diffraction = fft2_power_shifted(&image, n);
162
163        Ok(DiffractionPatternResult { diffraction, image })
164    }
165}
166
167/// 2-D FFT via row+column 1-D FFTs, return `|F|²` shifted so k=0 is at
168/// the grid centre.
169fn fft2_power_shifted(image: &Array2<F>, n: usize) -> Array2<F> {
170    let mut planner = FftPlanner::<F>::new();
171    let fft = planner.plan_fft_forward(n);
172
173    // Build a complex buffer.
174    let mut buf: Vec<RfComplex<F>> = vec![RfComplex::new(0.0, 0.0); n * n];
175    for i in 0..n {
176        for j in 0..n {
177            buf[i * n + j] = RfComplex::new(image[[i, j]], 0.0);
178        }
179    }
180
181    // Row pass: FFT every row of length `n`.
182    {
183        let mut scratch = vec![RfComplex::new(0.0, 0.0); n];
184        for i in 0..n {
185            scratch.copy_from_slice(&buf[i * n..(i + 1) * n]);
186            fft.process(&mut scratch);
187            buf[i * n..(i + 1) * n].copy_from_slice(&scratch);
188        }
189    }
190    // Column pass: gather column → FFT → scatter back.
191    {
192        let mut col = vec![RfComplex::new(0.0, 0.0); n];
193        for j in 0..n {
194            for i in 0..n {
195                col[i] = buf[i * n + j];
196            }
197            fft.process(&mut col);
198            for i in 0..n {
199                buf[i * n + j] = col[i];
200            }
201        }
202    }
203
204    // Magnitude squared + FFT-shift (swap upper-half and lower-half along
205    // both axes) so DC sits at (n/2, n/2).
206    let mut shifted = Array2::<F>::zeros((n, n));
207    let half = n / 2;
208    for i in 0..n {
209        for j in 0..n {
210            let si = (i + half) % n;
211            let sj = (j + half) % n;
212            shifted[[si, sj]] = buf[i * n + j].norm_sqr();
213        }
214    }
215    shifted
216}
217
218impl Compute for DiffractionPattern {
219    type Args<'a> = ();
220    type Output = Vec<DiffractionPatternResult>;
221
222    fn compute<'a, FA: FrameAccess + Sync + 'a>(
223        &self,
224        frames: &[&'a FA],
225        _: (),
226    ) -> Result<Vec<DiffractionPatternResult>, ComputeError> {
227        if frames.is_empty() {
228            return Err(ComputeError::EmptyInput);
229        }
230        let mut out = Vec::with_capacity(frames.len());
231        for f in frames {
232            out.push(self.one_frame(*f)?);
233        }
234        Ok(out)
235    }
236}
237
238#[cfg(test)]
239mod tests {
240    use super::*;
241    use molrs::Frame;
242    use molrs::spatial::region::simbox::SimBox;
243    use molrs::store::block::Block;
244    use ndarray::{Array1 as A1, array};
245
246    fn frame_with(positions: &[[F; 3]], box_len: F) -> Frame {
247        let x = A1::from_iter(positions.iter().map(|p| p[0]));
248        let y = A1::from_iter(positions.iter().map(|p| p[1]));
249        let z = A1::from_iter(positions.iter().map(|p| p[2]));
250        let mut block = Block::new();
251        block.insert("x", x.into_dyn()).unwrap();
252        block.insert("y", y.into_dyn()).unwrap();
253        block.insert("z", z.into_dyn()).unwrap();
254        let mut frame = Frame::new();
255        frame.insert("atoms", block);
256        frame.simbox =
257            Some(SimBox::cube(box_len, array![0.0 as F, 0.0 as F, 0.0 as F], [true; 3]).unwrap());
258        frame
259    }
260
261    #[test]
262    fn empty_frame_image_is_zero() {
263        let frame = frame_with(&[], 10.0);
264        let r = &DiffractionPattern::new(32, 0.3)
265            .unwrap()
266            .compute(&[&frame], ())
267            .unwrap()[0];
268        let s: F = r.image.iter().copied().sum();
269        assert_eq!(s, 0.0);
270    }
271
272    #[test]
273    fn single_particle_dc_component_dominant() {
274        // For a single Gaussian at the centre, |F(k=0)|² should be the
275        // largest entry in the shifted diffraction image.
276        let frame = frame_with(&[[5.0, 5.0, 5.0]], 10.0);
277        let r = &DiffractionPattern::new(32, 0.4)
278            .unwrap()
279            .compute(&[&frame], ())
280            .unwrap()[0];
281        let dc = r.diffraction[[16, 16]];
282        let mut max_off = 0.0_f64;
283        for ((i, j), &v) in r.diffraction.indexed_iter() {
284            if !(i == 16 && j == 16) && v > max_off {
285                max_off = v;
286            }
287        }
288        assert!(
289            dc > max_off,
290            "DC component {dc} must dominate; saw off-DC max {max_off}"
291        );
292    }
293
294    #[test]
295    fn periodic_lattice_has_bragg_peaks() {
296        // 4×4 square lattice in xy gives Bragg-like peaks at non-DC k.
297        let mut positions = Vec::new();
298        let a = 2.5_f64;
299        for ix in 0..4 {
300            for iy in 0..4 {
301                positions.push([ix as F * a, iy as F * a, 5.0]);
302            }
303        }
304        let frame = frame_with(&positions, a * 4.0);
305        let r = &DiffractionPattern::new(32, 0.2)
306            .unwrap()
307            .compute(&[&frame], ())
308            .unwrap()[0];
309        // Some off-DC element must be larger than typical background.
310        let dc = r.diffraction[[16, 16]];
311        let mut max_off = 0.0_f64;
312        for ((i, j), &v) in r.diffraction.indexed_iter() {
313            if !(i == 16 && j == 16) && v > max_off {
314                max_off = v;
315            }
316        }
317        // Lattices produce strong off-DC peaks comparable to the DC
318        // (within a factor of a few in this small system).
319        assert!(
320            max_off > 0.05 * dc,
321            "Bragg peak off-DC max ({max_off}) should be a non-negligible fraction of DC ({dc})"
322        );
323    }
324
325    #[test]
326    fn invalid_args_error() {
327        assert!(DiffractionPattern::new(0, 0.3).is_err());
328        assert!(DiffractionPattern::new(32, 0.0).is_err());
329        assert!(DiffractionPattern::new(32, -1.0).is_err());
330    }
331}
molrs_compute/diffraction/diffraction_pattern.rs

molrs_compute/diffraction/
diffraction_pattern.rs
