Simple RDKitConverter

package/MDAnalysis/coordinates/RDKit.py

@@ -55,11 +55,45 @@
 """
 
 import warnings
+import re
 
 import numpy as np
 
+from ..exceptions import NoDataError
+from ..topology.guessers import guess_atom_element
 from . import memory
-
+from . import base
+
+try:
+    from rdkit import Chem
+except ImportError:
+    pass
+else:
+    RDBONDTYPE = {
+        'AROMATIC': Chem.BondType.AROMATIC,
+        'SINGLE': Chem.BondType.SINGLE,
+        'DOUBLE': Chem.BondType.DOUBLE,
+        'TRIPLE': Chem.BondType.TRIPLE,
+    }
+    RDBONDORDER = {
+        1: Chem.BondType.SINGLE,
+        1.5: Chem.BondType.AROMATIC,
+        "ar": Chem.BondType.AROMATIC,
+        2: Chem.BondType.DOUBLE,
+        3: Chem.BondType.TRIPLE,
+    }
+    # add string version of the key for each bond
+    RDBONDORDER.update({str(key):value for key,value in RDBONDORDER.items()})
+    RDATTRIBUTES = {
+        "altLoc": "AltLoc",
+        "chainID": "ChainId",
+        "name": "Name",
+        "occupancy": "Occupancy",
+        "resname": "ResidueName",
+        "resid": "ResidueNumber",
+        "segid": "SegmentNumber",
+        "tempfactor": "TempFactor",
+    }
 
 class RDKitReader(memory.MemoryReader):
     """Coordinate reader for RDKit.
@@ -101,4 +135,113 @@ def __init__(self, filename, **kwargs):
             warnings.warn("No coordinates found in the RDKit molecule")
             coordinates = np.empty((1,n_atoms,3), dtype=np.float32)
             coordinates[:] = np.nan
-        super(RDKitReader, self).__init__(coordinates, order='fac', **kwargs)
+        super(RDKitReader, self).__init__(coordinates, order='fac', **kwargs)
+
+
+class RDKitConverter(base.ConverterBase):
+    """Convert MDAnalysis AtomGroup or Universe to `RDKit <https://www.rdkit.org/docs/source/rdkit.Chem.rdchem.html#rdkit.Chem.rdchem.Mol>`_ :class:`rdkit.Chem.rdchem.Mol`.
+
+    Example
+    -------
+
+    .. code-block:: python
+
+        import MDAnalysis as mda
+        from MDAnalysis.tests.datafiles import PDB_full
+        u = mda.Universe(PDB_full)
+        mol = u.select_atoms('resname DMS').convert_to('RDKIT')
+
+
+    .. versionadded:: 2.X.X
+    """
+
+    lib = 'RDKIT'
+    units = {'time': None, 'length': 'Angstrom'}
+
+    def convert(self, obj):
+        """Write selection at current trajectory frame to :class:`~rdkit.Chem.rdchem.Mol`.
+
+        Parameters
+        -----------
+        obj : AtomGroup or Universe
+        """
+        try:
+            from rdkit import Chem
+        except ImportError:
+            raise ImportError('RDKit is required for RDKitConverter but '
+                              'is not installed. Try installing it with \n'
+                              'conda install -c conda-forge rdkit')
+        try:
+            # make sure to use atoms (Issue 46)
+            ag = obj.atoms
+        except AttributeError as e:
+            raise TypeError("No `atoms` attribute in object of type {}, "
+                            "please use a valid AtomGroup or Universe".format(
+                            type(obj))) from e
+
+        mol = Chem.RWMol()
+        atom_mapper = {}
+
+        for atom in ag:
+            try:
+                element = atom.element
+            except NoDataError:
+                # guess atom element
+                # capitalize: transform CL to Cl and so on
+                element = guess_atom_element(atom.name).capitalize()
+            rdatom = Chem.Atom(element)
+            # add properties
+            mi = Chem.AtomPDBResidueInfo()
+            for attr, rdattr in RDATTRIBUTES.items():
+                try: # get value in MDA atom
+                    value = getattr(atom, attr)
+                except AttributeError:
+                    pass
+                else:
+                    if isinstance(value, np.generic):
+                        # convert numpy types to python standard types
+                        value = value.item()
+                    if attr == "segid":
+                        # RDKit needs segid to be an int
+                        try:
+                            value = int(value)
+                        except ValueError:
+                            # convert any string to int
+                            # can be mapped back with np.base_repr(x, 36)
+                            value = int(value, 36)
+                    elif attr == "name":
+                        # RDKit needs the name to be properly formated for a
+                        # PDB file (1 letter elements start at col 14)
+                        name = re.findall('(\D+|\d+)', value)
+                        if len(name) == 2:
+                            symbol, number = name
+                        else:
+                            symbol, number = name[0], ""
+                        value = "{:>2}".format(symbol) + "{:<2}".format(number)
+                    # set attribute value in RDKit MonomerInfo
+                    getattr(mi, "Set%s" % rdattr)(value)
+            rdatom.SetMonomerInfo(mi)
+            # TODO other properties (charges)
+            index = mol.AddAtom(rdatom)
+            # map index in universe to index in mol
+            atom_mapper[atom.ix] = index
+
+        try:
+            bonds = ag.bonds
+        except NoDataError:
+            ag.guess_bonds()
+            bonds = ag.bonds
+
+        for bond in bonds:
+            bond_indices = [atom_mapper[i] for i in bond.indices]
+            try:
+                bond_type = bond.type.upper()
+            except AttributeError:
+                # bond type can be a tuple for PDB files
+                bond_type = None
+            bond_type = RDBONDTYPE.get(bond_type, RDBONDORDER.get(
+                bond.order, Chem.BondType.SINGLE))
+            mol.AddBond(*bond_indices, bond_type)
+
+        Chem.SanitizeMol(mol)
+        return mol

testsuite/MDAnalysisTests/coordinates/test_rdkit.py

@@ -28,7 +28,7 @@
 from numpy.testing import (assert_equal,
                            assert_almost_equal)
 
-from MDAnalysisTests.datafiles import mol2_molecule
+from MDAnalysisTests.datafiles import mol2_molecule, PDB_full
 
 Chem = pytest.importorskip("rdkit.Chem")
 AllChem = pytest.importorskip("rdkit.Chem.AllChem")
@@ -76,3 +76,41 @@ def test_compare_mol2reader(self):
         assert_equal(universe.trajectory.ts.positions, 
                      mol2.trajectory.ts.positions)
 
+
+class TestRDKitConverter(object):
+    @pytest.fixture
+    def pdb(self):
+        return mda.Universe(PDB_full)
+
+    @pytest.fixture
+    def mol2(self):
+        return mda.Universe(mol2_molecule)
+
+    @pytest.mark.parametrize("sel_str", [
+        "resid 1",
+        "resname LYS and name NZ",
+        "resid 34 and altloc B",
+    ])
+    def test_monomer_info(self, pdb, sel_str):
+        rdmol = Chem.MolFromPDBFile(PDB_full)
+        sel = pdb.select_atoms(sel_str)
+        umol = sel.convert_to("RDKIT")
+        atom = umol.GetAtomWithIdx(0)
+        mi = atom.GetMonomerInfo()
+
+        for mda_attr, rd_attr in mda.coordinates.RDKit.RDATTRIBUTES.items():
+            if mda_attr == "occupancy":
+                mda_attr = "occupancie"
+            rd_value = getattr(mi, "Get%s" % rd_attr)()
+            mda_value = getattr(sel, "%ss" % mda_attr)[0]
+            if mda_attr == "name":
+                rd_value = rd_value.strip()
+            elif mda_attr == "segid":
+                rd_value = np.base_repr(rd_value, 36)
+            assert rd_value == mda_value
+
+    def test_identical_topology_mol2(self, mol2):
+        # no chirality check
+        rdmol = mol2_mol()
+        umol = mol2.atoms.convert_to("RDKIT")
+        assert rdmol.HasSubstructMatch(umol) and umol.HasSubstructMatch(rdmol)