Revamp Ellipsoid Class, add method to create bond constraints in a system's snapshot.

.github/workflows/pytest.yml

@@ -29,7 +29,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        os: [macos-14, ubuntu-24.04]
+        os: [macos-latest, ubuntu-latest]
         python-version: ['3.10', '3.11', '3.12']
 
     runs-on: ${{ matrix.os }}

codecov.yml

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

flowermd/base/simulation.py

@@ -52,7 +52,10 @@ class Simulation(hoomd.simulation.Simulation):
     thermostat : flowermd.utils.HOOMDThermostats, default
         HOOMDThermostats.MTTK
         The thermostat to use for the simulation.
-
+    constraint : hoomd.md.constrain object
+        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 = constraint
+        self._rigid_constraint = None
+        self._distance_constraint = None
+        if self.constraint:
+            if isinstance(self.constraint, hoomd.md.constrain.Rigid):
+                self._rigid_constraint = self.constraint
+            elif isinstance(self.constraint, hoomd.md.constrain.Distance):
+                self._distance_constraint = self.constraint
+            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,44 @@ 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", "T", "T"]
         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 + (self.bond_L / 2), 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

@@ -11,10 +11,14 @@
 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.library.forcefields import BeadSpring, EllipsoidForcefield
+from flowermd.library.polymers import EllipsoidChain, LJChain
 from flowermd.tests import BaseTest
-from flowermd.utils import create_rigid_body, get_target_box_mass_density
+from flowermd.utils import (
+    create_rigid_ellipsoid_chain,
+    get_target_box_mass_density,
+    set_bond_constraints,
+)
 
 
 class TestSimulate(BaseTest):
@@ -378,47 +382,63 @@ 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):
-        ellipsoid_chain = EllipsoidChain(
-            lengths=4,
-            num_mols=2,
-            lpar=0.5,
-            bead_mass=100,
-            bond_length=0.01,
+    def test_d_constrain_sim(self):
+        chains = LJChain(lengths=10, num_mols=1)
+        system = Pack(molecules=chains, density=0.001, base_units=dict())
+        snap, d = set_bond_constraints(
+            snapshot=system.hoomd_snapshot,
+            bond_types=["A-A"],
+            constraint_values=[1.0],
+        )
+        ff = BeadSpring(
+            beads={"A": dict(epsilon=1.0, sigma=1.0)},
+            angles={"A-A-A": dict(t0=2.2, k=100)},
+            r_cut=2.5,
+        )
+        sim = Simulation(
+            initial_state=snap, forcefield=ff.hoomd_forces, constraint=d
+        )
+        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
+
+    def test_ellipsoid_chain_sim(self):
+        chain = EllipsoidChain(
+            lengths=15, num_mols=15, bead_mass=1, lpar=1, bond_L=0.0
         )
         system = Pack(
-            molecules=ellipsoid_chain,
-            density=0.1,
-            base_units=dict(),
+            molecules=chain,
+            density=0.0005,
+            edge=5,
+            overlap=1,
             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,
+        rigid_snap, rigid = create_rigid_ellipsoid_chain(system.hoomd_snapshot)
+        forces = EllipsoidForcefield(
+            angle_k=25,
             angle_theta0=2.2,
-        )
-        rigid_frame, rigid = create_rigid_body(
-            system.hoomd_snapshot,
-            ellipsoid_chain.bead_constituents_types,
-            bead_name="R",
+            bond_r0=0.0,
+            lpar=1,
+            lperp=0.5,
+            epsilon=1,
+            r_cut=3.0,
         )
         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)
+            initial_state=rigid_snap,
+            constraint=rigid,
+            gsd_file_name="traj.gsd",
+            gsd_write_freq=100,
+            forcefield=forces.hoomd_forces,
+        )
+        assert isinstance(sim._rigid_constraint, hoomd.md.constrain.Rigid)
+        assert sim._distance_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)