Skip to main content

sci_form/forcefield/
builder.rs

1/// Build a full UFF force field for a given molecule.
2///
3/// Walks all bonds, angles, torsions, out-of-plane inversions, and non-bonded
4/// pairs and adds the appropriate [`ForceFieldContribution`] terms to a fresh
5/// [`MolecularForceField`].
6use crate::forcefield::atom_typer::{assign_uff_type, is_atom_aromatic};
7use crate::forcefield::params::{
8    get_uff_angle_force_constant, get_uff_bond_force_constant, get_uff_bond_length, get_uff_params,
9    get_uff_torsion_params,
10};
11use crate::forcefield::traits::MolecularForceField;
12use crate::forcefield::uff::{
13    UffAngleBend, UffHarmonicBondStretch, UffInversion, UffLennardJones, UffTorsion,
14};
15use crate::graph::{BondOrder, Hybridization, Molecule};
16use petgraph::graph::NodeIndex;
17use petgraph::visit::EdgeRef;
18
19/// Assign atom types and look up UFF params for every atom in `mol`.
20/// Returns a parallel `Vec<&'static str>` of atom type labels.
21fn atom_types(mol: &Molecule) -> Vec<&'static str> {
22    let n = mol.graph.node_count();
23    (0..n)
24        .map(|i| {
25            let node = NodeIndex::new(i);
26            let atom = &mol.graph[node];
27            let aromatic = is_atom_aromatic(mol, i);
28            assign_uff_type(atom.element, &atom.hybridization, aromatic)
29        })
30        .collect()
31}
32
33/// Build a topological distance matrix using BFS from each node.
34/// Returns `dist[i][j]` = number of bonds between atoms i and j (or 999 if not reachable).
35fn topological_distances(mol: &Molecule) -> Vec<Vec<u32>> {
36    let n = mol.graph.node_count();
37    let mut dist = vec![vec![999u32; n]; n];
38    for start in 0..n {
39        dist[start][start] = 0;
40        let mut queue = std::collections::VecDeque::new();
41        queue.push_back(start);
42        while let Some(u) = queue.pop_front() {
43            let node_u = NodeIndex::new(u);
44            for edge in mol.graph.edges(node_u) {
45                let v = edge.target().index();
46                if dist[start][v] == 999 {
47                    dist[start][v] = dist[start][u] + 1;
48                    queue.push_back(v);
49                }
50            }
51        }
52    }
53    dist
54}
55
56/// Build a complete UFF force field for `mol`, including:
57/// - Bond stretching (harmonic)
58/// - Angle bending (Fourier)
59/// - Torsions (for all i-j-k-l paths)
60/// - Out-of-plane inversions (sp2 atoms)
61/// - Non-bonded Lennard-Jones (12-6), skipping 1-2 / 1-3, scaling 1-4 by 0.5
62pub fn build_uff_force_field(mol: &Molecule) -> MolecularForceField {
63    let mut ff = MolecularForceField::new();
64    let types = atom_types(mol);
65    let n = mol.graph.node_count();
66    let topo = topological_distances(mol);
67
68    // ---- Bond stretching ----
69    for edge in mol.graph.edge_references() {
70        let i = edge.source().index();
71        let j = edge.target().index();
72        let bo = &mol.graph[edge.id()].order;
73
74        let (pi, pj) = match (get_uff_params(types[i]), get_uff_params(types[j])) {
75            (Some(a), Some(b)) => (a, b),
76            _ => continue,
77        };
78
79        let r0 = get_uff_bond_length(&pi, &pj, bo);
80        let k_b = get_uff_bond_force_constant(&pi, &pj, r0);
81
82        ff.insert_dynamic_term(Box::new(UffHarmonicBondStretch {
83            atom_i_idx: i,
84            atom_j_idx: j,
85            force_constant_kb: k_b,
86            equilibrium_r0: r0,
87        }));
88    }
89
90    // ---- Build adjacency list for angle/torsion iteration ----
91    let mut neighbors: Vec<Vec<usize>> = vec![Vec::new(); n];
92    for edge in mol.graph.edge_references() {
93        let i = edge.source().index();
94        let j = edge.target().index();
95        neighbors[i].push(j);
96        neighbors[j].push(i);
97    }
98
99    // ---- Angle bending (i–j–k, j = centre) ----
100    for j in 0..n {
101        let nb_j = &neighbors[j];
102        if nb_j.len() < 2 {
103            continue;
104        }
105        let pj = match get_uff_params(types[j]) {
106            Some(p) => p,
107            None => continue,
108        };
109
110        let theta0_deg = pj.theta0;
111        let theta0_rad = theta0_deg.to_radians();
112
113        // coordination_n: 0 = linear (sp), 3 = sp2, 4 = sp3
114        let coord_n = match mol.graph[NodeIndex::new(j)].hybridization {
115            Hybridization::SP => 0,
116            Hybridization::SP2 => 3,
117            _ => 4,
118        };
119
120        for ii in 0..nb_j.len() {
121            for kk in (ii + 1)..nb_j.len() {
122                let i = nb_j[ii];
123                let k = nb_j[kk];
124
125                let pi = match get_uff_params(types[i]) {
126                    Some(p) => p,
127                    None => continue,
128                };
129                let pk = match get_uff_params(types[k]) {
130                    Some(p) => p,
131                    None => continue,
132                };
133
134                let bo_ij = mol
135                    .graph
136                    .find_edge(NodeIndex::new(i), NodeIndex::new(j))
137                    .map(|e| mol.graph[e].order)
138                    .unwrap_or(BondOrder::Single);
139                let bo_jk = mol
140                    .graph
141                    .find_edge(NodeIndex::new(j), NodeIndex::new(k))
142                    .map(|e| mol.graph[e].order)
143                    .unwrap_or(BondOrder::Single);
144
145                let r_ij = get_uff_bond_length(&pi, &pj, &bo_ij);
146                let r_jk = get_uff_bond_length(&pj, &pk, &bo_jk);
147                let k_a = get_uff_angle_force_constant(&pi, &pk, r_ij, r_jk, theta0_rad);
148
149                ff.insert_dynamic_term(Box::new(UffAngleBend {
150                    atom_i_idx: i,
151                    atom_j_idx: j,
152                    atom_k_idx: k,
153                    force_constant_ka: k_a,
154                    equilibrium_theta0: theta0_rad,
155                    coordination_n: coord_n,
156                }));
157            }
158        }
159    }
160
161    // ---- Torsion (i–j–k–l) ----
162    for edge_jk in mol.graph.edge_references() {
163        let j = edge_jk.source().index();
164        let k = edge_jk.target().index();
165        let bo_jk = &mol.graph[edge_jk.id()].order;
166
167        let pj = match get_uff_params(types[j]) {
168            Some(p) => p,
169            None => continue,
170        };
171        let pk = match get_uff_params(types[k]) {
172            Some(p) => p,
173            None => continue,
174        };
175
176        let hyb_j = &mol.graph[NodeIndex::new(j)].hybridization;
177        let hyb_k = &mol.graph[NodeIndex::new(k)].hybridization;
178
179        let (v_barrier, period_n, cos_phi0) =
180            get_uff_torsion_params(hyb_j, hyb_k, pj.v_tors, pk.v_tors, bo_jk);
181
182        if v_barrier < 1e-10 {
183            continue;
184        }
185
186        for &i in &neighbors[j] {
187            if i == k {
188                continue;
189            }
190            for &l in &neighbors[k] {
191                if l == j || l == i {
192                    continue;
193                }
194                ff.insert_dynamic_term(Box::new(UffTorsion {
195                    atom_i_idx: i,
196                    atom_j_idx: j,
197                    atom_k_idx: k,
198                    atom_l_idx: l,
199                    force_constant_v: v_barrier,
200                    periodicity_n: period_n,
201                    cos_phi0,
202                }));
203            }
204        }
205    }
206
207    // ---- Out-of-plane inversions (sp2 centre, 3 substituents) ----
208    for j in 0..n {
209        if mol.graph[NodeIndex::new(j)].hybridization != Hybridization::SP2 {
210            continue;
211        }
212        let nb_j = &neighbors[j];
213        if nb_j.len() != 3 {
214            continue;
215        }
216
217        // UFF inversion parameters for sp2: K = 6.0 kcal/mol for C_R/N_R, else 6.0 flat
218        let k_inv = 6.0_f64;
219
220        // c0, c1, c2 for ideal planar psi=0: c0=1, c1=0, c2=-1 → E = K(1 - cos(2*psi))
221        let (c0, c1, c2) = (1.0, 0.0, -1.0);
222
223        let (i, k, l) = (nb_j[0], nb_j[1], nb_j[2]);
224        ff.insert_dynamic_term(Box::new(UffInversion {
225            idx_i: i,
226            idx_j: j,
227            idx_k: k,
228            idx_l: l,
229            k_inv,
230            c0,
231            c1,
232            c2,
233        }));
234    }
235
236    // ---- Non-bonded Lennard-Jones ----
237    // Skip 1-2, 1-3 pairs; scale 1-4 by 0.5; full weight for 1-5+
238    for i in 0..n {
239        let pi = match get_uff_params(types[i]) {
240            Some(p) => p,
241            None => continue,
242        };
243        for j in (i + 1)..n {
244            let d = topo[i][j];
245            if d <= 2 {
246                continue;
247            } // 1-2 and 1-3 excluded
248
249            let pj = match get_uff_params(types[j]) {
250                Some(p) => p,
251                None => continue,
252            };
253
254            let r_star = (pi.x1 + pj.x1) * 0.5;
255            let eps_full = (pi.d1 * pj.d1).sqrt();
256            let epsilon = if d == 3 { eps_full * 0.5 } else { eps_full };
257
258            ff.insert_dynamic_term(Box::new(UffLennardJones {
259                atom_i_idx: i,
260                atom_j_idx: j,
261                r_star,
262                epsilon,
263            }));
264        }
265    }
266
267    ff
268}