molrs_compute/environment/
angular_separation.rs

1//! Pairwise angular separation between unit quaternions.
2//!
3//! Mirrors `freud.environment.AngularSeparationGlobal` and
4//! `AngularSeparationNeighbor`
5//! ([source](https://github.com/glotzerlab/freud/blob/main/freud/environment/AngularSeparation.cc)).
6//!
7//! For two unit quaternions `q₁` and `q₂` the (rotational) angular
8//! distance is
9//!
10//! ```text
11//!   θ = 2 · arccos( |q₁ · q₂| )
12//! ```
13//!
14//! where the absolute value accounts for the double cover (`q` and `−q`
15//! represent the same rotation). The result is in radians, `0 ≤ θ ≤ π/2`
16//! (between rotations), or `0 ≤ θ ≤ π` if the user does *not* want the
17//! double-cover identification (`equivalent_orientations = false` in
18//! freud).
19//!
20//! Two flavours are provided:
21//!
22//! - [`AngularSeparationGlobal`]: dense `(N_query × N_global)` table of
23//!   angular distances between every query orientation and every reference
24//!   orientation.
25//! - [`AngularSeparationNeighbor`]: sparse, one angular distance per
26//!   neighbor pair, driven by a `NeighborList`.
27
28use molrs::spatial::neighbors::NeighborList;
29use molrs::store::frame_access::FrameAccess;
30use molrs::types::F;
31use ndarray::Array2;
32
33use crate::error::ComputeError;
34use crate::result::ComputeResult;
35use crate::traits::Compute;
36
37/// Quaternion (w, x, y, z), unit-normalised by convention.
38pub type Quat = [F; 4];
39
40#[inline]
41fn quat_norm(q: Quat) -> F {
42    (q[0] * q[0] + q[1] * q[1] + q[2] * q[2] + q[3] * q[3]).sqrt()
43}
44
45#[inline]
46fn quat_dot(a: Quat, b: Quat) -> F {
47    a[0] * b[0] + a[1] * b[1] + a[2] * b[2] + a[3] * b[3]
48}
49
50/// Angular distance between two unit quaternions, in radians.
51///
52/// `θ = 2 · arccos(|q₁ · q₂|)` when `equivalent_orientations = true`
53/// (the default), giving values in `[0, π/2]`. When `false`, the absolute
54/// value is dropped and the result lies in `[0, π]`.
55pub fn angular_distance(q1: Quat, q2: Quat, equivalent_orientations: bool) -> F {
56    let n1 = quat_norm(q1);
57    let n2 = quat_norm(q2);
58    if n1 == 0.0 || n2 == 0.0 {
59        return 0.0;
60    }
61    let mut d = quat_dot(q1, q2) / (n1 * n2);
62    if equivalent_orientations {
63        d = d.abs();
64    }
65    2.0 * d.clamp(-1.0, 1.0).acos()
66}
67
68// ---------------------------------------------------------------------------
69// Global: dense (N_query × N_global) table
70// ---------------------------------------------------------------------------
71
72#[derive(Debug, Clone, Default)]
73pub struct AngularSeparationGlobalResult {
74    /// `(n_query, n_global)` table of angular distances (radians).
75    pub angles: Array2<F>,
76}
77
78impl ComputeResult for AngularSeparationGlobalResult {}
79
80#[derive(Debug, Clone)]
81pub struct AngularSeparationGlobal {
82    equivalent_orientations: bool,
83}
84
85impl Default for AngularSeparationGlobal {
86    fn default() -> Self {
87        Self::new()
88    }
89}
90
91impl AngularSeparationGlobal {
92    pub fn new() -> Self {
93        Self {
94            equivalent_orientations: true,
95        }
96    }
97
98    pub fn with_equivalent_orientations(mut self, on: bool) -> Self {
99        self.equivalent_orientations = on;
100        self
101    }
102}
103
104/// Args for `AngularSeparationGlobal`: (per-particle query quaternions,
105/// global reference quaternions). The query set typically comes from the
106/// frame's atom orientations; the reference set is a small fixed list of
107/// canonical orientations (e.g. crystallographic point-group symmetries).
108pub struct AngularSeparationGlobalArgs<'a> {
109    pub query: &'a [Quat],
110    pub global: &'a [Quat],
111}
112
113impl Compute for AngularSeparationGlobal {
114    type Args<'a> = AngularSeparationGlobalArgs<'a>;
115    type Output = Vec<AngularSeparationGlobalResult>;
116
117    fn compute<'a, FA: FrameAccess + Sync + 'a>(
118        &self,
119        frames: &[&'a FA],
120        args: AngularSeparationGlobalArgs<'a>,
121    ) -> Result<Vec<AngularSeparationGlobalResult>, ComputeError> {
122        if frames.is_empty() {
123            return Err(ComputeError::EmptyInput);
124        }
125        if args.query.is_empty() || args.global.is_empty() {
126            return Err(ComputeError::EmptyInput);
127        }
128        let n_query = args.query.len();
129        let n_global = args.global.len();
130        let mut out = Vec::with_capacity(frames.len());
131        for _ in frames {
132            let mut a = Array2::<F>::zeros((n_query, n_global));
133            for i in 0..n_query {
134                for j in 0..n_global {
135                    a[[i, j]] = angular_distance(
136                        args.query[i],
137                        args.global[j],
138                        self.equivalent_orientations,
139                    );
140                }
141            }
142            out.push(AngularSeparationGlobalResult { angles: a });
143        }
144        Ok(out)
145    }
146}
147
148// ---------------------------------------------------------------------------
149// Neighbor: sparse, one angle per pair
150// ---------------------------------------------------------------------------
151
152#[derive(Debug, Clone, Default)]
153pub struct AngularSeparationNeighborResult {
154    /// One angle per neighbor pair, in radians; index matches the
155    /// underlying [`NeighborList`] pair index.
156    pub angles: Vec<F>,
157}
158
159impl ComputeResult for AngularSeparationNeighborResult {}
160
161#[derive(Debug, Clone)]
162pub struct AngularSeparationNeighbor {
163    equivalent_orientations: bool,
164}
165
166impl Default for AngularSeparationNeighbor {
167    fn default() -> Self {
168        Self::new()
169    }
170}
171
172impl AngularSeparationNeighbor {
173    pub fn new() -> Self {
174        Self {
175            equivalent_orientations: true,
176        }
177    }
178
179    pub fn with_equivalent_orientations(mut self, on: bool) -> Self {
180        self.equivalent_orientations = on;
181        self
182    }
183}
184
185pub struct AngularSeparationNeighborArgs<'a> {
186    pub nlists: &'a [NeighborList],
187    /// Per-frame query orientations (indexed by `query_point_indices`).
188    pub query_orientations: &'a [Vec<Quat>],
189    /// Per-frame point orientations (indexed by `point_indices`).
190    pub point_orientations: &'a [Vec<Quat>],
191}
192
193impl Compute for AngularSeparationNeighbor {
194    type Args<'a> = AngularSeparationNeighborArgs<'a>;
195    type Output = Vec<AngularSeparationNeighborResult>;
196
197    fn compute<'a, FA: FrameAccess + Sync + 'a>(
198        &self,
199        frames: &[&'a FA],
200        args: AngularSeparationNeighborArgs<'a>,
201    ) -> Result<Vec<AngularSeparationNeighborResult>, ComputeError> {
202        if frames.is_empty() {
203            return Err(ComputeError::EmptyInput);
204        }
205        let nf = frames.len();
206        if args.nlists.len() != nf
207            || args.query_orientations.len() != nf
208            || args.point_orientations.len() != nf
209        {
210            return Err(ComputeError::DimensionMismatch {
211                expected: nf,
212                got: args
213                    .nlists
214                    .len()
215                    .min(args.query_orientations.len())
216                    .min(args.point_orientations.len()),
217                what: "AngularSeparationNeighbor frame-aligned inputs",
218            });
219        }
220        let mut out = Vec::with_capacity(nf);
221        for k in 0..nf {
222            let nl = &args.nlists[k];
223            let q = &args.query_orientations[k];
224            let p = &args.point_orientations[k];
225            let i_idx = nl.query_point_indices();
226            let j_idx = nl.point_indices();
227            let mut angles = Vec::with_capacity(nl.n_pairs());
228            for pk in 0..nl.n_pairs() {
229                let i = i_idx[pk] as usize;
230                let j = j_idx[pk] as usize;
231                if i >= q.len() || j >= p.len() {
232                    return Err(ComputeError::DimensionMismatch {
233                        expected: i.max(j) + 1,
234                        got: q.len().min(p.len()),
235                        what: "AngularSeparationNeighbor orientations length",
236                    });
237                }
238                angles.push(angular_distance(q[i], p[j], self.equivalent_orientations));
239            }
240            out.push(AngularSeparationNeighborResult { angles });
241        }
242        Ok(out)
243    }
244}
245
246#[cfg(test)]
247mod tests {
248    use super::*;
249    use molrs::Frame;
250
251    const TOL: F = 1e-10;
252
253    fn frame() -> Frame {
254        Frame::new()
255    }
256
257    #[test]
258    fn identity_quaternion_zero_distance() {
259        let q = [1.0_f64, 0.0, 0.0, 0.0];
260        assert!(angular_distance(q, q, true).abs() < TOL);
261        assert!(angular_distance(q, q, false).abs() < TOL);
262    }
263
264    #[test]
265    fn antipodal_quaternions_zero_with_equivalence() {
266        let q = [1.0_f64, 0.0, 0.0, 0.0];
267        let neg_q = [-1.0_f64, 0.0, 0.0, 0.0];
268        // With equivalent_orientations = true, q and -q describe the same
269        // rotation → angular distance 0.
270        assert!(angular_distance(q, neg_q, true).abs() < TOL);
271        // Without equivalence the bare formula gives 2·acos(-1) = 2π.
272        let theta = angular_distance(q, neg_q, false);
273        assert!((theta - 2.0 * std::f64::consts::PI).abs() < TOL);
274    }
275
276    #[test]
277    fn ninety_degree_rotation_about_z() {
278        // Quaternion for 90° rotation about z: (cos(45°), 0, 0, sin(45°))
279        let q = [
280            std::f64::consts::FRAC_PI_4.cos(),
281            0.0,
282            0.0,
283            std::f64::consts::FRAC_PI_4.sin(),
284        ];
285        let identity = [1.0_f64, 0.0, 0.0, 0.0];
286        let theta = angular_distance(q, identity, true);
287        // Expected rotation angle: 90° = π/2.
288        assert!((theta - std::f64::consts::FRAC_PI_2).abs() < 1e-12);
289    }
290
291    #[test]
292    fn global_shape_and_values() {
293        let q = vec![[1.0_f64, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0]];
294        let g = vec![[1.0_f64, 0.0, 0.0, 0.0]];
295        let r = &AngularSeparationGlobal::new()
296            .compute(
297                &[&frame()],
298                AngularSeparationGlobalArgs {
299                    query: &q,
300                    global: &g,
301                },
302            )
303            .unwrap()[0];
304        assert_eq!(r.angles.dim(), (2, 1));
305        assert!(r.angles[[0, 0]].abs() < TOL);
306        // (0,1,0,0) is a 180° rotation about x; cf. equivalent_orientations
307        // = true → |dot| = 0 → angle = π.
308        assert!((r.angles[[1, 0]] - std::f64::consts::PI).abs() < 1e-12);
309    }
310
311    #[test]
312    fn empty_inputs_error() {
313        let err = AngularSeparationGlobal::new()
314            .compute(
315                &[&frame()],
316                AngularSeparationGlobalArgs {
317                    query: &[],
318                    global: &[[1.0, 0.0, 0.0, 0.0]],
319                },
320            )
321            .unwrap_err();
322        assert!(matches!(err, ComputeError::EmptyInput));
323    }
324}
molrs_compute/environment/angular_separation.rs

molrs_compute/environment/
angular_separation.rs
