diff --git a/constrainmol/constrainmol.py b/constrainmol/constrainmol.py index a90f75a..fffeb82 100644 --- a/constrainmol/constrainmol.py +++ b/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 diff --git a/constrainmol/tests/base_test.py b/constrainmol/tests/base_test.py index 61b526e..4a4d956 100644 --- a/constrainmol/tests/base_test.py +++ b/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], diff --git a/constrainmol/tests/test_constrainmol.py b/constrainmol/tests/test_constrainmol.py index b078def..9160aa0 100644 --- a/constrainmol/tests/test_constrainmol.py +++ b/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()