Rigid body simulation utils

[chrisjonesBSU]

This is still a WIP, but the goal is to port over most of the work Jenny did in the reproducibility project to prep a system to perform rigid body simulations.

I designed it to work with mBuild's built-in rigid body functionality (and it works well so far). Maybe we can change/update the functions to not require using mBuild.

The process might be a little clunky, but it essentially follows the process used in the reproducibility study. It works as follows:

1. Create mBuild system, and label rigid bodies
2. Pass in the mBuild compound into create_rigid_snapshot
3. Pass in the output of create_rigid_snapshot into create_hoomd_forcefield using the init_snap parameter
4. Pass in the snapshot returned by create_hoomd_forcefield into update_rigid_snapshot
5. Initialize the Hoomd simulation from the output of update_rigid_snapshot

Here is an example script to simulate rigid Benzene:

import mbuild as mb
from mbuild.utils.io import get_fn
from mbuild.formats.hoomd_forcefield import create_hoomd_forcefield
import foyer
import hoomd
from cmeutils.gsd_utils import create_rigid_snapshot, update_rigid_snapshot

benzene = mb.load(get_fn('benzene.mol2'))
benzene.name = "Benzene"
filled = mb.fill_box(benzene, n_compounds=20, box=[0,0,0, 4, 4, 4])
filled.label_rigid_bodies(discrete_bodies='Benzene')

init_rigid_snap = create_rigid_snapshot(filled)
opls = foyer.Forcefield(name="oplsaa")
typed_system = opls.apply(filled)

snapshot, forcefield, refs = create_hoomd_forcefield(
    typed_system,
    auto_scale=True,
    init_snap=init_rigid_snap,
    r_cut=2.5
)

snapshot, rigid = update_rigid_snapshot(
    snapshot=snapshot, mb_compound=filled
)

# Now we're ready for a Hoomd rigid body simulation:

for force in forcefield:
    if isinstance(force, hoomd.md.pair.LJ):
        for t in snapshot.particles.types:
            force.params[("R", t)] = dict(epsilon=0, sigma=0)
            force.r_cut[("R", t)] = 0

_all = hoomd.filter.Rigid(("center", "free"))

device = hoomd.device.auto_select()
sim = hoomd.Simulation(device=device, seed=42)
sim.create_state_from_snapshot(snapshot)
gsd_writer = hoomd.write.GSD(
    trigger=hoomd.trigger.Periodic(int(2e2)),
    filename="rigid.gsd",
    mode="wb"
)
sim.operations.writers.append(gsd_writer)

integrator = hoomd.md.Integrator(dt=0.001, integrate_rotational_dof=True)
integrator.rigid = rigid
integrator.forces = forcefield
integrator_method = hoomd.md.methods.NVT(filter=_all, kT=2.0, tau=0.01)
integrator.methods = [integrator_method]
sim.operations.integrator = integrator

sim.run(2e4)

[codecov]

Codecov Report

Merging #37 (cb08a1b) into master (f790902) will increase coverage by 0.11%.
The diff coverage is 100.00%.

Additional details and impacted files

@@            Coverage Diff             @@
##           master      #37      +/-   ##
==========================================
+ Coverage   98.96%   99.07%   +0.11%     
==========================================
  Files           7        7              
  Lines         385      431      +46     
==========================================
+ Hits          381      427      +46     
  Misses          4        4              

Impacted Files	Coverage Δ
cmeutils/geometry.py	97.50% <ø> (ø)
cmeutils/gsd_utils.py	98.44% <100.00%> (+0.85%)	⬆️

[chrisjonesBSU]

EDIT: I moved this part inside the update_rigid_snapshot function, and updated the workflow in the example above.

I think we could include the creation of the rigid object into the update_rigid_snapshot function and return rigid along with the updated snapshot. What do you think @jennyfothergill ?

Specifically, this bit:

rigid = hoomd.md.constrain.Rigid()
inds = mol_inds[0]
r_pos = snapshot.particles.position[0]
c_pos = snapshot.particles.position[inds]
c_pos -= r_pos
c_pos = [tuple(i) for i in c_pos]
c_types = [
    snapshot.particles.types[i] for i in snapshot.particles.typeid[inds]
]
c_orient = [tuple(i) for i in snapshot.particles.orientation[inds]]

c_charge = [i for i in snapshot.particles.charge[inds]]
c_diam = [i for i in snapshot.particles.diameter[inds]]
rigid.body["R"] = {
    "constituent_types": c_types,
    "positions": c_pos,
    "charges": c_charge,
    "orientations": c_orient,
    "diameters": c_diam,
}