Revamp Ellipsoid Class, use fixed bonds.

codecov.yml

@@ -11,4 +11,3 @@ ignore:
     - "flowermd/assets"
     - "flowermd/__version__.py"
     - "tutorials"
-    - "setup.py"

flowermd/base/simulation.py

@@ -13,7 +13,7 @@
 import numpy as np
 import unyt as u
 
-from flowermd.internal import validate_ref_value
+from flowermd.internal import check_return_iterable, validate_ref_value
 from flowermd.utils.actions import StdOutLogger, UpdateWalls
 from flowermd.utils.base_types import HOOMDThermostats
 
@@ -52,7 +52,10 @@ class Simulation(hoomd.simulation.Simulation):
     thermostat : flowermd.utils.HOOMDThermostats, default
         HOOMDThermostats.MTTK
         The thermostat to use for the simulation.
-
+    constraint : list of hoomd.md.constrain objects
+        Sets constraints for the simulation.
+        See flowermd.utils.constraints for built-in helpers
+        or see https://hoomd-blue.readthedocs.io/en/stable/hoomd/md/module-constrain.html
     """
 
     def __init__(
@@ -69,7 +72,7 @@ def __init__(
         log_write_freq=1e3,
         log_file_name="sim_data.txt",
         thermostat=HOOMDThermostats.MTTK,
-        rigid_constraint=None,
+        constraint=None,
     ):
         if not isinstance(forcefield, Iterable) or isinstance(forcefield, str):
             raise ValueError(
@@ -105,16 +108,20 @@ def __init__(
         self._dt = dt
         self._reference_values = dict()
         self._reference_values = reference_values
-        if rigid_constraint and not isinstance(
-            rigid_constraint, hoomd.md.constrain.Rigid
-        ):
-            raise ValueError(
-                "Invalid rigid constraint. Please provide a "
-                "hoomd.md.constrain.Rigid object."
-            )
-        self._rigid_constraint = rigid_constraint
+        self.constraint = check_return_iterable(constraint)
+        self._rigid_constraint = None
+        self._distance_constraint = None
+        for obj in self.constraint:
+            if isinstance(obj, hoomd.md.constrain.Rigid):
+                self._rigid_constraint = obj
+            elif isinstance(obj, hoomd.md.constrain.Distance):
+                self._distance_constraint = obj
+            else:
+                raise ValueError(
+                    "`constaint` must be an instance of hoomd.md.constrain."
+                )
         self._integrate_group = self._create_integrate_group(
-            rigid=True if rigid_constraint else False
+            rigid=True if self._rigid_constraint else False
         )
         self._wall_forces = dict()
         self._create_state(self.initial_state)
@@ -589,12 +596,12 @@ def set_integrator_method(self, integrator_method, method_kwargs):
         if not self.integrator:  # Integrator and method not yet created
             self.integrator = hoomd.md.Integrator(
                 dt=self.dt,
-                integrate_rotational_dof=(
-                    True if self._rigid_constraint else False
-                ),
+                integrate_rotational_dof=(True if self.constraint else False),
             )
             if self._rigid_constraint:
                 self.integrator.rigid = self._rigid_constraint
+            if self._distance_constraint:
+                self.integrator.constraints.append(self._distance_constraint)
             self.integrator.forces = self._forcefield
             self.operations.add(self.integrator)
             new_method = integrator_method(**method_kwargs)

flowermd/library/forcefields.py

@@ -584,14 +584,14 @@ class EllipsoidForcefield(BaseHOOMDForcefield):
         Semi-axis length of the ellipsoid along the minor axis.
     r_cut : float, required
         Cut off radius for pair interactions
-    bond_k : float, required
-        Spring constant in harmonic bond
-    bond_r0: float, required
-        Equilibrium tip-to-tip bond length.
     angle_k : float, required
         Spring constant in harmonic angle.
     angle_theta0: float, required
         Equilibrium angle between 2 consecutive beads.
+    bond_k : float, required
+        Spring constant in harmonic bond.
+    bond_r0: float, required
+        Equilibrium distance between 2 ellipsoid tips.
 
     """
 
@@ -601,48 +601,54 @@ def __init__(
         lpar,
         lperp,
         r_cut,
-        bond_k,
-        bond_r0,
         angle_k=None,
         angle_theta0=None,
+        bond_k=100,
+        bond_r0=0.1,
     ):
         self.epsilon = epsilon
         self.lperp = lperp
         self.lpar = lpar
         self.r_cut = r_cut
-        self.bond_k = bond_k
-        self.bond_r0 = bond_r0
         self.angle_k = angle_k
         self.angle_theta0 = angle_theta0
+        self.bond_k = bond_k
+        self.bond_r0 = bond_r0
         hoomd_forces = self._create_forcefield()
         super(EllipsoidForcefield, self).__init__(hoomd_forces)
 
     def _create_forcefield(self):
         forces = []
         # Bonds
         bond = hoomd.md.bond.Harmonic()
-        bond.params["A-A"] = dict(k=self.bond_k, r0=self.bond_r0)
-        bond.params["B-B"] = dict(k=0, r0=self.lpar)
+        bond.params["T-T"] = dict(k=self.bond_k, r0=self.bond_r0)
+        bond.params["A-X"] = dict(k=0, r0=0)
         forces.append(bond)
         # Angles
         if all([self.angle_k, self.angle_theta0]):
             angle = hoomd.md.angle.Harmonic()
-            angle.params["B-B-B"] = dict(k=self.angle_k, t0=self.angle_theta0)
+            angle.params["X-A-X"] = dict(k=self.angle_k, t0=self.angle_theta0)
+            angle.params["A-X-A"] = dict(k=0, t0=0)
             forces.append(angle)
         # Gay-Berne Pairs
-        nlist = hoomd.md.nlist.Cell(buffer=0.40)
+        nlist = hoomd.md.nlist.Cell(buffer=0.40, exclusions=["body"])
         gb = hoomd.md.pair.aniso.GayBerne(nlist=nlist, default_r_cut=self.r_cut)
-        gb.params[("R", "R")] = dict(
+        gb.params[("X", "X")] = dict(
             epsilon=self.epsilon, lperp=self.lperp, lpar=self.lpar
         )
         # Add zero pairs
         for pair in [
+            ("R", "R"),
+            ("T", "T"),
+            ("T", "R"),
             ("A", "A"),
-            ("B", "B"),
-            ("A", "B"),
+            ("A", "X"),
+            ("A", "T"),
             ("A", "R"),
-            ("B", "R"),
+            ("X", "R"),
+            ("X", "T"),
         ]:
             gb.params[pair] = dict(epsilon=0.0, lperp=0.0, lpar=0.0)
+            gb.params[pair].r_cut = 0.0
         forces.append(gb)
         return forces

flowermd/library/polymers.py

@@ -3,8 +3,8 @@
 import os
 
 import mbuild as mb
+import numpy as np
 from mbuild.coordinate_transform import z_axis_transform
-from mbuild.lib.recipes import Polymer as mbPolymer
 
 from flowermd import CoPolymer, Polymer
 from flowermd.assets import MON_DIR
@@ -288,8 +288,8 @@ class EllipsoidChain(Polymer):
 
     This is meant to be used with
     `flowermd.library.forcefields.EllipsoidForcefield`
-    and requires using `flowermd.utils.rigid_body` to set up
-    the rigid bodies correctly in HOOMD-Blue.
+    and requires using `flowermd.utils.constraints.set_bond_constraints` to set up
+    the fixed bonds correctly in HOOMD-Blue.
 
     Parameters
     ----------
@@ -301,53 +301,42 @@ class EllipsoidChain(Polymer):
         The semi-axis length of the ellipsoid bead along its major axis.
     bead_mass : float, required
         The mass of the ellipsoid bead.
-    bond_length : float, required
-        The bond length between connected beads.
-        This is used as the bond length between ellipsoid tips
-        rather than between ellipsoid centers.
-
     """
 
-    def __init__(self, lengths, num_mols, lpar, bead_mass, bond_length):
+    def __init__(self, lengths, num_mols, lpar, bead_mass, bond_L=0.1):
         self.bead_mass = bead_mass
-        self.bead_bond_length = bond_length
         self.lpar = lpar
+        self.bond_L = bond_L
         # get the indices of the particles in a rigid body
-        self.bead_constituents_types = ["A", "A", "B", "B"]
+        self.bead_constituents_types = ["X", "A"]
         super(EllipsoidChain, self).__init__(lengths=lengths, num_mols=num_mols)
 
     def _build(self, length):
         # Build up ellipsoid bead
         bead = mb.Compound(name="ellipsoid")
+        center = mb.Compound(pos=(0, 0, 0), name="X", mass=self.bead_mass / 4)
         head = mb.Compound(
-            pos=(0, 0, self.lpar), name="A", mass=self.bead_mass / 4
-        )
-        tail = mb.Compound(
-            pos=(0, 0, -self.lpar), name="A", mass=self.bead_mass / 4
-        )
-        head_mid = mb.Compound(
-            pos=(0, 0, self.lpar / 2), name="B", mass=self.bead_mass / 4
+            pos=(self.lpar, 0, 0), name="A", mass=self.bead_mass / 4
         )
-        tail_mid = mb.Compound(
-            pos=(0, 0, -self.lpar / 2), name="B", mass=self.bead_mass / 4
+        tether_head = mb.Compound(
+            pos=(self.lpar, 0, 0), name="T", mass=self.bead_mass / 4
         )
-        bead.add([head, tail, head_mid, tail_mid])
-        # Build the bead chain
-        chain = mbPolymer()
-        chain.add_monomer(
-            bead,
-            indices=[0, 1],
-            orientation=[[0, 0, 1], [0, 0, -1]],
-            replace=False,
-            separation=self.bead_bond_length,
-        )
-        chain.build(n=length, add_hydrogens=False)
-        # Generate bonds between the mid-particles.
-        # This is needed to use an angle potential between 2 beads.
-        chain.freud_generate_bonds(
-            name_a="B",
-            name_b="B",
-            dmin=self.lpar - 0.1,
-            dmax=self.lpar + self.bead_bond_length + 0.1,
+        tether_tail = mb.Compound(
+            pos=(-self.lpar, 0, 0), name="T", mass=self.bead_mass / 4
         )
+        bead.add([center, head, tether_head, tether_tail])
+        bead.add_bond([center, head])
+
+        chain = mb.Compound()
+        last_bead = None
+        for i in range(length):
+            translate_by = np.array([(i * self.lpar * 2) + self.bond_L, 0, 0])
+            this_bead = mb.clone(bead)
+            this_bead.translate(by=translate_by)
+            chain.add(this_bead)
+            if last_bead:
+                chain.add_bond([this_bead.children[0], last_bead.children[1]])
+                chain.add_bond([this_bead.children[3], last_bead.children[2]])
+            last_bead = this_bead
+
         return chain

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):
@@ -378,17 +378,16 @@ def test_pickle_ff(self, benzene_system):
             assert type(i) is type(j)
         os.remove("forcefield.pickle")
 
-    def test_bad_rigid(self, benzene_system):
+    def test_bad_constraint(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,25 @@ 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)
+        assert isinstance(sim._distance_constraint, hoomd.md.constrain.Distance)
+        assert sim._rigid_constraint is None
+        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