pub mod geometry;
pub mod angle;
pub mod bond;
pub mod dihedral;
pub mod improper;
pub mod kspace;
pub mod pair;
pub mod registry;
pub mod soft;
pub use registry::{KernelConstructor, KernelRegistry, lookup_kernel, register_kernel};
pub mod kernels {
pub use super::angle::harmonic::{AngleHarmonic, angle_harmonic_ctor};
pub use super::bond::harmonic::{BondHarmonic, bond_harmonic_ctor};
pub use super::pair::lj_cut::{PairLJCut as PairLJ126, pair_lj_cut_ctor as pair_lj126_ctor};
}
use std::borrow::Cow;
use std::collections::HashSet;
use ndarray::Array1;
use crate::ff::forcefield::{ForceField, Params, SpecialBonds};
use molrs::store::block::Block;
use molrs::store::frame::Frame;
use molrs::types::{F, U};
pub fn intramolecular_pairs(frame: &Frame) -> Block {
let n_atoms = frame.get("atoms").and_then(|b| b.nrows()).unwrap_or(0);
let excluded_12 = end_pairs(frame, "bonds", "atomi", "atomj");
let excluded_13 = end_pairs(frame, "angles", "atomi", "atomk");
let set_14 = end_pairs(frame, "dihedrals", "atomi", "atoml");
let mut pi: Vec<U> = Vec::new();
let mut pj: Vec<U> = Vec::new();
let mut p14: Vec<bool> = Vec::new();
for a in 0..n_atoms {
for b in (a + 1)..n_atoms {
let key = (a, b);
if excluded_12.contains(&key) || excluded_13.contains(&key) {
continue;
}
pi.push(a as U);
pj.push(b as U);
p14.push(set_14.contains(&key));
}
}
let mut pairs = Block::new();
if !pi.is_empty() {
pairs
.insert("atomi", Array1::from_vec(pi).into_dyn())
.expect("fresh pairs block");
pairs
.insert("atomj", Array1::from_vec(pj).into_dyn())
.expect("fresh pairs block");
pairs
.insert("is_14", Array1::from_vec(p14).into_dyn())
.expect("fresh pairs block");
}
pairs
}
fn end_pairs(frame: &Frame, block: &str, col_a: &str, col_b: &str) -> HashSet<(usize, usize)> {
let Some(b) = frame.get(block) else {
return HashSet::new();
};
let (Some(a_col), Some(b_col)) = (b.get_uint(col_a), b.get_uint(col_b)) else {
return HashSet::new();
};
a_col
.iter()
.zip(b_col.iter())
.map(|(&i, &j)| {
let (i, j) = (i as usize, j as usize);
if i < j { (i, j) } else { (j, i) }
})
.collect()
}
pub use crate::optimize::Potential;
pub struct Potentials {
inner: Vec<Box<dyn Potential>>,
n_atoms: usize,
}
impl std::fmt::Debug for Potentials {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Potentials")
.field("len", &self.inner.len())
.finish()
}
}
impl Potentials {
pub fn new() -> Self {
Self {
inner: Vec::new(),
n_atoms: 0,
}
}
pub fn push(&mut self, pot: Box<dyn Potential>) {
self.inner.push(pot);
}
pub fn len(&self) -> usize {
self.inner.len()
}
pub fn n_atoms(&self) -> usize {
self.n_atoms
}
pub fn set_n_atoms(&mut self, n_atoms: usize) {
self.n_atoms = n_atoms;
}
pub fn is_empty(&self) -> bool {
self.inner.is_empty()
}
pub fn calc_energy_forces(&self, coords: &[F]) -> (F, Vec<F>) {
let n = coords.len();
let mut total_e: F = 0.0;
let mut total_f = vec![0.0; n];
for p in &self.inner {
let (e, f) = p.calc_energy_forces(coords);
total_e += e;
for (t, fi) in total_f.iter_mut().zip(f.iter()) {
*t += fi;
}
}
(total_e, total_f)
}
pub fn calc_energy(&self, coords: &[F]) -> F {
self.calc_energy_forces(coords).0
}
pub fn calc_forces(&self, coords: &[F]) -> Vec<F> {
self.calc_energy_forces(coords).1
}
}
impl Default for Potentials {
fn default() -> Self {
Self::new()
}
}
impl Potential for Potentials {
fn calc_energy_forces(&self, coords: &[F]) -> (F, Vec<F>) {
Potentials::calc_energy_forces(self, coords)
}
}
impl crate::ff::forcefield::Style {
pub fn to_potential(
&self,
frame: &Frame,
special_bonds: &SpecialBonds,
) -> Result<Option<Box<dyn Potential>>, String> {
let category = self.category();
if category == "atom" {
return Ok(None);
}
let topo_block = match category {
"bond" => Some("bonds"),
"angle" => Some("angles"),
"dihedral" => Some("dihedrals"),
"improper" => Some("impropers"),
"pair" => Some("pairs"),
_ => None,
};
if let Some(block_name) = topo_block {
let rows = frame.get(block_name).and_then(|b| b.nrows()).unwrap_or(0);
if rows == 0 {
return Ok(None);
}
}
let type_params = self.defs.collect_type_params();
if type_params.is_empty() && category != "pair" {
return Err(format!(
"Style '{}' ({}) has no type definitions",
self.name, category
));
}
let type_refs: Vec<(&str, &Params)> = type_params
.iter()
.map(|(name, params)| (name.as_str(), params))
.collect();
let ctor = registry::lookup_kernel(category, &self.name).ok_or_else(|| {
format!(
"no kernel for style category '{}' name '{}'",
category, self.name
)
})?;
let params: Cow<Params> = if category == "pair" {
let mut p = self.params.clone();
p.set("lj14scale", special_bonds.lj_14());
p.set("coulomb14scale", special_bonds.coul_14());
Cow::Owned(p)
} else {
Cow::Borrowed(&self.params)
};
let pot = ctor(¶ms, &type_refs, frame)?;
Ok(Some(pot))
}
}
pub fn extract_coords(frame: &Frame) -> Result<Vec<F>, String> {
let atoms = frame
.get("atoms")
.ok_or_else(|| "Frame has no \"atoms\" block".to_string())?;
let (x, y, z) = (
atoms.get_float("x"),
atoms.get_float("y"),
atoms.get_float("z"),
);
let (Some(x), Some(y), Some(z)) = (x, y, z) else {
return Err("atoms block missing x/y/z float columns".into());
};
let xs: Vec<F> = x.iter().copied().collect();
let ys: Vec<F> = y.iter().copied().collect();
let zs: Vec<F> = z.iter().copied().collect();
let n = xs.len();
if ys.len() != n || zs.len() != n {
return Err("atoms x/y/z columns have mismatched lengths".into());
}
let mut coords = Vec::with_capacity(n * 3);
for i in 0..n {
coords.push(xs[i]);
coords.push(ys[i]);
coords.push(zs[i]);
}
Ok(coords)
}
impl ForceField {
pub fn to_potentials(&self, frame: &Frame) -> Result<Potentials, String> {
let mut pots = Potentials::new();
for style in self.styles() {
let block = match style.category() {
"bond" => Some("bonds"),
"angle" => Some("angles"),
"dihedral" => Some("dihedrals"),
"improper" => Some("impropers"),
"pair" => Some("pairs"),
_ => None,
};
if let Some(b) = block
&& frame.get(b).is_none()
{
continue;
}
if let Some(pot) = style.to_potential(frame, self.special_bonds())? {
pots.push(pot);
}
}
pots.set_n_atoms(frame.get("atoms").and_then(|b| b.nrows()).unwrap_or(0));
Ok(pots)
}
}
#[cfg(test)]
mod tests {
use super::*;
use molrs::store::block::Block;
use molrs::types::U;
use ndarray::Array1;
struct DummyPotential {
value: F,
}
impl Potential for DummyPotential {
fn calc_energy_forces(&self, coords: &[F]) -> (F, Vec<F>) {
(self.value, vec![self.value; coords.len()])
}
}
fn make_atoms_only_frame() -> Frame {
let mut frame = Frame::new();
let mut atoms = Block::new();
atoms
.insert("x", Array1::from_vec(vec![0.0 as F, 2.0 as F]).into_dyn())
.unwrap();
atoms
.insert("y", Array1::from_vec(vec![0.0 as F, 0.0 as F]).into_dyn())
.unwrap();
atoms
.insert("z", Array1::from_vec(vec![0.0 as F, 0.0 as F]).into_dyn())
.unwrap();
frame.insert("atoms", atoms);
frame
}
fn make_bond_frame() -> Frame {
let mut frame = make_atoms_only_frame();
let mut bonds = Block::new();
bonds
.insert("atomi", Array1::from_vec(vec![0 as U]).into_dyn())
.unwrap();
bonds
.insert("atomj", Array1::from_vec(vec![1 as U]).into_dyn())
.unwrap();
bonds
.insert("type", Array1::from_vec(vec!["A-A".to_string()]).into_dyn())
.unwrap();
frame.insert("bonds", bonds);
frame
}
fn make_lj_frame() -> Frame {
let mut frame = make_atoms_only_frame();
frame
.get_mut("atoms")
.unwrap()
.insert(
"type",
Array1::from_vec(vec!["A".to_string(), "A".to_string()]).into_dyn(),
)
.unwrap();
let pairs = intramolecular_pairs(&frame);
frame.insert("pairs", pairs);
frame
}
#[test]
fn test_potentials_collection() {
let mut pots = Potentials::new();
pots.push(Box::new(DummyPotential { value: 1.0 }));
pots.push(Box::new(DummyPotential { value: 2.0 }));
assert_eq!(pots.len(), 2);
let coords: Vec<F> = vec![0.0; 6];
assert!((pots.calc_energy(&coords) - 3.0).abs() < 1e-5);
let forces = pots.calc_forces(&coords);
for f in &forces {
assert!((*f - 3.0).abs() < 1e-5);
}
}
#[test]
fn unknown_style_kernel_is_error() {
let mut ff = ForceField::new("test");
ff.def_bondstyle("nonexistent")
.def_type("A-A", &[("k", 1.0)]);
let frame = make_bond_frame();
let err = ff.to_potentials(&frame).unwrap_err();
assert!(err.contains("no kernel"), "{err}");
}
#[test]
fn register_kernel_extends_dispatch() {
fn my_ctor(
_sp: &Params,
_tp: &[(&str, &Params)],
_f: &Frame,
) -> Result<Box<dyn Potential>, String> {
Ok(Box::new(DummyPotential { value: 42.0 }))
}
register_kernel("pair", "test/custom", my_ctor);
let mut ff = ForceField::new("test");
ff.def_pairstyle("test/custom", &[]).def_type("A", &[]);
let frame = make_lj_frame();
let pots = ff.to_potentials(&frame).unwrap();
let coords = extract_coords(&frame).unwrap();
assert!((pots.calc_energy(&coords) - 42.0).abs() < 1e-9);
}
#[test]
fn atom_style_is_skipped() {
let ff = ForceField::new("test").with_atomstyle("full");
let frame = make_atoms_only_frame();
let pots = ff.to_potentials(&frame).unwrap();
assert_eq!(pots.len(), 0);
}
#[test]
fn test_compile_requires_types() {
let mut ff = ForceField::new("test");
ff.def_bondstyle("harmonic");
let frame = make_bond_frame();
let err = ff
.to_potentials(&frame)
.expect_err("expected compile to fail");
assert!(err.contains("has no type definitions"));
}
#[test]
fn test_compile_energy() {
let mut ff = ForceField::new("test");
ff.def_pairstyle("lj/cut", &[("cutoff", 10.0)])
.def_type("A", &[("epsilon", 1.0), ("sigma", 1.0)]);
let frame = make_lj_frame();
let pots = ff.to_potentials(&frame).unwrap();
let coords = extract_coords(&frame).unwrap();
let (energy, _) = pots.calc_energy_forces(&coords);
let expected: F = 4.0 * (1.0 / 4096.0 - 1.0 / 64.0);
assert!((energy - expected).abs() < 1e-5);
}
#[test]
fn test_compile_forces() {
let mut ff = ForceField::new("test");
ff.def_pairstyle("lj/cut", &[("cutoff", 10.0)])
.def_type("A", &[("epsilon", 1.0), ("sigma", 1.0)]);
let frame = make_lj_frame();
let pots = ff.to_potentials(&frame).unwrap();
let coords = extract_coords(&frame).unwrap();
let (_, forces) = pots.calc_energy_forces(&coords);
for dim in 0..3 {
let sum = forces[dim] + forces[3 + dim];
assert!(sum.abs() < 1e-5);
}
}
#[test]
fn lj_cut_combines_distinct_types_lorentz_berthelot() {
let mut frame = make_atoms_only_frame();
frame
.get_mut("atoms")
.unwrap()
.insert("x", Array1::from_vec(vec![0.0 as F, 2.5 as F]).into_dyn())
.unwrap();
frame
.get_mut("atoms")
.unwrap()
.insert(
"type",
Array1::from_vec(vec!["A".to_string(), "B".to_string()]).into_dyn(),
)
.unwrap();
frame.insert("pairs", intramolecular_pairs(&frame));
let mut ff = ForceField::new("test");
ff.def_pairstyle("lj/cut", &[])
.def_type("A", &[("epsilon", 1.0), ("sigma", 1.0)])
.def_type("B", &[("epsilon", 4.0), ("sigma", 3.0)]);
let pots = ff.to_potentials(&frame).unwrap();
let coords = extract_coords(&frame).unwrap();
let (energy, _) = pots.calc_energy_forces(&coords);
let (eps, sigma, r) = (2.0_f64, 2.0_f64, 2.5_f64);
let sr6 = (sigma / r).powi(6);
let expected = 4.0 * eps * (sr6 * sr6 - sr6);
assert!(
(energy - expected).abs() < 1e-9,
"energy={energy} expected={expected}"
);
}
#[test]
fn lj_cut_applies_special_bonds_14_scaling() {
let mut frame = Frame::new();
let mut atoms = Block::new();
atoms
.insert(
"x",
Array1::from_vec(vec![0.0 as F, 1.0, 2.0, 3.0]).into_dyn(),
)
.unwrap();
for col in ["y", "z"] {
atoms
.insert(col, Array1::from_vec(vec![0.0 as F; 4]).into_dyn())
.unwrap();
}
atoms
.insert(
"type",
Array1::from_vec(vec!["A".to_string(); 4]).into_dyn(),
)
.unwrap();
frame.insert("atoms", atoms);
let mut bonds = Block::new();
bonds
.insert("atomi", Array1::from_vec(vec![0 as U, 1, 2]).into_dyn())
.unwrap();
bonds
.insert("atomj", Array1::from_vec(vec![1 as U, 2, 3]).into_dyn())
.unwrap();
frame.insert("bonds", bonds);
let mut angles = Block::new();
angles
.insert("atomi", Array1::from_vec(vec![0 as U, 1]).into_dyn())
.unwrap();
angles
.insert("atomk", Array1::from_vec(vec![2 as U, 3]).into_dyn())
.unwrap();
frame.insert("angles", angles);
let mut dihedrals = Block::new();
dihedrals
.insert("atomi", Array1::from_vec(vec![0 as U]).into_dyn())
.unwrap();
dihedrals
.insert("atoml", Array1::from_vec(vec![3 as U]).into_dyn())
.unwrap();
frame.insert("dihedrals", dihedrals);
let pairs = intramolecular_pairs(&frame);
assert_eq!(
pairs.nrows(),
Some(1),
"expected exactly the (0,3) 1-4 pair"
);
frame.insert("pairs", pairs);
let mut ff = ForceField::new("test");
ff.def_pairstyle("lj/cut", &[])
.def_type("A", &[("epsilon", 1.0), ("sigma", 1.0)]);
let mut sb = *ff.special_bonds();
sb.lj[2] = 0.5;
ff.set_special_bonds(sb);
let pots = ff.to_potentials(&frame).unwrap();
let coords = extract_coords(&frame).unwrap();
let (energy, _) = pots.calc_energy_forces(&coords);
let sr6 = (1.0_f64 / 3.0).powi(6);
let expected = 0.5 * 4.0 * (sr6 * sr6 - sr6);
assert!(
(energy - expected).abs() < 1e-12,
"energy={energy} expected={expected}"
);
}
#[test]
fn test_compile_empty_ff() {
let ff = ForceField::new("test");
let frame = make_lj_frame();
let pots = ff.to_potentials(&frame).unwrap();
let coords = extract_coords(&frame).unwrap();
let (energy, forces) = pots.calc_energy_forces(&coords);
assert!(energy.abs() < 1e-5);
assert_eq!(forces.len(), 6);
assert!(forces.iter().all(|x| x.abs() < 1e-5));
}
#[test]
fn test_compile_skips_absent_topology() {
let mut ff = ForceField::new("test");
ff.def_pairstyle("lj/cut", &[("cutoff", 10.0)])
.def_type("A", &[("epsilon", 1.0), ("sigma", 1.0)]);
let frame = make_atoms_only_frame();
let pots = ff.to_potentials(&frame).unwrap();
assert_eq!(pots.len(), 0);
let coords = extract_coords(&frame).unwrap();
assert!(pots.calc_energy(&coords).abs() < 1e-9);
}
}