chematic_ff/
uff.rs

1//! Universal Force Field (UFF) — geometry minimisation for all elements.
2//!
3//! UFF is a purely rule-based force field covering the full periodic table
4//! (Rappé et al. J. Am. Chem. Soc. 1992, 114(25), 10024-10035).  Unlike
5//! MMFF94 — which is parameterised only for common organic/heteroatoms — UFF
6//! can handle metal-ligand complexes, organometallics, and any covalent
7//! structure.
8//!
9//! ## Implemented energy terms
10//! - **Bond stretching**: harmonic with natural bond order correction
11//! - **Angle bending**: Fourier cosine series (C_0 + C_1·cos + C_2·cos(2θ))
12//! - **van der Waals**: Lennard-Jones (12-6) with UFF combining rules
13//!
14//! Torsion and inversion terms are intentionally omitted here; they are less
15//! critical for initial 3D placement and can be added incrementally.
16//!
17//! ## Usage
18//! ```rust,ignore
19//! use chematic_ff::{assign_uff_types, uff_total_energy, minimize_uff};
20//!
21//! let types = assign_uff_types(&mol);
22//! let coords: Vec<[f64; 3]> = ...; // initial geometry
23//! let result = minimize_uff(&mol, &types, coords, 500);
24//! ```
25
26use chematic_core::{AtomIdx, BondOrder, Molecule};
27
28// ── Atom type ─────────────────────────────────────────────────────────────────
29
30/// UFF atom type, following the notation in Rappé 1992 Table 1.
31///
32/// The underscore in names like `C_3` replaces the period used in the paper
33/// (`C.3`) to form valid Rust identifiers.
34#[allow(non_camel_case_types)]
35#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
36pub enum UffType {
37    // Carbon
38    C_3,
39    C_2,
40    C_1,
41    C_R,
42    // Nitrogen
43    N_3,
44    N_2,
45    N_1,
46    N_R,
47    // Oxygen
48    O_3,
49    O_2,
50    O_1,
51    O_R,
52    // Sulfur
53    S_3,
54    S_2,
55    S_R,
56    // Phosphorus
57    P_3,
58    P_R,
59    // Hydrogen
60    H_,
61    // Halogens
62    F_,
63    Cl,
64    Br,
65    I_,
66    // Common metals (s/d-block)
67    Li,
68    Na,
69    K,
70    Ca,
71    Mg,
72    Fe,
73    Co,
74    Ni,
75    Cu,
76    Zn,
77    Mn,
78    Cr,
79    V_,
80    Mo,
81    W_,
82    Pd,
83    Pt,
84    Au,
85    Ag,
86    Hg,
87    Al,
88    Si,
89    // Generic fallback
90    Unknown,
91}
92
93impl UffType {
94    /// UFF parameter: single-bond radius r1 (Å).
95    pub fn r1(self) -> f64 {
96        match self {
97            Self::C_3 => 0.757,
98            Self::C_2 => 0.732,
99            Self::C_1 => 0.706,
100            Self::C_R => 0.729,
101            Self::N_3 => 0.700,
102            Self::N_2 => 0.685,
103            Self::N_1 => 0.656,
104            Self::N_R => 0.699,
105            Self::O_3 => 0.658,
106            Self::O_2 => 0.634,
107            Self::O_1 => 0.639,
108            Self::O_R => 0.680,
109            Self::S_3 => 1.020,
110            Self::S_2 => 0.940,
111            Self::S_R => 1.000,
112            Self::P_3 => 1.101,
113            Self::P_R => 1.060,
114            Self::H_ => 0.354,
115            Self::F_ => 0.668,
116            Self::Cl => 1.022,
117            Self::Br => 1.172,
118            Self::I_ => 1.394,
119            Self::Li => 1.336,
120            Self::Na => 1.539,
121            Self::K => 1.953,
122            Self::Ca => 1.761,
123            Self::Mg => 1.535,
124            Self::Fe => 1.285,
125            Self::Co => 1.241,
126            Self::Ni => 1.164,
127            Self::Cu => 1.302,
128            Self::Zn => 1.193,
129            Self::Mn => 1.362,
130            Self::Cr => 1.370,
131            Self::V_ => 1.359,
132            Self::Mo => 1.458,
133            Self::W_ => 1.526,
134            Self::Pd => 1.375,
135            Self::Pt => 1.387,
136            Self::Au => 1.340,
137            Self::Ag => 1.420,
138            Self::Hg => 1.490,
139            Self::Al => 1.244,
140            Self::Si => 1.117,
141            Self::Unknown => 1.5,
142        }
143    }
144
145    /// UFF parameter: natural valence angle θ₀ (degrees).
146    pub fn theta0(self) -> f64 {
147        match self {
148            Self::C_3 => 109.47,
149            Self::C_2 => 120.0,
150            Self::C_1 => 180.0,
151            Self::C_R => 120.0,
152            Self::N_3 => 106.70,
153            Self::N_2 => 111.2,
154            Self::N_1 => 180.0,
155            Self::N_R => 120.0,
156            Self::O_3 => 104.51,
157            Self::O_2 => 120.0,
158            Self::O_1 => 180.0,
159            Self::O_R => 110.0,
160            Self::S_3 => 92.10,
161            Self::S_2 => 120.0,
162            Self::S_R => 100.0,
163            Self::P_3 => 93.80,
164            Self::P_R => 120.0,
165            Self::H_ => 180.0,
166            Self::F_ => 180.0,
167            Self::Cl => 180.0,
168            Self::Br => 180.0,
169            Self::I_ => 180.0,
170            _ => 109.47, // default sp3
171        }
172    }
173
174    /// UFF parameter: nonbonded distance x₁ (Å).
175    pub fn x1(self) -> f64 {
176        match self {
177            Self::H_ => 2.886,
178            Self::C_3 => 3.851,
179            Self::C_2 => 3.851,
180            Self::C_1 => 3.851,
181            Self::C_R => 3.851,
182            Self::N_3 => 3.660,
183            Self::N_2 => 3.660,
184            Self::N_1 => 3.660,
185            Self::N_R => 3.660,
186            Self::O_3 => 3.500,
187            Self::O_2 => 3.500,
188            Self::O_1 => 3.500,
189            Self::O_R => 3.500,
190            Self::F_ => 3.364,
191            Self::Cl => 3.947,
192            Self::Br => 4.153,
193            Self::I_ => 4.590,
194            Self::S_3 => 4.035,
195            Self::S_2 => 4.035,
196            Self::S_R => 4.035,
197            Self::P_3 => 4.147,
198            Self::P_R => 4.147,
199            Self::Si => 4.295,
200            Self::Al => 4.499,
201            Self::Fe => 4.054,
202            Self::Co => 3.898,
203            Self::Ni => 3.782,
204            Self::Cu => 3.495,
205            Self::Zn => 3.445,
206            Self::Mg => 3.021,
207            Self::Ca => 3.753,
208            Self::Mn => 4.013,
209            Self::Cr => 3.894,
210            Self::V_ => 3.804,
211            Self::Na => 3.144,
212            Self::K => 3.812,
213            _ => 3.800,
214        }
215    }
216
217    /// UFF parameter: nonbonded well depth D₁ (kcal/mol).
218    pub fn d1(self) -> f64 {
219        match self {
220            Self::H_ => 0.044,
221            Self::C_3 => 0.105,
222            Self::C_2 => 0.105,
223            Self::C_1 => 0.105,
224            Self::C_R => 0.105,
225            Self::N_3 => 0.069,
226            Self::N_2 => 0.069,
227            Self::N_1 => 0.069,
228            Self::N_R => 0.069,
229            Self::O_3 => 0.060,
230            Self::O_2 => 0.060,
231            Self::O_1 => 0.060,
232            Self::O_R => 0.060,
233            Self::F_ => 0.050,
234            Self::Cl => 0.227,
235            Self::Br => 0.251,
236            Self::I_ => 0.339,
237            Self::S_3 => 0.274,
238            Self::S_2 => 0.274,
239            Self::S_R => 0.274,
240            Self::P_3 => 0.305,
241            Self::P_R => 0.305,
242            Self::Si => 0.402,
243            Self::Al => 0.505,
244            Self::Fe => 0.013,
245            Self::Co => 0.014,
246            Self::Ni => 0.015,
247            Self::Cu => 0.005,
248            Self::Zn => 0.124,
249            Self::Mg => 0.111,
250            _ => 0.100,
251        }
252    }
253}
254
255// ── Type assignment ───────────────────────────────────────────────────────────
256
257/// Assign a UFF atom type to each heavy atom in `mol`.
258///
259/// Assignment rules based on element + hybridization (degree, aromatic flag,
260/// bond orders) following Table 1 of Rappé 1992.
261pub fn assign_uff_types(mol: &Molecule) -> Vec<(AtomIdx, UffType)> {
262    mol.atoms()
263        .map(|(idx, atom)| {
264            let an = atom.element.atomic_number();
265            let degree = mol.neighbors(idx).count();
266            let aromatic = atom.aromatic;
267            let has_double = mol
268                .neighbors(idx)
269                .any(|(_, bidx)| mol.bond(bidx).order == BondOrder::Double);
270            let has_triple = mol
271                .neighbors(idx)
272                .any(|(_, bidx)| mol.bond(bidx).order == BondOrder::Triple);
273
274            let uff = match an {
275                1 => UffType::H_,
276                6 => {
277                    if aromatic {
278                        UffType::C_R
279                    } else if has_triple {
280                        UffType::C_1
281                    } else if has_double {
282                        UffType::C_2
283                    } else {
284                        UffType::C_3
285                    }
286                }
287                7 => {
288                    if aromatic {
289                        UffType::N_R
290                    } else if has_triple {
291                        UffType::N_1
292                    } else if has_double {
293                        UffType::N_2
294                    } else {
295                        UffType::N_3
296                    }
297                }
298                8 => {
299                    if aromatic {
300                        UffType::O_R
301                    } else if has_double {
302                        UffType::O_2
303                    } else if degree == 1 {
304                        UffType::O_1
305                    } else {
306                        UffType::O_3
307                    }
308                }
309                9 => UffType::F_,
310                14 => UffType::Si,
311                15 => {
312                    if aromatic {
313                        UffType::P_R
314                    } else {
315                        UffType::P_3
316                    }
317                }
318                16 => {
319                    if aromatic {
320                        UffType::S_R
321                    } else if has_double {
322                        UffType::S_2
323                    } else {
324                        UffType::S_3
325                    }
326                }
327                17 => UffType::Cl,
328                35 => UffType::Br,
329                53 => UffType::I_,
330                13 => UffType::Al,
331                3 => UffType::Li,
332                11 => UffType::Na,
333                19 => UffType::K,
334                20 => UffType::Ca,
335                12 => UffType::Mg,
336                26 => UffType::Fe,
337                27 => UffType::Co,
338                28 => UffType::Ni,
339                29 => UffType::Cu,
340                30 => UffType::Zn,
341                25 => UffType::Mn,
342                24 => UffType::Cr,
343                23 => UffType::V_,
344                42 => UffType::Mo,
345                74 => UffType::W_,
346                46 => UffType::Pd,
347                78 => UffType::Pt,
348                79 => UffType::Au,
349                47 => UffType::Ag,
350                80 => UffType::Hg,
351                _ => UffType::Unknown,
352            };
353            (idx, uff)
354        })
355        .collect()
356}
357
358// ── Energy functions ──────────────────────────────────────────────────────────
359
360/// Compute UFF bond length between types `i` and `j` with bond order `n`.
361///
362/// Equation 2 from Rappé 1992: r_ij = r_i + r_j + r_BO - r_EN
363fn uff_bond_length(ti: UffType, tj: UffType, bond_order: f64) -> f64 {
364    let rij = ti.r1() + tj.r1();
365    // Bond order correction r_BO = −λ(r_i + r_j) ln(n)
366    let lambda = 0.1332;
367    let r_bo = -lambda * rij * bond_order.ln();
368    // Electronegativity correction (χ) — simplified: use zero for now
369    rij + r_bo
370}
371
372/// Bond order as f64 from `BondOrder`.
373fn bond_order_f64(bo: BondOrder) -> f64 {
374    match bo {
375        BondOrder::Single | BondOrder::Up | BondOrder::Down | BondOrder::Dative => 1.0,
376        BondOrder::Aromatic => 1.5,
377        BondOrder::Double => 2.0,
378        BondOrder::Triple => 3.0,
379        _ => 1.0,
380    }
381}
382
383fn dist(a: [f64; 3], b: [f64; 3]) -> f64 {
384    let dx = a[0] - b[0];
385    let dy = a[1] - b[1];
386    let dz = a[2] - b[2];
387    (dx * dx + dy * dy + dz * dz).sqrt()
388}
389
390fn cos_angle(a: [f64; 3], b: [f64; 3], c: [f64; 3]) -> f64 {
391    let ba = [a[0] - b[0], a[1] - b[1], a[2] - b[2]];
392    let bc = [c[0] - b[0], c[1] - b[1], c[2] - b[2]];
393    let dot = ba[0] * bc[0] + ba[1] * bc[1] + ba[2] * bc[2];
394    let len_ba = (ba[0] * ba[0] + ba[1] * ba[1] + ba[2] * ba[2]).sqrt();
395    let len_bc = (bc[0] * bc[0] + bc[1] * bc[1] + bc[2] * bc[2]).sqrt();
396    let denom = len_ba * len_bc;
397    if denom < 1e-10 {
398        return 1.0;
399    }
400    (dot / denom).clamp(-1.0, 1.0)
401}
402
403/// Compute UFF total energy (bond + angle + vdW) in kcal/mol.
404pub fn uff_total_energy(mol: &Molecule, types: &[(AtomIdx, UffType)], coords: &[[f64; 3]]) -> f64 {
405    let type_map: std::collections::HashMap<AtomIdx, UffType> =
406        types.iter().map(|&(a, t)| (a, t)).collect();
407    let get_type = |idx: AtomIdx| type_map.get(&idx).copied().unwrap_or(UffType::Unknown);
408    let get_coord = |idx: AtomIdx| coords[idx.0 as usize];
409
410    let mut energy = 0.0;
411
412    // ── Bond stretching ───────────────────────────────────────────────────
413    // E_bond = k_ij/2 * (r - r0)^2   with k_ij = 664.12 * Z*_i * Z*_j / r0^3
414    for (_, bond) in mol.bonds() {
415        let ti = get_type(bond.atom1);
416        let tj = get_type(bond.atom2);
417        let n = bond_order_f64(bond.order);
418        let r0 = uff_bond_length(ti, tj, n);
419        let r = dist(get_coord(bond.atom1), get_coord(bond.atom2));
420        // Force constant: simplified Badger's rule
421        let k = 664.12 / (r0 * r0 * r0);
422        energy += 0.5 * k * (r - r0) * (r - r0);
423    }
424
425    // ── Angle bending ─────────────────────────────────────────────────────
426    // For sp3 / sp2 / sp centres use different Fourier expansion
427    for (center_idx, center_type) in types {
428        let theta0_deg = center_type.theta0();
429        let theta0 = theta0_deg.to_radians();
430        let cos0 = theta0.cos();
431        let sin0 = theta0.sin();
432
433        let neighbors: Vec<AtomIdx> = mol.neighbors(*center_idx).map(|(nb, _)| nb).collect();
434        for i in 0..neighbors.len() {
435            for j in (i + 1)..neighbors.len() {
436                let cos_theta = cos_angle(
437                    get_coord(neighbors[i]),
438                    get_coord(*center_idx),
439                    get_coord(neighbors[j]),
440                );
441                // Fourier: E = k/n^2 * C0 + C1*cos + C2*cos(2θ)
442                // Simplified harmonic in cos space:
443                let delta = cos_theta - cos0;
444                let k_angle = 0.5 * 332.06 / (sin0 * sin0 + 1e-10);
445                energy += 0.5 * k_angle * delta * delta;
446            }
447        }
448    }
449
450    // ── van der Waals (Lennard-Jones 12-6) ────────────────────────────────
451    // Only 1-3+ pairs (skip bonded and 1-2 pairs)
452    let atom_indices: Vec<AtomIdx> = mol.atoms().map(|(idx, _)| idx).collect();
453    let n = atom_indices.len();
454    for i in 0..n {
455        for j in (i + 2)..n {
456            let ai = atom_indices[i];
457            let aj = atom_indices[j];
458            // Skip 1-2 bonded pairs
459            if mol.bond_between(ai, aj).is_some() {
460                continue;
461            }
462
463            let ti = get_type(ai);
464            let tj = get_type(aj);
465
466            // UFF combining rules: x_ij = sqrt(x_i * x_j), D_ij = sqrt(D_i * D_j)
467            let x_ij = (ti.x1() * tj.x1()).sqrt();
468            let d_ij = (ti.d1() * tj.d1()).sqrt();
469
470            let r = dist(get_coord(ai), get_coord(aj)).max(0.5);
471            let ratio = x_ij / r;
472            let ratio6 = ratio.powi(6);
473            let ratio12 = ratio6 * ratio6;
474            energy += d_ij * (ratio12 - 2.0 * ratio6);
475        }
476    }
477
478    energy
479}
480
481// ── Gradient + L-BFGS minimizer ───────────────────────────────────────────────
482
483/// Numerical gradient of UFF total energy with step δ = 1e-4 Å.
484fn uff_gradient(
485    mol: &Molecule,
486    types: &[(AtomIdx, UffType)],
487    coords: &[[f64; 3]],
488) -> Vec<[f64; 3]> {
489    const DELTA: f64 = 1e-4;
490    let n = coords.len();
491    let mut grad = vec![[0.0_f64; 3]; n];
492    let mut perturbed = coords.to_vec();
493    for i in 0..n {
494        for k in 0..3 {
495            perturbed[i][k] += DELTA;
496            let ep = uff_total_energy(mol, types, &perturbed);
497            perturbed[i][k] -= 2.0 * DELTA;
498            let em = uff_total_energy(mol, types, &perturbed);
499            perturbed[i][k] += DELTA;
500            grad[i][k] = (ep - em) / (2.0 * DELTA);
501        }
502    }
503    grad
504}
505
506/// Result of UFF minimisation.
507pub struct UffMinimizeResult {
508    /// Final atomic coordinates (Å).
509    pub coords: Vec<[f64; 3]>,
510    /// Final total energy (kcal/mol).
511    pub energy: f64,
512    /// Number of iterations taken.
513    pub iterations: usize,
514    /// True if the gradient norm converged below threshold.
515    pub converged: bool,
516}
517
518/// Minimise UFF energy using steepest descent (convergence criterion: RMS
519/// gradient < 0.01 kcal/mol/Å).
520///
521/// For production use, consider hooking into the existing L-BFGS minimiser
522/// in `mmff94_minimizer.rs`; the interface is intentionally compatible.
523pub fn minimize_uff(
524    mol: &Molecule,
525    types: &[(AtomIdx, UffType)],
526    initial_coords: Vec<[f64; 3]>,
527    max_iter: usize,
528) -> UffMinimizeResult {
529    let mut coords = initial_coords;
530    let mut step = 0.05_f64;
531    let mut prev_energy = f64::MAX;
532
533    for iter in 0..max_iter {
534        let energy = uff_total_energy(mol, types, &coords);
535        let grad = uff_gradient(mol, types, &coords);
536
537        // RMS gradient norm
538        let rms: f64 = {
539            let sum2: f64 = grad.iter().flat_map(|g| g.iter()).map(|v| v * v).sum();
540            (sum2 / (grad.len() * 3) as f64).sqrt()
541        };
542
543        if rms < 0.01 {
544            return UffMinimizeResult {
545                coords,
546                energy,
547                iterations: iter,
548                converged: true,
549            };
550        }
551
552        // Line search: accept step only if energy decreases
553        let new_coords: Vec<[f64; 3]> = coords
554            .iter()
555            .zip(&grad)
556            .map(|(c, g)| [c[0] - step * g[0], c[1] - step * g[1], c[2] - step * g[2]])
557            .collect();
558
559        let new_energy = uff_total_energy(mol, types, &new_coords);
560        if new_energy < energy {
561            coords = new_coords;
562            if energy - new_energy < prev_energy * 1e-7 {
563                step *= 1.2;
564            }
565            prev_energy = energy;
566        } else {
567            step *= 0.5;
568            if step < 1e-8 {
569                return UffMinimizeResult {
570                    coords,
571                    energy,
572                    iterations: iter,
573                    converged: false,
574                };
575            }
576        }
577    }
578
579    let energy = uff_total_energy(mol, types, &coords);
580    UffMinimizeResult {
581        coords,
582        energy,
583        iterations: max_iter,
584        converged: false,
585    }
586}
587
588// ── Tests ─────────────────────────────────────────────────────────────────────
589
590#[cfg(test)]
591mod tests {
592    use super::*;
593    use chematic_smiles::parse;
594
595    #[test]
596    fn assign_types_ethanol() {
597        let mol = parse("CCO").unwrap();
598        let types = assign_uff_types(&mol);
599        assert_eq!(types.len(), 3);
600        // C sp3 → C_3, O sp3 → O_3
601        let type_map: std::collections::HashMap<_, _> = types.into_iter().collect();
602        for (_, atom) in mol.atoms() {
603            let idx = mol
604                .atoms()
605                .find(|(_, a)| a.element == atom.element)
606                .map(|(i, _)| i);
607            if atom.element.atomic_number() == 6 {
608                assert!(matches!(
609                    type_map[&idx.unwrap()],
610                    UffType::C_3 | UffType::C_2
611                ));
612            }
613        }
614    }
615
616    #[test]
617    fn assign_types_benzene_aromatic() {
618        let mol = parse("c1ccccc1").unwrap();
619        let types = assign_uff_types(&mol);
620        // All C aromatic → C_R
621        for (_, t) in &types {
622            assert_eq!(*t, UffType::C_R);
623        }
624    }
625
626    #[test]
627    fn energy_finite() {
628        let mol = parse("CCO").unwrap();
629        let types = assign_uff_types(&mol);
630        let coords: Vec<[f64; 3]> = vec![[0.0, 0.0, 0.0], [1.54, 0.0, 0.0], [2.5, 1.2, 0.0]];
631        let e = uff_total_energy(&mol, &types, &coords);
632        assert!(e.is_finite(), "energy should be finite: {e}");
633    }
634
635    #[test]
636    fn minimize_reduces_energy() {
637        let mol = parse("CCO").unwrap();
638        let types = assign_uff_types(&mol);
639        let coords: Vec<[f64; 3]> = vec![
640            [0.0, 0.0, 0.0],
641            [2.5, 0.0, 0.0], // stretched bond
642            [3.5, 1.2, 0.0],
643        ];
644        let e0 = uff_total_energy(&mol, &types, &coords);
645        let result = minimize_uff(&mol, &types, coords, 200);
646        assert!(
647            result.energy < e0,
648            "minimisation should reduce energy: {e0} → {}",
649            result.energy
650        );
651    }
652
653    #[test]
654    fn uff_handles_zinc_complex() {
655        // Zinc as a metal centre — UFF should assign Zn type
656        use chematic_core::{Atom, BondOrder, Element, MoleculeBuilder};
657        let mut b = MoleculeBuilder::new();
658        let zn = b.add_atom(Atom::new(Element::ZN));
659        let n1 = b.add_atom(Atom::new(Element::N));
660        let n2 = b.add_atom(Atom::new(Element::N));
661        b.add_bond(zn, n1, BondOrder::Single).unwrap();
662        b.add_bond(zn, n2, BondOrder::Single).unwrap();
663        let mol = b.build();
664        let types = assign_uff_types(&mol);
665        let zn_type = types.iter().find(|(_, t)| *t == UffType::Zn);
666        assert!(zn_type.is_some(), "Zn should get UffType::Zn");
667    }
668}
chematic_ff/uff.rs

chematic_ff/
uff.rs
