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) } // Always use symbol_ for generic match
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            let threads = if num_threads > 0 { num_threads } else { 4 };
127            let pool = rayon::ThreadPoolBuilder::new().num_threads(threads).build().unwrap();
128            pool.install(|| self.compute_forces_parallel(cutoff))
129        } else {
130            self.compute_forces_serial(cutoff)
131        }
132    }
133
134    fn compute_forces_serial(&mut self, cutoff: f64) -> EnergyTerms {
135        let mut energy = EnergyTerms::default();
136        for atom in &mut self.atoms { atom.force = DVec3::ZERO; }
137        
138        let mut adj = vec![Vec::new(); self.atoms.len()];
139        for b in &self.bonds {
140            let (u, v) = b.atom_indices;
141            adj[u].push(v);
142            adj[v].push(u);
143        }
144
145        energy.bond = self.compute_bond_forces_sequential();
146        energy.angle = self.compute_angle_forces_sequential();
147        energy.torsion = self.compute_torsion_forces_sequential();
148        energy.non_bonded = self.compute_non_bonded_forces_sequential_cell_list(&adj, cutoff);
149        energy.total = energy.bond + energy.angle + energy.torsion + energy.non_bonded;
150        
151        energy
152    }
153
154    fn compute_forces_parallel(&mut self, cutoff: f64) -> EnergyTerms {
155        let mut energy = EnergyTerms::default();
156        for atom in &mut self.atoms { atom.force = DVec3::ZERO; }
157        
158        let mut adj = vec![Vec::new(); self.atoms.len()];
159        for b in &self.bonds {
160            let (u, v) = b.atom_indices;
161            adj[u].push(v);
162            adj[v].push(u);
163        }
164
165        energy.bond = self.compute_bond_forces_parallel();
166        energy.angle = self.compute_angle_forces_parallel();
167        energy.torsion = self.compute_torsion_forces_parallel();
168        energy.non_bonded = self.compute_non_bonded_forces_parallel_cell_list(&adj, cutoff);
169        energy.total = energy.bond + energy.angle + energy.torsion + energy.non_bonded;
170        
171        energy
172    }
173
174    pub(crate) fn get_cell_neighbors(&self, cl: &crate::spatial::CellList, pos: DVec3, _cutoff: f64) -> Vec<usize> {
175        let mut neighbors = Vec::new();
176        let rel = pos - cl.min_p;
177        let ix = (rel.x / cl.cell_size.x) as i32;
178        let iy = (rel.y / cl.cell_size.y) as i32;
179        let iz = (rel.z / cl.cell_size.z) as i32;
180
181        for dx in -1..=1 {
182            for dy in -1..=1 {
183                for dz in -1..=1 {
184                    let mut nx = ix + dx; let mut ny = iy + dy; let mut nz = iz + dz;
185                    
186                    // PBC wrap for cells
187                    if nx < 0 { nx += cl.dx as i32; } else if nx >= cl.dx as i32 { nx -= cl.dx as i32; }
188                    if ny < 0 { ny += cl.dy as i32; } else if ny >= cl.dy as i32 { ny -= cl.dy as i32; }
189                    if nz < 0 { nz += cl.dz as i32; } else if nz >= cl.dz as i32 { nz -= cl.dz as i32; }
190
191                    if nx >= 0 && nx < cl.dx as i32 && ny >= 0 && ny < cl.dy as i32 && nz >= 0 && nz < cl.dz as i32 {
192                        let idx = (nx as usize * cl.dy * cl.dz) + (ny as usize * cl.dz) + nz as usize;
193                        neighbors.extend(&cl.cells[idx]);
194                    }
195                }
196            }
197        }
198        
199        // Remove duplicates (e.g. if cell size is large or system is small)
200        neighbors.sort_unstable();
201        neighbors.dedup();
202        neighbors
203    }
204
205        pub(crate) fn get_exclusion_scale(&self, i: usize, j: usize, adj: &[Vec<usize>]) -> (bool, f64) {
206
207            if i == j { return (true, 0.0); }
208
209            
210
211            // 1-2 neighbors: Strictly excluded from LJ
212
213            for &n1 in &adj[i] {
214
215                if n1 == j { return (true, 0.0); }
216
217            }
218
219            
220
221            // 1-3 neighbors: Small scale (0.1) to prevent collapse
222
223            for &n1 in &adj[i] {
224
225                for &n2 in &adj[n1] {
226
227                    if n2 == j { return (false, 0.1); }
228
229                }
230
231            }
232
233            
234
235            // 1-4 neighbors: Standard scale 0.5
236
237            for &n1 in &adj[i] {
238
239                for &n2 in &adj[n1] {
240
241                    for &n3 in &adj[n2] {
242
243                        if n3 == j { return (false, 0.5); }
244
245                    }
246
247                }
248
249            }
250
251            
252
253            (false, 1.0)
254
255        }
256
257    }
258
259
uff_relax/forcefield/mod.rs

uff_relax/forcefield/
mod.rs
