'RDKit Optimize Geometry node' only uses 8 CPU cores

Dear all,

The ‘RDKIT Optimize Geometry’ node is one of the slowest nodes in KNIME. It is not a surprise, as it involves many iterations of reorganizing 3D coordinates, but still.

I observed it also consistently uses only 8 CPU cores, while many other RDKIT nodes do maximailize the use of CPU cores. It may be due to how the multiprocessing is being dealt with within the ‘RDKit Optimize Geometry’ node (and not in the other RDKIT nodes)?

When I perform the RDKIT geometry optimisation in a ‘Python Script’ node using ‘loky’ for multiprocessing, then all CPU cores can be used and you can run your large set of molecules signicifantly faster than with the RDKIT node.

I hope also this would also be an easy fix for the developers of the RDKIT Optimize Geometry node. Here’s the Python code used succesfully in the ‘Python Script node:

import os

os.environ["OMP_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"
os.environ["OPENBLAS_NUM_THREADS"] = "1"
os.environ["NUMEXPR_NUM_THREADS"] = "1"

import multiprocessing as mp

import pandas as pd
import knime.scripting.io as knio
from joblib import Parallel, delayed
from rdkit import Chem
from rdkit.Chem import AllChem


def parse_record(record):
  supplier = Chem.SDMolSupplier()
  supplier.SetData(record, sanitize=True, removeHs=False)
  return next(iter(supplier), None)


def optimise_molecule(mol, max_uff_iterations):
  mol = Chem.Mol(mol)

  # Match "Remove starting coordinates" in the KNIME RDKit Optimize Geometry node.
  mol.RemoveAllConformers()

  # Match the embedding procedure in the verified reference implementation.
  params = AllChem.ETKDG()
  params.randomSeed = 42
  if AllChem.EmbedMolecule(mol, params) != 0:
    raise ValueError("3D embedding failed")

  # A return value of 0 means that UFF converged.
  convergence_code = AllChem.UFFOptimizeMolecule(
    mol,
    maxIters=max_uff_iterations,
  )
  convered = convergence_code == 0

  ff = AllChem.UFFGetMoleculeForceField(mol)
  if ff is None:
    raise ValueError("UFF force field could not be created")

  uff_energy = float(ff.CalcEnergy())
  return mol, convered, uff_energy


def mol_to_sdf_record(mol):
  force_v3000 = mol.GetNumAtoms() > 999 or mol.GetNumBonds() > 999
  block = Chem.MolToMolBlock(
    mol,
    includeStereo=True,
    kekulize=True,
    forceV3000=force_v3000,
  ).rstrip()
  return block + "\n$$$$\n"


def process_one(task):
  row_index, record, max_uff_iterations = task
  mol = parse_record(record)
  if mol is None:
    return row_index, None, False, None

  try:
    mol, convered, uff_energy = optimise_molecule(
      mol,
      max_uff_iterations=max_uff_iterations,
    )
    return row_index, mol, convered, uff_energy
  except Exception:
    return row_index, None, False, None


def normalise_sdf_record(value):
  if value is None:
    return None
  try:
    if pd.isna(value):
      return None
  except (TypeError, ValueError):
    pass

  if isinstance(value, Chem.Mol):
    return mol_to_sdf_record(value)

  text = str(value).strip()
  if not text:
    return None
  if text.endswith("$$$$"):
    return text + "\n"
  return text + "\n$$$$\n"


def records_from_knime_table(table):
  if table.empty:
    raise RuntimeError("The KNIME input table contains no rows")

  for column in table.columns:
    records = [normalise_sdf_record(value) for value in table[column].tolist()]
    if any(record is None for record in records):
      continue
    try:
      if all(parse_record(record) is not None for record in records):
        return records
    except Exception:
      pass

  raise ValueError("No KNIME input column contains valid RDKit/SDF molecules")


def main():
  input_table = knio.input_tables[0].to_pandas()
  records = records_from_knime_table(input_table)

  logical = os.cpu_count() or 1
  workers = max(1, logical - 2)
  max_uff_iterations = 5000

  tasks = [
    (index, record, max_uff_iterations)
    for index, record in enumerate(records)
  ]

  results = Parallel(
    n_jobs=workers,
    backend="loky",
    batch_size=1,
    verbose=10,
  )(
    delayed(process_one)(task) for task in tasks
  )

  results.sort(key=lambda item: item[0])

  output = pd.DataFrame({
    "Molecule": pd.Series([item[1] for item in results], dtype=object),
    "convered": [item[2] for item in results],
    "UFF_kcal/mol": [item[3] for item in results],
  })

  knio.output_tables[0] = knio.Table.from_pandas(output)


mp.freeze_support()
main()