From 0834f32c4fde257d6eedb882cb81f3374a434e5e Mon Sep 17 00:00:00 2001 From: vandan revanur Date: Sat, 18 Jul 2026 00:18:15 +0200 Subject: [PATCH] fix: Improve handlng of heavy-atom-only inputs in RDKit conversion Fixes: Issue #6 [Problem] - Heavy-atom-only inputs (e.g. from PDB files) were causing failures in bond assignment. - Previous fallback methods returned disconnected atoms, leading to incorrect SMILES representations. [Solution] - Implemented a new method to create RDKit molecules from heavy-atom-only coordinates. - Added a retry mechanism using `DetermineConnectivity` when bond determination fails. --- steamroll/steamroll.py | 81 +++++++++++++++++++++++++++++++++++++++++ tests/test_steamroll.py | 76 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 157 insertions(+) diff --git a/steamroll/steamroll.py b/steamroll/steamroll.py index e93895b..e52869e 100644 --- a/steamroll/steamroll.py +++ b/steamroll/steamroll.py @@ -92,6 +92,74 @@ def _write_temp_xyz(atomic_numbers: list[int], coordinates: list[list[float]]) - return f.name +def _build_rwmol( + atomic_numbers: list[int], + coordinates: list[list[float]], +) -> Chem.RWMol: + """Create an ``RWMol`` with a conformer from atomic numbers and coordinates.""" + n = len(atomic_numbers) + xyz_pos = np.array(coordinates) + rwmol = Chem.RWMol() + conf = Chem.Conformer(n) + for i, z in enumerate(atomic_numbers): + rwmol.AddAtom(Chem.Atom(z)) + conf.SetAtomPosition(i, Point3D(*xyz_pos[i].tolist())) + rwmol.AddConformer(conf, assignId=True) + return rwmol + + +def _from_coords_rdkit( + atomic_numbers: list[int], + coordinates: list[list[float]], + charge: int = 0, +) -> Chem.rdchem.Mol: + """Build an RDKit mol from 3D coordinates using ``rdDetermineBonds``. + + First attempts ``DetermineBonds`` (Hückel theory, includes bond orders). + Heavy-atom-only inputs (e.g. from PDB files, no explicit H) often cause + ``DetermineBonds`` to raise a charge-balance error, so on failure the + method retries with ``DetermineConnectivity``, which assigns correct + topology (single bonds only) without requiring the valence balance + constraint to hold. Either result is far better than the obabel fallback + that returns atoms with no bonds at all. + + Args: + atomic_numbers: atomic numbers for each atom. + coordinates: Cartesian coordinates for each atom, in Å. + charge: total molecular charge. + + Returns: + RDKit molecule; bond orders present when ``DetermineBonds`` succeeded, + otherwise single-bond connectivity from ``DetermineConnectivity``. + + Raises: + Exception: if both methods fail. + """ + rwmol = _build_rwmol(atomic_numbers, coordinates) + try: + rdDetermineBonds.DetermineBonds(rwmol, charge=charge) + logger.debug("rdDetermineBonds.DetermineBonds succeeded") + return rwmol.GetMol() + except Exception as exc: + logger.debug( + f"DetermineBonds failed ({exc}); retrying with DetermineConnectivity" + ) + + # DetermineBonds raises when it cannot satisfy the charge constraint + # (common for heavy-atom-only inputs). DetermineConnectivity does not + # apply that constraint, so it reliably gives the correct connectivity. + rwmol2 = _build_rwmol(atomic_numbers, coordinates) + rdDetermineBonds.DetermineConnectivity(rwmol2) + # DetermineConnectivity sets noImplicit=True on every atom so that the + # conformer is taken as-is; clear the flag before sanitizing so that + # implicit H counts are computed correctly from valence. + for atom in rwmol2.GetAtoms(): + atom.SetNoImplicit(False) + Chem.SanitizeMol(rwmol2) + logger.debug("rdDetermineBonds.DetermineConnectivity succeeded") + return rwmol2.GetMol() + + def _from_smiles_and_coords( smiles: str, atomic_numbers: list[int], @@ -343,6 +411,19 @@ def _topology_ok(mol: Chem.rdchem.Mol) -> bool: "provide a SMILES string or fix the geometry" ) + # rdDetermineBonds fallback — handles heavy-atom-only inputs (e.g. from PDB files) + # where xyz2mol fails because it relies on explicit H atoms for bond-order assignment. + if rdkm is None: + try: + candidate = _from_coords_rdkit(atomic_numbers, coords, charge=charge) + if _topology_ok(candidate): + rdkm = candidate + logger.debug("rdDetermineBonds succeeded") + else: + logger.debug("rdDetermineBonds produced wrong topology, trying obabel") + except Exception as exc: + logger.debug(f"rdDetermineBonds failed, trying obabel: {exc}") + if rdkm is None: # Geometry-only fallback via obabel — no bond orders, last resort. try: diff --git a/tests/test_steamroll.py b/tests/test_steamroll.py index 04b213d..2ba87c9 100644 --- a/tests/test_steamroll.py +++ b/tests/test_steamroll.py @@ -332,6 +332,82 @@ def test_from_smiles_and_coords_charges_preserved() -> None: assert -1 in charges_by_element[8], "one O should have -1 formal charge" +_BZU_PDB_BLOCK = """\ +HETATM 2 C9 BZU A 555 -3.394 5.205 13.019 1.00 0.00 C +HETATM 3 C10 BZU A 555 -3.055 3.867 13.710 1.00 0.00 C +HETATM 4 C11 BZU A 555 -3.349 7.408 13.935 1.00 0.00 C +HETATM 5 C12 BZU A 555 -2.076 7.426 14.764 1.00 0.00 C +HETATM 6 C14 BZU A 555 -1.915 9.769 15.059 1.00 0.00 C +HETATM 7 C15 BZU A 555 -0.878 10.903 15.178 1.00 0.00 C +HETATM 8 O1A BZU A 555 -5.425 0.994 18.287 1.00 0.00 O +HETATM 9 O2A BZU A 555 -7.398 0.817 16.995 1.00 0.00 O +HETATM 10 N21 BZU A 555 -5.486 -0.545 16.461 1.00 0.00 N +HETATM 11 S1 BZU A 555 -5.966 0.750 16.997 1.00 0.00 S +HETATM 12 C4 BZU A 555 -4.321 2.152 15.292 1.00 0.00 C +HETATM 13 C5 BZU A 555 -4.165 3.345 14.642 1.00 0.00 C +HETATM 14 C6 BZU A 555 -5.216 4.200 14.931 1.00 0.00 C +HETATM 15 S2 BZU A 555 -6.379 3.501 15.979 1.00 0.00 S +HETATM 16 S7 BZU A 555 -5.295 5.827 14.368 1.00 0.00 S +HETATM 17 O3B BZU A 555 -6.160 5.873 13.244 1.00 0.00 O +HETATM 18 O4B BZU A 555 -5.858 6.680 15.373 1.00 0.00 O +HETATM 19 N8 BZU A 555 -3.848 6.065 14.090 1.00 0.00 N +HETATM 20 N16 BZU A 555 -2.776 2.895 12.649 1.00 0.00 N +HETATM 21 O13 BZU A 555 -1.307 8.596 14.533 1.00 0.00 O +HETATM 22 C17 BZU A 555 -1.564 3.011 11.841 1.00 0.00 C +HETATM 23 C18 BZU A 555 -1.463 1.927 10.754 1.00 0.00 C +""" + + +def test_pdb_heavy_atom_only_gets_bonds() -> None: + """Regression test: heavy-atom-only input (e.g. from a PDB file) must yield + a connected molecule with proper bond orders, not isolated atoms. + + Before the fix, xyz2mol would silently fail on inputs with no explicit + hydrogens and the obabel fallback would return atoms with no bonds, giving + a SMILES like ``C.C.C.C...``. After the fix, rdDetermineBonds.DetermineBonds + handles the heavy-atom-only case and assigns connectivity + bond orders. + """ + pdb_mol = Chem.MolFromPDBBlock(_BZU_PDB_BLOCK, removeHs=False) + assert pdb_mol is not None, "RDKit could not parse the PDB block" + + crds = pdb_mol.GetConformer(0).GetPositions().tolist() + atomic_nums = [atom.GetAtomicNum() for atom in pdb_mol.GetAtoms()] + + # Confirm the input really is heavy-atom-only (no hydrogens) + assert 1 not in atomic_nums, "Expected no explicit H atoms in this PDB input" + + rd_mol = to_rdkit(atomic_nums, crds, charge=0) + + # All 22 heavy atoms must be present + assert rd_mol.GetNumAtoms() == 22 + + # Every atom must have at least one bond — the core regression check + isolated = [ + atom.GetIdx() + for atom in rd_mol.GetAtoms() + if atom.GetDegree() == 0 + ] + assert isolated == [], f"Atoms with no bonds (isolated): {isolated}" + + # The SMILES must not be the all-disconnected pattern produced by the bug + smiles = Chem.MolToSmiles(rd_mol) + assert "." not in smiles or smiles.count(".") < rd_mol.GetNumAtoms() - 1, ( + f"Molecule looks disconnected: {smiles}" + ) + + # Canonical heavy-atom SMILES must match the topology reported by RDKit's own + # PDB parser. We compare non-isomeric SMILES after stripping any remaining + # explicit Hs so that bookkeeping differences (e.g. "[SH]" vs "S" with an + # implicit H) do not cause spurious mismatches. + ref_smiles = Chem.MolToSmiles( + Chem.RemoveHs(Chem.RWMol(pdb_mol).GetMol()), isomericSmiles=False + ) + got_smiles = Chem.MolToSmiles(rd_mol, isomericSmiles=False) + assert got_smiles == ref_smiles, ( + f"SMILES mismatch:\n expected (PDB parser): {ref_smiles}\n got (steamroll): {got_smiles}" + ) + + def test_tmc_conformer_preserved() -> None: """to_rdkit preserves 3D coordinates for transition metal complexes.