diff --git a/flowermd/library/polymers.py b/flowermd/library/polymers.py index bb544c9d..d9952038 100644 --- a/flowermd/library/polymers.py +++ b/flowermd/library/polymers.py @@ -308,7 +308,7 @@ 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): @@ -316,7 +316,9 @@ def _build(self, length): 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 diff --git a/flowermd/tests/utils/test_constraints.py b/flowermd/tests/utils/test_constraints.py index 9eb72611..33048a4a 100644 --- a/flowermd/tests/utils/test_constraints.py +++ b/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"], ) diff --git a/flowermd/utils/constraints.py b/flowermd/utils/constraints.py index cb33ec76..83ed2813 100644 --- a/flowermd/utils/constraints.py +++ b/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