Add export to RDKIT to allow for chemdraw style visualization.

mbuild/compound.py

@@ -1189,9 +1189,16 @@
         for b1, b2 in self.root.bond_graph.edges(self):
             yield b2
 
-    def bonds(self):
+    def bonds(self, return_bond_order=False):
         """Return all bonds in the Compound and sub-Compounds.
 
+        Parameters
+        ----------
+        return_bond_order : bool, optional, default=False
+            Return the bond order of the bond as the 3rd argument in the tuple.
+            bond order is returned as a dictionary with 'bo' as the key.
+            If bond order is not set, the value will be set to 'default'.
+
         Yields
         ------
         tuple of mb.Compound
@@ -1204,9 +1211,11 @@
         """
         if self.root.bond_graph:
             if self.root == self:
-                return self.root.bond_graph.edges()
+                return self.root.bond_graph.edges(data=return_bond_order)
             else:
-                return self.root.bond_graph.subgraph(self.particles()).edges()
+                return self.root.bond_graph.subgraph(self.particles()).edges(
+                    data=return_bond_order
+                )
         else:
             return iter(())
 
@@ -1246,18 +1255,23 @@
             )
         return sum(1 for _ in self.bonds())
 
-    def add_bond(self, particle_pair):
+    def add_bond(self, particle_pair, bond_order=None):
         """Add a bond between two Particles.
 
         Parameters
         ----------
         particle_pair : indexable object, length=2, dtype=mb.Compound
             The pair of Particles to add a bond between
+        bond_order : float, optional, default=None
+            Bond order of the bond.
         """
         if self.root.bond_graph is None:
             self.root.bond_graph = BondGraph()
-
-        self.root.bond_graph.add_edge(particle_pair[0], particle_pair[1])
+        if bond_order is None:
+            bond_order = "default"
+        self.root.bond_graph.add_edge(
+            particle_pair[0], particle_pair[1], bo=bond_order
+        )
 
     def generate_bonds(self, name_a, name_b, dmin, dmax):
         """Add Bonds between all pairs of types a/b within [dmin, dmax].
@@ -2654,7 +2668,7 @@
                     particle._element = element_from_symbol(particle.name)
                 except ElementError:
                     try:
-                        element_from_name(particle.name)
+                        particle._element = element_from_name(particle.name)
                     except ElementError:
                         raise MBuildError(
                             "No element assigned to {}; element could not be"
@@ -3228,6 +3242,58 @@
             infer_hierarchy=infer_hierarchy,
         )
 
+    def to_rdkit(self):
+        rdkit = import_("rdkit")
+        from rdkit import Chem
+        from rdkit.Chem import AllChem
+
+        for particle in self.particles():
+            if particle.element is None:
+                try:
+                    particle._element = element_from_symbol(particle.name)
+                except ElementError:
+                    try:
+                        particle._element = element_from_name(particle.name)
+                    except ElementError:
+                        raise MBuildError(
+                            "No element assigned to {}; element could not be"
+                            "inferred from particle name {}. Cannot perform"
+                            "an energy minimization.".format(
+                                particle, particle.name
+                            )
+                        )
+
+        temp_mol = Chem.RWMol()
+        p_dict = {particle: i for i, particle in enumerate(self.particles())}
+
+        bo_dict = {
+            1.0: Chem.BondType.SINGLE,
+            2.0: Chem.BondType.DOUBLE,
+            3.0: Chem.BondType.TRIPLE,
+            1.5: Chem.BondType.AROMATIC,
+            0.0: Chem.BondType.UNSPECIFIED,
+            "default": Chem.BondType.SINGLE,
+        }
+
+        for particle in self.particles():
+            temp_atom = Chem.Atom(particle.element.atomic_number)
+
+            # this next line is necessary to prevent rdkit from adding hydrogens
+            # this will also set the label to be the lement with particle index
+            temp_atom.SetProp(
+                "atomLabel", f"{temp_atom.GetSymbol()}:{p_dict[particle]}"
+            )
+
+            temp_mol.AddAtom(temp_atom)
+
+        for bond in self.bonds(return_bond_order=True):
+            bond_indices = (p_dict[bond[0]], p_dict[bond[1]])
+            temp_mol.AddBond(*bond_indices)
+            rdkit_bond = temp_mol.GetBondBetweenAtoms(*bond_indices)
+            rdkit_bond.SetBondType(bo_dict[bond[2]["bo"]])
+
+        return temp_mol
+
     def to_parmed(
         self,
         box=None,
@@ -3552,9 +3618,12 @@
         newone.bond_graph = BondGraph()
         for particle in self.particles():
             newone.bond_graph.add_node(clone_of[particle])
-        for c1, c2 in self.bonds():
+        for c1, c2, data in self.bonds(return_bond_order=True):
             try:
-                newone.add_bond((clone_of[c1], clone_of[c2]))
+                # bond order is added to the data dictionary as 'bo'
+                newone.add_bond(
+                    (clone_of[c1], clone_of[c2]), bond_order=data["bo"]
+                )
             except KeyError:
                 raise MBuildError(
                     "Cloning failed. Compound contains bonds to "

mbuild/conversion.py

@@ -884,10 +884,18 @@ def from_rdkit(rdkit_mol, compound=None, coords_only=False, smiles_seed=0):
         part_list.append(part)
 
     comp.add(part_list)
+    bo_dict = {
+        Chem.BondType.SINGLE: 1.0,
+        Chem.BondType.DOUBLE: 2.0,
+        Chem.BondType.TRIPLE: 3.0,
+        Chem.BondType.AROMATIC: 1.5,
+        Chem.BondType.UNSPECIFIED: 0.0,
+    }
 
     for bond in mymol.GetBonds():
         comp.add_bond(
-            [comp[bond.GetBeginAtomIdx()], comp[bond.GetEndAtomIdx()]]
+            [comp[bond.GetBeginAtomIdx()], comp[bond.GetEndAtomIdx()]],
+            bond_order=bo_dict[bond.GetBondType()],
         )
 
     return comp

mbuild/tests/test_compound.py

@@ -115,6 +115,53 @@ def test_direct_bonds(self, methane):
         for H in methane.particles_by_name("H"):
             assert H in bond_particles
 
+    def test_bond_order(self, methane):
+        # test the default behavior
+        bonds = [bond for bond in methane.bonds(return_bond_order=True)]
+        assert bonds[0][2]["bo"] == "default"
+
+        # test setting bond order when adding a bond
+        carbon = mb.Compound(pos=[0, 0, 0], name="C")
+        carbon_clone = mb.clone(carbon)
+        C2 = mb.Compound([carbon, carbon_clone])
+        C2.add_bond([carbon, carbon_clone], bond_order=2.0)
+        bonds = [bond for bond in C2.bonds(return_bond_order=True)]
+
+        assert bonds[0][2]["bo"] == 2.0
+
+        # ensure we propagate bond order information when we clone a compound
+        C2_clone = mb.clone(C2)
+        bonds = [bond for bond in C2_clone.bonds(return_bond_order=True)]
+
+        assert bonds[0][2]["bo"] == 2.0
+
+    @pytest.mark.skipif(not has_rdkit, reason="RDKit is not installed")
+    def test_convert_to_rdkit(self, methane):
+        # check basic conversion
+        rdkit = import_("rdkit")
+        rdmol = methane.to_rdkit()
+        assert isinstance(rdmol, rdkit.Chem.rdchem.RWMol)
+
+        from rdkit import Chem
+
+        atoms = [atom for atom in rdmol.GetAtoms()]
+        bonds = [bond for bond in rdmol.GetBonds()]
+
+        bo_dict = {
+            Chem.BondType.SINGLE: 1.0,
+            Chem.BondType.DOUBLE: 2.0,
+            Chem.BondType.TRIPLE: 3.0,
+            Chem.BondType.AROMATIC: 1.5,
+            Chem.BondType.UNSPECIFIED: 0.0,
+        }
+        assert len(atoms) == 5
+        assert len(bonds) == 4
+
+        assert bo_dict[bonds[0].GetBondType()] == 1.0
+        assert bo_dict[bonds[1].GetBondType()] == 1.0
+        assert bo_dict[bonds[2].GetBondType()] == 1.0
+        assert bo_dict[bonds[3].GetBondType()] == 1.0
+
     def test_n_direct_bonds_parent(self, ethane):
         with pytest.raises(MBuildError):
             ethane.n_direct_bonds