Unit tests for Bias and Umbrella, umbrella sampling over n windows

mltrain/__init__.py

@@ -4,6 +4,7 @@
 from mltrain.system import System
 from mltrain.box import Box
 from mltrain.bias import Bias
+from mltrain.umbrella import UmbrellaSampling
 from mltrain import md
 from mltrain import potentials
 from mltrain.training import selection
@@ -17,7 +18,9 @@
            'Molecule',
            'System',
            'Box',
+           'Bias',
+           'UmbrellaSampling',
            'md',
            'selection',
-           'potentials',
-           'Bias']
+           'potentials'
+           ]

mltrain/bias.py

@@ -62,7 +62,7 @@ def __init__(self,
             {to_add, to_subtract, to_average}: (list) Indices of the atoms
             which are combined in some way to define the reaction rxn_coord
         """
-        self.s = reference
+        self.ref = reference
         self.kappa = kappa
 
         if len(kwargs) != 1:
@@ -90,14 +90,14 @@ def __init__(self,
     def __call__(self, atom_pair_list, atoms):
         """Value of the bias for set of atom pairs in atoms"""
 
-        return 0.5 * self.kappa * (self.func(atom_pair_list, atoms) - self.s)**2
+        return 0.5 * self.kappa * (self.func(atom_pair_list, atoms) - self.ref)**2
 
     def grad(self, atom_pair_list, atoms):
         """Gradient of the biasing potential a set of atom pairs in atoms"""
 
         return (self.kappa
                 * self.func.grad(atom_pair_list, atoms)
-                * (self.func(atom_pair_list, atoms) - self.s))
+                * (self.func(atom_pair_list, atoms) - self.ref))
 
     def adjust_potential_energy(self, atoms):
         """Adjust the energy of a set of atoms using the bias function"""
@@ -135,12 +135,15 @@ def __call__(self, atom_pair_list, atoms):
         return np.mean(euclidean_dists)
 
     def grad(self, atom_pair_list, atoms):
-        """Gradient of the average distance between atom pairs
+        """Gradient of the average distance between atom pairs. Each component
+        of the gradient is calculated using ∇B_i,m:
 
-        .. math::
-
-            LaTeX..
+        ∇B_i,m = (1/M) * (r_i,m - r_i,m') / ||r_m||
 
+        ∇B_i,m:  Gradient of bias for atom m along component i for pair m, m'
+        M:       Number of atom pairs
+        r_i,m:   i (= x, y or z) position of atom in pair m, m'
+        ||r_m||: Euclidean distance between atoms in pair m, m'
         """
 
         derivative = np.zeros(shape=(len(atoms), 3))

mltrain/md.py

@@ -82,7 +82,9 @@ def run_mlp_md(configuration: 'mltrain.Configuration',
 
     ase_atoms = configuration.ase_atoms
     ase_atoms.set_calculator(mlp.ase_calculator)
-    ase_atoms.set_constraint(bias)
+
+    if bias is not None:
+        ase_atoms.set_constraint(bias)
 
     _set_momenta(ase_atoms,
                  temp=init_temp if init_temp is not None else temp,

mltrain/tests/data/h2_traj.xyz

@@ -0,0 +1,84 @@
+2
+Lattice="100.000000 0.000000 0.000000 0.000000 100.000000 0.000000 0.000000 0.000000 100.000000" 
+H 24.76963 25.26878 25.00000 
+H 25.23038 24.73122 25.00000 
+2
+Lattice="100.000000 0.000000 0.000000 0.000000 100.000000 0.000000 0.000000 0.000000 100.000000" 
+H 24.07381 25.93476 25.08863 
+H 25.94379 24.05945 24.91863 
+2
+Lattice="100.000000 0.000000 0.000000 0.000000 100.000000 0.000000 0.000000 0.000000 100.000000" 
+H 24.44639 25.53838 25.06493 
+H 25.57345 24.45285 24.94070 
+2
+Lattice="100.000000 0.000000 0.000000 0.000000 100.000000 0.000000 0.000000 0.000000 100.000000" 
+H 24.52384 25.36621 25.11032 
+H 25.49658 24.62412 24.89383 
+2
+Lattice="100.000000 0.000000 0.000000 0.000000 100.000000 0.000000 0.000000 0.000000 100.000000" 
+H 23.89950 25.79419 25.26704 
+H 26.12117 24.19476 24.73532 
+2
+Lattice="100.000000 0.000000 0.000000 0.000000 100.000000 0.000000 0.000000 0.000000 100.000000" 
+H 24.66819 25.21790 25.09581 
+H 25.35108 24.77000 24.90832 
+2
+Lattice="100.000000 0.000000 0.000000 0.000000 100.000000 0.000000 0.000000 0.000000 100.000000" 
+H 24.09053 25.49411 25.32562 
+H 25.92872 24.49672 24.67665 
+2
+Lattice="100.000000 0.000000 0.000000 0.000000 100.000000 0.000000 0.000000 0.000000 100.000000" 
+H 24.11118 25.45973 25.32734 
+H 25.90996 24.52988 24.67234 
+2
+Lattice="100.000000 0.000000 0.000000 0.000000 100.000000 0.000000 0.000000 0.000000 100.000000" 
+H 24.61501 25.11814 25.16434 
+H 25.40722 24.87117 24.83571 
+2
+Lattice="100.000000 0.000000 0.000000 0.000000 100.000000 0.000000 0.000000 0.000000 100.000000" 
+H 23.77229 25.28780 25.54274 
+H 26.25134 24.70113 24.45674 
+2
+Lattice="100.000000 0.000000 0.000000 0.000000 100.000000 0.000000 0.000000 0.000000 100.000000" 
+H 24.46641 25.08991 25.25368 
+H 25.55701 24.89956 24.74784 
+2
+Lattice="100.000000 0.000000 0.000000 0.000000 100.000000 0.000000 0.000000 0.000000 100.000000" 
+H 24.28908 24.89950 25.41426 
+H 25.73620 25.09140 24.58697 
+2
+Lattice="100.000000 0.000000 0.000000 0.000000 100.000000 0.000000 0.000000 0.000000 100.000000" 
+H 23.93955 24.79223 25.63082 
+H 26.08527 25.19938 24.37127 
+2
+Lattice="100.000000 0.000000 0.000000 0.000000 100.000000 0.000000 0.000000 0.000000 100.000000" 
+H 24.70199 24.85212 25.21266 
+H 25.32230 25.13983 24.78881 
+2
+Lattice="100.000000 0.000000 0.000000 0.000000 100.000000 0.000000 0.000000 0.000000 100.000000" 
+H 24.18112 24.32190 25.67205 
+H 25.84437 25.67090 24.33231 
+2
+Lattice="100.000000 0.000000 0.000000 0.000000 100.000000 0.000000 0.000000 0.000000 100.000000" 
+H 24.48644 24.52555 25.45356 
+H 25.53874 25.46889 24.54878 
+2
+Lattice="100.000000 0.000000 0.000000 0.000000 100.000000 0.000000 0.000000 0.000000 100.000000" 
+H 24.73808 24.57492 25.35570 
+H 25.28781 25.41922 24.64488 
+2
+Lattice="100.000000 0.000000 0.000000 0.000000 100.000000 0.000000 0.000000 0.000000 100.000000" 
+H 24.46477 24.03101 25.78700 
+H 25.56175 25.96104 24.21417 
+2
+Lattice="100.000000 0.000000 0.000000 0.000000 100.000000 0.000000 0.000000 0.000000 100.000000" 
+H 24.85595 24.64596 25.27895 
+H 25.16951 25.34776 24.72335 
+2
+Lattice="100.000000 0.000000 0.000000 0.000000 100.000000 0.000000 0.000000 0.000000 100.000000" 
+H 24.83496 24.17036 25.62398 
+H 25.18796 25.82224 24.37982 
+2
+Lattice="100.000000 0.000000 0.000000 0.000000 100.000000 0.000000 0.000000 0.000000 100.000000" 
+H 24.86401 24.14410 25.64057 
+H 25.16169 25.84892 24.36302 

mltrain/tests/test_bias.py

@@ -0,0 +1,161 @@
+import os
+import numpy as np
+import mltrain as mlt
+from autode.atoms import Atom
+from mltrain.potentials._base import MLPotential
+from ase.calculators.calculator import Calculator
+from mltrain.utils import work_in_tmp_dir
+mlt.Config.n_cores = 1
+here = os.path.abspath(os.path.dirname(__file__))
+
+
+def _h2():
+    """Dihydrogen molecule"""
+    atoms = [Atom('H', -0.80952, 2.49855, 0.), Atom('H', -0.34877, 1.961, 0.)]
+    return mlt.Molecule(atoms=atoms, charge=0, mult=1)
+
+
+class HarmonicPotential(Calculator):
+
+    def get_potential_energy(self, atoms):
+
+        r = atoms.get_distance(0, 1)
+
+        return (r - 0.7)**2
+
+    def get_forces(self, atoms):
+
+        derivative = np.zeros((len(atoms), 3))
+
+        r = atoms.get_distance(0, 1)
+
+        x_dist, y_dist, z_dist = [atoms[0].position[j] - atoms[1].position[j]
+                                  for j in range(3)]
+
+        x_i, y_i, z_i = (x_dist / r), (y_dist / r), (z_dist / r)
+
+        derivative[0][:] = [x_i, y_i, z_i]
+        derivative[1][:] = [-x_i, -y_i, -z_i]
+
+        force = - 2 * derivative * (r-0.7)
+
+        return force
+
+
+class TestPotential(MLPotential):
+
+    __test__ = False
+
+    def __init__(self,
+                 name: str,
+                 system=None):
+
+        super().__init__(name=name, system=system)
+
+    @property
+    def ase_calculator(self):
+
+        potential = HarmonicPotential()
+
+        return potential
+
+    def _train(self) -> None:
+        """Train this potential on self._training_data"""
+
+    def requires_atomic_energies(self) -> bool:
+        """Does this potential need E_0s for each atom to be specified"""
+
+    def requires_non_zero_box_size(self) -> bool:
+        """Can this potential be run in a box with side lengths = 0"""
+
+
+@work_in_tmp_dir()
+def test_bias():
+
+    system = mlt.System(_h2(), box=[50, 50, 50])
+    pot = TestPotential('1D')
+
+    config = system.random_configuration()
+
+    bias = mlt.Bias(to_average=[[0, 1]], reference=2, kappa=10)
+
+    assert bias.ref is not None
+    assert bias.kappa is not None
+    assert bias.rxn_coord == [[0, 1]]
+
+    new_pos = [[0, 0, 0],
+               [0, 0, 1]]
+
+    ase_atoms = config.ase_atoms
+    ase_atoms.set_positions(new_pos, apply_constraint=False)
+
+    bias_energy = bias.__call__([[0, 1]], ase_atoms)
+
+    assert bias_energy == 5  # (kappa / 2) * (1-2)^2
+
+    bias_force = bias.grad([[0, 1]], ase_atoms)
+
+    assert bias_force[0][2] == - bias_force[1][[2]]
+    assert bias_force[0][2] == 10
+
+    trajectory = mlt.md.run_mlp_md(configuration=config,
+                                   mlp=pot,
+                                   fs=100,
+                                   temp=300,
+                                   dt=0.5,
+                                   interval=10,
+                                   bias=bias)
+
+    trajectory.save_xyz('tmp.xyz')
+
+    assert os.path.exists('tmp.xyz')
+
+
+@work_in_tmp_dir()
+def test_window_umbrella():
+
+    charge, mult = -1, 1
+
+    system = mlt.System(_h2(), box=[50, 50, 50])
+    pot = TestPotential('1D')
+
+    config = system.random_configuration()
+    new_pos = [[0, 0, 0],
+               [0, 0, 1]]
+
+    ase_atoms = config.ase_atoms
+    ase_atoms.set_positions(new_pos, apply_constraint=False)
+
+    mean_distances = mlt.umbrella._get_rxn_coords(ase_atoms, [[0, 1]])
+
+    assert mean_distances == 1.0
+
+    umbrella = mlt.UmbrellaSampling(to_average=[[0, 1]], kappa=10)
+
+    assert umbrella.kappa is not None
+    assert umbrella.num_pairs == 1
+    assert umbrella.refs is None
+
+    traj = mlt.ConfigurationSet()
+    traj.load_xyz(os.path.join(here, 'data', 'h2_traj.xyz'),
+                  charge=charge,
+                  mult=mult)
+
+    _ = umbrella._get_window_frames(traj, num_windows=10,
+                                    init_ref=0.7, final_ref=2)
+
+    assert np.alltrue(umbrella.refs == np.linspace(0.7, 2, 10))
+
+    umbrella.run_umbrella_sampling(traj,
+                                   pot,
+                                   temp=300,
+                                   interval=5,
+                                   dt=0.5,
+                                   num_windows=3,
+                                   fs=500)
+
+    assert umbrella.bias.kappa == 10
+    assert umbrella.refs is not None
+
+    assert os.path.exists('combined_windows.xyz')
+    assert os.path.exists('fitted_data.pdf')

mltrain/training/tests/test_selection.py

@@ -35,6 +35,11 @@ def test_selection_on_structures():
 
     configs = mlt.ConfigurationSet()
 
+    try:
+        import dscribe
+    except ModuleNotFoundError:
+        return
+
     file_path = os.path.join(here, 'data', 'methane.xyz')
     configs.load_xyz(filename=file_path,
                      charge=0, mult=1, box=None)