Source code for boltz_data.rdkit._from_ccd

from rdkit import Chem

from boltz_data.ccd import ChemicalComponent

from ._constants import RDKIT_BOND_TYPE_FROM_ORDER


def add_chemical_component_to_rdmol(
    *,
    rdmol: Chem.RWMol,
    chemical_component: ChemicalComponent,
    excluded_atom_names: set[str] | None = None,
    atom_id_to_atom_idx: dict[str, int] | None = None,
    residue_number: int | None = None,
) -> None:
    atom_id_to_idx: dict[str, int] = {}
    for atom in chemical_component.atoms.values():
        if excluded_atom_names and atom.atom_id in excluded_atom_names:
            continue
        rdkit_atom = Chem.Atom(atom.element.capitalize())

        info = Chem.AtomPDBResidueInfo()
        info.SetName(atom.atom_id)
        info.SetResidueName(chemical_component.comp_id)
        if residue_number is not None:
            info.SetResidueNumber(residue_number)
        rdkit_atom.SetPDBResidueInfo(info)

        rdkit_atom.SetFormalCharge(atom.charge)
        atom_idx = rdmol.AddAtom(rdkit_atom)
        atom_id_to_idx[atom.atom_id] = atom_idx

    for bond in chemical_component.bonds.values():
        if bond.atom_id_1 not in atom_id_to_idx or bond.atom_id_2 not in atom_id_to_idx:
            continue
        rdmol.AddBond(
            beginAtomIdx=atom_id_to_idx[bond.atom_id_1],
            endAtomIdx=atom_id_to_idx[bond.atom_id_2],
            order=RDKIT_BOND_TYPE_FROM_ORDER[bond.order],
        )

    if atom_id_to_atom_idx is not None:
        atom_id_to_atom_idx.update(atom_id_to_idx)




def rdmol_from_chemical_component(chemical_component: ChemicalComponent, /) -> Chem.Mol:
    """Create an RDKit Mol object from a ChemicalComponent."""
    rdkit_mol = Chem.RWMol()

    add_chemical_component_to_rdmol(rdmol=rdkit_mol, chemical_component=chemical_component)

    rdmol = rdkit_mol.GetMol()
    Chem.SanitizeMol(rdmol)
    return rdmol