"""Generated nanostructures as species of a document (#137).
A nanotube, graphene flake, nanoribbon or fullerene is added as a species
whose first conformer holds the generated atoms and whose ``nanostructure``
record holds the parameters, the bonds and an identity key in place of an
InChIKey. It is a finite cluster model: writers take it like any species
with a 3D structure, with its charge and multiplicity checked against the
electron count. It is not drawn on the canvas, so a redraw keeps it
(:func:`vibesketch.document.derive.derive`); its depiction is #138.
"""
from __future__ import annotations
from typing import Any
from pydantic import TypeAdapter, ValidationError
from vibesketch.chem.nanostructures import (
AtomModel,
NanostructureRefused,
build,
wall_properties,
)
from vibesketch.document.model import (
Conformer,
Document,
NanostructureParameters,
Nanotube,
Species,
)
from vibesketch.elements import ATOMIC_NUMBER
__all__ = [
"NanostructureRefused",
"attach_to_picture",
"masses",
"pictures",
"summary",
"with_nanostructure",
"with_picture",
]
_PARAMETERS: TypeAdapter[NanostructureParameters] = TypeAdapter(NanostructureParameters)
def _parameters(parameters: NanostructureParameters | dict[str, Any]) -> NanostructureParameters:
if not isinstance(parameters, dict):
return parameters
try:
return _PARAMETERS.validate_python(parameters)
except ValidationError as error:
problems = "; ".join(
f"{'.'.join(str(p) for p in e['loc']) or 'parameters'}: {e['msg']}"
for e in error.errors()
)
raise NanostructureRefused(problems) from None
def _free_id(document: Document, prefix: str) -> str:
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),
}
number = 1
while f"{prefix}{number}" in taken:
number += 1
return f"{prefix}{number}"
[docs]
def masses(symbols: list[str]) -> tuple[float, float]:
"""Average and monoisotopic mass of the atoms, from RDKit's periodic table."""
from rdkit import Chem
table = Chem.GetPeriodicTable()
counts: dict[str, int] = {}
for symbol in symbols:
counts[symbol] = counts.get(symbol, 0) + 1
average = sum(float(table.GetAtomicWeight(s)) * n for s, n in counts.items())
mono = sum(float(table.GetMostCommonIsotopeMass(s)) * n for s, n in counts.items())
return round(average, 4), round(mono, 6)
[docs]
def summary(species: Species) -> dict[str, Any]:
"""What a client shows of a nanostructure species: identity, formula, atom count."""
made = species.nanostructure
if made is None:
raise ValueError(f"species {species.id} is not a generated nanostructure")
return {
"species": species.id,
"key": made.key,
"description": made.description,
"formula": made.formula,
"atoms": len(species.conformers[0].symbols),
"bonds": len(made.bonds),
"charge": species.charge,
"multiplicity": species.multiplicity,
"kind": made.parameters.kind,
"periodic": made.periodic.model_dump() if made.periodic is not None else None,
**(
{"walls": [wall_properties(w) for w in made.parameters.walls]}
if isinstance(made.parameters, Nanotube)
else {}
),
}
[docs]
def with_nanostructure(
document: Document,
parameters: NanostructureParameters | dict[str, Any],
*,
species_id: str | None = None,
label: str | None = None,
charge: int = 0,
multiplicity: int | None = None,
picture_style: str | None = None,
picture_label: bool = False,
) -> tuple[Document, dict[str, Any]]:
"""``document`` with a generated nanostructure added as a new species, and its summary.
The multiplicity defaults to the lowest the electron count allows
(singlet or doublet); one that is impossible is refused, as is any
parameter the generator cannot build, with the reason. Raises
:class:`NanostructureRefused`. Optional ``picture_style`` also adds the
depiction; any refusal returns neither model nor picture.
"""
resolved = _parameters(parameters)
resolved, model = build(resolved)
electrons = sum(ATOMIC_NUMBER[s] for s in model.symbols) - charge
if multiplicity is None:
multiplicity = 1 if electrons % 2 == 0 else 2
identifier = species_id or _free_id(document, "nano")
species_data = _species_data(identifier, model, resolved, label, charge, multiplicity)
data = document.to_json_data()
data["species"] = [*data.get("species", []), species_data]
try:
updated = Document.model_validate(data)
except ValidationError as error:
raise NanostructureRefused(
"; ".join(str(e["msg"]).removeprefix("Value error, ") for e in error.errors())
) from None
if picture_style is not None:
updated = with_picture(updated, identifier, picture_style, label=picture_label)
return updated, summary(updated.species_by_id(identifier))
def _species_data(
identifier: str,
model: AtomModel,
parameters: NanostructureParameters,
label: str | None,
charge: int,
multiplicity: int,
) -> dict[str, Any]:
conformer = Conformer(
method="nanostructure", symbols=model.symbols, xyz_angstrom=model.positions
)
data: dict[str, Any] = {
"id": identifier,
"charge": charge,
"multiplicity": multiplicity,
"conformers": [conformer.model_dump(mode="json", exclude_none=True)],
"nanostructure": model.record(parameters).model_dump(mode="json", exclude_none=True),
}
if label is not None:
data["label"] = label
return data
[docs]
def pictures(document: Document) -> list[dict[str, Any]]:
"""The nanostructure pictures on the canvas: index, species, style, box and anchors (#138)."""
from vibesketch.render.nanostructures import anchor_points, depiction_of, image_box
found = []
canvas: Any = document.canvas.data
for index, node in enumerate(canvas.get("root", {}).get("nodes", [])):
if not isinstance(node, dict) or (seen := depiction_of(node)) is None:
continue
found.append(
{
"node": index,
"species": seen[0].species,
"style": seen[0].style,
"box": list(image_box(node)),
"anchors": {k: list(v) for k, v in anchor_points(node).items()},
}
)
from vibesketch.render.surfaces import anchor_points as surface_anchors
from vibesketch.render.surfaces import surface_of
for index, node in enumerate(canvas.get("root", {}).get("nodes", [])):
if isinstance(node, dict) and (surface := surface_of(node)) is not None:
found.append(
{
"node": index,
"species": surface.id,
"style": surface.kind,
"box": list(image_box(node)),
"anchors": {k: list(v) for k, v in surface_anchors(node).items()},
}
)
return found
[docs]
def with_picture(
document: Document,
species_id: str,
style: str = "schematic",
*,
label: bool = False,
rear: str = "light",
brk: str = "auto",
width: float = 8.0,
at: tuple[float, float] | None = None,
) -> Document:
"""``document`` with a picture of a nanostructure species on its canvas, as one object.
The picture goes to the right of the drawing unless ``at`` (its top-left
corner, KET units) is given; ``width`` is in bond lengths. Raises
:class:`NanostructureRefused` for an unknown species, style or option.
"""
from vibesketch.document.formats import _points
from vibesketch.render.nanostructures import (
STYLES,
Depiction,
DepictionRefused,
canvas_image,
)
try:
species = document.species_by_id(species_id)
except KeyError:
raise NanostructureRefused(f"there is no species {species_id!r}") from None
if species.nanostructure is None:
raise NanostructureRefused(f"species {species_id} is not a generated nanostructure")
if style not in STYLES:
raise NanostructureRefused(f"no style {style!r}; styles: {', '.join(STYLES)}")
if rear not in ("light", "omit") or brk not in ("auto", "on", "off"):
raise NanostructureRefused("rear is light or omit, and the break auto, on or off")
if not 1.0 <= width <= 100.0:
raise NanostructureRefused("a picture is 1 to 100 bond lengths wide")
# The style, rear and break were checked against their choices above.
options: dict[str, Any] = {"style": style, "label": label, "rear": rear, "brk": brk}
depiction = Depiction(species_id, species.nanostructure.key, **options)
canvas: dict[str, Any] = dict(document.canvas.data) or {"root": {"nodes": []}}
if at is None:
points = _points(canvas)
at = (max(p[0] for p in points) + 2.0, max(p[1] for p in points)) if points else (0.0, 0.0)
try:
node = canvas_image(species, depiction, at, width)
except DepictionRefused as error:
raise NanostructureRefused(str(error)) from None
root: dict[str, Any] = dict(canvas.get("root") or {})
root["nodes"] = [*root.get("nodes", []), node]
root.setdefault("connections", [])
root.setdefault("templates", [])
canvas["root"] = root
canvas.setdefault("ket_version", "2.0.0")
data = document.to_json_data()
data["canvas"] = {"format": "ket", "data": canvas}
return Document.model_validate(data)
[docs]
def attach_to_picture(
document: Document, generic_id: str, species_id: str, anchor: str = "upper-end"
) -> Document:
"""``document`` with a labelled structure moved onto a nanostructure picture's anchor.
The structure's one attachment atom (#139) is put on the anchor and the
structure turned so its bond leaves the picture outwards; figures then
print the two as one structure. Raises :class:`NanostructureRefused`.
"""
import copy
import math
from vibesketch.document.derive import derive
from vibesketch.document.generics import attachment_labels, cxsmiles
from vibesketch.document.ket import read_scene
generic = next((g for g in document.generics if g.id == generic_id), None)
if generic is None or generic.kind != "labelled":
raise NanostructureRefused(
f"{generic_id!r} is not a labelled structure (one bonded to CNT, Au, ...)"
)
if len(generic.labels) != 1:
raise NanostructureRefused(
f"{generic_id} has {len(generic.labels)} labels; attach needs one"
)
shown = [p for p in pictures(document) if p["species"] == species_id]
if not shown:
raise NanostructureRefused(f"no picture of {species_id} is on the canvas")
anchors = shown[0]["anchors"]
if anchor not in anchors:
raise NanostructureRefused(
f"the picture has no anchor {anchor!r}; anchors: {', '.join(sorted(anchors))}"
)
ax, ay, dx, dy = anchors[anchor]
canvas: Any = copy.deepcopy(document.canvas.data)
for fragment in read_scene(canvas).fragments:
figure = fragment.figure
if figure is None or cxsmiles(figure) != generic.cxsmiles:
continue
if not attachment_labels(figure):
continue
label = next(a for a in figure.GetAtoms() if a.HasProp("attachmentLabel"))
neighbour = label.GetNeighbors()[0]
atoms = canvas[fragment.node]["atoms"]
start = atoms[fragment.atoms[label.GetIdx()]]["location"]
towards = atoms[fragment.atoms[neighbour.GetIdx()]]["location"]
turn = math.atan2(dy, dx) - math.atan2(towards[1] - start[1], towards[0] - start[0])
cos, sin = math.cos(turn), math.sin(turn)
sx, sy = float(start[0]), float(start[1])
for index in fragment.atoms:
location = atoms[index]["location"]
x, y = float(location[0]) - sx, float(location[1]) - sy
location[0] = round(ax + cos * x - sin * y, 6)
location[1] = round(ay + sin * x + cos * y, 6)
break
else:
raise NanostructureRefused(f"{generic_id} is not drawn on the canvas")
moved = document.model_copy(
update={"canvas": document.canvas.model_copy(update={"data": canvas})}
)
return derive(Document.model_validate(moved.to_json_data())).document