"""Clean up a drawn structure: standard bond lengths and angles, the drawing kept.
The reference-table objective in :mod:`vibesketch.chem.depiction` refines a
standard seed with FIRE. Every seed's objective, and the gate that compares
it with the legacy result, take complex ring systems from one shared list of
reference layouts: the legacy result first, then each standard layout.
Convergence, geometry and overlap gates retain the legacy result on refusal.
The seed/fallback algorithm is described below.
A hand-drawn structure has bonds of uneven length and angles that are almost
120°. The seed layouts standardize lengths and angles, with complex-ring
exceptions, while preserving the drawing:
1. **Standard layouts.** RDKit's depictor, RDKit's depictor told to mimic the
drawing's inter-atom distances (``Compute2DCoordsMimicDistmat``), and
CoordGen each propose a layout with exact bond lengths and angles.
2. **The drawing's choices.** In two dimensions an acyclic single bond can
only be drawn one of two ways: the neighbours of its two atoms on the same
side of the bond (a U-turn) or on opposite sides (a zig-zag). For every
such bond, each layout is made to agree with the drawing by mirroring the
smaller side across the bond. A mirror keeps every length and angle, and
no other bond's choice depends on it, so the bonds are settled one at a
time and in any order.
3. **Placement.** Each layout is rotated, or mirrored, onto the drawing by
least squares (W. Kabsch, *Acta Cryst.* A32, 922-923 (1976)) and put where
the drawing was. The layout with the fewest overlapping atoms, then the
smallest deviation from the drawing, wins.
Stereochemistry is not this module's concern: chiral tags and double-bond
geometry are inputs to the layouts, and the caller re-derives the wedges and
checks the result (:mod:`vibesketch.document.clean`).
"""
from __future__ import annotations
import math
from dataclasses import dataclass
from typing import Any
import numpy as np
from numpy.typing import NDArray
__all__ = ["MAX_CLEAN_ATOMS", "Cleaned", "clean_2d", "legacy_clean_2d"]
#: Atoms closer than this, in bond lengths, overlap; standard layouts keep them apart.
_OVERLAP = 0.55
# CoordGen's macrocycle search has no interrupt/budget API and can take minutes
# for a 50-membered ring (#78). Limit that optional candidate to perceived
# rings of at most eight atoms; RDKit/mimic still handle larger rings. This is
# an execution policy, not a chemical size limit or an IUPAC requirement.
_COORDGEN_MAX_RING = 8
# The layout candidates' cost grows faster than the molecule (#68). Measured on
# zigzag chains (RDKit 2026.09, Apple M4 Pro), seconds for rdkit / mimic / coordgen:
# 500 atoms 0.10 / 0.08 / 0.13, 1,000 atoms 0.72 / 0.47 / 0.60, 2,000 atoms
# 5.2 / 1.8 / 2.1, and mimic 3.5 s at 3,000 and 10 s at 5,000. Above
# _ALL_CANDIDATES_ATOMS only the distance-matrix mimic runs: the fastest there,
# and the most faithful to the drawing (on the corpus's 2,000-atom chain all three
# tie on overlaps and deviation, and cleanup chose mimic). Above MAX_CLEAN_ATOMS a
# molecule is left as drawn, with the reason. An execution policy, not a chemical
# limit; the maintainer judged the 2,000-atom cleanup over budget (2026-10-22).
_ALL_CANDIDATES_ATOMS = 500
#: The most atoms a molecule may have and still be cleaned.
MAX_CLEAN_ATOMS = 3_000
_Points = NDArray[np.float64]
[docs]
@dataclass(frozen=True)
class Cleaned:
"""New 2D positions for a molecule's atoms, in its atom order."""
xy: list[tuple[float, float]]
#: Root-mean-square distance from the drawing, in bond lengths, after alignment.
deviation: float
#: Seed name, or ``fire-v1:<seed>`` for an accepted converged refinement.
layout: str
def _points(mol: Any) -> _Points:
conformer = mol.GetConformer()
return np.array(
[
[conformer.GetAtomPosition(i).x, conformer.GetAtomPosition(i).y]
for i in range(mol.GetNumAtoms())
],
dtype=np.float64,
)
def _bond_lengths(mol: Any, points: _Points) -> list[float]:
return [
float(np.linalg.norm(points[b.GetBeginAtomIdx()] - points[b.GetEndAtomIdx()]))
for b in mol.GetBonds()
]
def _normalized(mol: Any, points: _Points) -> _Points | None:
"""``points`` scaled to a median bond length of one; ``None`` if degenerate."""
if not np.isfinite(points).all():
return None
lengths = sorted(_bond_lengths(mol, points))
median = lengths[len(lengths) // 2] if lengths else 0.0
if median < 1e-6:
return None
return points / median
def _layouts(mol: Any, drawn: _Points) -> dict[str, _Points]:
"""Standard layouts of ``mol``, each with bond length one."""
from rdkit import Chem, rdBase
from rdkit.Chem import rdCoordGen, rdDepictor
from vibesketch.chem._rdkit import untyped
depictor = untyped(rdDepictor)
count = mol.GetNumAtoms()
# RDKit lays out at 1.5 per bond; the distances to mimic use the same scale.
# The lower triangle row by row, in RDKit's order, without an n x n x 2 temporary.
distances = (
np.concatenate([np.linalg.norm(drawn[:k] - drawn[k], axis=1) for k in range(1, count)])
* 1.5
if count > 1
else np.zeros(0)
)
found: dict[str, _Points] = {}
candidates = ["rdkit", "mimic"]
if all(len(ring) <= _COORDGEN_MAX_RING for ring in mol.GetRingInfo().AtomRings()):
candidates.append("coordgen")
if count > _ALL_CANDIDATES_ATOMS:
candidates = ["mimic"]
blocker = rdBase.BlockLogs()
try:
for name in candidates:
work = Chem.Mol(mol)
try:
if name == "rdkit":
depictor.Compute2DCoords(work, forceRDKit=True)
elif name == "mimic":
depictor.Compute2DCoordsMimicDistmat(
work, distances, weightDistMat=0.5, sampleSeed=7, forceRDKit=True
)
else:
rdCoordGen.AddCoords(work)
except (RuntimeError, ValueError):
continue
points = _normalized(work, _points(work))
if points is not None:
found[name] = points
finally:
del blocker
return found
def _mirror_sides(mol: Any) -> list[tuple[int, int, list[int]]]:
"""Per acyclic single bond between non-terminal atoms: ``(i, j, atoms on j's side)``.
``j`` is the end with the smaller side, so the mirror moves fewer atoms.
"""
from rdkit import Chem
found: list[tuple[int, int, list[int]]] = []
for bond in mol.GetBonds():
if bond.GetBondType() != Chem.BondType.SINGLE or bond.IsInRing():
continue
i, j = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()
if mol.GetAtomWithIdx(i).GetDegree() < 2 or mol.GetAtomWithIdx(j).GetDegree() < 2:
continue
cut = Chem.RWMol(mol)
cut.RemoveBond(i, j)
pieces = Chem.GetMolFrags(cut)
side = next(piece for piece in pieces if j in piece)
if 2 * len(side) > mol.GetNumAtoms():
i, j = j, i
side = next(piece for piece in pieces if j in piece)
found.append((i, j, list(side)))
return found
def _cross(origin: _Points, a: _Points, b: _Points) -> float:
return float((a[0] - origin[0]) * (b[1] - origin[1]) - (a[1] - origin[1]) * (b[0] - origin[0]))
def _agreement(mol: Any, drawn: _Points, layout: _Points, i: int, j: int) -> float:
"""How far the layout agrees with the drawing on which side of bond i-j the neighbours sit.
Positive when they agree, negative when the layout needs the mirror.
"""
total = 0.0
for left in mol.GetAtomWithIdx(i).GetNeighbors():
a = left.GetIdx()
if a == j:
continue
for right in mol.GetAtomWithIdx(j).GetNeighbors():
b = right.GetIdx()
if b == i:
continue
seen = _cross(drawn[i], drawn[j], drawn[a]) * _cross(drawn[i], drawn[j], drawn[b])
laid = _cross(layout[i], layout[j], layout[a]) * _cross(layout[i], layout[j], layout[b])
if abs(seen) > 1e-3 and abs(laid) > 1e-3:
total += math.copysign(1.0, seen) * math.copysign(1.0, laid)
return total
def _mirrored(points: _Points, i: int, j: int, side: list[int]) -> _Points:
"""``points`` with ``side`` mirrored across the line through atoms i and j."""
result = points.copy()
direction = (points[j] - points[i]) / np.linalg.norm(points[j] - points[i])
offsets = result[side] - points[i]
result[side] = points[i] + 2 * (offsets @ direction)[:, None] * direction - offsets
return result
def _aligned(layout: _Points, drawn: _Points) -> _Points:
"""``layout`` rotated or mirrored onto ``drawn`` by least squares (Kabsch 1976)."""
layout_centre, drawn_centre = layout.mean(axis=0), drawn.mean(axis=0)
covariance = (layout - layout_centre).T @ (drawn - drawn_centre)
u, _, vt = np.linalg.svd(covariance)
rotation = vt.T @ u.T # a reflection when its determinant is -1, which 2D allows
aligned: _Points = (layout - layout_centre) @ rotation.T + drawn_centre
return aligned
def _overlaps(mol: Any, points: _Points) -> int:
"""Pairs of atoms closer than :data:`_OVERLAP` that no bond joins.
Row by row, so memory stays linear in the atoms rather than an n x n
matrix (about 130 MB at 2,000 atoms, #68).
"""
bonded = {
(max(b.GetBeginAtomIdx(), b.GetEndAtomIdx()), min(b.GetBeginAtomIdx(), b.GetEndAtomIdx()))
for b in mol.GetBonds()
}
total = 0
for k in range(1, len(points)):
close = np.flatnonzero(np.linalg.norm(points[:k] - points[k], axis=1) < _OVERLAP)
total += sum(1 for j in close if (k, int(j)) not in bonded)
return total
[docs]
def legacy_clean_2d(mol: Any, *, bond_length: float = 1.0) -> Cleaned | None:
"""Standard coordinates for ``mol`` that follow its drawn 2D conformer.
``mol`` is one connected molecule with its drawing as conformer 0 and its
stereochemistry assigned. The result targets bonds ``bond_length`` long
and sits where the drawing was. ``None`` when there is nothing to clean
(a lone atom) or no layout could be made.
"""
return _legacy(mol, bond_length)[0]
def _legacy(mol: Any, bond_length: float) -> tuple[Cleaned | None, dict[str, _Points]]:
""":func:`legacy_clean_2d`'s result, and every layout it placed on the drawing."""
if mol.GetNumAtoms() < 2 or mol.GetNumBonds() == 0 or mol.GetNumAtoms() > MAX_CLEAN_ATOMS:
return None, {}
original = _points(mol)
drawn = _normalized(mol, original)
if drawn is None:
return None, {}
sides = _mirror_sides(mol)
best: tuple[int, float, str, _Points] | None = None
placed_layouts: dict[str, _Points] = {}
for name, layout in _layouts(mol, drawn).items():
for i, j, side in sides:
if _agreement(mol, drawn, layout, i, j) < 0:
layout = _mirrored(layout, i, j, side)
aligned = _aligned(layout, drawn)
placed_layouts[name] = aligned
deviation = float(np.sqrt(((aligned - drawn) ** 2).sum(axis=1).mean()))
rank = (_overlaps(mol, layout), deviation, name, aligned)
if best is None or rank[:2] < best[:2]:
best = rank
if best is None:
return None, placed_layouts
_, deviation, name, aligned = best
placed = (aligned - drawn.mean(axis=0)) * bond_length + original.mean(axis=0)
return Cleaned([(float(x), float(y)) for x, y in placed], deviation, name), placed_layouts
[docs]
def clean_2d(mol: Any, *, bond_length: float = 1.0) -> Cleaned | None:
"""Reference-table/FIRE cleanup, with the legacy depictor as a fallback.
Only converged, non-regressing refinements are accepted. Complex ring
seeds may come from RDKit's public template API; no template data or
global template settings are changed. The document caller verifies
stereo and retries the legacy layout on refusal.
"""
from rdkit import Chem, rdBase
from rdkit.Chem import rdDepictor
from vibesketch.chem._rdkit import untyped
from vibesketch.chem.depiction import DepictionField, _table
if not math.isfinite(bond_length) or bond_length <= 0:
raise ValueError("drawing bond length must be finite and positive")
fallback, placed_layouts = _legacy(mol, bond_length)
if fallback is None or mol.GetNumAtoms() > _table()["numerics"]["max_atoms"]:
return fallback
original = _points(mol)
drawn = _normalized(mol, original)
if drawn is None:
return fallback
baseline = np.asarray(fallback.xy) / bond_length
baseline -= baseline.mean(axis=0)
baseline += drawn.mean(axis=0)
seeds = {fallback.layout: baseline}
sides = _mirror_sides(mol)
if mol.GetRingInfo().NumRings():
work = Chem.Mol(mol)
blocker = rdBase.BlockLogs()
try:
untyped(rdDepictor).Compute2DCoords(work, forceRDKit=True, useRingTemplates=True)
template = _normalized(work, _points(work))
if template is not None:
for i, j, side in sides:
if _agreement(mol, drawn, template, i, j) < 0:
template = _mirrored(template, i, j, side)
seeds["rdkit-templates"] = _aligned(template, drawn)
except (RuntimeError, ValueError):
pass
finally:
del blocker
# One list of reference layouts for every field, the baseline's first, so
# each ring system gets the same targets whichever seed is refined.
shapes = [baseline, *placed_layouts.values()] + [
layout for name, layout in seeds.items() if name != fallback.layout
]
baseline_field = DepictionField.build(mol, baseline, shapes)
best = fallback
best_rank: tuple[int, float, float] | None = None
for name, seed in seeds.items():
field = DepictionField.build(mol, seed, shapes)
# A different template must also pass the common quality gate before
# spending iterations on it. All comparisons use fixed targets.
if not field.acceptable(seed, baseline):
continue
try:
refined = field.refine()
except (ValueError, RuntimeError, FloatingPointError):
continue
if refined is None or refined.shape != baseline.shape or not np.isfinite(refined).all():
continue
aligned = _aligned(refined, drawn)
if not field.acceptable(aligned, baseline) or not baseline_field.acceptable(
aligned, baseline
):
continue
if any(_agreement(mol, baseline, aligned, i, j) < 0 for i, j, _ in sides):
continue
deviation = float(np.sqrt(((aligned - drawn) ** 2).sum(axis=1).mean()))
metrics = field.metrics(aligned)
rank = (metrics[2], metrics[3], deviation)
if best_rank is not None and rank >= best_rank:
continue
placed = (aligned - drawn.mean(axis=0)) * bond_length + original.mean(axis=0)
best = Cleaned([(float(x), float(y)) for x, y in placed], deviation, f"fire-v1:{name}")
best_rank = rank
return best