use crate::forcefield::*;
use crate::math::Vec3;
use crate::traits::{AtomTrait, BondTrait};
use crate::vsepr::{calculate_steric_number, Geometry};
pub struct VseprOptimizer {
pub iterations: usize,
pub force_constant: f64,
}
impl Default for VseprOptimizer {
fn default() -> Self {
Self {
iterations: 1500,
force_constant: 0.15,
}
}
}
impl VseprOptimizer {
pub fn new() -> Self {
Self::default()
}
pub fn optimize<A: AtomTrait, B: BondTrait>(&self, atoms: &mut [A], bonds: &[B]) {
let n = atoms.len();
if n == 0 { return; }
let mut adj = vec![Vec::new(); n];
for (idx, bond) in bonds.iter().enumerate() {
let (i, j) = bond.get_atom_indices();
if i < n && j < n {
adj[i].push((j, idx));
adj[j].push((i, idx));
}
}
let mut positions: Vec<Vec3> = atoms
.iter()
.enumerate()
.map(|(i, a)| {
let p = Vec3::from(a.get_position());
let noise_scale = 0.01;
let dx = ((i * 1337) % 100) as f64 / 100.0 - 0.5;
let dy = ((i * 7331) % 100) as f64 / 100.0 - 0.5;
let dz = ((i * 9973) % 100) as f64 / 100.0 - 0.5;
p.add(Vec3(dx, dy, dz).mul(noise_scale))
})
.collect();
for i in 0..n {
if positions[i].length_squared() < 0.001 {
positions[i] = Vec3(
(i as f64 * 1.1).cos() * 1.0,
(i as f64 * 1.3).sin() * 1.0,
(i as f64 * 1.7).cos() * 1.0,
);
}
}
let geometries: Vec<Geometry> = (0..n)
.map(|i| Geometry::from_steric_number(calculate_steric_number(i, atoms, bonds)))
.collect();
for iter in 0..self.iterations {
let mut forces = vec![Vec3::ZERO; n];
let damping = (-3.0 * (iter as f64 / self.iterations as f64)).exp();
apply_bond_constraints(atoms, bonds, &positions, &mut forces);
apply_angle_constraints(n, &adj, &geometries, &positions, &mut forces);
apply_planarity_constraints(n, &adj, &geometries, &positions, &mut forces);
apply_torsion_constraints(n, &adj, bonds, &geometries, &positions, &mut forces);
apply_repulsion_constraints(n, &adj, &positions, &mut forces);
for i in 0..n {
positions[i] = positions[i].add(forces[i].mul(self.force_constant * damping));
}
}
for i in 0..n {
atoms[i].set_position(positions[i].into());
}
}
}