Add some topology slicing methods.

gmso/core/topology.py

@@ -1256,6 +1256,21 @@
 
         return index
 
+    def calculate_charges(self, method="am1bcc", slice_by="Molecule"):
+        print("Not implemented yet. We are working on it.")
+        # Get all of the moleucles types in the topology
+        molecules = set()
+        for site in self.sites:
+            molecules.add(site.molecule.name)
+
+        # Figure out how to handle the tags
+        for mol in molecules:
+            mol_slice = slice_topology_by_molecule(topology=self, molecule_tag=0)
+            openff_mol = to_openff_molecule(mol_slice)
+            openff_mol.assign_partial_charges(method=method)
+            for atom in openff_mol.atoms:
+                print(atom.partial_charge)
+
     def write_forcefield(self, filename, overwrite=False):
         """Save an xml file for all parameters found in the topology.
 
@@ -1508,7 +1523,7 @@
                 if getattr(site, key) == value:
                     yield site
 
-    def iter_sites_by_residue(self, residue_tag):
+    def iter_sites_by_residue(self, residue_tag, residue_number=None):
         """Iterate through this topology's sites which contain this specific residue name.
 
         See Also
@@ -1518,12 +1533,29 @@
         """
         if isinstance(residue_tag, str):
             for site in self._sites:
-                if site.residue and getattr(site, "residue").name == residue_tag:
-                    yield site
+                if residue_tag and residue_number is None:
+                    if site.molecule and getattr(site, "residue").name == residue_tag:
+                        yield site
+                elif residue_number and residue_tag is None:
+                    if (
+                        site.molecule
+                        and getattr(site, "residue").number == residue_number
+                    ):
+                        yield site
+                else:
+                    if all(
+                        [
+                            site.molecule
+                            and getattr(site, "residue").name == residue_tag,
+                            site.molecule
+                            and getattr(site, "residue").number == residue_number,
+                        ]
+                    ):
+                        yield site
         else:
             return self.iter_sites("residue", residue_tag)
 
-    def iter_sites_by_molecule(self, molecule_tag):
+    def iter_sites_by_molecule(self, molecule_tag, molecule_number=None):
         """Iterate through this topology's sites which contain this specific molecule name.
 
         See Also
@@ -1533,8 +1565,25 @@
         """
         if isinstance(molecule_tag, str):
             for site in self._sites:
-                if site.molecule and getattr(site, "molecule").name == molecule_tag:
-                    yield site
+                if molecule_tag and molecule_number is None:
+                    if site.molecule and getattr(site, "molecule").name == molecule_tag:
+                        yield site
+                elif molecule_number and molecule_tag is None:
+                    if (
+                        site.molecule
+                        and getattr(site, "molecule").number == molecule_number
+                    ):
+                        yield site
+                else:
+                    if all(
+                        [
+                            site.molecule
+                            and getattr(site, "molecule").name == molecule_tag,
+                            site.molecule
+                            and getattr(site, "molecule").number == molecule_number,
+                        ]
+                    ):
+                        yield site
         else:
             return self.iter_sites("molecule", molecule_tag)
 

gmso/external/convert_openmm.py

@@ -4,6 +4,7 @@
 """Convert to and from an OpenMM Topology or System object."""
 
 import unyt as u
+from openff.toolkit import Molecule as openFF_Molecule
 
 from gmso.utils.io import has_openmm, has_openmm_unit, import_
 
@@ -13,6 +14,49 @@
     from openmm.app import *
 
 
+def to_openff_molecule(topology):
+    """Converts a GMSO topology into an OpenFF Molecule."""
+    mol = openFF_Molecule()
+    # Bond order not available at the site level
+    # But, it can be set during openFF's add_atom'
+    # Build up a list of sites first to track their bond orders later
+    sites = []
+    for site in topology.sites:
+        sites.append(site)
+
+    # Index-to-index mapping with sites list
+    is_aromatic = [False for i in sites]
+    bond_pairs = []
+    bond_orders = []
+
+    for bond in topology.bonds:
+        bond_pair = []
+        bond_orders.append(bond.bond_order)
+        for site in bond.connection_members:
+            site_index = sites.index(site)
+            if bond.bond_order == 1.5:
+                is_aromatic[site_index] = True
+            bond_pair.append(site_index)
+        bond_pairs.append(bond_pair)
+    # Now that we have parsed all sites for aromaticity, add each to Molecule
+    for site, aromatic in zip(sites, is_aromatic):
+        mol.add_atom(
+            atomic_number=site.element.atomic_number,
+            is_aromatic=aromatic,
+            formal_charge=0,
+        )
+
+    for bond, border in zip(bond_pairs, bond_orders):
+        mol.add_bond(
+            atom1=bond[0],
+            atom2=bond[1],
+            bond_order=int(border),
+            fractional_bond_order=border,
+            is_aromatic=True if border == 1.5 else False,
+        )
+    return mol
+
+
 def to_openmm(topology, openmm_object="topology"):
     """Convert an untyped topology object to an untyped OpenMM modeller or topology.
 

gmso/tests/test_topology.py

@@ -886,6 +886,17 @@ def test_iter_sites_by_molecule(self, labeled_top):
             for site in labeled_top.iter_sites_by_molecule(molecule_name):
                 assert site.molecule.name == molecule_name
 
+    def test_iter_sites_by_molecule_tag_and_number(self, labeled_top):
+        molecules = labeled_top.unique_site_labels("molecule", name_only=False)
+        for molecule in molecules:
+            for site in labeled_top.iter_sites_by_molecule(molecule):
+                assert site.residue == molecule
+
+        molecule_names = labeled_top.unique_site_labels("molecule", name_only=True)
+        for molecule_name in molecule_names:
+            for site in labeled_top.iter_sites_by_molecule(molecule_name):
+                assert site.molecule.name == molecule_name
+
     @pytest.mark.parametrize(
         "connections",
         ["bonds", "angles", "dihedrals", "impropers"],

gmso/tests/test_utils.py

@@ -1,3 +1,4 @@
+import mbuild as mb
 import numpy as np
 import pytest
 import unyt as u
@@ -7,6 +8,10 @@
 from gmso.utils.geometry import moment_of_inertia
 from gmso.utils.io import run_from_ipython
 from gmso.utils.misc import unyt_to_hashable
+from gmso.utils.slicing import (
+    slice_topology_by_molecule,
+    slice_topology_by_residue,
+)
 from gmso.utils.sorting import sort_connection_members, sort_connection_strings
 
 
@@ -81,3 +86,37 @@ def test_moment_of_inertia():
         masses=np.array([1.0 for i in xyz]),
     )
     assert np.array_equal(tensor, np.array([1, 1, 2]))
+
+
+def test_slice_by_molecule():
+    benzene = mb.load("c1ccccc1", smiles=True)
+    benzene.name = "Benzene"
+    ethane = mb.load("CC", smiles=True)
+    ethane.name = "Ethane"
+
+    system = mb.fill_box(compound=[benzene, ethane], n_compounds=[2, 2], box=[2, 2, 2])
+    topology = system.to_gmso()
+    topology.identify_connections()
+
+    single_benzene_top = slice_topology_by_molecule(topology, "Benzene", 0)
+    assert single_benzene_top.n_sites == 12
+
+    all_benzene_top = slice_topology_by_molecule(topology, "Benzene")
+    assert all_benzene_top.n_sites == 24
+
+
+def test_slice_by_residue():
+    benzene = mb.load("c1ccccc1", smiles=True)
+    benzene.name = "Benzene"
+    ethane = mb.load("CC", smiles=True)
+    ethane.name = "Ethane"
+
+    system = mb.fill_box(compound=[benzene, ethane], n_compounds=[2, 2], box=[2, 2, 2])
+    topology = system.to_gmso()
+    topology.identify_connections()
+
+    single_ethane_top = slice_topology_by_residue(topology, "Ethane", 0)
+    assert single_ethane_top.n_sites == 8
+
+    all_ethane_top = slice_topology_by_residue(topology, "Ethane")
+    assert all_ethane_top.n_sites == 16

gmso/utils/charges.py

@@ -0,0 +1,21 @@
+"""Calculate charges for a gmso Topology."""
+
+from gmso.external.convert_openmm import to_openff_molecule
+from gmso.utils.slicing import slice_topology_by_molecule
+
+
+def calculate_charges(topology, method="am1bcc", slice_by="Molecule"):
+    # Get all of the moleucles types in the topology
+    charges_dict = dict()
+
+    molecules = set()
+    for site in topology.sites:
+        molecules.add(site.molecule.name)
+
+    # Figure out how to handle the tags
+    for mol in molecules:
+        mol_slice = slice_topology_by_molecule(topology=topology, molecule_tag=0)
+        openff_mol = to_openff_molecule(mol_slice)
+        openff_mol.assign_partial_charges(partial_charge_method=method)
+        for atom in openff_mol.atoms:
+            print(atom.partial_charge)

gmso/utils/slicing.py

@@ -0,0 +1,81 @@
+import gmso
+
+
+def slice_topology_by_molecule(topology, molecule_tag, molecule_number=None):
+    """Create a Topology that contains a subset of molecules from another Topology.
+
+    Parameters
+    ----------
+    topology : gmso.core.topology.Topology
+        The gmso Topology to perform the slice on
+    molecule_tag : str
+        The name of the gmso.abstract_site.Molecule object to include in the slice
+    molecule_number : int, default None
+        If given, only include a single molecule's sites
+        If None, then all sites in every molecule matching `molecule_tag` are included
+        in the sliced topology.
+
+    Returns
+    -------
+    gmso.core.topology.Topology
+        A new Topology instance containing only sites and connections from matching molecules.
+    """
+    sites = [
+        s
+        for s in topology.iter_sites_by_molecule(
+            molecule_tag=molecule_tag, molecule_number=molecule_number
+        )
+    ]
+    return slice_by_sites(topology=topology, sites=sites)
+
+
+def slice_topology_by_residue(topology, residue_tag, residue_number=None):
+    """Create a Topology that contains a subset of residues from another Topology.
+
+    Parameters
+    ----------
+    topology : gmso.core.topology.Topology
+        The gmso Topology to perform the slice on.
+    residue_tag : str
+        The name of the gmso.abstract_site.Residue object to include in the slice.
+    residue_number : int, default None
+        If given, only include a single residue's sites
+        If None, then all sites in every residue matching `residue_tag` are included
+        in the sliced topology.
+
+    Returns
+    -------
+    gmso.core.topology.Topology
+        A new Topology instance containing only sites and connections from matching molecules.
+    """
+    sites = [
+        s
+        for s in topology.iter_sites_by_residue(
+            residue_tag=residue_tag, residue_number=residue_number
+        )
+    ]
+    return slice_by_sites(topology=topology, sites=sites)
+
+
+def slice_by_sites(topology, sites):
+    """Used by slice_topology_by_molecule() and slice_topology_by_residue()
+
+    Parameters
+    ----------
+    sites : list of gmso.core.atom.Atom
+        List of sites to include in the sub-topology.
+    topology : gmso.core.topology.Topology
+        The topology being sliced.
+    """
+    new_topology = gmso.core.topology.Topology()
+
+    connections = set()
+    for site in sites:
+        new_topology.add_site(site)
+        for connection in topology.iter_connections_by_site(site):
+            connections.add(connection)
+
+    for connection in connections:
+        new_topology.add_connection(connection)
+
+    return new_topology