gosh-elastic 0.1.2

// [[file:../elastic.note::ca6990d3][ca6990d3]]
use super::*;
use gchemol::{Bond, Molecule};

use std::collections::HashMap;
// ca6990d3 ends here

// [[file:../elastic.note::23b9ecbf][23b9ecbf]]
/// An elastic spring attached between two atoms
#[derive(Clone, Debug, Copy, Serialize, Deserialize)]
pub struct ElasticSpring {
    /// atom `u`
    pub u: usize,
    /// atom `v`
    pub v: usize,
    /// spring constant
    pub k: Option<f64>,
    /// natural length. if not set, should use the distance between `u` and `v`
    /// in current structure
    pub l: Option<f64>,
    /// hook length: when in range l..l+r, there is no elastic force. The default
    /// is zero.
    pub r: Option<f64>,
}

impl ElasticSpring {
    pub fn spring_constant(&self) -> f64 {
        let k = self.k.unwrap_or(1.0);
        assert!(k.is_sign_positive(), "invalid spring constant {}", k);
        k
    }
    pub fn natural_length(&self) -> f64 {
        let l = self.l.expect("natural length not set");
        assert!(l.is_sign_positive(), "invalid spring natural length {}", l);
        l
    }

    pub fn hook_length(&self) -> f64 {
        let r = self.r.unwrap_or_default();
        assert!(r.is_sign_positive(), "invalid spring hook length {}", r);
        r
    }

    /// The atom with a smaller serial number
    pub fn atom1(&self) -> usize {
        if self.u > self.v {
            self.v
        } else {
            self.u
        }
    }

    /// The atom with a larger serial number
    pub fn atom2(&self) -> usize {
        if self.u > self.v {
            self.u
        } else {
            self.v
        }
    }

    fn elastic_energy_between(&self, pij: Vector3f) -> f64 {
        let d = pij.norm();
        let l = self.natural_length();
        let k = self.spring_constant();
        let r = self.hook_length();
        info!(
            "distance({}--{}) = {:4.2}, spring (l= {}, r= {}, k= {})",
            self.u, self.v, d, l, r, k
        );
        let vij = if d > l && d < l + r {
            0.0
        } else {
            0.5 * k * (l - d).powi(2)
        };
        vij
    }

    fn elastic_force_between(&self, pij: Vector3f) -> Vector3f {
        let d = pij.norm();
        let l = self.natural_length();
        let k = self.spring_constant();
        let r = self.hook_length();
        let k = self.spring_constant();
        let r = self.hook_length();
        let fij = -k * (l - d) * pij / l;
        if d > l && d < l + r {
            fij * 0.0
        } else {
            fij
        }
    }
}
// 23b9ecbf ends here

// [[file:../elastic.note::88d69ea4][88d69ea4]]
/// Elastic model based spring potential
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub struct ElasticModel {
    /// constraints that can be serialized
    pub constraints: HashMap<String, ElasticSpring>,
}
// 88d69ea4 ends here

// [[file:../elastic.note::f02d0fe2][f02d0fe2]]
use vecfx::*;

impl ElasticModel {
    fn get_elastic_energy_between(&self, mol: &Molecule, i: usize, j: usize, b: &Bond) -> Option<f64> {
        let sij = self.get_spring_between(i, j)?;
        let pij = get_pij_vector(mol, i, j)?;
        let vij = sij.elastic_energy_between(pij);
        Some(vij)
    }

    fn get_elastic_force_between(&self, mol: &Molecule, i: usize, j: usize) -> Option<Vector3f> {
        let sij = mol.get_bond(i, j).and_then(|b| self.get_spring_between(i, j))?;
        let pij = get_pij_vector(mol, i, j)?;
        sij.elastic_force_between(pij).into()
    }

    // FIXME: only using the matched spring found for the first time
    fn get_spring_between(&self, u: usize, v: usize) -> Option<ElasticSpring> {
        let (u, v) = if u > v { (v, u) } else { (u, v) };
        self.constraints.iter().find_map(|(_, &s)| {
            if s.atom1() == u && s.atom2() == v {
                Some(s)
            } else {
                None
            }
        })
    }

    /// Add a dummy bonds between constrained atom pairs
    fn add_dummy_bonds(&self, mol: &mut Molecule) {
        use gchemol::Bond;

        for (_, s) in &self.constraints {
            mol.add_bond(s.atom1(), s.atom2(), Bond::dummy());
        }
    }
}

fn get_pij_vector(mol: &Molecule, i: usize, j: usize) -> Option<Vector3f> {
    let pi: Vector3f = mol.get_atom(i)?.position().into();
    let pj: Vector3f = mol.get_atom(j)?.position().into();

    if let Some(lat) = mol.lattice {
        lat.apply_mic(pj - pi)
    } else {
        pj - pi
    }
    .into()
}
// f02d0fe2 ends here

// [[file:../elastic.note::b9afbe6a][b9afbe6a]]
use gosh_model::*;
use gut::config::*;

impl ElasticModel {
    /// Apply elastic constraints in Molecule `mol` by setting dummy
    /// bonds between pair of atoms.
    pub fn apply_constraints(&self, mol: &mut Molecule) {
        self.add_dummy_bonds(mol);
    }

    /// Compute elastic energy from `mol` with elastic constraints.
    pub fn elastic_energy(&self, mol: &Molecule) -> f64 {
        mol.bonds()
            .map(|(i, j, b)| self.get_elastic_energy_between(mol, i, j, b).unwrap_or_default())
            .sum()
    }

    /// Compute elastic force from `mol` with elastic constraints.
    pub fn elastic_force(&self, mol: &Molecule) -> Vec<[f64; 3]> {
        mol.atoms()
            .map(|(i, a)| {
                // if there is no neighbor, returns a zero vector
                let fi: Vector3f = mol
                    .connected(i)
                    .map(|j| self.get_elastic_force_between(mol, i, j).unwrap_or_default())
                    .sum();
                fi.into()
            })
            .collect()
    }
}

/// Allow to use in `Chemicalmodel` interface for constrained optimization.
impl ChemicalModel for ElasticModel {
    fn compute(&mut self, mol: &Molecule) -> Result<ModelProperties> {
        let mut mp = ModelProperties::default();
        mp.set_energy(self.elastic_energy(&mol));
        mp.set_forces(self.elastic_force(&mol));
        Ok(mp)
    }
}
// b9afbe6a ends here