Initial addition of angle constraints. Tests not passing.

constrainmol/constrainmol.py

@@ -6,7 +6,7 @@
 
 class ConstrainedMolecule(object):
 
-    def __init__(self, structure):
+    def __init__(self, structure, constrain_angles=False):
         """Initialize the ConstrainedMolecule from a parmed.Structure
 
         Parameters
@@ -28,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 = deepcopy(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
         -------
@@ -59,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)
 
@@ -102,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 +
@@ -133,6 +157,31 @@ def _calc_bond_length(m, i, j):
         else:
             return pyo.Constraint.Skip
 
+    @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):
         """Solve the pyomo model to find the constrained coordinates
 
@@ -151,10 +200,17 @@ 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 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

@@ -33,9 +33,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],

constrainmol/tests/test_constrainmol.py

@@ -48,10 +48,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()
@@ -107,6 +111,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.deg2rad(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()
@@ -118,6 +142,25 @@ def test_dimethylether(self, dimehtylether_oplsaa):
             dist = np.sqrt(np.sum((xyz[idx2] - xyz[idx1])**2))
             assert np.allclose(dist, bond.type.req)
 
+        constrain_mol = ConstrainedMolecule(dimehtylether_oplsaa, 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.deg2rad(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()