"""Clean up the drawing: standard bond lengths and angles, everything else kept.
Cleaning works on the canvas (KET) in place. Each molecule gets the
coordinates :func:`vibesketch.chem.clean.clean_2d` computes, with bonds of
the length Ketcher draws with; arrows, plus signs, text and every atom and
bond property stay as they were, and molecules stay where they were drawn.
Moving atoms can change what a wedge means, so the wedges are re-derived and
every molecule is checked: its canonical isomeric SMILES, read back from the
cleaned canvas the way the rest of vibe-sketch reads it, must equal the one
before. The wedges are tried in this order:
1. as drawn, on the same bonds;
2. all wedges and hashes swapped, for a layout that came out mirrored;
3. swapped only at the stereocentres whose CIP label changed;
4. placed afresh by RDKit (``WedgeMolBonds``).
A molecule for which none of them keeps the stereochemistry is left as drawn
and reported, as is anything that is not a plain structure (query atoms,
R-groups, polymer brackets).
"""
from __future__ import annotations
import copy
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from vibesketch.chem.clean import Cleaned
from vibesketch.document.derive import Derivation, derive
from vibesketch.document.ket import BOND_LENGTH, molecule_reason, molfile
from vibesketch.document.model import Canvas, Document
__all__ = ["CleanReport", "clean", "clean_ket"]
#: KET bond stereo values for a wedge (up) and a hash (down).
_WEDGE, _HASH = 1, 6
[docs]
@dataclass
class CleanReport:
"""What cleaning did."""
#: Molecules that got standard coordinates.
cleaned: int = 0
#: Why molecules were left as drawn, one line each.
left: list[str] = field(default_factory=list)
def _read(atoms: list[dict[str, Any]], bonds: list[dict[str, Any]]) -> Any:
from rdkit import Chem, rdBase
blocker = rdBase.BlockLogs()
mol = Chem.MolFromMolBlock(molfile(atoms, bonds), sanitize=True, removeHs=False)
del blocker
return mol
def _smiles(mol: Any) -> str | None:
from rdkit import Chem
return None if mol is None else str(Chem.MolToSmiles(mol))
def _cip_codes(mol: Any) -> dict[int, str]:
return {
atom.GetIdx(): str(atom.GetProp("_CIPCode"))
for atom in mol.GetAtoms()
if atom.HasProp("_CIPCode")
}
def _swapped(bonds: list[dict[str, Any]], starts: set[int] | None = None) -> list[dict[str, Any]]:
"""``bonds`` with wedges and hashes swapped, all of them or those starting at ``starts``."""
result = []
for bond in bonds:
stereo = bond.get("stereo")
if stereo in (_WEDGE, _HASH) and (starts is None or bond["atoms"][0] in starts):
bond = {**bond, "stereo": _HASH if stereo == _WEDGE else _WEDGE}
result.append(bond)
return result
def _rdkit_wedges(
mol: Any, xy: list[tuple[float, float]], bonds: list[dict[str, Any]]
) -> list[dict[str, Any]]:
"""``bonds`` with the wedges RDKit places for ``mol`` at positions ``xy``."""
from rdkit import Chem
from rdkit.Geometry import Point3D
work = Chem.Mol(mol)
conformer = work.GetConformer()
for index, (x, y) in enumerate(xy):
conformer.SetAtomPosition(index, Point3D(x, y, 0.0))
for rdkit_bond in work.GetBonds():
rdkit_bond.SetBondDir(Chem.BondDir.NONE)
Chem.WedgeMolBonds(work, conformer)
# The molfile writer puts each wedge's stereocentre first, as KET wants it.
placed: dict[frozenset[int], tuple[list[int], int]] = {}
lines = Chem.MolToMolBlock(work).splitlines()
count_atoms, count_bonds = int(lines[3][0:3]), int(lines[3][3:6])
for line in lines[4 + count_atoms : 4 + count_atoms + count_bonds]:
first, second, stereo = int(line[0:3]) - 1, int(line[3:6]) - 1, int(line[9:12])
if stereo in (_WEDGE, _HASH):
placed[frozenset((first, second))] = ([first, second], stereo)
result: list[dict[str, Any]] = []
for bond in bonds:
entry = {key: value for key, value in bond.items() if key != "stereo"}
if bond.get("stereo") not in (None, _WEDGE, _HASH):
entry["stereo"] = bond["stereo"] # "either" bonds keep their meaning
found = placed.get(frozenset(int(i) for i in bond["atoms"]))
if found is not None:
entry["atoms"], entry["stereo"] = list(found[0]), found[1]
result.append(entry)
return result
def _clean_piece(
atoms: list[dict[str, Any]], bonds: list[dict[str, Any]]
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]] | None:
"""Cleaned atoms and bonds of one connected molecule, or ``None`` to leave it."""
from vibesketch.chem.clean import clean_2d, legacy_clean_2d
before = _read(atoms, bonds)
target = _smiles(before)
if before is None or target is None:
return None
for layout in (clean_2d, legacy_clean_2d):
cleaned = layout(before, bond_length=BOND_LENGTH)
if cleaned is not None:
outcome = _verified_layout(before, target, atoms, bonds, cleaned)
if outcome is not None:
return outcome
return None
def _verified_layout(
before: Any,
target: str,
atoms: list[dict[str, Any]],
bonds: list[dict[str, Any]],
cleaned: Cleaned,
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]] | None:
"""Accept only coordinates whose re-read chemistry exactly matches input."""
moved = [
{**atom, "location": [round(x, 6), round(y, 6), 0.0]}
for atom, (x, y) in zip(atoms, cleaned.xy, strict=True)
]
def by_cip(trial: list[dict[str, Any]]) -> list[dict[str, Any]]:
after = _read(moved, trial)
if after is None:
return trial
old, new = _cip_codes(before), _cip_codes(after)
changed = {index for index, code in old.items() if new.get(index) not in (None, code)}
return _swapped(trial, changed)
attempts: list[Callable[[], list[dict[str, Any]]]] = [
lambda: bonds,
lambda: _swapped(bonds),
lambda: by_cip(bonds),
lambda: _rdkit_wedges(before, cleaned.xy, bonds),
]
for attempt in attempts:
trial = attempt()
if _smiles(_read(moved, trial)) == target:
return moved, trial
return None
def _pieces(atoms: list[dict[str, Any]], bonds: list[dict[str, Any]]) -> list[list[int]]:
"""The atom indices of each connected molecule in a KET molecule node."""
parent = list(range(len(atoms)))
def root(index: int) -> int:
while parent[index] != index:
parent[index] = parent[parent[index]]
index = parent[index]
return index
for bond in bonds:
first, second = (int(i) for i in bond["atoms"])
parent[root(first)] = root(second)
groups: dict[int, list[int]] = {}
for index in range(len(atoms)):
groups.setdefault(root(index), []).append(index)
return list(groups.values())
[docs]
def clean_ket(ket: dict[str, Any]) -> tuple[dict[str, Any], CleanReport]:
"""A copy of ``ket`` with every molecule cleaned, and what was done."""
from vibesketch.chem.clean import MAX_CLEAN_ATOMS
result = copy.deepcopy(ket)
report = CleanReport()
for node in result.get("root", {}).get("nodes", []):
reference = node.get("$ref") if isinstance(node, dict) else None
molecule = result.get(reference) if isinstance(reference, str) else None
if not isinstance(molecule, dict) or molecule.get("type", "molecule") != "molecule":
continue
atoms: list[dict[str, Any]] = list(molecule.get("atoms", []))
bonds: list[dict[str, Any]] = list(molecule.get("bonds", []))
reason = molecule_reason(molecule)
if reason is not None:
report.left.append(f"a molecule was left as drawn: {reason}")
continue
for piece in _pieces(atoms, bonds):
if len(piece) < 2:
continue
if len(piece) > MAX_CLEAN_ATOMS:
report.left.append(
f"a molecule of {len(piece):,} atoms was left as drawn: cleanup lays out "
f"at most {MAX_CLEAN_ATOMS:,}"
)
continue
local = {old: new for new, old in enumerate(piece)}
piece_bonds = [
(position, {**bond, "atoms": [local[int(i)] for i in bond["atoms"]]})
for position, bond in enumerate(bonds)
if int(bond["atoms"][0]) in local
]
outcome = _clean_piece([atoms[i] for i in piece], [b for _, b in piece_bonds])
if outcome is None:
report.left.append(
"a molecule was left as drawn: cleaning it would have changed its "
"stereochemistry, or it could not be laid out"
)
continue
new_atoms, new_bonds = outcome
for local_index, atom in enumerate(new_atoms):
atoms[piece[local_index]] = atom
for (position, _), bond in zip(piece_bonds, new_bonds, strict=True):
bonds[position] = {**bond, "atoms": [piece[int(i)] for i in bond["atoms"]]}
report.cleaned += 1
molecule["atoms"], molecule["bonds"] = atoms, bonds
return result, report
[docs]
def clean(document: Document) -> tuple[Derivation, CleanReport]:
"""``document`` with its drawing cleaned; the normalized layer is re-derived.
The structures do not change, so species keep their ids, names and
conformers.
"""
if not document.canvas.data:
return Derivation(document), CleanReport()
ket, report = clean_ket(document.canvas.data)
updated = document.model_copy(update={"canvas": Canvas(format="ket", data=ket)})
return derive(updated), report