diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index be42b1aa..12fbc18d 100644 --- a/.pre-commit-config.yaml +++ b/.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 diff --git a/flowermd/base/simulation.py b/flowermd/base/simulation.py index c6225937..04ba2d44 100644 --- a/flowermd/base/simulation.py +++ b/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. diff --git a/flowermd/base/system.py b/flowermd/base/system.py index 4fbfd75f..6416be53 100644 --- a/flowermd/base/system.py +++ b/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 diff --git a/flowermd/library/__init__.py b/flowermd/library/__init__.py index 7fa65b37..64ea7c0d 100644 --- a/flowermd/library/__init__.py +++ b/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, diff --git a/flowermd/library/forcefields.py b/flowermd/library/forcefields.py index 8f071f25..3721ef8b 100644 --- a/flowermd/library/forcefields.py +++ b/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 diff --git a/flowermd/library/polymers.py b/flowermd/library/polymers.py index f287b12c..2ac9038f 100644 --- a/flowermd/library/polymers.py +++ b/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 diff --git a/flowermd/tests/base/test_simulation.py b/flowermd/tests/base/test_simulation.py index e0b99c51..a756f808 100644 --- a/flowermd/tests/base/test_simulation.py +++ b/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") diff --git a/flowermd/tests/library/test_forcefields.py b/flowermd/tests/library/test_forcefields.py index 3ae9478a..30ae0920 100644 --- a/flowermd/tests/library/test_forcefields.py +++ b/flowermd/tests/library/test_forcefields.py @@ -11,6 +11,7 @@ OPLS_AA_DIMETHYLETHER, OPLS_AA_PPS, BeadSpring, + EllipsoidForcefield, FF_from_file, KremerGrestBeadSpring, TableForcefield, @@ -104,6 +105,31 @@ def test_BeadSpring(self): assert ff.hoomd_forces[3].params[param]["d"] == -1 assert ff.hoomd_forces[3].params[param]["n"] == 1 + def test_ellipsoid_ff(self): + ellipsoid_ff = EllipsoidForcefield( + epsilon=1.0, + lperp=0.5, + lpar=1.0, + r_cut=3, + bond_k=500, + bond_r0=0.1, + angle_k=100, + angle_theta0=1.57, + ) + assert len(ellipsoid_ff.hoomd_forces) == 3 + assert isinstance( + ellipsoid_ff.hoomd_forces[-1], hoomd.md.pair.aniso.GayBerne + ) + assert ("R", "R") in list( + dict(ellipsoid_ff.hoomd_forces[-1].params).keys() + ) + assert ("B", "R") in list( + dict(ellipsoid_ff.hoomd_forces[-1].params).keys() + ) + assert ellipsoid_ff.hoomd_forces[-1].params["A", "A"]["epsilon"] == 0.0 + assert ellipsoid_ff.hoomd_forces[-1].params["A", "A"]["lperp"] == 0.0 + assert ellipsoid_ff.hoomd_forces[-1].params["A", "A"]["lpar"] == 0.0 + class TestTableForcefield: def test_from_txt_file(self): diff --git a/flowermd/tests/library/test_polymers.py b/flowermd/tests/library/test_polymers.py index 20f3b8ed..12b5e9f0 100644 --- a/flowermd/tests/library/test_polymers.py +++ b/flowermd/tests/library/test_polymers.py @@ -5,6 +5,7 @@ PEEK, PEKK, PPS, + EllipsoidChain, LJChain, PEKK_meta, PEKK_para, @@ -106,3 +107,16 @@ def test_lj_chain_sequence_bad_mass(self): def test_copolymer(self): pass + + def test_ellipsoid_chain(self): + ellipsoid_chain = EllipsoidChain( + lengths=4, + num_mols=2, + lpar=0.5, + bead_mass=100, + bond_length=0.01, + ) + assert ellipsoid_chain.n_particles == 32 + assert ellipsoid_chain.molecules[0].mass == 400 + assert ellipsoid_chain.molecules[0].n_particles == 16 + assert ellipsoid_chain.molecules[0].n_bonds == 10 diff --git a/flowermd/tests/utils/test_rigid_body.py b/flowermd/tests/utils/test_rigid_body.py new file mode 100644 index 00000000..08ed20d5 --- /dev/null +++ b/flowermd/tests/utils/test_rigid_body.py @@ -0,0 +1,70 @@ +import numpy as np + +from flowermd.base import Pack +from flowermd.library.polymers import EllipsoidChain +from flowermd.tests import BaseTest +from flowermd.utils import create_rigid_body + + +class TestRigidBody(BaseTest): + def test_ellipsoid_create_rigid_body(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, + ) + + rigid_frame, rigid = create_rigid_body( + system.hoomd_snapshot, + ellipsoid_chain.bead_constituents_types, + bead_name="R", + ) + assert rigid_frame.particles.N == 8 + ellipsoid_chain.n_particles + assert ( + rigid_frame.particles.types + == ["R"] + system.hoomd_snapshot.particles.types + ) + assert rigid_frame.particles.mass[0] == 100 + assert np.all( + np.isclose( + rigid_frame.particles.position[0], + np.mean(system.hoomd_snapshot.particles.position[:4], axis=0), + ) + ) + + points = ( + rigid_frame.particles.position[0] + - system.hoomd_snapshot.particles.position[:4] + ) + + x = points[:, 0] + y = points[:, 1] + z = points[:, 2] + I_xx = np.sum((y**2 + z**2) * system.hoomd_snapshot.particles.mass[:4]) + I_yy = np.sum((x**2 + z**2) * system.hoomd_snapshot.particles.mass[:4]) + I_zz = np.sum((x**2 + y**2) * system.hoomd_snapshot.particles.mass[:4]) + assert np.all( + np.isclose( + rigid_frame.particles.moment_inertia[0], + np.array((I_xx, I_yy, I_zz)), + ) + ) + + assert np.all(rigid_frame.particles.body[:8] == np.arange(8)) + assert np.all(rigid_frame.particles.body[8:12] == [0]) + + assert rigid_frame.bonds.N == 20 + assert rigid_frame.bonds.types == system.hoomd_snapshot.bonds.types + assert rigid_frame.angles.N == 12 + assert rigid_frame.angles.types == system.hoomd_snapshot.angles.types + + assert "R" in list(rigid.body.keys()) + assert rigid.body["R"]["constituent_types"] == ["A", "A", "B", "B"] diff --git a/flowermd/utils/__init__.py b/flowermd/utils/__init__.py index 65e69050..f3e56cb6 100644 --- a/flowermd/utils/__init__.py +++ b/flowermd/utils/__init__.py @@ -1,5 +1,6 @@ from .actions import * from .base_types import HOOMDThermostats +from .rigid_body import create_rigid_body from .utils import ( _calculate_box_length, get_target_box_mass_density, diff --git a/flowermd/utils/rigid_body.py b/flowermd/utils/rigid_body.py new file mode 100644 index 00000000..4484b1e6 --- /dev/null +++ b/flowermd/utils/rigid_body.py @@ -0,0 +1,186 @@ +import gsd +import gsd.hoomd +import hoomd +import numpy as np +from cmeutils.geometry import moit + + +def create_rigid_body( + snapshot, + bead_constituents_types, + bead_name="R", + initial_orientation=[1, 0, 0, 0], +): + """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. + + Examples + -------- + This example demonstrates how to create a rigid body snapshot from a + snapshot of a system of ellipsoids. The ellipsoids are created using + the EllipsoidChain class from the `flowermd.library.polymers`. + The `rigid_frame` snapshot contains all rigid body centers and the + constituent particles along with information about the positions of + center of mass, orientations and moment of inertia. The + `rigid_constrain` object is used by hoomd in the simulation class to + constrain the rigid bodies. + + :: + + from flowermd.library.polymers import EllipsoidChain + from flowermd.library.forcefields import EllipsoidForcefield + from flowermd.base import Pack + from flowermd.base import Simulation + + ellipsoid_chain = EllipsoidChain(lengths=9, num_mols=20, + bead_length=1, bead_mass=100, + bond_length=0.01) + + system = Pack(molecules=ellipsoid_chain, density=0.1, + fix_orientation=True) + + rigid_frame, rigid = create_rigid_body(system.hoomd_snapshot, + ellipsoid_chain.bead_constituents_types) + + + ellipsoid_ff = EllipsoidForcefield(epsilon= 1.0, lperp=0.5 , + lpar=1.0, bead_length=1, + r_cut=3, bond_k=500, + bond_r0=0.1) + + simulation = Simulation(initial_state=rigid_frame, + forcefield=ellipsoid_ff.hoomd_forces, + rigid_constraint=rigid) + + simulation.run_NVT(n_steps=10000, kT=3.5, tau_kt=1) + + """ + # 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 = [initial_orientation] * ( + n_rigid + snapshot.particles.N + ) + rigid_frame.particles.body = np.concatenate( + ( + np.arange(n_rigid), + np.arange(n_rigid).repeat(rigid_const_idx.shape[1]), + ) + ) + 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 = 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": [initial_orientation] * 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) + 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