"""Derive the normalized layer of a document from its canvas.
The canvas (KET) is what the user drew; the normalized layer is what the
compute side reads. Each time the canvas changes, this module rebuilds the
species and reactions from it:
* **Species** are the distinct molecules on the canvas, identified by
InChIKey. A species that was there before keeps its id, label, name,
conformers and a multiplicity someone set, so jobs and edits survive a
redraw. Charge is the sum of the formal charges; multiplicity defaults from
the radical count, is checked against the electron count, and is flagged
for review when the species contains a transition metal.
* **Reactions** come from the straight arrows; curved arrows (electron
pushing in a mechanism) and crossed-out ones are not reactions. For arrows
pointing any way, each arrow takes the
molecules beside it as reagents, the nearest molecule ahead of its head as
its product and the nearest behind its tail as its reactant, each with
what a plus sign joins to it and what is drawn close beside it; the steps
follow the molecules they share, then the reading order. Text beside an
arrow becomes its condition lines; a short number under a molecule becomes
its compound label.
Molecules that cannot be a single structure (query features, polymer
brackets, valence errors) stay on the canvas and are reported, never
guessed.
"""
from __future__ import annotations
import math
import re
from dataclasses import dataclass, field
from typing import Any
from vibesketch.chem.species import Identity, identity
from vibesketch.document.ket import KetArrow, KetFragment, KetScene, KetText, read_scene
from vibesketch.document.model import Document, Job, Reaction, Species
__all__ = ["Derivation", "arrow_kind", "derive", "is_reaction_arrow"]
#: A compound label as schemes print it: 1, 2a, (3), 4'.
_LABEL = re.compile(r"^\(?\d{1,3}[a-z]?'?\)?$")
#: How far from an arrow a text or a reagent may sit, in bond lengths.
_ARROW_REACH = 2.5
#: How far below a molecule its label may sit, in bond lengths.
_LABEL_REACH = 2.0
[docs]
@dataclass
class Derivation:
document: Document
warnings: list[str] = field(default_factory=list)
#: What a finished QVF job container says about its calculation
#: (:mod:`vibesketch.document.qvf`), when the document came from one.
result: dict[str, Any] | None = None
[docs]
def arrow_kind(mode: str) -> str:
"""The document's arrow for a KET arrow mode."""
if "retrosynthetic" in mode:
return "retrosynthetic"
if "equilibrium" in mode:
return "equilibrium"
if "both-ends" in mode:
return "resonance"
return "forward"
def _along(arrow: KetArrow, point: tuple[float, float]) -> tuple[float, float]:
"""Position along the arrow (0 at the tail, 1 at the head) and distance from it."""
dx, dy = arrow.head[0] - arrow.tail[0], arrow.head[1] - arrow.tail[1]
length2 = dx * dx + dy * dy or 1.0
px, py = point[0] - arrow.tail[0], point[1] - arrow.tail[1]
t = (px * dx + py * dy) / length2
distance = abs(px * dy - py * dx) / math.sqrt(length2)
return t, distance
def _labels(scene: KetScene) -> tuple[dict[int, str], set[int]]:
"""Compound labels found under fragments, and which texts were used."""
texts = _label_texts(scene)
found = {fragment: scene.texts[index].text.strip("()") for fragment, index in texts.items()}
return found, set(texts.values())
def _label_texts(scene: KetScene) -> dict[int, int]:
"""For each fragment with a compound label under it, the label's text index."""
found: dict[int, int] = {}
for index, text in enumerate(scene.texts):
if not _LABEL.match(text.text):
continue
best: tuple[float, int] | None = None
for position, fragment in enumerate(scene.fragments):
xmin, ymin, xmax, _ = fragment.bounds
x, y = text.centre
if xmin - 0.75 <= x <= xmax + 0.75 and ymin - _LABEL_REACH <= y < ymin:
gap = ymin - y
if best is None or gap < best[0]:
best = (gap, position)
if best is not None and best[1] not in found:
found[best[1]] = index
return found
#: Per reaction arrow: the arrow, its reactant, reagent and product fragments, its texts.
_Step = tuple[KetArrow, list[int], list[int], list[int], list[KetText]]
[docs]
def is_reaction_arrow(mode: str) -> bool:
"""Whether a KET arrow mode draws a reaction step: not a curved (electron-pushing)
arrow, not a crossed-out one."""
return not mode.startswith("elliptical-arc") and mode != "failed"
#: Lateral room, in bond lengths, beyond the 45° cone ahead of a head or behind a tail.
_CONE_SLACK = 1.0
#: Molecules this close to a side, in bond lengths, belong to it.
_SIDE_GAP = 1.5
#: A lone hand-drawn row often omits plus signs and leaves several bond lengths
#: between components. Allow up to 3.5, still bounded instead of consuming
#: every disconnected structure on that side of the canvas.
_SINGLE_SIDE_GAP = 3.5
def _box_distance(box: tuple[float, float, float, float], point: tuple[float, float]) -> float:
dx = max(box[0] - point[0], 0.0, point[0] - box[2])
dy = max(box[1] - point[1], 0.0, point[1] - box[3])
return math.hypot(dx, dy)
def _box_gap(a: tuple[float, float, float, float], b: tuple[float, float, float, float]) -> float:
dx = max(a[0] - b[2], 0.0, b[0] - a[2])
dy = max(a[1] - b[3], 0.0, b[1] - a[3])
return math.hypot(dx, dy)
def _plus_groups(scene: KetScene) -> list[int]:
"""For each fragment, the representative of the fragments plus signs join it to."""
parent = list(range(len(scene.fragments)))
def root(index: int) -> int:
while parent[index] != index:
parent[index] = parent[parent[index]]
index = parent[index]
return index
for plus in scene.pluses:
near = sorted(
(_box_distance(f.bounds, plus), i)
for i, f in enumerate(scene.fragments)
if _box_distance(f.bounds, plus) <= _ARROW_REACH
)
best: tuple[float, int, int] | None = None
for first, (gap_a, a) in enumerate(near):
for gap_b, b in near[first + 1 :]:
ca, cb = scene.fragments[a].centre, scene.fragments[b].centre
opposite = (ca[0] - plus[0]) * (cb[0] - plus[0]) + (ca[1] - plus[1]) * (
cb[1] - plus[1]
) < 0
if opposite and (best is None or gap_a + gap_b < best[0]):
best = (gap_a + gap_b, a, b)
if best is not None:
parent[root(best[1])] = root(best[2])
return [root(i) for i in range(len(parent))]
def _side(
scene: KetScene,
arrow: KetArrow,
end: str,
groups: list[int],
excluded: set[int],
gap: float = _SIDE_GAP,
) -> list[int]:
"""The fragments ahead of an arrow's head (``end="head"``) or behind its tail."""
tail, head = arrow.tail, arrow.head
length = math.dist(tail, head) or 1.0
ux, uy = (head[0] - tail[0]) / length, (head[1] - tail[1]) / length
origin, sign = (head, 1.0) if end == "head" else (tail, -1.0)
candidates: dict[int, float] = {}
for index, fragment in enumerate(scene.fragments):
if index in excluded:
continue
vx, vy = fragment.centre[0] - origin[0], fragment.centre[1] - origin[1]
along = sign * (vx * ux + vy * uy)
lateral = abs(vx * uy - vy * ux)
if along > 0 and lateral <= along + _CONE_SLACK:
candidates[index] = _box_distance(fragment.bounds, origin)
if not candidates:
return []
nearest = min(candidates, key=lambda i: candidates[i])
side = {i for i in range(len(scene.fragments)) if groups[i] == groups[nearest]} - excluded
grown = True
while grown:
grown = False
for index in candidates:
if index in side:
continue
if (
min(
_box_gap(scene.fragments[index].bounds, scene.fragments[s].bounds) for s in side
)
<= gap
):
side |= {i for i in range(len(scene.fragments)) if groups[i] == groups[index]}
side -= excluded
grown = True
# Reading order: top to bottom in bands of a bond length and a half, then left to right.
return sorted(
side,
key=lambda i: (-round(scene.fragments[i].centre[1] / 1.5), scene.fragments[i].centre[0]),
)
def _roles(
scene: KetScene,
) -> list[_Step]:
"""Per reaction arrow, in step order: reactant, reagent and product fragments."""
from vibesketch.document.electrochem import widened
# Coefficients and states drawn beside molecules belong to them (#141).
scene = widened(scene)
arrows = [a for a in scene.arrows if is_reaction_arrow(a.mode)]
groups = _plus_groups(scene)
gap = _SINGLE_SIDE_GAP if len(arrows) == 1 else _SIDE_GAP
beside = [
[
position
for position, fragment in enumerate(scene.fragments)
if 0.0 <= _along(arrow, fragment.centre)[0] <= 1.0
and _along(arrow, fragment.centre)[1] <= _ARROW_REACH
]
for arrow in arrows
]
# A molecule beside an arrow is that arrow's; it is no other step's product.
excluded = {position for positions in beside for position in positions}
steps: list[_Step] = [
(
arrow,
_side(scene, arrow, "tail", groups, excluded, gap),
reagents,
_side(scene, arrow, "head", groups, excluded, gap),
[],
)
for arrow, reagents in zip(arrows, beside, strict=True)
]
return _in_step_order(steps)
def _in_step_order(
steps: list[_Step],
) -> list[_Step]:
"""Steps that share a molecule follow each other; otherwise top to bottom, left to right."""
def reading(
step: _Step,
) -> tuple[float, float]:
arrow = step[0]
middle = ((arrow.tail[0] + arrow.head[0]) / 2, (arrow.tail[1] + arrow.head[1]) / 2)
return (-round(middle[1]), middle[0])
pending = sorted(range(len(steps)), key=lambda i: reading(steps[i]))
feeds = {
i: [j for j in pending if j != i and set(steps[i][3]) & set(steps[j][1])] for i in pending
}
fed = {j for targets in feeds.values() for j in targets}
order: list[int] = []
def follow(index: int) -> None:
if index in order:
return
order.append(index)
for target in feeds[index]:
follow(target)
for index in pending:
if index not in fed:
follow(index)
for index in pending: # a cycle with no start
follow(index)
return [steps[i] for i in order]
def _condition_texts(scene: KetScene, arrow: KetArrow, used: set[int]) -> list[int]:
"""The texts beside ``arrow`` not yet ``used``, in reading order; marks them used."""
found: list[tuple[float, str, int]] = []
for index, text in enumerate(scene.texts):
if index in used:
continue
t, distance = _along(arrow, text.centre)
if -0.15 <= t <= 1.15 and distance <= _ARROW_REACH:
used.add(index)
# Above the arrow first, then below; top to bottom within each.
found.append((-text.centre[1], text.text, index))
return [index for _, _, index in sorted(found)]
def _texts_for(scene: KetScene, arrow: KetArrow, used: set[int]) -> list[str]:
return [
line
for index in _condition_texts(scene, arrow, used)
for line in scene.texts[index].text.splitlines()
if line
]
def _next_id(prefix: str, taken: set[str]) -> str:
number = 1
while f"{prefix}{number}" in taken:
number += 1
taken.add(f"{prefix}{number}")
return f"{prefix}{number}"
def _species_for(
found: Identity,
previous: dict[str, Species],
taken: set[str],
label: str | None,
warnings: list[str],
) -> Species:
multiplicity = found.multiplicity
old = previous.get(found.key)
if old is not None and (found.electrons - (old.multiplicity - 1)) % 2 == 0:
multiplicity = old.multiplicity
if (found.electrons - (multiplicity - 1)) % 2:
multiplicity += 1
warnings.append(
f"{found.smiles}: {found.electrons} electrons cannot be a "
f"multiplicity {multiplicity - 1}; using {multiplicity}"
)
data: dict[str, Any] = {
"id": old.id if old is not None else _next_id("m", taken),
"smiles": found.smiles,
"inchikey": found.inchikey,
"charge": found.charge,
"multiplicity": multiplicity,
}
if label is not None:
data["label"] = label
elif old is not None and old.label is not None:
data["label"] = old.label
if old is not None:
if old.name is not None:
data["name"] = old.name
if (
old.charge == found.charge
and old.multiplicity == multiplicity
and old.smiles == found.smiles
):
# Mapped SMILES can change canonical atom order without changing InChI.
data["conformers"] = old.conformers
if old.sequence is not None:
# A sequence is kept only for the structure it builds (#98).
if old.sequence.inchikey == found.inchikey:
data["sequence"] = old.sequence
else:
warnings.append(
f"{data['id']}: its drawing no longer is the structure its sequence "
"builds, so the sequence was dropped"
)
if found.has_transition_metal:
warnings.append(
f"{data['id']} ({found.smiles}) contains a transition metal: its spin state cannot "
"be read from the drawing, so set its multiplicity explicitly"
)
return Species.model_validate(data)
[docs]
def derive(document: Document) -> Derivation:
"""Rebuild ``document``'s species and reactions from its canvas."""
warnings: list[str] = []
scene = read_scene(document.canvas.data) if document.canvas.data else KetScene()
from vibesketch.document.solids import solid_data, with_solids
# Solids written as formulas take part like molecules (#142).
scene, solids = with_solids(scene)
previous = {
("solid:" + s.solid.key if s.solid else s.inchikey or s.smiles or s.id): s
for s in document.species
}
from vibesketch.render.surfaces import surface_of
# Canvas metadata is authoritative on fresh imports, before surfaces have
# been derived. Reserve its IDs alongside every existing document record.
canvas: Any = document.canvas.data
surface_ids = {
surface.id
for node in canvas.get("root", {}).get("nodes", [])
if isinstance(node, dict) and (surface := surface_of(node)) is not None
}
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),
*surface_ids,
}
labels, used = _labels(scene)
species: dict[str, Species] = {}
by_fragment: dict[int, str] = {}
for position, fragment in enumerate(scene.fragments):
if position in solids:
formula = solids[position]
key = "solid:" + formula.key()
if key not in species:
old = previous.get(key)
record: dict[str, Any] = {
"id": old.id if old is not None else _next_id("m", taken),
"solid": solid_data(formula),
}
label = labels.get(position) or (old.label if old is not None else None)
if label is not None:
record["label"] = label
species[key] = Species.model_validate(record)
by_fragment[position] = species[key].id
continue
if fragment.mol is None:
warnings.append(f"a molecule on the canvas is not a species: {fragment.unsupported}")
continue
found = identity(fragment.mol)
if found.key not in species:
species[found.key] = _species_for(
found, previous, taken, labels.get(position), warnings
)
by_fragment[position] = species[found.key].id
# Generated nanostructures are not drawn as atoms; a redraw keeps them.
for kept in document.species:
if kept.nanostructure is not None:
species[kept.nanostructure.key + kept.id] = kept
from vibesketch.document.electrochem import quantities, read_half_reaction
reactions: list[Reaction] = []
previous_reactions = {
(frozenset(r.reactants), frozenset(r.products)): r for r in document.reactions
}
taken_reactions = {r.id for r in document.reactions} | set(taken)
for step, (arrow, reactant_idx, reagent_idx, product_idx, _) in enumerate(_roles(scene)):
def ids(indices: list[int]) -> list[str]:
return list(dict.fromkeys(by_fragment[i] for i in indices if i in by_fragment))
reactants, reagents, products = ids(reactant_idx), ids(reagent_idx), ids(product_idx)
# Electrons, states and coefficients are read before the condition lines (#141).
half, notes = read_half_reaction(
scene, arrow, (reactant_idx, reagent_idx, product_idx), by_fragment, used
)
warnings.extend(notes)
conditions = _texts_for(scene, arrow, used)
if not reactants or not products:
warnings.append("an arrow without molecules on both sides is not a reaction yet")
continue
old_reaction = previous_reactions.get((frozenset(reactants), frozenset(products)))
if old_reaction is None and step < len(document.reactions):
# Same arrow, different molecules: keep the reaction's id and yield.
candidate = document.reactions[step]
if candidate.id not in {r.id for r in reactions}:
old_reaction = candidate
data: dict[str, Any] = {
"id": old_reaction.id if old_reaction else _next_id("r", taken_reactions),
"reactants": reactants,
"reagents": [r for r in reagents if r not in reactants and r not in products],
"products": products,
"arrow": arrow_kind(arrow.mode),
"conditions": conditions,
**half,
}
typed = quantities(conditions)
if typed:
data["quantities"] = [q.model_dump(exclude_none=True) for q in typed]
# A solid's state is part of its formula, (s) for C(s): it is the participant's state.
stated = {e["species"] for e in data.get("states", [])}
solid_states = [
{"species": s.id, "state": s.solid.state}
for s in species.values()
if s.solid is not None
and s.solid.state is not None
and s.id not in stated
and s.id in {*reactants, *data["reagents"], *products}
]
if solid_states:
data["states"] = [*data.get("states", []), *solid_states]
if old_reaction is not None and old_reaction.yield_ is not None:
data["yield"] = old_reaction.yield_
if old_reaction is not None and old_reaction.amounts:
# Entered amounts stay with the species still in the reaction.
present = {*reactants, *data["reagents"], *products}
data["amounts"] = [
a.model_dump(exclude_none=True)
for a in old_reaction.amounts
if a.species in present
]
dropped = [a.species for a in old_reaction.amounts if a.species not in present]
if dropped:
warnings.append(
f"reaction {data['id']}: the amounts of {', '.join(dropped)} were dropped, "
"as they are no longer in it"
)
# A state or coefficient drawn for a species that is a reagent here is not kept.
present = {*data["reactants"], *data["reagents"], *data["products"]}
for entries in ("states", "coefficients"):
if entries in data:
data[entries] = [e for e in data[entries] if e["species"] in present]
reactions.append(Reaction.model_validate(data))
from vibesketch.document.generics import derive_generics
generics, rgroups = derive_generics(
document,
scene,
labels,
{*taken_reactions, *(r.id for r in document.recipes), *(g.id for g in document.generics)},
warnings,
)
# A fragment capped on request stays while its labelled structure is drawn (#139).
labelled = {g.cxsmiles: g.id for g in generics if g.kind == "labelled"}
for kept in document.species:
if kept.capped is None:
continue
generic = labelled.get(kept.capped.cxsmiles)
if generic is None:
warnings.append(
f"{kept.id} was removed: the labelled structure it caps is no longer drawn"
)
continue
capping = kept.capped.model_copy(update={"generic": generic})
species["capped:" + kept.id] = kept.model_copy(update={"capped": capping})
present = {s.id for s in species.values()}
jobs: list[Job] = []
for job in document.jobs:
if job.species in present:
jobs.append(job)
else:
warnings.append(f"job {job.id} was removed: its species is no longer on the canvas")
highlights = [h for h in document.highlights if h.species in present]
if len(highlights) != len(document.highlights):
warnings.append("highlights of species no longer on the canvas were removed")
rebuilt = document.model_copy(
update={
"species": list(species.values()),
"reactions": reactions,
"jobs": jobs,
"highlights": highlights,
"generics": generics,
"rgroups": rgroups,
}
)
from vibesketch.document.surfaces import records
rebuilt = rebuilt.model_copy(update={"surfaces": records(rebuilt)})
# model_copy skips validation; validate the whole document once.
return Derivation(Document.model_validate(rebuilt.to_json_data()), warnings)
def fragments(document: Document) -> list[KetFragment]:
"""The canvas's fragments, for callers that need coordinates."""
return read_scene(document.canvas.data).fragments if document.canvas.data else []