Source code for boltz_data.mol._to._to_rdkit

from rdkit import Chem

from boltz_data.mol._mol import BZBioMol, BZMol
from boltz_data.rdkit._constants import RDKIT_BOND_TYPE_FROM_ORDER




def rdmol_from_bzmol(bzmol: BZMol | BZBioMol, /) -> Chem.RWMol:
    """Convert a BZMol or BZBioMol to an RDKit molecule."""
    rdmol = Chem.RWMol()
    atom_idx_to_rdmol_idx: list[int] = []

    conformer = Chem.Conformer(bzmol.num_atoms)

    # Check if this is a BZBioMol with residue information
    if isinstance(bzmol, BZBioMol):
        for element, name, residue, charge, coordinates in zip(
            bzmol.atom_element,
            bzmol.atom_name,
            bzmol.atom_residue,
            bzmol.atom_charge,
            bzmol.atom_coordinates,
            strict=True,
        ):
            atom = Chem.Atom(int(element))
            atom.SetFormalCharge(int(charge))
            residue_name = bzmol.residue_name[residue]
            atom.SetPDBResidueInfo(
                Chem.AtomPDBResidueInfo(atomName=name, residueName=str(residue_name), residueNumber=int(residue + 1))
            )
            atom.SetProp("name", name)
            atom_idx_to_rdmol_idx.append(rdmol.AddAtom(atom))
            conformer.SetAtomPosition(atom_idx_to_rdmol_idx[-1], coordinates.tolist())
    else:
        # Plain BZMol without residue information
        for element, charge, coordinates, name in zip(
            bzmol.atom_element, bzmol.atom_charge, bzmol.atom_coordinates, bzmol.atom_name, strict=True
        ):
            atom = Chem.Atom(int(element))
            atom.SetFormalCharge(int(charge))
            atom.SetProp("name", name)
            atom_idx_to_rdmol_idx.append(rdmol.AddAtom(atom))
            conformer.SetAtomPosition(atom_idx_to_rdmol_idx[-1], coordinates.tolist())

    for (atom1, atom2), order in zip(bzmol.bond_atoms, bzmol.bond_order, strict=True):
        rdmol.AddBond(atom_idx_to_rdmol_idx[atom1], atom_idx_to_rdmol_idx[atom2], RDKIT_BOND_TYPE_FROM_ORDER[order]) - 1

    if bzmol.atom_resolved.all():
        rdmol.AddConformer(conformer, assignId=True)

    Chem.SanitizeMol(rdmol)

    return rdmol