Add simple finegraining for UA systems

grits/finegrain.py

@@ -2,11 +2,89 @@
 __all__ = ["backmap"]
 
 import itertools as it
+import tempfile
 from collections import defaultdict
 
+import gsd.hoomd
+import numpy as np
+from ele import element_from_atomic_number
 from mbuild import Compound, Particle, load
+from openbabel import openbabel, pybel
 
-from grits.utils import align, get_hydrogen, get_index
+from grits.utils import (
+    align,
+    comp_from_snapshot,
+    get_hydrogen,
+    get_index,
+    snap_molecules,
+)
+
+
+def _add_hydrogens(snap, conversion_dict=None, scale=1.0):
+    """Add hydrogens to a united atom gsd snapshot.
+
+    Parameters
+    ----------
+    snap : gsd.hoomd.Snapshot
+        The united-atom gsd snapshot
+    conversion_dict : dict, default None
+        Dictionary to map particle types to their element.
+    scale : float, default 1.0
+        Factor by which to scale length values.
+
+    Returns
+    -------
+    gsd.hoomd.Snapshot
+    """
+    # Use the conversion dictionary to map particle type to element symbol
+    if conversion_dict is not None:
+        snap.particles.types = [
+            conversion_dict[i].symbol for i in snap.particles.types
+        ]
+    # Break apart the snapshot into separate molecules
+    molecules = snap_molecules(snap)
+    mol_inds = []
+    for i in range(max(molecules) + 1):
+        mol_inds.append(np.where(molecules == i)[0])
+
+    types = []
+    coords = []
+    bonds = []
+    # Read each molecule into a compound
+    # Then convert to pybel to add hydrogens
+    for inds in mol_inds:
+        l = len(inds)
+        compound = comp_from_snapshot(
+            snap, inds, scale=scale, shift_coords=False
+        )
+        mol = compound.to_pybel()
+        mol.OBMol.PerceiveBondOrders()
+
+        # Add hydrogens
+        # This is a goofy work around necessary for the aromaticity
+        # to be set correctly.
+        with tempfile.NamedTemporaryFile() as f:
+            mol.write(format="mol2", filename=f.name, overwrite=True)
+            mol = list(pybel.readfile("mol2", f.name))[0]
+
+        mol.addh()
+        for atom in mol.atoms:
+            coords.append(np.array(atom.coords) / 10)
+            types.append(element_from_atomic_number(atom.atomicnum).symbol)
+        for bond in openbabel.OBMolBondIter(mol.OBMol):
+            bonds.append([bond.GetBeginAtomIdx() - 1, bond.GetEndAtomIdx() - 1])
+    alltypes = list(set(types))
+    typeids = [alltypes.index(i) for i in types]
+    new_snap = gsd.hoomd.Snapshot()
+    new_snap.configuration.box = snap.configuration.box
+    new_snap.particles.N = len(types)
+    new_snap.particles.types = alltypes
+    new_snap.particles.typeid = typeids
+    new_snap.particles.position = np.array(coords)
+    new_snap.bonds.N = len(bonds)
+    new_snap.bonds.group = np.array(bonds)
+    new_snap.validate()
+    return new_snap
 
 
 def backmap(cg_compound):

grits/utils.py

@@ -57,8 +57,11 @@ def comp_from_snapshot(snapshot, indices, length_scale=1.0, mass_scale=1.0):
         lengths=box[:3] * length_scale, tilt_factors=box[3:]
     )
 
-    # GSD and HOOMD snapshots center their boxes on the origin (0,0,0)
-    shift = np.array(comp.box.lengths) / 2
+    if shift_coords:
+        # GSD and HOOMD snapshots center their boxes on the origin (0,0,0)
+        shift = np.array(comp.box.lengths) / 2
+    else:
+        shift = np.zeros(3)
     particle_dict = {}
     # Add particles
     for i in range(n_atoms):