diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 674aa5e1db..8768e58eff 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -13,6 +13,8 @@ Next release (`#2136 `__). * Add thermodynamically consistent frictional contact forces: ``hoomd.md.pair.friction`` (`#2116 `__). +* ``create_bodies`` method takes optional ``masses`` argument to set masses + (`#2169 `__). *Fixed* diff --git a/hoomd/md/ForceComposite.cc b/hoomd/md/ForceComposite.cc index 092f55f288..a6d700f81f 100644 --- a/hoomd/md/ForceComposite.cc +++ b/hoomd/md/ForceComposite.cc @@ -434,15 +434,10 @@ void ForceComposite::validateRigidBodies() m_particles_added_removed = false; } -void ForceComposite::pyCreateRigidBodies(pybind11::dict charges) +void ForceComposite::pyCreateRigidBodies(pybind11::dict charges, pybind11::dict masses) { - if (pybind11::len(charges) == 0) - { - createRigidBodies(std::unordered_map>()); - return; - } - std::unordered_map> charges_map; + std::unordered_map> masses_map; { ArrayHandle h_body_len(m_body_len, access_location::host, access_mode::read); for (const auto& item : charges) @@ -464,13 +459,33 @@ void ForceComposite::pyCreateRigidBodies(pybind11::dict charges) } charges_map.insert({type, charges_vector}); } + for (const auto& item : masses) + { + const auto type = m_pdata->getTypeByName(item.first.cast()); + if (h_body_len.data[type] == 0) + { + throw std::runtime_error("Masses provided for non-central particle type."); + } + const auto masses_list = item.second.cast(); + if (pybind11::len(masses_list) != h_body_len.data[type]) + { + throw std::runtime_error("Masses provided not consistent with rigid body size."); + } + std::vector masses_vector; + for (auto& mass : masses_list) + { + masses_vector.emplace_back(mass.cast()); + } + masses_map.insert({type, masses_vector}); + } } - createRigidBodies(charges_map); + createRigidBodies(charges_map, masses_map); } void ForceComposite::createRigidBodies( - const std::unordered_map> charges) + const std::unordered_map> charges, + const std::unordered_map> masses) { SnapshotParticleData snap; const BoxDim& global_box = m_pdata->getGlobalBox(); @@ -573,7 +588,10 @@ void ForceComposite::createRigidBodies( snap.charge[constituent_particle_tag] = charges.at(body_type)[current_body_index]; } - + if (!masses.empty()) + { + snap.mass[constituent_particle_tag] = masses.at(body_type)[current_body_index]; + } // Set position and orientation of constituents vec3 body_position(snap.pos[particle_tag]); quat body_orientation(snap.orientation[particle_tag]); diff --git a/hoomd/md/ForceComposite.h b/hoomd/md/ForceComposite.h index f0697486ea..442933a802 100644 --- a/hoomd/md/ForceComposite.h +++ b/hoomd/md/ForceComposite.h @@ -99,11 +99,12 @@ class PYBIND11_EXPORT ForceComposite : public MolecularForceCompute virtual void validateRigidBodies(); //! Create rigid body constituent particles - void pyCreateRigidBodies(pybind11::dict charges); + void pyCreateRigidBodies(pybind11::dict charges, pybind11::dict masses); //! Create rigid body constituent particles virtual void - createRigidBodies(const std::unordered_map> charges); + createRigidBodies(const std::unordered_map> charges, + const std::unordered_map> masses); /// Construct from a Python dictionary void setBody(std::string typ, pybind11::object v) diff --git a/hoomd/md/constrain.py b/hoomd/md/constrain.py index e222837812..53d11d8878 100644 --- a/hoomd/md/constrain.py +++ b/hoomd/md/constrain.py @@ -331,7 +331,7 @@ def __init__(self): self._add_typeparam(body) self.body.default = None - def create_bodies(self, state, charges=None): + def create_bodies(self, state, charges=None, masses=None): r"""Create rigid bodies from central particles in state. Args: @@ -339,6 +339,9 @@ def create_bodies(self, state, charges=None): charges (dict[str, list[float]]): (optional) The charges for each of the constituent particles, defaults to ``None``. If ``None``, all charges are zero. The keys should be the central particles. + masses (dict[str, list[float]]): (optional) The masses for each of + the constituent particles, defaults to ``None``. If ``None``, + all masses are one. The keys should be the central particles. `create_bodies` removes any existing constituent particles and adds new ones based on the body definitions in `body`. It overwrites all existing @@ -347,7 +350,9 @@ def create_bodies(self, state, charges=None): if self._attached: raise RuntimeError("Cannot call create_bodies after running simulation.") super()._attach(state._simulation) - self._cpp_obj.createRigidBodies({} if charges is None else charges) + self._cpp_obj.createRigidBodies( + {} if charges is None else charges, {} if masses is None else masses + ) # Restore previous state self._detach() diff --git a/hoomd/md/pytest/test_rigid.py b/hoomd/md/pytest/test_rigid.py index 94933b0df0..6d69f44d09 100644 --- a/hoomd/md/pytest/test_rigid.py +++ b/hoomd/md/pytest/test_rigid.py @@ -66,7 +66,7 @@ def test_body_setting(valid_body_definition): current_body_definition[key] = valid_body_definition[key] -def check_bodies(snapshot, definition, charges=None): +def check_bodies(snapshot, definition, charges=None, masses=None): """Non-general assumes a snapshot from two_particle_snapshot_factory. This is just to prevent duplication of code from test_create_bodies and @@ -86,6 +86,12 @@ def check_bodies(snapshot, definition, charges=None): assert snapshot.particles.charge[i + 2] == charges[i] assert snapshot.particles.charge[i + 6] == charges[i] + # check masses + if masses is not None: + for i in range(4): + assert snapshot.particles.mass[i + 2] == masses[i] + assert snapshot.particles.mass[i + 6] == masses[i] + particle_one = (snapshot.particles.position[0], snapshot.particles.orientation[0]) particle_two = (snapshot.particles.position[1], snapshot.particles.orientation[1]) @@ -131,11 +137,18 @@ def check_orientation( @skip_rowan +@pytest.mark.parametrize( + "include_charge,include_mass", + [(False, False), (True, False), (False, True), (True, True)], + ids=["default", "optional-change", "optional-mass", "optional-charge-mass"], +) def test_create_bodies( simulation_factory, two_particle_snapshot_factory, lattice_snapshot_factory, valid_body_definition, + include_charge, + include_mass, ): rigid = md.constrain.Rigid() rigid.body["A"] = valid_body_definition @@ -145,11 +158,20 @@ def test_create_bodies( initial_snapshot.particles.types = ["A", "B"] sim = simulation_factory(initial_snapshot) - charges = [1.0, 2.0, 3.0, 4.0] - rigid.create_bodies(sim.state, charges={"A": charges}) + optional_kwargs = {} + if include_charge: + optional_kwargs["charges"] = {"A": [1.0, 2.0, 3.0, 4.0]} + if include_mass: + optional_kwargs["masses"] = {"A": [5.0, 6.0, 7.0, 8.0]} + rigid.create_bodies(sim.state, **optional_kwargs) snapshot = sim.state.get_snapshot() if snapshot.communicator.rank == 0: - check_bodies(snapshot, valid_body_definition, charges) + check_bodies( + snapshot, + valid_body_definition, + optional_kwargs["charges"]["A"] if include_charge else None, + optional_kwargs["masses"]["A"] if include_mass else None, + ) sim.operations.integrator = hoomd.md.Integrator(dt=0.005, rigid=rigid) # Ensure validate bodies passes @@ -262,12 +284,13 @@ def test_running_simulation( sim.seed = 5 charges = [1.0, 2.0, 3.0, 4.0] - rigid.create_bodies(sim.state, charges={"A": charges}) + masses = [5.0, 6.0, 7.0, 8.0] + rigid.create_bodies(sim.state, charges={"A": charges}, masses={"A": masses}) sim.operations += integrator sim.run(5) snapshot = sim.state.get_snapshot() if sim.device.communicator.rank == 0: - check_bodies(snapshot, valid_body_definition, charges) + check_bodies(snapshot, valid_body_definition, charges, masses) autotuned_kernel_parameter_check(instance=rigid, activate=lambda: sim.run(1)) @@ -587,16 +610,7 @@ def test_velocity_constituents_constant_angmom( rigid = md.constrain.Rigid() rigid.body["A"] = body_definition - rigid.create_bodies(sim.state) - - intermed_snapshot = sim.state.get_snapshot() - if intermed_snapshot.communicator.rank == 0: - flags = ( - intermed_snapshot.particles.typeid - == intermed_snapshot.particles.types.index("B") - ) - intermed_snapshot.particles.mass[flags] = mass[flags] - sim.state.set_snapshot(intermed_snapshot) + rigid.create_bodies(sim.state, masses={"A": mass[1:]}) constvol = md.methods.ConstantVolume(filter=hoomd.filter.Rigid()) integrator = md.Integrator( diff --git a/hoomd/mpcd/pytest/test_collide.py b/hoomd/mpcd/pytest/test_collide.py index d76cc452d0..3bc4852f9b 100644 --- a/hoomd/mpcd/pytest/test_collide.py +++ b/hoomd/mpcd/pytest/test_collide.py @@ -225,29 +225,18 @@ def test_rigid_collide( initial_snap.particles.velocity[:] = [velo_rigid] initial_snap.particles.angmom[:] = [angmom_rigid] + # place the mpcd particles on top of constituents + initial_snap.mpcd.N = N_mpcd + initial_snap.mpcd.types = ["C"] + initial_snap.mpcd.position[:] = def_rigid["positions"] + initial_snap.mpcd.velocity[:] = velo_mpcd + sim = simulation_factory(initial_snap) sim.seed = 5 rigid = hoomd.md.constrain.Rigid() rigid.body["A"] = def_rigid - rigid.create_bodies(sim.state) - - intermed_snap = sim.state.get_snapshot() - if intermed_snap.communicator.rank == 0: - # add mass of constituents - flags = ( - intermed_snap.particles.typeid - == intermed_snap.particles.types.index("B") - ) - intermed_snap.particles.mass[flags] = properties_rigid["mass"][flags] - intermed_snap.wrap() - - # place the mpcd particles on top of constituents - intermed_snap.mpcd.N = N_mpcd - intermed_snap.mpcd.types = ["C"] - intermed_snap.mpcd.position[:] = intermed_snap.particles.position[flags] - intermed_snap.mpcd.velocity[:] = velo_mpcd - sim.state.set_snapshot(intermed_snap) + rigid.create_bodies(sim.state, masses={"A": properties_rigid["mass"][1:]}) sim.operations.integrator = hoomd.mpcd.Integrator( dt=0, integrate_rotational_dof=True, rigid=rigid @@ -364,31 +353,18 @@ def test_rigid_collide_free( initial_snap.particles.angmom[:] = [[0, 0, 0, 0], [0, 0, 0, 0]] initial_snap.particles.typeid[:] = [0, 2] + # place the mpcd particles on top of constituents + initial_snap.mpcd.N = N_mpcd + initial_snap.mpcd.types = ["C"] + initial_snap.mpcd.position[:] = def_rigid["positions"] + initial_snap.mpcd.velocity[:] = velo_mpcd + sim = simulation_factory(initial_snap) sim.seed = 5 rigid = hoomd.md.constrain.Rigid() rigid.body["A"] = def_rigid - rigid.create_bodies(sim.state) - - intermed_snap = sim.state.get_snapshot() - if intermed_snap.communicator.rank == 0: - # add mass of constituents - flags = np.logical_or( - intermed_snap.particles.typeid - == intermed_snap.particles.types.index("B"), - intermed_snap.particles.typeid - == intermed_snap.particles.types.index("D"), - ) - intermed_snap.particles.mass[flags] = properties_rigid["mass"][1:] - intermed_snap.wrap() - - # place the mpcd particles on top of constituents - intermed_snap.mpcd.N = N_mpcd - intermed_snap.mpcd.types = ["C"] - intermed_snap.mpcd.position[:] = intermed_snap.particles.position[flags] - intermed_snap.mpcd.velocity[:] = velo_mpcd - sim.state.set_snapshot(intermed_snap) + rigid.create_bodies(sim.state, masses={"A": properties_rigid["mass"][1:]}) sim.operations.integrator = hoomd.mpcd.Integrator( dt=0, integrate_rotational_dof=True, rigid=rigid @@ -538,17 +514,7 @@ def test_rigid_mass_errors( rigid = hoomd.md.constrain.Rigid() rigid.body["A"] = def_rigid - rigid.create_bodies(sim.state) - - # add mass of constituents - intermed_snap = sim.state.get_snapshot() - if intermed_snap.communicator.rank == 0: - flags = intermed_snap.particles.typeid == intermed_snap.particles.types.index( - "B" - ) - intermed_snap.particles.mass[flags] = masses[1:] - intermed_snap.wrap() - sim.state.set_snapshot(intermed_snap) + rigid.create_bodies(sim.state, masses={"A": masses[1:]}) sim.operations.integrator = hoomd.mpcd.Integrator(dt=0, rigid=rigid) sim.operations.integrator.collision_method = (