"""Generic structures in the drawing as document entries: R-group cores, their
definitions and bracketed groups (#90).
They are derived from the drawing like species, and written as CXSMILES
without coordinates: R-group atoms carry their labels (``$_R1$``), a
definition's attachment points ``_AP1``, and brackets their S-groups
(``Sg:n:...``). The drawing stays the source of truth; nothing here is
edited by clients (maintainer decisions, 2026-10-24).
"""
from __future__ import annotations
from typing import Any
from vibesketch.document.ket import KetFragment, KetScene
from vibesketch.document.model import Document, Generic, RGroup, Species
__all__ = [
"CapRefused",
"attachment_labels",
"cap_labelled",
"cxsmiles",
"derive_generics",
"generics_report",
]
[docs]
def cxsmiles(mol: Any) -> str:
"""``mol`` as canonical CXSMILES without coordinates, its R-group atoms as ``$_R1$``."""
from rdkit import Chem
copy = Chem.Mol(mol)
for atom in copy.GetAtoms():
if atom.GetAtomicNum() == 0 and atom.HasProp("attachmentLabel"):
atom.SetProp("atomLabel", atom.GetProp("attachmentLabel"))
for name in ("dummyLabel", "attachmentLabel", "molFileAlias"):
atom.ClearProp(name)
elif atom.GetAtomicNum() == 0 and atom.HasProp("dummyLabel"):
label = atom.GetProp("dummyLabel")
if label.startswith("R"):
atom.SetProp("atomLabel", f"_{label}")
atom.ClearProp("dummyLabel")
fields = Chem.CXSmilesFields.CX_ALL_BUT_COORDS ^ Chem.CXSmilesFields.CX_ATOM_PROPS
text = str(Chem.MolToCXSmiles(copy, Chem.SmilesWriteParams(), fields))
# RDKit CXSMILES omits MUL groups. Retain their count and canonical atom
# membership in an ordinary CX atom property, so different drawings do
# not collapse to one generic entry merely because their parent is equal.
multiples = [g for g in Chem.GetMolSubstanceGroups(copy) if g.GetProp("TYPE") == "MUL"]
if multiples:
import json
Chem.MolToSmiles(copy)
order = json.loads(copy.GetProp("_smilesAtomOutputOrder"))
positions = {atom: index for index, atom in enumerate(order)}
markers = sorted(
f"{g.GetProp('LABEL')}:"
+ ",".join(str(positions[a]) for a in sorted(g.GetAtoms(), key=lambda a: positions[a]))
for g in multiples
)
for atom in copy.GetAtoms():
for name in atom.GetPropNames():
if name != "atomLabel":
atom.ClearProp(name)
copy.GetAtomWithIdx(order[0]).SetProp("vibesketchMUL", ";".join(markers))
text = str(
Chem.MolToCXSmiles(
copy, Chem.SmilesWriteParams(), fields | Chem.CXSmilesFields.CX_ATOM_PROPS
)
)
return text
[docs]
class CapRefused(ValueError):
"""A labelled structure that cannot be capped as asked; the message says why."""
[docs]
def cap_labelled(document: Document, generic_id: str, cap: str) -> tuple[Document, Species]:
"""``document`` with a species made by capping a labelled structure, and that species.
Every attachment atom of the labelled generic structure ``generic_id``
(CNT, Au, ...) is replaced by ``cap``, ``H`` or ``CH3``: the organic part
as a molecule a calculation can take, which the user asked for by
choosing the cap. The species records the generic structure, its labels
and the cap. Raises :class:`CapRefused`.
"""
from rdkit import Chem
from vibesketch.chem.species import identity
from vibesketch.document.derive import _next_id
if cap not in ("H", "CH3"):
raise CapRefused(f"the cap is H or CH3, not {cap!r}")
generic = next((g for g in document.generics if g.id == generic_id), None)
if generic is None:
raise CapRefused(f"there is no generic structure {generic_id!r}")
if generic.kind != "labelled":
raise CapRefused(
f"{generic_id} is a {generic.kind} structure; only one bonded to labelled "
"attachment atoms (CNT, Au, ...) is capped"
)
for existing in document.species:
if existing.capped is not None and (existing.capped.cxsmiles, existing.capped.cap) == (
generic.cxsmiles,
cap,
):
raise CapRefused(f"{generic_id} is already capped with {cap} as {existing.id}")
mol = Chem.RWMol(Chem.MolFromSmiles(generic.cxsmiles))
for atom in mol.GetAtoms():
if atom.GetAtomicNum() == 0:
for name in list(atom.GetPropNames()):
atom.ClearProp(name)
atom.SetAtomicNum(1 if cap == "H" else 6)
atom.SetNoImplicit(False)
atom.SetIsotope(0)
capped = Chem.RemoveHs(mol.GetMol())
Chem.SanitizeMol(capped)
found = identity(capped)
taken = {
*(s.id for s in document.species),
*(r.id for r in document.reactions),
*(j.id for j in document.jobs),
*(r.id for r in document.recipes),
*(g.id for g in document.generics),
*(s.id for s in document.surfaces),
}
species = Species.model_validate(
{
"id": _next_id("m", taken),
"smiles": found.smiles,
"inchikey": found.inchikey,
"charge": found.charge,
"multiplicity": found.multiplicity,
"capped": {
"generic": generic.id,
"cxsmiles": generic.cxsmiles,
"cap": cap,
"labels": generic.labels,
},
}
)
data = document.to_json_data()
data["species"] = [*data.get("species", []), species.model_dump(mode="json", exclude_none=True)]
return Document.model_validate(data), species
[docs]
def generics_report(document: Document) -> dict[str, list[dict[str, Any]]]:
"""The document's generic structures and R-group definitions, as JSON data.
A labelled structure also lists the species capped from it (#139).
"""
def entry(generic: Generic) -> dict[str, Any]:
data = generic.model_dump(mode="json", exclude_none=True)
if generic.kind == "labelled":
data["capped"] = [
{"species": s.id, "cap": s.capped.cap, "smiles": s.smiles}
for s in document.species
if s.capped is not None and s.capped.generic == generic.id
]
else:
data.pop("labels", None)
return data
return {
"generics": [entry(g) for g in document.generics],
"rgroups": [r.model_dump(mode="json", exclude_none=True) for r in document.rgroups],
}
def _has_r_atom(mol: Any) -> bool:
"""Whether ``mol`` has an R-group atom, numbered or not (an unnumbered one is R)."""
return any(
atom.GetAtomicNum() == 0
and not atom.HasProp("attachmentLabel")
and atom.HasProp("dummyLabel")
and str(atom.GetProp("dummyLabel")).startswith("R")
for atom in mol.GetAtoms()
)
def _numbers(mol: Any) -> list[int]:
"""The R-group numbers a core's labelled atoms use, R1 as 1, in order."""
found = set()
for atom in mol.GetAtoms():
label = atom.GetProp("dummyLabel") if atom.HasProp("dummyLabel") else ""
if atom.GetAtomicNum() == 0 and label[:1] == "R" and label[1:].isdigit():
found.add(int(label[1:]))
return sorted(found)
[docs]
def attachment_labels(mol: Any) -> list[str]:
"""The labels of a structure's attachment atoms (CNT, Au, ...), in atom order (#139)."""
return [
str(atom.GetProp("attachmentLabel"))
for atom in mol.GetAtoms()
if atom.HasProp("attachmentLabel")
]
def _definition(fragment: KetFragment, node: dict[str, Any]) -> RGroup:
from rdkit import Chem
logic = node.get("rlogic") or {}
substituents = sorted(
cxsmiles(piece) for piece in Chem.GetMolFrags(fragment.figure, asMols=True)
)
occurrence = str(logic.get("range") or "").strip()
return RGroup(
number=int(logic["number"]),
substituents=substituents,
occurrence=occurrence or None,
rest_h=logic.get("resth") is True,
if_then=int(logic["ifthen"]) if logic.get("ifthen") else None,
)
[docs]
def derive_generics(
document: Document,
scene: KetScene,
labels: dict[int, str],
taken: set[str],
warnings: list[str],
) -> tuple[list[Generic], list[RGroup]]:
"""The drawing's R-group cores, bracketed groups and R-group definitions.
Only what figures can draw becomes an entry; a structure the drawing
cannot express stays a note, as before. A generic structure keeps its id
and compound number while its CXSMILES is the same. ``taken`` holds the
identifiers already used, and gains the new ones.
"""
from vibesketch.document.derive import _next_id
previous: dict[tuple[str, str], Generic] = {(g.kind, g.cxsmiles): g for g in document.generics}
generics: dict[tuple[str, str], Generic] = {}
rgroups: dict[int, RGroup] = {}
for position, fragment in enumerate(scene.fragments):
if fragment.mol is not None or fragment.figure is None:
continue
node = document.canvas.data.get(fragment.node)
if isinstance(node, dict) and node.get("type") == "rgroup":
definition = _definition(fragment, node)
if definition.number in rgroups:
warnings.append(f"R{definition.number} is defined twice; the first is kept")
else:
rgroups[definition.number] = definition
continue
numbers = _numbers(fragment.figure)
attached = attachment_labels(fragment.figure)
if numbers or _has_r_atom(fragment.figure):
kind = "rgroup-core"
else:
kind = "labelled" if attached else "bracketed"
text = cxsmiles(fragment.figure)
if (kind, text) in generics:
continue
old = previous.get((kind, text))
data: dict[str, Any] = {
"id": old.id if old is not None else _next_id("g", taken),
"kind": kind,
"cxsmiles": text,
"rgroups": numbers,
}
if kind == "labelled":
data["labels"] = attached
label = labels.get(position) or (old.label if old is not None else None)
if label is not None:
data["label"] = label
generics[(kind, text)] = Generic.model_validate(data)
return list(generics.values()), [rgroups[n] for n in sorted(rgroups)]