Initial addition of angle constraints

azure-pipelines.yml

@@ -6,14 +6,14 @@ pr:
   autoCancel: true
   branches:
     include:
-      - master
+      - main
 
 schedules:
 - cron: "0 0 * * *"
   displayName: Daily midnight build for master
   branches:
     include:
-      - master
+      - main
 
 jobs:
 - job: Stable

constrainmol/constrainmol.py

@@ -1,11 +1,12 @@
 import pyomo.environ as pyo
 import parmed
 import numpy as np
+from warnings import warn
 
 
 class ConstrainedMolecule(object):
 
-    def __init__(self, structure):
+    def __init__(self, structure, constrain_angles=False):
         """Initialize the ConstrainedMolecule from a parmed.Structure
 
         Parameters
@@ -27,27 +28,43 @@ def __init__(self, structure):
                 f"structure: {structure} contains no bonds"
             )
 
-        # Extract coordinates and bonds
+        # Extract coordinates
         xyz = structure.coordinates
-        constraints = {}
+
+        # Extract bond constraints
+        bond_constraints = {}
         for bond in structure.bonds:
             idx1 = bond.atom1.idx
             idx2 = bond.atom2.idx
-            constraints[(idx1, idx2)] = bond.type.req
+            bond_constraints[(idx1, idx2)] = bond.type.req
+
+        # Extract angle constraints if desired
+        if constrain_angles:
+            angle_constraints = {}
+            for angle in structure.angles:
+                idx1 = angle.atom1.idx
+                idx2 = angle.atom2.idx
+                idx3 = angle.atom3.idx
+                angle_constraints[(idx1, idx2, idx3)] = angle.type.theteq
+        else:
+            angle_constraints = None
 
+        # Copy the pmd.structure
         self.structure = parmed.structure.copy(structure)
-        self.model = self._create_model(xyz, constraints)
+        self.model = self._create_model(xyz, bond_constraints, angle_constraints)
         self.model_solved = False
 
-    def _create_model(self, xyz, constraints):
-        """Create a pyomo model to make xyz satisfy bond constraints
+    def _create_model(self, xyz, bond_constraints, angle_constraints):
+        """Create a pyomo model to make xyz satisfy bond and angle constraints
 
         Parameters
         ----------
         xyz: np.ndarray, shape=(n_atoms, 3)
             initial xyz coordinates
-        constraints: dict, keys=(idx1, idx2), values=bond length
+        bond_constraints: dict, keys=(idx1, idx2), values=bond length
             dictionary with bond length constraints
+        angle_constraints: dict, keys=(idx1, idx2, idx2), values=angle
+            dictionary with bond angle constraints
 
         Returns
         -------
@@ -58,7 +75,7 @@ def _create_model(self, xyz, constraints):
         # Create pyomo model
         m = pyo.ConcreteModel()
 
-        # Create list of atom indicies
+        # Create list of atom indices
         ids = range(xyz.shape[0])
         m.atom_ids = pyo.Set(initialize=ids)
 
@@ -101,14 +118,22 @@ def _create_model(self, xyz, constraints):
 
         # Create bond length parameter
         m.bond_lengths = pyo.Param(
-            m.atom_ids, m.atom_ids, initialize=constraints
+            m.atom_ids, m.atom_ids, initialize=bond_constraints
         )
 
         # Add bond length constraints
-        m.eq_calc_bound_length = pyo.Constraint(
+        m.eq_calc_bond_length = pyo.Constraint(
             m.atom_ids, m.atom_ids, rule=self._calc_bond_length
         )
 
+        if angle_constraints is not None:
+            m.bond_angles = pyo.Param(
+                m.atom_ids, m.atom_ids, m.atom_ids, initialize=angle_constraints
+            )
+            m.eq_calc_angle = pyo.Constraint(
+                m.atom_ids, m.atom_ids, m.atom_ids, rule=self._calc_angle
+            )
+
         m.obj = pyo.Objective(
             expr=sum(
                 (m.x[i] - m.x_start[i])**2 +
@@ -132,14 +157,40 @@ def _calc_bond_length(m, i, j):
         else:
             return pyo.Constraint.Skip
 
-    def solve(self, verbose=False):
+    @staticmethod
+    def _calc_angle(m, i, j, k):
+        if (i, j, k) in m.bond_angles.keys():
+            ji = [
+                m.x[i] - m.x[j],
+                m.y[i] - m.y[j],
+                m.z[i] - m.z[j],
+            ]
+            jk = [
+                m.x[k] - m.x[j],
+                m.y[k] - m.y[j],
+                m.z[k] - m.z[j],
+            ]
+            norm_ji = pyo.sqrt(ji[0]**2 + ji[1]**2 + ji[2]**2)
+            norm_jk = pyo.sqrt(jk[0]**2 + jk[1]**2 + jk[2]**2)
+
+            dot = (
+                ji[0] * jk[0] + ji[1]*jk[1] + ji[2]*jk[2]
+            )
+            cos_angle = dot / (norm_ji * norm_jk)
+            angle = pyo.acos(cos_angle) * 180 / np.pi
+            return angle == m.bond_angles[(i, j, k)]
+        else:
+            return pyo.Constraint.Skip
+
+    def solve(self, verbose=False, ignore_errors=False):
         """Solve the pyomo model to find the constrained coordinates
 
         Parameters
         ----------
         verbose : boolean, optional, default=False
             print the solver output to screen
-
+        ignore_errors : boolean, optional, default=False
+            ignore solver errors and return
         Notes
         -----
         Updates ConstrainedMolecule.structure.coordinates if solve is successful
@@ -150,10 +201,24 @@ def solve(self, verbose=False):
                 str(result["Solver"][0]["Termination condition"]) == 'optimal'
         )
         if not success:
-            raise ValueError(
-                "Optimal solution not found. You may want to run "
-                "'solve' with 'verbose=True' to see the solver output."
-            )
+            if ignore_errors:
+                warn(
+                    "Optimal solution not found. Your constraints may not "
+                    "be satisfied. Run 'solve' with 'verbose=True' "
+                    "to see the solver output."
+                )
+            else:
+                if not verbose:
+                    raise ValueError(
+                        "Optimal solution not found. You may want to run "
+                        "'solve' with 'verbose=True' to see the solver output."
+                    )
+                else:
+                    raise ValueError(
+                        "Optimal solution not found. See the solver output "
+                        "for details."
+                    )
+
         constrained_xyz = np.zeros((len(self.model.atom_ids), 3))
         for idx in self.model.atom_ids:
             constrained_xyz[idx, 0] = self.model.x[idx].value

constrainmol/tests/base_test.py

@@ -4,6 +4,8 @@
 import parmed
 import numpy as np
 
+from os.path import join, split, abspath
+
 
 class BaseTest:
 
@@ -33,9 +35,13 @@ def propane_ua(self):
         bond_type = parmed.topologyobjects.BondType(1.5, 1.5)
         b1 = parmed.topologyobjects.Bond(a1, a2, type=bond_type)
         b2 = parmed.topologyobjects.Bond(a2, a3, type=bond_type)
+        angle_type = parmed.topologyobjects.AngleType(1.0, 120)
+        ang1 = parmed.topologyobjects.Angle(a1, a2, a3, type=angle_type)
         propane.bonds.append(b1)
         propane.bonds.append(b2)
+        propane.angles.append(ang1)
         propane.bond_types.append(bond_type)
+        propane.angle_types.append(angle_type)
         propane.coordinates = np.array(
             [
                 [0.0, 0.1, 0.2],
@@ -45,6 +51,14 @@ def propane_ua(self):
         )
         return propane
 
+    @pytest.fixture
+    def water_spce(self):
+        ff_file = get_fn("spce.xml")
+        ff = foyer.Forcefield(ff_file)
+        water = mbuild.load("O", smiles=True)
+        water = ff.apply(water)
+        return water
+
     @pytest.fixture
     def dimehtylether_oplsaa(self):
         ff = foyer.forcefields.load_OPLSAA()
@@ -66,3 +80,19 @@ def diethylether_box(self):
         dee_ff = ff.apply(dee)
         box = mbuild.fill_box(dee, 50, density=600)
         return dee_ff, box
+
+
+def get_fn(filename):
+    """Gets the full path of the file name for a particular test file.
+
+    Parameters
+    ----------
+    filename : str
+        name of the file to get
+
+    Returns
+    -------
+    path: str
+        path to the file
+    """
+    return join(split(abspath(__file__))[0], "files", filename)

constrainmol/tests/files/spce.xml

@@ -0,0 +1,16 @@
+<ForceField>
+ <AtomTypes>
+  <Type name="spce_o" class="OW" element="O" mass="15.999" def="O" desc="SPC/E oxygen" doi="10.1021/j100308a038"/>
+  <Type name="spce_h" class="HW" element="H" mass="1.008" def="H" desc="SPC/E hydrogen" doi="10.1021/j100308a038"/>
+ </AtomTypes>
+ <HarmonicBondForce>
+  <Bond class1="OW" class2="HW" length="0.1" k="502416.0"/>
+ </HarmonicBondForce>
+ <HarmonicAngleForce>
+  <Angle class1="HW" class2="OW" class3="HW" angle="1.91061193" k="383.0"/>
+ </HarmonicAngleForce>
+ <NonbondedForce coulomb14scale="0.5" lj14scale="0.5">
+  <Atom type="spce_o" charge="-0.8476" sigma="0.316557" epsilon="0.650194"/>
+  <Atom type="spce_h" charge="0.4238" sigma="0.0" epsilon="0.0"/>
+ </NonbondedForce>
+</ForceField>

constrainmol/tests/test_constrainmol.py

@@ -47,10 +47,14 @@ def test_model_xyz(self, ethane_ua):
         assert np.allclose(constrain_mol.model.z[0].value, ethane_ua.coordinates[0, 2])
         assert np.allclose(constrain_mol.model.z[1].value, ethane_ua.coordinates[1, 2])
 
-    def test_model_constraints(self, ethane_ua):
+    def test_model_bond_constraints(self, ethane_ua):
         constrain_mol = ConstrainedMolecule(ethane_ua)
         assert np.allclose(constrain_mol.model.bond_lengths[(0, 1)], ethane_ua.bonds[0].type.req)
 
+    def test_model_angle_constraints(self, propane_ua):
+        constrain_mol = ConstrainedMolecule(propane_ua, constrain_angles=True)
+        assert np.allclose(constrain_mol.model.bond_angles[(0, 1, 2)], propane_ua.angles[0].type.theteq)
+
     def test_solved_model(self, ethane_ua):
         constrain_mol = ConstrainedMolecule(ethane_ua)
         constrain_mol.solve()
@@ -106,6 +110,26 @@ def test_resolve_model(self, propane_ua):
         assert constrain_mol.model_solved is True
         assert np.allclose(propane_solved.coordinates, constrain_mol.structure.coordinates)
 
+    def test_propane_angles(self, propane_ua):
+        constrain_mol = ConstrainedMolecule(propane_ua, constrain_angles=True)
+        constrain_mol.solve(verbose=True)
+        optimized = constrain_mol.structure
+        xyz = optimized.coordinates
+        for bond in optimized.bonds:
+            idx1 = bond.atom1.idx
+            idx2 = bond.atom2.idx
+            dist = np.sqrt(np.sum((xyz[idx2] - xyz[idx1])**2))
+            assert np.allclose(dist, bond.type.req)
+        for angle in optimized.angles:
+            idx1 = angle.atom1.idx
+            idx2 = angle.atom2.idx
+            idx3 = angle.atom3.idx
+            ji = xyz[idx1] - xyz[idx2]
+            jk = xyz[idx3] - xyz[idx2]
+            cos_angle = np.dot(ji, jk) / (np.linalg.norm(ji) * np.linalg.norm(jk))
+            calc_angle = np.rad2deg(np.arccos(cos_angle))
+            assert np.allclose(calc_angle, angle.type.theteq)
+
     def test_dimethylether(self, dimehtylether_oplsaa):
         constrain_mol = ConstrainedMolecule(dimehtylether_oplsaa)
         constrain_mol.solve()
@@ -117,6 +141,35 @@ def test_dimethylether(self, dimehtylether_oplsaa):
             dist = np.sqrt(np.sum((xyz[idx2] - xyz[idx1])**2))
             assert np.allclose(dist, bond.type.req)
 
+        with pytest.raises(ValueError):
+            constrain_mol = ConstrainedMolecule(dimehtylether_oplsaa, constrain_angles=True)
+            constrain_mol.solve(verbose=True)
+
+        constrain_mol = ConstrainedMolecule(dimehtylether_oplsaa, constrain_angles=True)
+        constrain_mol.solve(ignore_errors=True)
+
+    def test_spce_water(self, water_spce):
+
+        constrain_mol = ConstrainedMolecule(water_spce, constrain_angles=True)
+        constrain_mol.solve(verbose=True)
+        optimized = constrain_mol.structure
+        xyz = optimized.coordinates
+        for bond in optimized.bonds:
+            idx1 = bond.atom1.idx
+            idx2 = bond.atom2.idx
+            dist = np.sqrt(np.sum((xyz[idx2] - xyz[idx1])**2))
+            assert np.allclose(dist, bond.type.req)
+        for angle in optimized.angles:
+            idx1 = angle.atom1.idx
+            idx2 = angle.atom2.idx
+            idx3 = angle.atom3.idx
+            ji = xyz[idx1] - xyz[idx2]
+            jk = xyz[idx3] - xyz[idx2]
+            cos_angle = np.dot(ji, jk) / (np.linalg.norm(ji) * np.linalg.norm(jk))
+            calc_angle = np.rad2deg(np.arccos(cos_angle))
+            assert np.allclose(calc_angle, angle.type.theteq)
+
+
     def test_benzene(self, benzene_oplsaa):
         constrain_mol = ConstrainedMolecule(benzene_oplsaa)
         constrain_mol.solve()