Source code for vibesketch.chem.layout

"""2D coordinates for molecules that have none, with every coordinate finite."""

from __future__ import annotations

import math
from typing import Any

__all__ = ["compute_2d"]


def _positions(mol: Any) -> list[list[float]]:
    """Every atom's coordinates, read in one call: one Point3D per atom is far slower."""
    return list(mol.GetConformer().GetPositions().tolist())


def _finite(mol: Any) -> bool:
    return all(math.isfinite(value) for point in _positions(mol) for value in point)


def _separate_fragments(mol: Any) -> None:
    """Pack newly drawn components without changing their internal geometry.

    Depictor atom extents omit the visible hydrogen/charge labels of isolated
    ions. Leave two bond lengths between components to accommodate them.
    Existing authored conformers never pass through this layout function.
    """
    from rdkit import Chem
    from rdkit.Geometry import Point3D

    fragments = Chem.GetMolFrags(mol)
    if len(fragments) < 2:
        return
    conformer = mol.GetConformer()
    positions = _positions(mol)
    lengths = [
        math.dist(positions[b.GetBeginAtomIdx()], positions[b.GetEndAtomIdx()])
        for b in mol.GetBonds()
    ]
    bond_length = sum(lengths) / len(lengths) if lengths else 1.5
    gap = 2.0 * (bond_length if bond_length > 1e-6 else 1.5)
    cursor = 0.0
    for fragment in fragments:
        points = [conformer.GetAtomPosition(i) for i in fragment]
        xmin, xmax = min(p.x for p in points), max(p.x for p in points)
        middle = (min(p.y for p in points) + max(p.y for p in points)) / 2
        for index, point in zip(fragment, points, strict=True):
            conformer.SetAtomPosition(index, Point3D(point.x - xmin + cursor, point.y - middle, 0))
        cursor += xmax - xmin + gap




def compute_2d(mol: Any) -> None:
    """Give ``mol`` a 2D conformer in place.

    CoordGen first (it lays out rings and macrocycles better), then RDKit's
    own depictor, because CoordGen returns NaN for a molecule made only of
    hydrogen atoms (H+, H2). If both fail, the atoms go on a line. The
    process-wide CoordGen preference is restored afterwards.
    """
    from rdkit.Chem import rdDepictor
    from rdkit.Geometry import Point3D

    previous = rdDepictor.GetPreferCoordGen()
    try:
        for prefer in (True, False):
            rdDepictor.SetPreferCoordGen(prefer)
            mol.RemoveAllConformers()
            rdDepictor.Compute2DCoords(mol)
            if _finite(mol):
                _separate_fragments(mol)
                return
    finally:
        rdDepictor.SetPreferCoordGen(previous)
    conformer = mol.GetConformer()
    for index in range(mol.GetNumAtoms()):
        conformer.SetAtomPosition(index, Point3D(1.5 * index, 0.0, 0.0))
    _separate_fragments(mol)