Source code for boltz_data.mol._op._conformer

import numpy as np
from rdkit.Chem import rdDepictor

from boltz_data.mol._mol import BZMol
from boltz_data.mol._to import rdmol_from_bzmol
from boltz_data.rdkit._conformer import calculate_conformer

from ._transform import transform_bzmol



[docs]
def generate_conformer(bzmol: BZMol, /, *, seed: int) -> BZMol:
    """Generate a 3D conformer for a BZMol using RDKit."""
    rdmol = rdmol_from_bzmol(bzmol)

    rdmol = calculate_conformer(rdmol, seed=seed)

    coordinates = np.array(rdmol.GetConformer(-1).GetPositions())[: bzmol.num_atoms]

    return transform_bzmol(
        bzmol=bzmol, atom_coordinates=coordinates, atom_resolved=np.ones(bzmol.num_atoms, dtype=bool)
    )




[docs]
def generate_depiction(bzmol: BZMol, /, *, match_3d: bool = True) -> BZMol:
    """Generate 2D coordinates for a BZMol using RDKit."""
    rdmol = rdmol_from_bzmol(bzmol)

    rdDepictor.SetPreferCoordGen(val=True)

    if match_3d and bzmol.atom_resolved.any():
        rdDepictor.GenerateDepictionMatching3DStructure(rdmol, rdmol)
    else:
        rdDepictor.Compute2DCoords(rdmol)

    rdDepictor.NormalizeDepiction(rdmol)
    rdDepictor.StraightenDepiction(rdmol, minimizeRotation=True)

    coordinates = np.array([rdmol.GetConformer(-1).GetAtomPosition(i) for i in range(bzmol.num_atoms)])

    return transform_bzmol(
        bzmol=bzmol, atom_coordinates=coordinates, atom_resolved=np.ones(bzmol.num_atoms, dtype=bool)
    )