"""Atoms and bonds marked in colour, on the canvas and in figures (#91).
The canvas engine keeps no highlight in its drawing, so the document does:
a :class:`~vibesketch.document.model.Highlight` names a species, a colour
and the canonical ranks of its marked heavy atoms (RDKit's
``CanonicalRankAtoms`` with ties broken), which every drawing of the species
maps back to its own atoms. Bonds between two marked atoms of one highlight
are marked too, so marking a ring's atoms marks the ring. On a symmetric
molecule a mark may land on an equivalent atom when the structure is drawn
anew, since such atoms are told apart only by their order.
The colours are Okabe and Ito's palette for colour-vision deficiency
("Color Universal Design", 2008, https://jfly.uni-koeln.de/color/), printed
as an opaque 40 % tint on white so the line art stays legible in every
journal preset and in SVG, PNG and PDF alike.
"""
from __future__ import annotations
from collections.abc import Sequence
from typing import Any, cast
from vibesketch.document.model import Document, Highlight, HighlightColour
__all__ = [
"PALETTE",
"HighlightRefused",
"highlight_points",
"marks_for",
"tint",
"with_highlight",
"with_selection_highlighted",
"without_highlights",
]
#: Okabe and Ito's colours by name.
PALETTE: dict[str, str] = {
"orange": "#E69F00",
"sky-blue": "#56B4E9",
"green": "#009E73",
"blue": "#0072B2",
"vermillion": "#D55E00",
"purple": "#CC79A7",
}
#: Share of the colour in the printed tint; the rest is white.
_STRENGTH = 0.4
#: How far, in bond lengths, a selected point may be from the atom it names.
_TOLERANCE = 0.05
[docs]
class HighlightRefused(ValueError):
"""A highlight that cannot be made as asked."""
[docs]
def tint(colour: str) -> tuple[float, float, float]:
"""The printed fill of ``colour``: an opaque tint, as RDKit takes colours (0 to 1)."""
value = PALETTE[colour].lstrip("#")
rgb = (int(value[i : i + 2], 16) / 255 for i in (0, 2, 4))
red, green, blue = (_STRENGTH * c + (1 - _STRENGTH) for c in rgb)
return (red, green, blue)
def _ranks(mol: Any) -> dict[int, int]:
"""Each heavy atom's canonical rank among the heavy atoms, by its index in ``mol``."""
from rdkit import Chem
heavy = Chem.Mol(mol)
for atom in heavy.GetAtoms():
atom.SetIntProp("vibesketch_index", atom.GetIdx())
heavy = Chem.RemoveHs(heavy, sanitize=False)
ranks = Chem.CanonicalRankAtoms(heavy, breakTies=True)
return {
atom.GetIntProp("vibesketch_index"): int(rank)
for atom, rank in zip(heavy.GetAtoms(), ranks, strict=True)
}
[docs]
def marks_for(document: Document, species: str, mol: Any) -> list[tuple[list[int], str]]:
"""The species' highlights as atom indices of ``mol``, a drawing of it, with colours."""
wanted = [h for h in document.highlights if h.species == species]
if not wanted:
return []
by_rank = {rank: index for index, rank in _ranks(mol).items()}
return [
(sorted(by_rank[r] for r in h.atoms if r in by_rank), h.colour)
for h in wanted
if any(r in by_rank for r in h.atoms)
]
def _check_colour(colour: str) -> HighlightColour:
if colour not in PALETTE:
raise HighlightRefused(f"colour must be one of {', '.join(PALETTE)}, not {colour!r}")
return cast(HighlightColour, colour)
def _merged(document: Document, species: str, ranks: set[int], colour: HighlightColour) -> Document:
"""``document`` with ``ranks`` of ``species`` in ``colour``, taken out of any other colour."""
kept: list[Highlight] = []
for highlight in document.highlights:
if highlight.species != species:
kept.append(highlight)
continue
atoms = set(highlight.atoms) - ranks
if highlight.colour == colour:
ranks |= atoms
elif atoms:
kept.append(highlight.model_copy(update={"atoms": sorted(atoms)}))
kept.append(Highlight(species=species, atoms=sorted(ranks), colour=colour))
return document.model_copy(update={"highlights": kept})
[docs]
def with_highlight(
document: Document, species: str, atoms: Sequence[int], colour: str = "orange"
) -> Document:
"""``document`` with ``atoms`` of ``species`` marked, numbered from 0 as in its report."""
from vibesketch.document.formats import species_mols
chosen = _check_colour(colour)
mol = species_mols(document).get(species)
if mol is None:
raise HighlightRefused(f"the document has no drawn species {species!r}")
ranks = _ranks(mol)
unknown = [a for a in atoms if a not in ranks]
if unknown or not atoms:
count = mol.GetNumAtoms()
raise HighlightRefused(
f"mark heavy atoms of {species} by number, 1 to {count}; not "
+ ", ".join(str(a + 1) for a in unknown or atoms)
)
return _merged(document, species, {ranks[a] for a in atoms}, chosen)
[docs]
def with_selection_highlighted(
document: Document, ket: dict[str, Any], points: Sequence[tuple[float, float]], colour: str
) -> Document:
"""``document`` with the atoms drawn at ``points`` in ``ket``, the canvas's own copy, marked."""
from vibesketch.chem.species import identity
from vibesketch.document.ket import read_scene
chosen = _check_colour(colour)
by_smiles = {s.smiles: s.id for s in document.species}
found: dict[str, set[int]] = {}
for fragment in read_scene(ket).fragments:
if fragment.mol is None:
continue
conformer = fragment.mol.GetConformer()
hits = [
atom.GetIdx()
for atom in fragment.mol.GetAtoms()
if any(
abs(conformer.GetAtomPosition(atom.GetIdx()).x - x) < _TOLERANCE
and abs(conformer.GetAtomPosition(atom.GetIdx()).y - y) < _TOLERANCE
for x, y in points
)
]
if not hits:
continue
species = by_smiles.get(identity(fragment.mol).smiles)
if species is None:
raise HighlightRefused("a selected structure is not a species of the document yet")
ranked = _ranks(fragment.mol)
found.setdefault(species, set()).update(ranked[h] for h in hits if h in ranked)
if not any(found.values()):
raise HighlightRefused("select atoms of a structure on the canvas first")
for species, marked in found.items():
if marked:
document = _merged(document, species, marked, chosen)
return document
[docs]
def without_highlights(document: Document, species: str | None = None) -> Document:
"""``document`` without its highlights, or without those of ``species``."""
kept = [h for h in document.highlights if species is not None and h.species != species]
return document.model_copy(update={"highlights": kept})
[docs]
def highlight_points(document: Document, ket: dict[str, Any]) -> list[dict[str, Any]]:
"""Where the marked atoms are drawn in ``ket``, by colour, for the canvas to show."""
from vibesketch.chem.species import identity
from vibesketch.document.ket import read_scene
by_smiles = {s.smiles: s.id for s in document.species}
groups: dict[str, list[dict[str, float]]] = {}
for fragment in read_scene(ket).fragments:
if fragment.mol is None:
continue
species = by_smiles.get(identity(fragment.mol).smiles)
if species is None:
continue
conformer = fragment.mol.GetConformer()
for atoms, colour in marks_for(document, species, fragment.mol):
for index in atoms:
point = conformer.GetAtomPosition(index)
groups.setdefault(colour, []).append({"x": point.x, "y": point.y})
return [{"colour": c, "color": PALETTE[c], "atoms": atoms} for c, atoms in groups.items()]