uff_relax/forcefield/
mod.rs

1pub mod interactions;
2pub mod parallel;
3pub mod sequential;
4
5use crate::atom::{Atom, Bond, UffAtomType};
6use crate::cell::UnitCell;
7use crate::params::element_symbol;
8use glam::DVec3;
9
10const PARALLEL_THRESHOLD: usize = 1000;
11
12#[derive(Debug, Default, Clone, Copy)]
13pub struct EnergyTerms {
14    pub bond: f64,
15    pub angle: f64,
16    pub torsion: f64,
17    pub non_bonded: f64,
18    pub total: f64,
19}
20
21/// Represents a molecular system consisting of atoms, bonds, and an optional unit cell.
22pub struct System {
23    /// List of atoms in the system.
24    pub atoms: Vec<Atom>,
25    /// List of chemical bonds.
26    pub bonds: Vec<Bond>,
27    /// Unit cell for periodic boundary conditions.
28    pub cell: UnitCell,
29}
30
31impl System {
32    /// Creates a new molecular system and automatically assigns UFF atom types.
33    ///
34    /// # Arguments
35    /// * `atoms` - Initial atom positions and elements.
36    /// * `bonds` - Connectivity and bond orders.
37    /// * `cell` - Boundary conditions (use `UnitCell::new_none()` for gas phase).
38    pub fn new(atoms: Vec<Atom>, bonds: Vec<Bond>, cell: UnitCell) -> Self {
39        let mut system = Self { atoms, bonds, cell };
40        system.auto_assign_uff_types();
41        system
42    }
43
44    /// Automatically infers UFF atom types based on element, connectivity, and bond orders.
45    pub fn auto_assign_uff_types(&mut self) {
46        let n = self.atoms.len();
47        let mut adj = vec![Vec::new(); n];
48        for bond in &self.bonds {
49            adj[bond.atom_indices.0].push(bond);
50            adj[bond.atom_indices.1].push(bond);
51        }
52
53        for i in 0..n {
54            let z = self.atoms[i].element;
55            let symbol = element_symbol(z);
56            let neighbors = &adj[i];
57            let n_neighbors = neighbors.len();
58            let has_order_1_5 = neighbors.iter().any(|b| (b.order - 1.5).abs() < 0.1);
59            let has_order_2_0 = neighbors.iter().any(|b| (b.order - 2.0).abs() < 0.1);
60            let bond_order_sum: f32 = neighbors.iter().map(|b| b.order).sum();
61
62            let label = match z {
63                1 => "H_".to_string(),
64                6 => { // Carbon
65                    match n_neighbors {
66                        4 => "C_3".to_string(),
67                        3 => if has_order_1_5 || has_order_2_0 { "C_R".to_string() } else { "C_2".to_string() },
68                        2 => "C_1".to_string(),
69                        _ => "C_3".to_string(),
70                    }
71                }
72                7 => { // Nitrogen
73                    match n_neighbors {
74                        3 => if has_order_1_5 { "N_R".to_string() } else { "N_3".to_string() },
75                        2 => "N_2".to_string(),
76                        1 => "N_1".to_string(),
77                        _ => "N_3".to_string(),
78                    }
79                }
80                8 => { // Oxygen
81                    if has_order_1_5 { "O_R".to_string() }
82                    else if n_neighbors == 1 && has_order_2_0 { "O_2".to_string() }
83                    else { "O_3".to_string() }
84                }
85                9 => "F_".to_string(),
86                15 => { // Phosphorus
87                    if n_neighbors >= 4 || bond_order_sum > 4.0 { "P_3+5".to_string() }
88                    else { "P_3+3".to_string() }
89                }
90                16 => { // Sulfur
91                    if has_order_1_5 { "S_R".to_string() }
92                    else if has_order_2_0 && n_neighbors == 1 { "S_2".to_string() }
93                    else if n_neighbors == 3 || (bond_order_sum > 3.0 && bond_order_sum < 5.0) { "S_3+4".to_string() }
94                    else if n_neighbors >= 4 || bond_order_sum >= 5.0 { "S_3+6".to_string() }
95                    else { "S_3+2".to_string() }
96                }
97                17 => "Cl".to_string(),
98                35 => "Br".to_string(),
99                53 => "I_".to_string(),
100                _ => {
101                    if n_neighbors == 0 { format!("{}_", symbol) } 
102                    else { format!("{}_{}", symbol, n_neighbors) }
103                }
104            };
105            self.atoms[i].uff_type = UffAtomType(label);
106        }
107    }
108
109    /// Computes forces and total energy breakdown.
110    pub fn compute_forces(&mut self) -> EnergyTerms {
111        self.compute_forces_with_threads(0, 6.0) // Default auto, cutoff 6.0
112    }
113
114    pub fn compute_forces_with_threads(&mut self, num_threads: usize, cutoff: f64) -> EnergyTerms {
115        if num_threads == 1 {
116            return self.compute_forces_serial(cutoff);
117        }
118
119        let use_parallel = if num_threads > 1 {
120            true
121        } else {
122            self.atoms.len() >= PARALLEL_THRESHOLD
123        };
124
125        if use_parallel {
126            if num_threads > 1 {
127                let pool = rayon::ThreadPoolBuilder::new().num_threads(num_threads).build().unwrap();
128                pool.install(|| self.compute_forces_parallel(cutoff))
129            } else {
130                crate::init_parallelism(None);
131                self.compute_forces_parallel(cutoff)
132            }
133        } else {
134            self.compute_forces_serial(cutoff)
135        }
136    }
137
138    fn compute_forces_serial(&mut self, cutoff: f64) -> EnergyTerms {
139        let mut energy = EnergyTerms::default();
140        for atom in &mut self.atoms { atom.force = DVec3::ZERO; }
141        
142        let mut adj = vec![Vec::new(); self.atoms.len()];
143        for b in &self.bonds {
144            let (u, v) = b.atom_indices;
145            adj[u].push(v);
146            adj[v].push(u);
147        }
148
149        energy.bond = self.compute_bond_forces_sequential();
150        energy.angle = self.compute_angle_forces_sequential();
151        energy.torsion = self.compute_torsion_forces_sequential();
152        energy.non_bonded = self.compute_non_bonded_forces_sequential_cell_list(&adj, cutoff);
153        energy.total = energy.bond + energy.angle + energy.torsion + energy.non_bonded;
154        
155        energy
156    }
157
158    fn compute_forces_parallel(&mut self, cutoff: f64) -> EnergyTerms {
159        let mut energy = EnergyTerms::default();
160        for atom in &mut self.atoms { atom.force = DVec3::ZERO; }
161        
162        let mut adj = vec![Vec::new(); self.atoms.len()];
163        for b in &self.bonds {
164            let (u, v) = b.atom_indices;
165            adj[u].push(v);
166            adj[v].push(u);
167        }
168
169        energy.bond = self.compute_bond_forces_parallel();
170        energy.angle = self.compute_angle_forces_parallel();
171        energy.torsion = self.compute_torsion_forces_parallel();
172        energy.non_bonded = self.compute_non_bonded_forces_parallel_cell_list(&adj, cutoff);
173        energy.total = energy.bond + energy.angle + energy.torsion + energy.non_bonded;
174        
175        energy
176    }
177
178    pub(crate) fn get_cell_neighbors(&self, cl: &crate::spatial::CellList, pos: DVec3, _cutoff: f64) -> Vec<usize> {
179        let mut neighbors = Vec::new();
180        let rel = pos - cl.min_p;
181        let ix = (rel.x / cl.cell_size.x) as i32;
182        let iy = (rel.y / cl.cell_size.y) as i32;
183        let iz = (rel.z / cl.cell_size.z) as i32;
184
185        for dx in -1..=1 {
186            for dy in -1..=1 {
187                for dz in -1..=1 {
188                    let nx = ix + dx; let ny = iy + dy; let nz = iz + dz;
189                    if nx >= 0 && nx < cl.dx as i32 && ny >= 0 && ny < cl.dy as i32 && nz >= 0 && nz < cl.dz as i32 {
190                        let idx = (nx as usize * cl.dy * cl.dz) + (ny as usize * cl.dz) + nz as usize;
191                        neighbors.extend(&cl.cells[idx]);
192                    }
193                }
194            }
195        }
196        neighbors
197    }
198
199    pub(crate) fn get_exclusion_scale(&self, i: usize, j: usize, adj: &[Vec<usize>]) -> (bool, f64) {
200        for &n1 in &adj[i] {
201            if n1 == j { return (true, 0.0); }
202        }
203        for &n1 in &adj[i] {
204            for &n2 in &adj[n1] {
205                if n2 == j { return (true, 0.0); }
206            }
207        }
208        for &n1 in &adj[i] {
209            for &n2 in &adj[n1] {
210                for &n3 in &adj[n2] {
211                    if n3 == j { return (false, 0.5); }
212                }
213            }
214        }
215        (false, 1.0)
216    }
217}
uff_relax/forcefield/mod.rs

uff_relax/forcefield/
mod.rs
