Merge pull request PolusAI#229 from misterbrandonwalker/patch_diffdock_sanitize

add timeout for rdkit conformer generation

Author: Nicholas-Schaub

examples/diffdock/sanitize_ligand.py

@@ -1,10 +1,13 @@
+# mypy: disable-error-code="attr-defined"
 # pylint: disable=E0401,E1101,I1101
 """Filter molecules with kekulization errors.
 Handle molecules with sanitization errors by assign formal charge based on valence."""
 # https://depth-first.com/articles/2020/02/10/a-comprehensive-treatment-of-aromaticity-in-the-smiles-language/
 import argparse
 from pathlib import Path
-from typing import Dict, Tuple
+from typing import Dict, Tuple, Optional
+from types import FrameType
+import signal  # only unix supported
 
 import rdkit
 from rdkit import Chem
@@ -16,6 +19,24 @@
 input_small_mol_ligand = Path(args.input_small_mol_ligand).name
 
 
+class TimeoutException(Exception):
+    pass
+
+
+def timeout(seconds: int, error_message: str = 'Timeout') -> None:
+    """Time out for conformer generation, only works for unix. For windows have to ignore mypy error.
+
+    Args:
+        seconds (int): The time in seconds to wait before timing out
+        error_message (str, optional): The error message to display. Defaults to 'Timeout'.
+    """
+    def handle_timeout(signum: int, frame: Optional[FrameType]) -> None:
+        raise TimeoutException(error_message)
+
+    signal.signal(signal.SIGALRM, handle_timeout)
+    signal.alarm(seconds)
+
+
 def adjust_formal_charges(molecule: Chem.SDMolSupplier) -> Chem.SDMolSupplier:
     """Sometimes input structures do not have correct formal charges corresponding
     to bond order topology. So choose to trust bond orders assigned and generate formal
@@ -59,14 +80,19 @@ def adjust_formal_charges(molecule: Chem.SDMolSupplier) -> Chem.SDMolSupplier:
         atom.SetFormalCharge(chg)
     return molecule
 
+# for some ligands can take up to minutes to generate conformer!!
+
 
-def generate_conformer(molecule: Chem.SDMolSupplier) -> None:
+def generate_conformer(molecule: Chem.SDMolSupplier, timeout_value: int = 5) -> None:
     """ Generate conformer for molecule. Sometimes rdkit embedding can fail,
       so use random coordinates if failed at first.
 
     Args:
-        molecule (Chem.SDMolSupplier): _description_
+        molecule (Chem.SDMolSupplier): The rdkit molecule object
+        timeout_value (float): The time in seconds to wait before timing out
     """
+    # Set up the signal handler for the timeout
+    timeout(timeout_value)
     ps = AllChem.ETKDGv2()
     confid = AllChem.EmbedMolecule(molecule, ps)
     if confid == -1:
@@ -75,6 +101,8 @@ def generate_conformer(molecule: Chem.SDMolSupplier) -> None:
         # only want to catch error Bad Conformation Id,  dont need to spend time optimizing
         AllChem.MMFFOptimizeMolecule(molecule, confId=0, maxIters=1)
 
+    signal.alarm(0)
+
 
 def is_valid_ligand(molecule: Chem.SDMolSupplier) -> Tuple[bool, bool, Chem.SDMolSupplier]:
     """Check for sanitization errors, attempt to fix formal charges/valence consistency errors.