refactor: code simplification pass on nvMolKit conformer worker

- Extract _rdkit_generate_chunk() to eliminate 3x duplicated blocks
- Consolidate max_failures/max_iterations into single max_attempts
- Extract _conf_has_finite_coords predicate in worker CLI
- Compact setup script: str.removeprefix(), f-strings, shorter docstrings
- Remove 8 obvious/redundant comments

Author: SergeiNikolenko

matcha/utils/preprocessing.py

@@ -383,11 +383,10 @@ def generate_multiple_conformers(orig_mol, num_conformers):
     mol = copy.deepcopy(orig_mol)
     ps = AllChem.ETKDGv3()
     failures, ids = 0, []
-    max_failures = 3
-    max_iterations = max_failures  # Prevent infinite loops
+    max_attempts = 3
 
     iteration = 0
-    while mol.GetNumConformers() < num_conformers and iteration < max_iterations:
+    while mol.GetNumConformers() < num_conformers and iteration < max_attempts:
         current_count = mol.GetNumConformers()
         needed = num_conformers - current_count
 
@@ -405,20 +404,18 @@ def generate_multiple_conformers(orig_mol, num_conformers):
 
         ids = [id for id in ids if id != -1]
 
-        # Manually add each new conformer to the main molecule
         added_count = 0
         for conf_id in ids:
-                conf = temp_mol.GetConformer(conf_id)
-                mol.AddConformer(conf, assignId=True)
-                added_count += 1
-        
+            conf = temp_mol.GetConformer(conf_id)
+            mol.AddConformer(conf, assignId=True)
+            added_count += 1
+
         new_count = mol.GetNumConformers()
 
         if added_count == 0:
-            # No new conformers were added
-            logger.debug(f"No new conformers added. Retrying {iteration + 1}/{max_iterations}")
+            logger.debug(f"No new conformers added. Retrying {iteration + 1}/{max_attempts}")
             failures += 1
-            if failures >= max_failures:
+            if failures >= max_attempts:
                 break
         else:
             # Successfully added some conformers, reset failure counter
@@ -452,9 +449,9 @@ def generate_multiple_conformers(orig_mol, num_conformers):
     if mol.GetNumConformers() == 0:
         logger.warning("No conformers generated, using original molecule")
         return orig_mol
-    else:
-        logger.debug(f"Generated {mol.GetNumConformers()} conformers")
-        return mol
+
+    logger.debug(f"Generated {mol.GetNumConformers()} conformers")
+    return mol
 
 
 def generate_conformer_mols(orig_mol, num_conformers, backend: str | None = None):
@@ -501,7 +498,6 @@ def generate_conformer_mols_batch(
 
     results = [None] * len(mols)
 
-    # Prepare molecules: keep originals intact for safety.
     prepared = []
     for mol in mols:
         init = copy.deepcopy(mol)
@@ -520,6 +516,17 @@ def generate_conformer_mols_batch(
     use_worker = backend == "worker" or (backend == "auto" and worker_cmd_configured)
     worker_disabled = False
 
+    def _rdkit_generate_chunk(chunk_mols):
+        result = []
+        for m in chunk_mols:
+            m_with_confs = generate_multiple_conformers(m, int(confs_per_mol))
+            m_with_confs = _remove_hs_safe(m_with_confs)
+            result.append(_split_single_conformer_mols(m_with_confs, int(confs_per_mol)))
+        if "rdkit" not in _CONFORMER_BACKEND_LOGGED:
+            logger.info("Conformer backend: RDKit")
+            _CONFORMER_BACKEND_LOGGED.add("rdkit")
+        return result
+
     for start in range(0, len(prepared), chunk_size):
         end = min(len(prepared), start + chunk_size)
         chunk = prepared[start:end]
@@ -538,29 +545,13 @@ def generate_conformer_mols_batch(
                     logger.info("Conformer backend: worker")
                     _CONFORMER_BACKEND_LOGGED.add("worker")
             else:
-                chunk_results = []
-                for mol in chunk:
-                    mol_with_confs = generate_multiple_conformers(mol, int(confs_per_mol))
-                    mol_with_confs = _remove_hs_safe(mol_with_confs)
-                    chunk_results.append(
-                        _split_single_conformer_mols(mol_with_confs, int(confs_per_mol))
-                    )
-                if "rdkit" not in _CONFORMER_BACKEND_LOGGED:
-                    logger.info("Conformer backend: RDKit")
-                    _CONFORMER_BACKEND_LOGGED.add("rdkit")
+                chunk_results = _rdkit_generate_chunk(chunk)
         except Exception as exc:
             if backend == "worker":
                 raise RuntimeError(f"Worker conformer generation failed: {exc}") from exc
             worker_disabled = True
             logger.warning(f"Worker conformer backend failed, falling back to RDKit: {exc}")
-            chunk_results = []
-            for mol in chunk:
-                mol_with_confs = generate_multiple_conformers(mol, int(confs_per_mol))
-                mol_with_confs = _remove_hs_safe(mol_with_confs)
-                chunk_results.append(_split_single_conformer_mols(mol_with_confs, int(confs_per_mol)))
-            if "rdkit" not in _CONFORMER_BACKEND_LOGGED:
-                logger.info("Conformer backend: RDKit")
-                _CONFORMER_BACKEND_LOGGED.add("rdkit")
+            chunk_results = _rdkit_generate_chunk(chunk)
 
         for local_idx, conformer_mols in enumerate(chunk_results):
             processed = []
@@ -582,7 +573,6 @@ def generate_conformer_mols_batch(
                 processed = _split_single_conformer_mols(fallback, 1)
             results[start + local_idx] = processed
 
-    # Satisfy type checker: results is fully populated.
     return [r if r is not None else [copy.deepcopy(mols[i])] for i, r in enumerate(results)]
 
 
@@ -615,8 +605,7 @@ def safe_index(items, element):
 
 
 def parse_receptor(pdbid, pdbbind_dir, dataset_type):
-    rec = parsePDB(pdbid, pdbbind_dir, dataset_type)
-    return rec
+    return parsePDB(pdbid, pdbbind_dir, dataset_type)
 
 
 def parsePDB(pdbid, pdbbind_dir, dataset_type):

packages/matcha_nvmolkit_worker/src/matcha_nvmolkit_worker/cli.py

@@ -119,7 +119,6 @@ def _read_input_molecules(input_sdf: Path):
     molecules = []
     names = []
     valid_mask = []
-    invalid_count = 0
     for idx, mol in enumerate(supplier):
         if mol is None:
             continue
@@ -133,41 +132,39 @@ def _read_input_molecules(input_sdf: Path):
             current.RemoveAllConformers()
             valid_mask.append(True)
         except Exception:
-            invalid_count += 1
             valid_mask.append(False)
         molecules.append(current)
         names.append(uid)
+    invalid_count = valid_mask.count(False)
     return molecules, names, valid_mask, invalid_count
 
 
 def _ensure_minimum_conformers(mol, confs_per_mol: int, seed: int | None):
-    target = max(1, int(confs_per_mol))
-    if mol.GetNumConformers() >= target:
+    if mol.GetNumConformers() >= max(1, int(confs_per_mol)):
+        return False
+    if mol.GetNumConformers() > 0:
         return False
 
-    before = mol.GetNumConformers()
+    conf = Chem.Conformer(mol.GetNumAtoms())
+    for atom_idx in range(mol.GetNumAtoms()):
+        conf.SetAtomPosition(atom_idx, Point3D(0.0, 0.0, 0.0))
+    mol.AddConformer(conf, assignId=True)
+    return True
 
-    if mol.GetNumConformers() == 0:
-        conf = Chem.Conformer(mol.GetNumAtoms())
-        for atom_idx in range(mol.GetNumAtoms()):
-            conf.SetAtomPosition(atom_idx, Point3D(0.0, 0.0, 0.0))
-        mol.AddConformer(conf, assignId=True)
 
-    return mol.GetNumConformers() != before
+def _conf_has_finite_coords(conf) -> bool:
+    for atom_idx in range(conf.GetNumAtoms()):
+        pos = conf.GetAtomPosition(atom_idx)
+        if not (math.isfinite(pos.x) and math.isfinite(pos.y) and math.isfinite(pos.z)):
+            return False
+    return True
 
 
 def _drop_non_finite_conformers(mol):
-    bad_ids = []
-    for cid in range(mol.GetNumConformers()):
-        conf = mol.GetConformer(cid)
-        ok = True
-        for atom_idx in range(conf.GetNumAtoms()):
-            pos = conf.GetAtomPosition(atom_idx)
-            if not (math.isfinite(pos.x) and math.isfinite(pos.y) and math.isfinite(pos.z)):
-                ok = False
-                break
-        if not ok:
-            bad_ids.append(cid)
+    bad_ids = [
+        cid for cid in range(mol.GetNumConformers())
+        if not _conf_has_finite_coords(mol.GetConformer(cid))
+    ]
     for cid in reversed(bad_ids):
         mol.RemoveConformer(cid)
     return len(bad_ids)