fix tests, add doc strings

flowermd/base/simulation.py

@@ -115,6 +115,10 @@ def __init__(
             self._rigid_constraint = constraint
         elif constraint and isinstance(constraint, hoomd.md.constrain.Distance):
             self._distance_constraint = constraint
+        elif constraint:
+            raise ValueError(
+                "`constaint` must be an instance of hoomd.md.constrain."
+            )
         self._integrate_group = self._create_integrate_group(
             rigid=True if self._rigid_constraint else False
         )

flowermd/tests/base/test_simulation.py

@@ -14,7 +14,7 @@
 from flowermd.library.forcefields import EllipsoidForcefield
 from flowermd.library.polymers import EllipsoidChain
 from flowermd.tests import BaseTest
-from flowermd.utils import create_rigid_body, get_target_box_mass_density
+from flowermd.utils import get_target_box_mass_density, set_bond_constraints
 
 
 class TestSimulate(BaseTest):
@@ -380,15 +380,14 @@ def test_pickle_ff(self, benzene_system):
 
     def test_bad_rigid(self, benzene_system):
         with pytest.raises(ValueError):
-            Simulation.from_system(benzene_system, rigid_constraint="A")
+            Simulation.from_system(benzene_system, constraint="A")
 
-    def test_rigid_sim(self):
+    def test_d_constrain_sim(self):
         ellipsoid_chain = EllipsoidChain(
             lengths=4,
             num_mols=2,
-            lpar=0.5,
+            lpar=1.0,
             bead_mass=100,
-            bond_length=0.01,
         )
         system = Pack(
             molecules=ellipsoid_chain,
@@ -397,28 +396,23 @@ def test_rigid_sim(self):
             fix_orientation=True,
         )
         ellipsoid_ff = EllipsoidForcefield(
-            lpar=0.5,
-            lperp=0.25,
+            lpar=1,
+            lperp=0.5,
             epsilon=1.0,
-            r_cut=2.0,
-            bond_k=500,
-            bond_r0=0.01,
-            angle_k=250,
+            r_cut=2.5,
+            angle_k=50,
             angle_theta0=2.2,
         )
-        rigid_frame, rigid = create_rigid_body(
-            system.hoomd_snapshot,
-            ellipsoid_chain.bead_constituents_types,
-            bead_name="R",
+        constrain_snap, d_constraint = set_bond_constraints(
+            system.hoomd_snapshot, constrain_value=1.0, bond_type="_C-_H"
         )
         sim = Simulation(
-            initial_state=rigid_frame,
+            initial_state=constrain_snap,
             forcefield=ellipsoid_ff.hoomd_forces,
-            rigid_constraint=rigid,
+            constraint=d_constraint,
         )
-        sim.run_NVT(n_steps=0, kT=1.0, tau_kt=sim.dt * 100)
+        sim.run_NVT(n_steps=10, kT=1.0, tau_kt=sim.dt * 100)
         assert sim.integrator.integrate_rotational_dof is True
-        assert sim.mass_reduced == 800.0
 
     def test_save_restart_gsd(self, benzene_system):
         sim = Simulation.from_system(benzene_system)

flowermd/tests/library/test_forcefields.py

@@ -111,24 +111,27 @@ def test_ellipsoid_ff(self):
             lperp=0.5,
             lpar=1.0,
             r_cut=3,
-            bond_k=500,
-            bond_r0=0.1,
             angle_k=100,
             angle_theta0=1.57,
         )
-        assert len(ellipsoid_ff.hoomd_forces) == 3
+        assert len(ellipsoid_ff.hoomd_forces) == 2
         assert isinstance(
             ellipsoid_ff.hoomd_forces[-1], hoomd.md.pair.aniso.GayBerne
         )
-        assert ("R", "R") in list(
+        assert ("_C", "_C") in list(
             dict(ellipsoid_ff.hoomd_forces[-1].params).keys()
         )
-        assert ("B", "R") in list(
+        assert ("_C", "_H") in list(
             dict(ellipsoid_ff.hoomd_forces[-1].params).keys()
         )
-        assert ellipsoid_ff.hoomd_forces[-1].params["A", "A"]["epsilon"] == 0.0
-        assert ellipsoid_ff.hoomd_forces[-1].params["A", "A"]["lperp"] == 0.0
-        assert ellipsoid_ff.hoomd_forces[-1].params["A", "A"]["lpar"] == 0.0
+        assert (
+            ellipsoid_ff.hoomd_forces[-1].params["_C", "_H"]["epsilon"] == 0.0
+        )
+        assert (
+            ellipsoid_ff.hoomd_forces[-1].params["_H", "_H"]["epsilon"] == 0.0
+        )
+        assert ellipsoid_ff.hoomd_forces[-1].params["_C", "_H"]["lperp"] == 0.0
+        assert ellipsoid_ff.hoomd_forces[-1].params["_H", "_H"]["lpar"] == 0.0
 
 
 class TestTableForcefield:

flowermd/tests/library/test_polymers.py

@@ -114,9 +114,8 @@ def test_ellipsoid_chain(self):
             num_mols=2,
             lpar=0.5,
             bead_mass=100,
-            bond_length=0.01,
         )
-        assert ellipsoid_chain.n_particles == 32
+        assert ellipsoid_chain.n_particles == 16
         assert ellipsoid_chain.molecules[0].mass == 400
-        assert ellipsoid_chain.molecules[0].n_particles == 16
-        assert ellipsoid_chain.molecules[0].n_bonds == 10
+        assert ellipsoid_chain.molecules[0].n_particles == 8
+        assert ellipsoid_chain.molecules[0].n_bonds == 7

flowermd/utils/constraints.py

@@ -6,7 +6,38 @@
 
 
 def set_bond_constraints(snapshot, bond_type, constrain_value, tolerance=1e-5):
-    """"""
+    """Helper method to add fixed bond constraints to a gsd.hoomd.Frame.
+
+    Parameters
+    ----------
+    snapshot : gsd.hoomd.Frame, required
+        Snapshot of complete topology that will have constraints added.
+    bond_type : str, required
+        The bond type to add constraints for. Must match snapshot.bonds.types.
+    constrain_value : float, required
+        The value to use for the constrained bond length.
+        Must be close to the exisitng bond lenghts in snapshot.bonds
+    tolerance : float, default 1e-5
+        Used to compare actual bond lengths vs `constraint_value`
+        Sets the tolerance property in hoomd.md.constrain.Distance
+
+    Returns
+    -------
+    snapshot : gsd.hoomd.Frame
+        The modified snapshot with populated snapshot.constraints
+    d : gsd.hoomd.constrain.Distance
+        Used when initializing a simulation in flowermd.base.Simulation
+
+    Notes
+    -----
+    This method was added as a helper function to be used with the
+    ellipsoid chain module. See flowermd.library.polymer.EllipsoidChain
+    and flowermd.library.forcefields.EllipsoidForcefield.
+
+    Pass the snapshot and constraint object into flowermd.base.Simulation
+    in order for the fixed bond lengths to take effect.
+
+    """
     constraint_values = []
     constraint_groups = []
     bond_type_id = snapshot.bonds.types.index(bond_type)