Merge 48c8fc5 into 5679d70

foyer/topology_graph.py

@@ -1,4 +1,6 @@
 """Module to represent chemical systems as graph structures."""
+import enum
+
 import networkx as nx
 from parmed import Structure
 from parmed import periodic_table as pt
@@ -35,6 +37,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.
 
@@ -79,7 +99,7 @@ def add_atom(self, index, name, atomic_number=None, element=None, **kwargs):
         atom_data = AtomData(index, name, atomic_number, element, **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
@@ -88,8 +108,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."""
@@ -140,9 +170,21 @@ def from_parmed(cls, structure: Structure):
                 atomic_number=atomic_number,
                 element=element,
             )
-
+        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
 
@@ -183,11 +225,22 @@ def from_openff_topology(cls, openff_topology):
                 element=element,
             )
 
+        bond_type_dict = {
+            "Single": BondOrder.SINGLE,
+            "Double": BondOrder.DOUBLE,
+            "Triple": BondOrder.TRIPLE,
+            "Aromatic": BondOrder.AROMATIC,
+            "Amide": BondOrder.AMIDE,
+        }
         for top_bond in openff_topology.topology_bonds:
             atoms_indices = [
                 atom.topology_atom_index for atom in top_bond.atoms
             ]
-            top_graph.add_bond(atoms_indices[0], atoms_indices[1])
+            top_graph.add_bond(
+                atoms_indices[0],
+                atoms_indices[1],
+                bond_type=bond_type_dict.get(top_bond.type, BondOrder.UNKNOWN),
+            )
 
         return top_graph
 
@@ -233,12 +286,14 @@ def from_gmso_topology(cls, gmso_topology):
                         atomic_number=atom.element.atomic_number,
                         element=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