1use 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
19fn 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
33fn 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
56pub 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 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 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 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 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 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 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 let k_inv = 6.0_f64;
219
220 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 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 } 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}