Add classes for simulating ellipsoid polymer systems. (#102)

* create EllipsoidChain class

* remove blank space

* add class for ellipsoid chain FF

* add FF class to library

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* update particle mass, add bead particle types sequence

* update class parent name

* add func to create rigid frame

* fix formatting, remove extra param

* fix numpy bug, add box to config

* fix rigid body sim bugs

* add bonds between B particles after building the chain

* add fix_orientation variable to pack

* add bond, angle and dihedral

* remove create_rigid_body from sim

* add body tags to rigid frame

* add new functions to __init__

* add unit tests

* move private func to bottom

* add doc strings

* add example code to docstring

* replace bead_length with lpar in EllipsoidChain class

* add rigid body handling of mass

* add overlap to Pack

* remove bead_length from Ellipsoid FF class, fix unit test

* fix rigid unit test

* add 2 unit tests for rigid simulations

* fix check for body tags in mass, add assertion for mass

* ignore E203 in precommit

* remove unused var

* use >= when checking for ascending next body tag

* actaully, we shouldn't use >= in body tag count

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix rigid body import

* Switch to z-axis for ellipsoid chain alignment, add param for initial orientaiton to rigid body util

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix test

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: marjanalbouye <[email protected]>

.pre-commit-config.yaml

@@ -37,6 +37,7 @@ repos:
       - id: flake8
         args:
           - --max-line-length=80
+          - --extend-ignore=E203
         exclude: '__init__.py'
 
   - repo: https://github.com/pycqa/pydocstyle

flowermd/base/simulation.py

@@ -67,6 +67,7 @@ def __init__(
         log_write_freq=1e3,
         log_file_name="sim_data.txt",
         thermostat=HOOMDThermostats.MTTK,
+        rigid_constraint=None,
     ):
         if not isinstance(forcefield, Iterable) or isinstance(forcefield, str):
             raise ValueError(
@@ -103,7 +104,17 @@ def __init__(
         self._dt = dt
         self._reference_values = dict()
         self._reference_values = reference_values
-        self._integrate_group = hoomd.filter.All()
+        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._integrate_group = self._create_integrate_group(
+            rigid=True if rigid_constraint else False
+        )
         self._wall_forces = dict()
         self._create_state(self.initial_state)
         # Add a gsd and thermo props logger to sim operations
@@ -289,7 +300,16 @@ def volume(self):
     def mass_reduced(self):
         """The total mass of the system in reduced units."""
         with self.state.cpu_local_snapshot as snap:
-            return sum(snap.particles.mass)
+            if self._rigid_constraint:
+                last_body_tag = -1
+                for body_tag in snap.particles.body:
+                    if body_tag > last_body_tag:
+                        last_body_tag += 1
+                    else:
+                        break
+                return sum(snap.particles.mass[last_body_tag + 1 :])
+            else:
+                return sum(snap.particles.mass)
 
     @property
     def mass(self):
@@ -529,7 +549,14 @@ 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)
+            self.integrator = hoomd.md.Integrator(
+                dt=self.dt,
+                integrate_rotational_dof=(
+                    True if self._rigid_constraint else False
+                ),
+            )
+            if self._rigid_constraint:
+                self.integrator.rigid = self._rigid_constraint
             self.integrator.forces = self._forcefield
             self.operations.add(self.integrator)
             new_method = integrator_method(**method_kwargs)
@@ -1098,6 +1125,11 @@ def _lj_force(self):
             ][0]
         return lj_force
 
+    def _create_integrate_group(self, rigid):
+        if rigid:
+            return hoomd.filter.Rigid(("center", "free"))
+        return hoomd.filter.All()
+
     def _create_state(self, initial_state):
         """Create the simulation state.
 

flowermd/base/system.py

@@ -624,6 +624,8 @@ class Pack(System):
         The factor by which to expand the box for packing.
     edge : float, default 0.2
         The space (nm) between the edge of the box and the molecules.
+    overlap : float, default 0.2
+        Minimum separation (nm) between particles of different molecules.
 
 
     .. warning::
@@ -653,6 +655,7 @@ def __init__(
         packing_expand_factor=5,
         edge=0.2,
         overlap=0.2,
+        fix_orientation=False,
     ):
         if not isinstance(density, u.array.unyt_quantity):
             self.density = density * u.Unit("g") / u.Unit("cm**3")
@@ -665,6 +668,7 @@ def __init__(
         self.packing_expand_factor = packing_expand_factor
         self.edge = edge
         self.overlap = overlap
+        self.fix_orientation = fix_orientation
         super(Pack, self).__init__(molecules=molecules, base_units=base_units)
 
     def _build_system(self):
@@ -692,6 +696,7 @@ def _build_system(self):
             box=list(target_box * self.packing_expand_factor),
             overlap=self.overlap,
             edge=self.edge,
+            fix_orientation=self.fix_orientation,
         )
         return system
 

flowermd/library/__init__.py

@@ -9,6 +9,7 @@
     BaseHOOMDForcefield,
     BaseXMLForcefield,
     BeadSpring,
+    EllipsoidForcefield,
     FF_from_file,
     KremerGrestBeadSpring,
     TableForcefield,
@@ -17,6 +18,7 @@
     PEEK,
     PEKK,
     PPS,
+    EllipsoidChain,
     LJChain,
     PEKK_meta,
     PEKK_para,

flowermd/library/forcefields.py

@@ -557,3 +557,92 @@ def _check_widths(self):
                         "number of points for table energies and forces."
                     )
         return bond_width, angle_width, dih_width
+
+
+class EllipsoidForcefield(BaseHOOMDForcefield):
+    """A forcefield for modeling anisotropic bead polymers.
+
+    Notes
+    -----
+    This is designed to be used with `flowermd.library.polymers.EllipsoidChain`
+    and uses ghost particles of type "A" and "B" for intra-molecular
+    interactions of bonds and two-body angles.
+    Ellipsoid centers (type "R") are used in inter-molecular pair interations.
+
+    The set of interactions are:
+    1. `hoomd.md.bond.Harmonic`: Models ellipsoid bonds as tip-to-tip bonds
+    2. `hoomd.md.angle.Harmonic`: Models angles of two neighboring ellipsoids.
+    3. `hoomd.md.pair.aniso.GayBerne`" Model pair interactions between beads.
+
+    Parameters
+    ----------
+    epsilon : float, required
+        energy
+    lpar: float, required
+        Semi-axis length of the ellipsoid along the major axis.
+    lperp : float, required
+        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.
+
+    """
+
+    def __init__(
+        self,
+        epsilon,
+        lpar,
+        lperp,
+        r_cut,
+        bond_k,
+        bond_r0,
+        angle_k=None,
+        angle_theta0=None,
+    ):
+        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
+        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)
+        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)
+            forces.append(angle)
+        # Gay-Berne Pairs
+        nlist = hoomd.md.nlist.Cell(buffer=0.40)
+        gb = hoomd.md.pair.aniso.GayBerne(nlist=nlist, default_r_cut=self.r_cut)
+        gb.params[("R", "R")] = dict(
+            epsilon=self.epsilon, lperp=self.lperp, lpar=self.lpar
+        )
+        # Add zero pairs
+        for pair in [
+            ("A", "A"),
+            ("B", "B"),
+            ("A", "B"),
+            ("A", "R"),
+            ("B", "R"),
+        ]:
+            gb.params[pair] = dict(epsilon=0.0, lperp=0.0, lpar=0.0)
+        forces.append(gb)
+        return forces

flowermd/library/polymers.py

@@ -4,6 +4,7 @@
 
 import mbuild as mb
 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
@@ -272,3 +273,81 @@ def _build(self, length):
                     chain.add_bond([next_bead, last_bead])
                 last_bead = next_bead
         return chain
+
+
+class EllipsoidChain(Polymer):
+    """Create an ellipsoid polymer chain.
+
+    This is a coarse-grained molecule where each monomer is modeled
+    as an anisotropic bead (i.e. ellipsoid).
+
+    Notes
+    -----
+    In order to form chains of connected ellipsoids, "ghost"
+    particles of types "A" and "B" are used.
+
+    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.
+
+    Parameters
+    ----------
+    lengths : int, required
+        The number of monomer repeat units in the chain.
+    num_mols : int, required
+        The number of chains to create.
+    lpar : float, required
+        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):
+        self.bead_mass = bead_mass
+        self.bead_bond_length = bond_length
+        self.lpar = lpar
+        # get the indices of the particles in a rigid body
+        self.bead_constituents_types = ["A", "A", "B", "B"]
+        super(EllipsoidChain, self).__init__(lengths=lengths, num_mols=num_mols)
+
+    def _build(self, length):
+        # Build up ellipsoid bead
+        bead = mb.Compound(name="ellipsoid")
+        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
+        )
+        tail_mid = mb.Compound(
+            pos=(0, 0, -self.lpar / 2), name="B", 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,
+        )
+        return chain

flowermd/tests/base/test_simulation.py

@@ -9,11 +9,12 @@
 import unyt as u
 
 from flowermd import Simulation
+from flowermd.base import Pack
 from flowermd.library import OPLS_AA_PPS
+from flowermd.library.forcefields import EllipsoidForcefield
+from flowermd.library.polymers import EllipsoidChain
 from flowermd.tests import BaseTest
-from flowermd.utils import (  # get_target_box_number_density,
-    get_target_box_mass_density,
-)
+from flowermd.utils import create_rigid_body, get_target_box_mass_density
 
 
 class TestSimulate(BaseTest):
@@ -294,6 +295,48 @@ def test_pickle_ff(self, benzene_system):
             assert type(i) is type(j)
         os.remove("forcefield.pickle")
 
+    def test_bad_rigid(self, benzene_system):
+        with pytest.raises(ValueError):
+            Simulation.from_system(benzene_system, rigid_constraint="A")
+
+    def test_rigid_sim(self):
+        ellipsoid_chain = EllipsoidChain(
+            lengths=4,
+            num_mols=2,
+            lpar=0.5,
+            bead_mass=100,
+            bond_length=0.01,
+        )
+        system = Pack(
+            molecules=ellipsoid_chain,
+            density=0.1,
+            base_units=dict(),
+            fix_orientation=True,
+        )
+        ellipsoid_ff = EllipsoidForcefield(
+            lpar=0.5,
+            lperp=0.25,
+            epsilon=1.0,
+            r_cut=2.0,
+            bond_k=500,
+            bond_r0=0.01,
+            angle_k=250,
+            angle_theta0=2.2,
+        )
+        rigid_frame, rigid = create_rigid_body(
+            system.hoomd_snapshot,
+            ellipsoid_chain.bead_constituents_types,
+            bead_name="R",
+        )
+        sim = Simulation(
+            initial_state=rigid_frame,
+            forcefield=ellipsoid_ff.hoomd_forces,
+            rigid_constraint=rigid,
+        )
+        sim.run_NVT(n_steps=0, 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)
         sim.save_restart_gsd("restart.gsd")