update tests

flowermd/library/polymers.py

@@ -308,15 +308,17 @@ def __init__(self, lengths, num_mols, lpar, bead_mass, bond_L=0.1):
         self.lpar = lpar
         self.bond_L = bond_L
         # get the indices of the particles in a rigid body
-        self.bead_constituents_types = ["X", "A"]
+        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=(self.lpar, 0, 0), name="A", mass=self.bead_mass / 4
+            pos=(self.lpar + (self.bond_L / 2), 0, 0),
+            name="A",
+            mass=self.bead_mass / 4,
         )
         tether_head = mb.Compound(
             pos=(self.lpar, 0, 0), name="T", mass=self.bead_mass / 4

flowermd/tests/utils/test_constraints.py

@@ -2,50 +2,63 @@
 import pytest
 
 from flowermd.base import Pack
+from flowermd.library import LJChain
 from flowermd.library.polymers import EllipsoidChain
 from flowermd.tests import BaseTest
 from flowermd.utils import create_rigid_ellipsoid_chain, set_bond_constraints
 
 
 class TestBondConstraint(BaseTest):
-    def test_ellipsoid_fixed_bonds(self):
-        ellipsoid_chain = EllipsoidChain(
-            lengths=4,
-            num_mols=2,
-            lpar=1,
-            bead_mass=50,
-        )
-        system = Pack(
-            molecules=ellipsoid_chain,
-            density=0.1,
-            base_units=dict(),
-        )
+    def test_single_fixed_bonds(self):
+        chains = LJChain(lengths=10, num_mols=1)
+        system = Pack(molecules=chains, density=0.001)
         snap, d = set_bond_constraints(
-            system.hoomd_snapshot, constraint_values=[1.0], bond_types=["_C-_H"]
+            snapshot=system.hoomd_snapshot,
+            bond_types=["A-A"],
+            constraint_values=[1.0],
         )
-        assert snap.constraints.N == (4 * 2 * 2) - 2
+        assert snap.constraints.N == 9
         for group in snap.bonds.group:
             assert group in snap.constraints.group
         assert d.tolerance == 1e-5
         assert all([val == 1.0 for val in snap.constraints.value])
 
-    def test_ellipsoid_fixed_bonds_bad_val(self):
-        ellipsoid_chain = EllipsoidChain(
-            lengths=4,
-            num_mols=2,
-            lpar=1,
-            bead_mass=50,
+    def test_multiple_fixed_bonds(self):
+        chains = LJChain(
+            lengths=15,
+            num_mols=1,
+            bead_sequence=["A", "B", "B"],
+            bead_mass={"A": 1.0, "B": 0.7},
+            bond_lengths={"A-B": 1.0, "B-B": 0.80},
         )
-        system = Pack(
-            molecules=ellipsoid_chain,
-            density=0.1,
-            base_units=dict(),
+        system = Pack(molecules=chains, density=0.0001)
+        snap, d = set_bond_constraints(
+            snapshot=system.hoomd_snapshot,
+            bond_types=["A-B", "B-B"],
+            constraint_values=[1.0, 0.8],
         )
+        assert snap.constraints.N == 44
+        for group in snap.bonds.group:
+            assert group in snap.constraints.group
+
+        ab_index = snap.bonds.types.index("A-B")
+        bb_index = snap.bonds.types.index("B-B")
+        for group, val in zip(snap.constraints.group, snap.constraints.value):
+            group_index = snap.bonds.group.index(group)
+            bond_id = snap.bonds.typeid[group_index]
+            if bond_id == bb_index:
+                assert np.allclose(val, 0.80)
+            elif bond_id == ab_index:
+                assert np.allclose(val, 1.0)
+
+    def test_ellipsoid_fixed_bonds_bad_val(self):
+        chains = LJChain(lengths=10, num_mols=1)
+        system = Pack(molecules=chains, density=0.001)
         with pytest.raises(ValueError):
             set_bond_constraints(
                 system.hoomd_snapshot,
                 constraint_values=[2.0],
-                bond_types=["_C-_H"],
+                bond_types=["A-A"],
             )
 
 

flowermd/utils/constraints.py

@@ -2,7 +2,8 @@
 import gsd.hoomd
 import hoomd
 import numpy as np
-from cmeutils.geometry import moit
+
+from flowermd.internal import check_return_iterable
 
 
 def set_bond_constraints(
@@ -14,9 +15,9 @@ def set_bond_constraints(
     ----------
     snapshot : gsd.hoomd.Frame, required
         Snapshot of complete topology that will have constraints added.
-    bond_type : list of str, required
+    bond_type : str or list of str, required
         The bond type to add constraints for. Must match snapshot.bonds.types.
-    constrain_value : list of float, required
+    constrain_value : float or list of 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
@@ -69,6 +70,8 @@ def set_bond_constraints(
             sim.flush_writers()
 
     """
+    bond_types = check_return_iterable(bond_types)
+    constraint_values = check_return_iterable(constraint_values)
     constraint_values_list = []
     constraint_groups_list = []
     for b_type, val in zip(bond_types, constraint_values):
@@ -123,7 +126,7 @@ def create_rigid_ellipsoid_chain(snapshot):
         The rigid body constrain object.
 
     """
-    bead_len = 4
+    bead_len = 4  # Number of particles belonging to 1 rigid body
     typeids = snapshot.particles.typeid.reshape(-1, bead_len)
     matches = np.where((typeids == typeids))
     rigid_const_idx = (matches[0] * bead_len + matches[1]).reshape(-1, bead_len)
@@ -136,10 +139,9 @@ def create_rigid_ellipsoid_chain(snapshot):
     for idx in rigid_const_idx:
         mass = np.sum(np.array(snapshot.particles.mass)[idx])
         pos = snapshot.particles.position[idx][0]
-        moi = [0, 2, 2]
         rigid_masses.append(mass)
         rigid_pos.append(pos)
-        rigid_moi.append(moi)
+        rigid_moi.append([0, 2, 2])
 
     rigid_frame = gsd.hoomd.Frame()
     rigid_frame.particles.types = ["R"] + snapshot.particles.types
@@ -183,13 +185,6 @@ def create_rigid_ellipsoid_chain(snapshot):
         rigid_frame.angles.group = [
             list(np.add(g, n_rigid)) for g in snapshot.angles.group
         ]
-    # set up constraints
-    if snapshot.constraints.N > 0:
-        rigid_frame.constraints.N = snapshot.constraints.N
-        rigid_frame.constraints.value = snapshot.constraints.value
-        rigid_frame.constraints.group = [
-            list(np.add(g, n_rigid)) for g in snapshot.constraints.group
-        ]
 
     # find local coordinates of the particles in the first rigid body
     # only need to find the local coordinates for the first rigid body
@@ -204,29 +199,3 @@ def create_rigid_ellipsoid_chain(snapshot):
         "orientations": [[1, 0, 0, 0]] * len(local_coords),
     }
     return rigid_frame, rigid_constrain
-
-
-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)[idx][0]
-        constituents_pos = np.array(snapshot.particles.position)[idx]
-        total_mass = np.sum(constituents_mass)
-        com_mass.append(total_mass)
-        com = (
-            np.sum(
-                constituents_pos * constituents_mass,
-                axis=0,
-            )
-            / total_mass
-        )
-        com_position.append(com)
-        moi = moit(
-            points=constituents_pos,
-            masses=constituents_mass,
-            center=com,
-        )
-        com_moi.append(moi)
-    return com_mass, com_position, com_moi