diff --git a/gmso/core/topology.py b/gmso/core/topology.py index 7248da7da..1306b3b52 100644 --- a/gmso/core/topology.py +++ b/gmso/core/topology.py @@ -1256,6 +1256,21 @@ def get_index(self, member): 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 @@ def iter_sites(self, key, value): 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 @@ def iter_sites_by_residue(self, residue_tag): """ 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 @@ def iter_sites_by_molecule(self, molecule_tag): """ 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) diff --git a/gmso/external/convert_openmm.py b/gmso/external/convert_openmm.py index b3ea139a6..b1a5c65c9 100644 --- a/gmso/external/convert_openmm.py +++ b/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. diff --git a/gmso/tests/test_topology.py b/gmso/tests/test_topology.py index bde0de76f..fbd0a1b62 100644 --- a/gmso/tests/test_topology.py +++ b/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"], diff --git a/gmso/tests/test_utils.py b/gmso/tests/test_utils.py index c23b7ef0a..e2e5c5e9a 100644 --- a/gmso/tests/test_utils.py +++ b/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 diff --git a/gmso/utils/charges.py b/gmso/utils/charges.py new file mode 100644 index 000000000..e336f72a1 --- /dev/null +++ b/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) diff --git a/gmso/utils/slicing.py b/gmso/utils/slicing.py new file mode 100644 index 000000000..d94b73eb2 --- /dev/null +++ b/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