Add classes for simulating ellipsoid polymer systems.

flowermd/base/simulation.py

@@ -73,6 +73,7 @@ def __init__(
         log_write_freq=1e3,
         log_file_name="sim_data.txt",
         thermostat=HOOMDThermostats.MTTK,
+        rigid_constraint=None,
     ):
         super(Simulation, self).__init__(device, seed)
         self.initial_state = initial_state
@@ -98,9 +99,22 @@ 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)
+        if rigid_constraint:
+            # place constituent particles in the simulation state
+            rigid_constraint.create_bodies(self.state)
         # Add a gsd and thermo props logger to sim operations
         self._add_hoomd_writers()
         self._thermostat = thermostat
@@ -524,7 +538,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)
@@ -1088,6 +1109,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

@@ -729,9 +729,11 @@ def __init__(
         base_units=dict(),
         packing_expand_factor=5,
         edge=0.2,
+        fix_orientation=False,
     ):
         self.packing_expand_factor = packing_expand_factor
         self.edge = edge
+        self.fix_orientation = fix_orientation
         super(Pack, self).__init__(
             molecules=molecules,
             density=density,
@@ -746,6 +748,7 @@ def _build_system(self):
             box=list(self._target_box * self.packing_expand_factor),
             overlap=0.2,
             edge=self.edge,
+            fix_orientation=self.fix_orientation,
         )
         return system
 

flowermd/library/__init__.py

@@ -15,6 +15,7 @@
     PEEK,
     PEKK,
     PPS,
+    EllipsoidChain,
     LJChain,
     PEKK_meta,
     PEKK_para,

flowermd/library/forcefields.py

@@ -481,3 +481,61 @@ 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."""
+
+    def __init__(
+        self,
+        epsilon,
+        lperp,
+        lpar,
+        bead_length,
+        r_cut,
+        bond_k,
+        bond_r0,
+        angle_k=None,
+        angle_theta0=None,
+    ):
+        self.epsilon = epsilon
+        self.lperp = lperp
+        self.lpar = lpar
+        self.bead_length = bead_length
+        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.bead_length / 2)
+        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

@@ -3,6 +3,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
@@ -261,3 +262,66 @@ 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.
+
+    Parameters
+    ----------
+    lengths : int, required
+        The number of monomer repeat units in the chain.
+    num_mols : int, required
+        The number of chains to create.
+    bead_length : float, required
+        The length of the ellipsoid bead.
+    bead_mass : float, required
+        The mass of the ellipsoid bead.
+    bond_length : float, required
+        The bond length between connected beads.
+
+    """
+
+    def __init__(self, lengths, num_mols, bead_length, bead_mass, bond_length):
+        self.bead_mass = bead_mass
+        self.bead_bond_length = bond_length
+        self.bead_length = bead_length
+        super(EllipsoidChain, self).__init__(lengths=lengths, num_mols=num_mols)
+        # get the indices of the particles in a rigid body
+        self.bead_constituents_types = ["A", "A", "B", "B"]
+
+    def _build(self, length):
+        # Build up ellipsoid bead
+        bead = mb.Compound(name="ellipsoid")
+        head = mb.Compound(
+            pos=(self.bead_length / 2, 0, 0), name="A", mass=self.bead_mass / 4
+        )
+        tail = mb.Compound(
+            pos=(-self.bead_length / 2, 0, 0), name="A", mass=self.bead_mass / 4
+        )
+        head_mid = mb.Compound(
+            pos=(self.bead_length / 4, 0, 0), name="B", mass=self.bead_mass / 4
+        )
+        tail_mid = mb.Compound(
+            pos=(-self.bead_length / 4, 0, 0), name="B", mass=self.bead_mass / 4
+        )
+        bead.add([head, tail, head_mid, tail_mid])
+
+        chain = mbPolymer()
+        chain.add_monomer(
+            bead,
+            indices=[0, 1],
+            orientation=[[1, 0, 0], [-1, 0, 0]],
+            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.bead_length / 2 - 0.1,
+            dmax=self.bead_length / 2 + self.bead_bond_length + 0.1,
+        )
+        return chain

flowermd/utils/rigid_body.py

@@ -0,0 +1,134 @@
+import gsd
+import gsd.hoomd
+import hoomd
+import numpy as np
+from cmeutils.geometry import moit
+
+
+def _get_com_mass_pos_moi(snapshot, rigid_const_idx):
+    com_mass = []
+    com_position = []
+    com_moi = []
+    for idx in rigid_const_idx:
+        constituents_mass = np.array(snapshot.particles.mass)
+        constituents_pos = np.array(snapshot.particles.position)
+        total_mass = np.sum(constituents_mass[idx])
+        com_mass.append(total_mass)
+        com = (
+            np.sum(
+                constituents_pos[idx] * constituents_mass[idx, np.newaxis],
+                axis=0,
+            )
+            / total_mass
+        )
+        com_position.append(com)
+        com_moi.append(
+            moit(
+                points=constituents_pos[idx],
+                masses=constituents_mass[idx],
+                center=com,
+            )
+        )
+    return com_mass, com_position, com_moi
+
+
+def create_rigid_body(snapshot, bead_constituents_types, bead_name="R"):
+    """Create rigid bodies from a snapshot.
+
+    Parameters
+    ----------
+    snapshot : gsd.hoomd.Snapshot; required
+        The snapshot of the system.
+    bead_constituents_types : list of particle types; required
+        The list of particle types that make up a rigid body.
+
+    Returns
+    -------
+    rigid_frame : gsd.hoomd.Frame
+        The snapshot of the rigid bodies.
+    rigid_constrain : hoomd.md.constrain.Rigid
+        The rigid body constrain object.
+    """
+    # find typeid sequence of the constituent particles types in a rigid bead
+    p_types = np.array(snapshot.particles.types)
+    constituent_type_ids = np.where(
+        p_types[:, None] == bead_constituents_types
+    )[0]
+
+    # find indices that matches the constituent particle types
+    typeids = snapshot.particles.typeid
+    bead_len = len(bead_constituents_types)
+    matches = np.where((typeids.reshape(-1, bead_len) == constituent_type_ids))
+    if len(matches[0]) == 0:
+        raise ValueError(
+            "No matches found between particle types in the system"
+            " and bead constituents particles"
+        )
+    rigid_const_idx = (matches[0] * bead_len + matches[1]).reshape(-1, bead_len)
+
+    n_rigid = rigid_const_idx.shape[0]  # number of rigid bodies
+
+    # calculate center of mass and its position for each rigid body
+    com_mass, com_position, com_moi = _get_com_mass_pos_moi(
+        snapshot, rigid_const_idx
+    )
+
+    rigid_frame = gsd.hoomd.Frame()
+    rigid_frame.particles.types = [bead_name] + snapshot.particles.types
+    rigid_frame.particles.N = n_rigid + snapshot.particles.N
+    rigid_frame.particles.typeid = np.concatenate(
+        (([0] * n_rigid), snapshot.particles.typeid + 1)
+    )
+    rigid_frame.particles.mass = np.concatenate(
+        (com_mass, snapshot.particles.mass)
+    )
+    rigid_frame.particles.position = np.concatenate(
+        (com_position, snapshot.particles.position)
+    )
+    rigid_frame.particles.moment_inertia = np.concatenate(
+        (com_moi, np.zeros((snapshot.particles.N, 3)))
+    )
+    rigid_frame.particles.orientation = [(1.0, 0.0, 0.0, 0.0)] * (
+        n_rigid + snapshot.particles.N
+    )
+    rigid_frame.configuration.box = snapshot.configuration.box
+
+    # set up bonds
+    if snapshot.bonds.N > 0:
+        rigid_frame.bonds.N = snapshot.bonds.N
+        rigid_frame.bonds.types = snapshot.bonds.types
+        rigid_frame.bonds.typeid = snapshot.bonds.typeid
+        rigid_frame.bonds.group = [
+            list(np.add(g, n_rigid)) for g in snapshot.bonds.group
+        ]
+    # set up angles
+    if snapshot.angles.N > 0:
+        rigid_frame.angles.N = not snapshot.angles.N
+        rigid_frame.angles.types = snapshot.angles.types
+        rigid_frame.angles.typeid = snapshot.angles.typeid
+        rigid_frame.angles.group = [
+            list(np.add(g, n_rigid)) for g in snapshot.angles.group
+        ]
+
+    # set up dihedrals
+    if snapshot.dihedrals.N > 0:
+        rigid_frame.dihedrals.N = snapshot.dihedrals.N
+        rigid_frame.dihedrals.types = snapshot.dihedrals.types
+        rigid_frame.dihedrals.typeid = snapshot.dihedrals.typeid
+        rigid_frame.dihedrals.group = [
+            list(np.add(g, n_rigid)) for g in snapshot.dihedrals.group
+        ]
+
+    # find local coordinates of the particles in the first rigid body
+    # only need to find the local coordinates for the first rigid body
+    local_coords = (
+        snapshot.particles.position[rigid_const_idx[0]] - com_position[0]
+    )
+
+    rigid_constrain = hoomd.md.constrain.Rigid()
+    rigid_constrain.body["R"] = {
+        "constituent_types": bead_constituents_types,
+        "positions": local_coords,
+        "orientations": [(1.0, 0.0, 0.0, 0.0)] * len(local_coords),
+    }
+    return rigid_frame, rigid_constrain