Copying molecule copies constraints and forcefields

moldesign/_tests/test_copies.py

@@ -166,3 +166,44 @@ def test_chain_rename(pdb3aid):
     assert newmol.chains[0].name == 'A'
     assert newmol.chains[1].name == 'B'
 
+
+@pytest.mark.parametrize('molkey', ["create_parameterize_am1bcc",
+                                    "create_protein_default_amber_forcefield"])
+def test_forcefield_copied_with_molecule(molkey, request):
+    mol = request.getfixturevalue(molkey)
+    m2 = mol.copy()
+
+    assert isinstance(m2.ff, mol.ff.__class__)
+    assert isinstance(m2.ff.parmed_obj, mol.ff.parmed_obj.__class__)
+    assert m2.ff.parmed_obj is not mol.ff.parmed_obj
+
+    p1 = m2.ff.parmed_obj
+    p2 = m2.ff.parmed_obj
+    assert m2.ff.parmed_obj.LJ_depth == mol.ff.parmed_obj.LJ_depth
+    assert p1.bond_types == p2.bond_types
+    assert p1.angle_types == p2.angle_types
+    assert p1.dihedral_types == p2.dihedral_types
+    assert p1.improper_types == p2.improper_types
+
+
+def test_constraints_copied_with_molecule(parameterize_zeros):
+    mol = parameterize_zeros
+
+    mol.constrain_distance(*mol.atoms[:2])
+    mol.constrain_angle(*mol.atoms[:3])
+    mol.constrain_dihedral(*mol.atoms[:4])
+    mol.constrain_atom(mol.atoms[0])
+    mol.constrain_hbonds()
+
+    mcpy = mol.copy()
+    assert mcpy.constraints[0].desc == 'distance'
+    assert mcpy.constraints[1].desc == 'angle'
+    assert mcpy.constraints[2].desc == 'dihedral'
+    assert mcpy.constraints[3].desc == 'position'
+    assert mcpy.constraints[4].desc == 'hbonds'
+
+    for constraint in mcpy.constraints:
+        assert constraint.mol is mcpy
+        if constraint.desc != 'hbonds':
+            for atom in constraint.atoms:
+                assert atom.molecule is mcpy

moldesign/geom/constraints.py

@@ -16,6 +16,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import copy
 
 import moldesign as mdt
 from .. import units as u
@@ -53,6 +54,24 @@ def __init__(self, atoms, value=None, tolerance=None, force_constant=None):
             assert atom.molecule is self.mol
         self.value = mdt.utils.if_not_none(value, self.current())
 
+    def copy(self, mol=None):
+        """ Copy this constraint, optionally relinking to a new molecule
+
+        Args:
+            mol (moldesign.Molecule): optional new molecule to track.
+
+        Returns:
+            GeometryConstraint: new constraint instance
+        """
+        if mol is None:
+            mol = self.mol
+        newatoms = [mol.atoms[atom.index] for atom in self.atoms]
+
+        # Note: this is the call signature for most subclasses, different than this base class's
+        return self.__class__(*newatoms,
+                              value=self.value, tolerance=self.tolerance,
+                              force_constant=self.force_constant)
+
     def current(self):
         """
         Return the current value of the constrained quantity
@@ -254,8 +273,6 @@ class HBondsConstraint(GeometryConstraint):
     desc = 'hbonds'
 
     def __init__(self, mol):
-        from ..interfaces.openmm import simtk2pint
-
         self.mol = mol
         self.bonds = []
         self.subconstraints = []
@@ -266,6 +283,12 @@ def __init__(self, mol):
                                                               value=bond.ff.equilibrium_length))
         self.values = [c.current() for c in self.subconstraints]
 
+    def copy(self, mol=None):
+        if mol is None:
+            return super().copy()
+        else:
+            return self.__class__(mol)
+
     @property
     def tolerance(self):
         return len(self.bonds) * DIST_TOLERANCE**2
@@ -336,6 +359,22 @@ def __init__(self, atom, vector, value=None,
         super().__init__([atom], value=self.value,
                          tolerance=tolerance, force_constant=force_constant)
 
+    def copy(self, mol=None):
+        """ Copy this constraint, optionally relinking to a new molecule
+
+        Args:
+            mol (moldesign.Molecule): optional new molecule to track.
+
+        Returns:
+            GeometryConstraint: new constraint instance
+        """
+        if mol is None:
+            mol = self.mol
+        newatom = mol.atoms[self.atom.index]
+        return self.__class__(newatom, self.vector.copy(),
+                              value=self.value, tolerance=self.tolerance,
+                              force_constant=self.force_constant)
+
     def current(self):
         return self.atom.position.dot(self.vector)
     current.__doc__ = GeometryConstraint.current.__doc__

moldesign/molecules/molecule.py

@@ -417,7 +417,7 @@ class MolTopologyMixin(object):
 
     Note:
         This is a mixin class designed only to be mixed into the :class:`Molecule` class. Routines
-        are separated are here for code organization only - they could be included in the main
+        are here for code organization only - they could be included in the main
         Atom class without changing any functionality
     """
     def copy(self, name=None):
@@ -440,6 +440,9 @@ def copy(self, name=None):
         newintegrator = self._copy_method(newmol, 'integrator')
         if newintegrator is not None:
             newmol.set_integrator(newintegrator)
+        if self.ff is not None:
+            self.ff.copy_to(newmol)
+        newmol.constraints = [c.copy(newmol) for c in self.constraints]
         return newmol
 
     def _copy_method(self, newmol, methodname):