Skip to main content

sci_form/ani/
aev_params.rs

1//! ANI-2x symmetry function parameter tables.
2//!
3//! Defines the radial (η, Rs) and angular (η, Rs, ζ, θs) parameters
4//! that characterize the Behler-Parrinello atomic environment vectors.
5
6use serde::{Deserialize, Serialize};
7
8/// Supported chemical elements for ANI potentials.
9pub const ANI_ELEMENTS: [u8; 7] = [1, 6, 7, 8, 9, 16, 17]; // H,C,N,O,F,S,Cl
10
11/// Number of supported element types.
12pub const N_SPECIES: usize = ANI_ELEMENTS.len();
13
14/// AEV hyperparameters.
15#[derive(Debug, Clone, Serialize, Deserialize)]
16pub struct AevParams {
17    /// Cutoff radius for radial terms (Å).
18    pub radial_cutoff: f64,
19    /// Cutoff radius for angular terms (Å).
20    pub angular_cutoff: f64,
21    /// η values for radial symmetry functions.
22    pub radial_eta: Vec<f64>,
23    /// Shifted center Rs values for radial symmetry functions.
24    pub radial_rs: Vec<f64>,
25    /// η values for angular symmetry functions.
26    pub angular_eta: Vec<f64>,
27    /// Shifted center Rs values for angular symmetry functions.
28    pub angular_rs: Vec<f64>,
29    /// ζ values for angular symmetry functions.
30    pub angular_zeta: Vec<f64>,
31    /// θs angular shift values (radians).
32    pub angular_theta_s: Vec<f64>,
33}
34
35impl AevParams {
36    /// Dimension of the radial sub-vector per species pair.
37    pub fn radial_length(&self) -> usize {
38        self.radial_eta.len() * self.radial_rs.len()
39    }
40
41    /// Dimension of the angular sub-vector per species-pair combination.
42    pub fn angular_length(&self) -> usize {
43        self.angular_eta.len()
44            * self.angular_rs.len()
45            * self.angular_zeta.len()
46            * self.angular_theta_s.len()
47    }
48
49    /// Total AEV length for one atom.
50    pub fn total_aev_length(&self) -> usize {
51        // Radial: N_SPECIES components × radial_length
52        // Angular: N_SPECIES*(N_SPECIES+1)/2 components × angular_length
53        let n_rad = N_SPECIES * self.radial_length();
54        let n_ang = N_SPECIES * (N_SPECIES + 1) / 2 * self.angular_length();
55        n_rad + n_ang
56    }
57}
58
59/// Map atomic number to species index for AEV computation.
60pub fn species_index(z: u8) -> Option<usize> {
61    ANI_ELEMENTS.iter().position(|&e| e == z)
62}
63
64/// Default ANI-2x-like parameters.
65///
66/// These are representative values inspired by the ANI-2x publication:
67/// Devereux et al., JCTC 16 (2020) 4192.
68pub fn default_ani2x_params() -> AevParams {
69    use std::f64::consts::PI;
70
71    // 8 radial functions
72    let radial_eta = vec![19.7; 8];
73    let radial_rs: Vec<f64> = (0..8).map(|i| 0.8 + 0.5625 * i as f64).collect();
74
75    // 4 angular shells
76    let angular_eta = vec![12.5; 4];
77    let angular_rs: Vec<f64> = (0..4).map(|i| 0.8 + 0.95 * i as f64).collect();
78    let angular_zeta = vec![14.1; 1];
79    let angular_theta_s: Vec<f64> = (0..8).map(|i| PI * i as f64 / 8.0).collect();
80
81    AevParams {
82        radial_cutoff: 5.2,
83        angular_cutoff: 3.5,
84        radial_eta,
85        radial_rs,
86        angular_eta,
87        angular_rs,
88        angular_zeta,
89        angular_theta_s,
90    }
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96
97    #[test]
98    fn test_species_index() {
99        assert_eq!(species_index(1), Some(0)); // H
100        assert_eq!(species_index(6), Some(1)); // C
101        assert_eq!(species_index(8), Some(3)); // O
102        assert_eq!(species_index(26), None); // Fe not supported
103    }
104
105    #[test]
106    fn test_aev_dimensions() {
107        let params = default_ani2x_params();
108        assert!(params.radial_length() > 0);
109        assert!(params.angular_length() > 0);
110        assert!(params.total_aev_length() > 0);
111    }
112}