diff --git a/foyer/topology_graph.py b/foyer/topology_graph.py index 2fef4776..27836754 100644 --- a/foyer/topology_graph.py +++ b/foyer/topology_graph.py @@ -1,4 +1,5 @@ """Module to represent chemical systems as graph structures.""" +import enum from typing import TYPE_CHECKING, Optional import networkx as nx @@ -43,6 +44,24 @@ def __init__(self, index, name, atomic_number=None, element=None, **kwargs): setattr(self, key, value) +class BondOrder(enum.Enum): + """Enum to represent various bond orders from multiple sources.""" + + SINGLE = "1" + DOUBLE = "2" + TRIPLE = "3" + AMIDE = "am" + AROMATIC = "ar" + DUMMY = "du" + NOTCONNECTED = "nc" + UNKNOWN = "un" + + @classmethod + def _missing_(cls, value): + """If value cannot be found, default to unknown.""" + return cls.UNKNOWN + + class TopologyGraph(nx.Graph): """A general TopologyGraph. @@ -94,7 +113,7 @@ def add_atom( atom_data = AtomData(index, name, atomic_number, symbol, **kwargs) self.add_node(index, atom_data=atom_data) - def add_bond(self, atom_1_index, atom_2_index): + def add_bond(self, atom_1_index, atom_2_index, bond_type=BondOrder.SINGLE): """Add a bond(edge) between two atoms in this TopologyGraph. Parameters @@ -103,8 +122,18 @@ def add_bond(self, atom_1_index, atom_2_index): The index of the first atom that forms this bond atom_2_index: int The index of the second atom that forms this bond + bond_type: str, default = BondOrder.SINGLE + The type of bond being added, can be: + "1": single, + "2": double, + "3": triple, + "am": amide, + "ar": aromatic, + "un": unknown, + "du": dummy + "nc": not connected """ - self.add_edge(atom_1_index, atom_2_index) + self.add_edge(atom_1_index, atom_2_index, {"bond_type": bond_type}) def atoms(self, data=False): """Iterate through atoms in the TopologyGraph.""" @@ -155,9 +184,21 @@ def from_parmed(cls, structure: Structure): atomic_number=atomic_number, symbol=symbol, ) - + bond_type_dict = { + round(1, 2): BondOrder.SINGLE, + round(2, 2): BondOrder.DOUBLE, + round(3, 2): BondOrder.TRIPLE, + round(1.5, 2): BondOrder.AROMATIC, + round(1.25, 2): BondOrder.AMIDE, + } for bond in structure.bonds: - topology_graph.add_bond(bond.atom1.idx, bond.atom2.idx) + topology_graph.add_bond( + bond.atom1.idx, + bond.atom2.idx, + bond_type=bond_type_dict.get( + round(bond.order, 2), BondOrder.UNKNOWN + ), + ) return topology_graph @@ -204,11 +245,23 @@ def from_openff_topology(cls, openff_topology: "OpenFFTopology"): symbol=element_symbol, ) + bond_type_dict = { + "Single": BondOrder.SINGLE, + "Double": BondOrder.DOUBLE, + "Triple": BondOrder.TRIPLE, + "Aromatic": BondOrder.AROMATIC, + "Amide": BondOrder.AMIDE, + } + for bond in openff_topology.bonds: atoms_indices = [ openff_topology.atom_index(atom) for atom in bond.atoms ] - top_graph.add_bond(*atoms_indices) + top_graph.add_bond( + atoms_indices[0], + atoms_indices[1], + bond_type=bond_type_dict.get(top_bond.type, BondOrder.UNKNOWN), + ) return top_graph @@ -257,12 +310,14 @@ def from_gmso_topology(cls, gmso_topology: "gmso.Topology"): atomic_number=atom.element.atomic_number, symbol=atom.element.symbol, ) - + # until gmso has bond orders, assume single bonds for top_bond in gmso_topology.bonds: atoms_indices = [ gmso_topology.get_index(atom) for atom in top_bond.connection_members ] - top_graph.add_bond(atoms_indices[0], atoms_indices[1]) + top_graph.add_bond( + atoms_indices[0], atoms_indices[1], bond_type=BondOrder.SINGLE + ) return top_graph