How to speed up tautomer enumeration on a large compound library

[NiBoyang]

Hi all, I have a large compound library of about 4 million compounds, and I want to perform tautomer enumeration on each compound. I've already tried using multiprocessing, but the time taken is similar to a direct loop without any noticeable speed improvement. Could you help me understand the possible reasons for this lack of speedup and suggest any viable methods to accelerate the process?

My original code:

from rdkit.Chem.MolStandardize.rdMolStandardize import TautomerEnumerator

mols = MyLibraryCompounds
te = TautomerEnumerator()
tautomers = [te.Enumerate(mol) for mol in mols[0:10000]]

The multiprocessing code using mpire (similar to the native multiprocessing):

from rdkit.Chem.MolStandardize.rdMolStandardize import TautomerEnumerator
from mpire import WorkerPool

te = TautomerEnumerator()

def enumerate_tautomers(mol):
    return [tau for tau in te.Enumerate(mol)]

with WorkerPool(n_jobs=4, start_method='threading') as pool:
    tautomers = pool.map(enumerate_tautomers, mols[0:10000], progress_bar=True)

[DavidACosgrove]

Hi,
This is a classic Python problem. Due to the Global Interpreter Lock (GIL) only a single Python thread can execute at a time. This means that any multi-processing code using threads rarely gives a performance improvement. I don’t know how mpire works, but most such libraries allow for multi-process running rather than multi- threaded. This runs a separate copy of the code in each process so the GIL is not relevant. There are memory implications though.

Dave

[bp-kelley]

We also can do another sweep of functions that should release the GIL. I think standardization functions tautomer enumeration are good candidates for this.

[bp-kelley]

We do it in a few places in the standardize library, but not everywhere:

rdkit/Code/GraphMol/MolStandardize/Wrap/rdMolStandardize.cpp

Line 190 in 8dae48b

NOGIL gil;

[NiBoyang]

Hi both, many thanks for you responses.

I've tried using Multiprocessing, and it took about 1 min to complete.

from rdkit.Chem.MolStandardize.rdMolStandardize import TautomerEnumerator
import multiprocessing

te = TautomerEnumerator()

def enumerate_tautomers(mol):
    return [tau for tau in te.Enumerate(mol)]

with multiprocessing.pool() as pool:
    tautomers = list(pool.map(enumerate_tautomers, mols[0:10000]))

The original 'loop' also took about 1 min to complete

from rdkit.Chem.MolStandardize.rdMolStandardize import TautomerEnumerator

mols = MyLibraryCompounds
te = TautomerEnumerator()
tautomers = [te.Enumerate(mol) for mol in mols[0:10000]]

I wonder if I was using the multiprocessing in a wrong way? Moreover, the multiprocessing seems not able to pickle what the TautomerEnumerator returns, any solutions to this? Sorry for these questions as I'm new to this area. Many thanks!

[bp-kelley]

Your code is essentially correct, here is a sample I'm using with a data file in the rdkit

from rdkit.rdBase import BlockLogs
from rdkit.Chem.MolStandardize.rdMolStandardize import TautomerEnumerator
import multiprocessing
import os
import time

te = TautomerEnumerator()

def enumerate_tautomers(mol):
    try:
        return [tau for tau in te.Enumerate(mol)]
    except Exception as e:
        print(str(e))
        return None

def doit():
    blocker = BlockLogs()
    t1 = time.time()
    mols = [Chem.MolFromSmiles(line.strip()) for line in
        open(os.path.join(os.environ["RDBASE"], "Data/NCI/first_5K.smi"))]
    # remove empties
    mols = [mol for mol in mols if mol]
    t2 = time.time()
    print(f'{t2-t1} seconds to read molecules')

    t1 = time.time()
    tautomers = [enumerate_tautomers(mol) for mol in mols]
    t2 = time.time()
    print(f"{t2-t1} seconds to process")
    
    t1 = time.time()
    with multiprocessing.Pool() as pool:
        tautomers = list(pool.map(enumerate_tautomers, mols))
    t2 = time.time()
    print(f"{t2-t1} seconds to process multiproc")

if __name__ == "__main__":
    doit()

results

0.18940377235412598 seconds to read molecules
16.685894012451172 seconds to process
4.813543796539307 seconds to process multiproc

[NiBoyang]

Many thanks it worked now!

Sorry I have another question to post. I'm now generating conformers for each of the generated tautomers, when I multiprocess it, the process always gets killed in halfway. I think it's because of the memory usage as I see kswapd() raised when I run that.

I was wondering if you could take time to have a look of my conformer generation functions to see what leads to this memory issue.

from rdkit import Chem
from rdkit.Chem import AllChem
from rdkit.Chem import MolStandardize
from rdkit.Chem import Crippen, Lipinski
from rdkit.Chem import rdMolTransforms
from rdkit.Chem import rdmolops
from rdkit.Chem.MolStandardize.rdMolStandardize import TautomerEnumerator
import itertools, zlib
import multiprocessing

def _get_flanking_atoms(
    bonds: Chem.rdchem._ROBondSeq, atom_1_idx: int, atom_2_idx: int
) -> tuple:
    """Search for two atoms connecting to either atom in a rotatable bond

    A private method to identify two atoms flanking atoms in a rotatable bond

    Parameters
    ----------
    bonds : Chem.rdchem._ROBondSeq
        The Rdkit molecule bond object that contains the indexes of both begin and end atoms in a bond.
    atom_1_idx : int
        The index of the first atom in a rotatable bond.
    atom_2_inx : int
        The index of the second atom in a rotatable bond.

    Returns
    -------
    tuple
        A tuple containing the indexes of two flanking atoms for the given atoms of a rotatable bond.
    """
    bound_to_atom_1 = None
    bound_to_atom_2 = None

    for bond in bonds:
        bond_indexes = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()
        if atom_1_idx not in bond_indexes and atom_2_idx not in bond_indexes:
            continue

        if atom_1_idx in bond_indexes:  # Atom 1 in bond indexes
            if atom_2_idx in bond_indexes:
                continue
            if atom_1_idx == bond_indexes[0]:
                bound_to_atom_1 = bond_indexes[1]
                continue
            if atom_1_idx == bond_indexes[1]:
                bound_to_atom_1 = bond_indexes[0]
                continue
        else:  # Atom 2 in bond indexes
            if atom_2_idx == bond_indexes[0]:
                bound_to_atom_2 = bond_indexes[1]
                continue
            if atom_2_idx == bond_indexes[1]:
                bound_to_atom_2 = bond_indexes[0]
                continue

            if bound_to_atom_1 is not None and bound_to_atom_2 is not None:
                return bound_to_atom_1, bound_to_atom_2
    return bound_to_atom_1, bound_to_atom_2

def _get_dihedrals(mol: Chem.rdchem.Mol) -> tuple:
    """Identify dihedrals in order to obtain rotatable bonds in a molecule

    Identify dihedrals using flanking atoms of rotatable bonds

    Parameters
    ----------
    mol : Chem.rdchem.Mol
        The Rdkit mol object of the input molecule.

    Returns
    -------
    tuple
        A tuple containing all identified dihedrals with the index of their four defined atoms.
    """
    dihedrals = []
    for atom_j, atom_k in mol.GetSubstructMatches(rotatable_bond_from_smarts):
        atom_i, atom_l = _get_flanking_atoms(mol.GetBonds(), atom_j, atom_k)
        if atom_i is not None and atom_l is not None:
            dihedrals.append((atom_i, atom_j, atom_k, atom_l))
    return dihedrals

def _sample_bond_states(n_rot: int, seed: int) -> list:
    """Sample a set of conformers with different rotation angles

    A private method used to generate a set of bond angle multipliers (0 to
    3) using all rotatable bonds within a molecule. Up to 1024 conformers
    are sampled for a molecule if n_rot is greater than five, otherwise all
    n_rot^4 are returned.

    Parameters
    ----------
    n_nor : int
        The number of rotatable bonds in a molecule.
    seed : int
        Seed for random sampling of rotamer space. Typically the hash of
        molecule coords.
    Returns
    -------
    List
        A list containing the sampled bond states (rotation angles) for all
        of the rotatable bonds of a molecule.
    """
    if n_rot <= 5:
        return list(itertools.product(range(4), repeat=n_rot))

    np_rng = np.random.default_rng(seed=seed)

    rotation_list = set(
        tuple(np_rng.choice(range(4), size=n_rot)) for counter in range(1024)
    )

    while len(rotation_list) < 1024:
        rotation_list.add(tuple(np_rng.choice(range(4), n_rot)))
    return list(rotation_list)

def _generate_conf(
    conformer: Chem.rdchem.Conformer,
    dihedrals: tuple,
    starting_angles: tuple,
    bond_state: tuple,
) -> None:
    """Rotate the assigned rotatable bonds

    Change conformers by rotating the assigned rotatable bond based on a set
    of dihedral angles defined by four flanking atoms.

    Parameters
    ----------
    conformer : Chem.rdchem.Conformer
        The Rdkit conformer object.
    dihedrals : tuple
        A tuple containing all identified dihedrals with the index of their
        four defined atoms.
    starting_angles : tuple
        A tuple containing the orignal states (dihedral angles) of all the
        rotatable bond before rotating.
    bond_state : tuple
        A tuple containing a specific bond state (a combination of various
        rotation angles) for all rotatable bonds of a molecule.
    """
    for dihedral_atoms, torsion_angle_multiplier, orig_torsion_angle in zip(
        dihedrals, bond_state, starting_angles
    ):
        rdMolTransforms.SetDihedralDeg(
            conformer,
            *dihedral_atoms,
            orig_torsion_angle + torsion_angle_multiplier * 90.0,
        )

The main generate_conformers function:

def generate_conformers(mol: Chem.rdchem.Mol, random_state: int = 42) -> list:
    """Generate conformers with rotatable bonds

    Generate conformers for a molecule, enumerating rotatable bonds over 90
    degree angles. This 90 degree increment was deemed opimal in literature.

    Parameters
    ----------
    mol : Chem.rdchem.Mol
        The Rdkit mol object of the input molecule.
    random_state : int
        Integer to use as a random state when seeding the random number
        generator.  By default 42.


    Returns
    -------
    List
        A list containing mol objects of different conformers with different
        angles of rotetable bonds

    """

    try:
        # mol = rdmolops.AddHs(mol)
        params = AllChem.ETKDGv3()
        params.useSmallRingTorsions = True
        params.randomSeed = random_state
        original_conformer = mol.GetConformer(AllChem.EmbedMolecule(mol, params))
    except ValueError:
        params = AllChem.ETKDGv2()
        id = AllChem.EmbedMolecule(mol, params)
        if id == -1:
            logging.warning(
                "Coords could not be generated without using random coords. using random coords now"
            )
            params.useRandomCoords = True
        try:
            original_conformer = mol.GetConformer(
                AllChem.EmbedMolecule(mol, params)
            )
        except ValueError:
            logging.warning("Conformer embedding failed")
            return []
    dihedrals = _get_dihedrals(mol)
    starting_angles = [
        rdMolTransforms.GetDihedralDeg(original_conformer, *dihedral_atoms)
        for dihedral_atoms in dihedrals
    ]
    bond_states = _sample_bond_states(
        len(dihedrals), zlib.crc32(mol.GetConformer(0).GetPositions().tobytes())
    )

    new_conf_mol_list = []
    for bond_state in bond_states:
        _generate_conf(
            original_conformer, dihedrals, starting_angles, bond_state
        )
        new_conf_mol_list.append(Chem.Mol(mol))

    return new_conf_mol_list

I have a list of lists of tautomers named tautomers

I defined a function to be multiprocessed:

def process_tautomer_list(tautomer_list):
    molecule_conformers = []
    for tautomer in tautomer_list:
        conformers = generate_conformers(tautomer)
        molecule_conformers.extend(conformers)
    return molecule_conformers

And multiprocess using multiprocessing:

with multiprocessing.Pool() as pool:
    results = list((pool.map(process_tautomer_list, tautomers)))