"""Same-orientation renders of a recipe campaign's stages, from its archived artifacts (#42).
For one species of a compiled campaign this gathers the starting geometry
(the sketch's prepared conformer, from the campaign's pinned sources) and
the geometry every stage's validator accepted (from vibe-queue's verified
artifacts), with each optimisation's trajectory where its result carries
one. Nothing is taken on trust: every artifact is checked against the
manifest digest vibe-queue's journal records for the producing attempt, and
every file against its manifest entry, before it is read. A stage that was
not accepted is listed and never rendered.
**One frame for all.** The start geometry is put in its principal-axes frame
(centroid at the origin, axes by decreasing spread, signs fixed by the
coordinates' skew, a proper rotation), and every accepted geometry and every
trajectory frame is superposed on it by the rotation that minimises the RMSD
(W. Kabsch, *Acta Cryst.* A32, 922-923 (1976); A34, 827-828 (1978)), never a
reflection, so chirality is kept. A rigid drift of the optimiser therefore
does not turn the molecule between figures, and the RMSD to the start says
how far each stage moved it.
Each geometry is written as a QVF archive of its own (``<species>-<stage>.qvf``,
the aligned structure and, for an optimisation, the aligned trajectory),
beside ``renders.json``, which records for every figure the artifact it came
from: node, attempt, manifest digest and result digest. The images are
vibe-view's: ``vibe-view capture`` draws each archive with one fixed
isometric camera and ``vibe-view animate --kind trajectory`` each
trajectory, so vibe-sketch prepares the geometry and imports no rendering
stack (the shared scope decision). Without vibe-view the archives and the
record are still written, and the warning gives the command to run.
"""
from __future__ import annotations
import json
import os
import subprocess
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any
import numpy as np
from numpy.typing import NDArray
from vibesketch.document.model import Conformer, Document
__all__ = ["RendersRefused", "StageRenders", "aligned_frame", "render_stages", "superpose"]
_Points = NDArray[np.float64]
_SCHEMA = "vibe-sketch.stage-renders/1"
_RENDER_TIMEOUT = 600.0
[docs]
class RendersRefused(ValueError):
"""The renders cannot be made as asked; the reasons say why."""
def __init__(self, reasons: list[str]) -> None:
self.reasons = reasons
super().__init__("; ".join(reasons))
[docs]
@dataclass
class StageRenders:
"""What was written: the record, the files, and what needs attention."""
directory: Path
record: dict[str, Any]
files: list[str] = field(default_factory=list)
warnings: list[str] = field(default_factory=list)
# ── Geometry ────────────────────────────────────────────────────────────
[docs]
def aligned_frame(points: _Points) -> tuple[_Points, _Points]:
"""``points`` in their principal-axes frame, and the centroid they were moved from.
The axes run by decreasing spread; each axis points where the
coordinates are skewed (the larger tail), ties broken by the first atom
that is off the plane; the third axis completes a right-handed frame.
"""
centroid = points.mean(axis=0)
centred = points - centroid
_, vectors = np.linalg.eigh(centred.T @ centred)
axes = vectors[:, ::-1].copy()
for column in range(2):
projection = centred @ axes[:, column]
skew = float(np.sum(projection**3))
if abs(skew) < 1e-9:
off = next((float(p) for p in projection if abs(p) > 1e-6), 0.0)
skew = off
if skew < 0:
axes[:, column] = -axes[:, column]
axes[:, 2] = np.cross(axes[:, 0], axes[:, 1])
return centred @ axes, centroid
[docs]
def superpose(reference: _Points, mobile: _Points) -> tuple[_Points, float]:
"""``mobile`` turned and moved onto ``reference``, and the RMSD left, in Å.
The proper rotation that minimises the RMSD (Kabsch 1976, with the
determinant correction of Kabsch 1978): a reflection is never used.
"""
reference_centre = reference.mean(axis=0)
mobile_centre = mobile.mean(axis=0)
a = mobile - mobile_centre
b = reference - reference_centre
u, _, vt = np.linalg.svd(a.T @ b)
sign = 1.0 if np.linalg.det(u @ vt) >= 0 else -1.0
rotation = u @ np.diag([1.0, 1.0, sign]) @ vt
moved = a @ rotation + reference_centre
rmsd = float(np.sqrt(np.mean(np.sum((moved - reference) ** 2, axis=1))))
return moved, rmsd
# ── Archived artifacts ──────────────────────────────────────────────────
def _sha256(path: Path) -> str:
import hashlib
digest = hashlib.sha256()
with path.open("rb") as handle:
for chunk in iter(lambda: handle.read(1 << 20), b""):
digest.update(chunk)
return digest.hexdigest()
def _verified(run_dir: Path, journal: dict[str, Any], node: str, output: str) -> dict[str, Any]:
"""The latest succeeded attempt's verified ``output`` of ``node``, checked here.
Returns ``{"path", "attempt", "manifest_sha256"}``; ``{}`` when the node
has no such artifact. Raises :class:`RendersRefused` when the bytes on
disk are not the ones vibe-queue verified.
"""
from vibesketch.document.campaign import _canonical
attempts = [a for a in journal.get("attempts", []) if a.get("node") == node]
if not attempts or attempts[-1].get("state") != "succeeded":
return {}
latest = attempts[-1]
record = (latest.get("artifacts") or {}).get(output) or {}
if record.get("status") != "verified":
return {}
base = run_dir / "artifacts" / node / str(latest["attempt"])
manifest_path = base / f"{output}.manifest.json"
try:
manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
except (OSError, ValueError) as error:
raise RendersRefused([f"{node}: its {output} manifest cannot be read: {error}"]) from None
if _canonical(manifest) != record.get("manifest_sha256"):
raise RendersRefused(
[f"{node}: {manifest_path} is not the manifest the journal recorded for attempt "
f"{latest['attempt']}"]
) # fmt: skip
path = base / output
problems = []
if manifest.get("kind") == "file":
if not path.is_file() or _sha256(path) != manifest.get("sha256"):
problems.append(f"{node}: {path} differs from its verified artifact")
else:
for entry in manifest.get("entries", []):
if entry.get("type") != "file":
continue
target = path / entry["path"]
if not target.is_file() or _sha256(target) != entry.get("sha256"):
problems.append(f"{node}: {target} differs from its verified artifact")
if problems:
raise RendersRefused(problems)
return {
"path": path,
"attempt": int(latest["attempt"]),
"manifest_sha256": record["manifest_sha256"],
}
def _trajectory(result: Path, symbols: list[str]) -> tuple[_Points, list[float] | None] | None:
"""The first ``trajectory`` of a result container: coordinates and energies, or ``None``."""
from vibesketch.document.qvf import _modules
_, qvf_reader = _modules()
with qvf_reader.QvfReader(str(result)) as reader:
sections = [s for s in reader.sections if s.get("kind") == "trajectory"]
if not sections:
return None
section = sections[0]
metadata = reader.read_json(section["members"]["metadata"])
coords = np.asarray(reader.read_array(section["members"]["coords"]), dtype=np.float64)
named = metadata.get("symbols")
if coords.ndim != 3 or coords.shape[1:] != (len(symbols), 3) or (named and named != symbols):
return None
energies = metadata.get("energies")
if not (isinstance(energies, list) and len(energies) == coords.shape[0]):
energies = None
return coords, energies
# ── Writing ─────────────────────────────────────────────────────────────
def _archive(
document: Document,
species_id: str,
title: str,
positions: _Points,
trajectory: tuple[_Points, list[float] | None] | None,
) -> bytes:
from vibesketch import __version__
from vibesketch.document.qvf import _atoms, _bonds, _identity, _modules
qvf_writer, _ = _modules()
species = document.species_by_id(species_id)
conformer = Conformer(
method="aligned",
symbols=list(species.conformers[0].symbols),
xyz_angstrom=[[round(float(v), 8) for v in row] for row in positions],
)
writer = qvf_writer.QvfWriter(program="vibe-sketch", version=__version__, calculation=title)
writer.add_structure(
_atoms(conformer),
bonds=_bonds(species, conformer),
label=title,
section_id="structure",
**_identity(species),
)
if trajectory is not None:
coords, energies = trajectory
metadata: dict[str, Any] = {"symbols": list(conformer.symbols)}
if energies is not None:
metadata["energies"] = [float(e) for e in energies]
writer.add_trajectory(
metadata, coords, label=f"{title}: trajectory", section_id="trajectory"
)
return bytes(writer.to_bytes())
def _vibe_view(command: list[str], directory: Path) -> tuple[bool, str]:
environment = {**os.environ, "PYVISTA_OFF_SCREEN": "True"}
try:
done = subprocess.run(
command, cwd=directory, env=environment, capture_output=True, text=True,
timeout=_RENDER_TIMEOUT, check=False,
) # fmt: skip
except (OSError, subprocess.TimeoutExpired) as error:
return False, str(error)
tail = (done.stderr.strip() or done.stdout.strip()).splitlines()[-1:] or [""]
return done.returncode == 0, tail[0]
def _render_stages(
document: Document,
recipe_id: str,
campaign_id: str,
species_id: str,
directory: str | os.PathLike[str],
*,
size: tuple[int, int] = (900, 600),
publication: bool = False,
images: bool = True,
) -> StageRenders:
""":func:`render_stages` into ``target``; the caller turns file errors into refusals."""
from vibesketch.document.campaign import UnknownCampaign, _record, inside, source_digest
try:
recipe, record = _record(document, recipe_id, campaign_id)
except (KeyError, UnknownCampaign) as error:
raise RendersRefused([str(error)]) from None
if species_id not in recipe.species:
raise RendersRefused([f"recipe {recipe_id} does not run species {species_id!r}"])
target = Path(directory).expanduser().resolve()
if target.exists() and (not target.is_dir() or any(target.iterdir())):
raise RendersRefused([f"{target} is not an empty directory"])
if size[0] < 16 or size[1] < 16:
raise RendersRefused(["an image is at least 16x16 pixels"])
campaign_dir = Path(record.directory)
campaign = json.loads((campaign_dir / "campaign.json").read_text(encoding="utf-8"))
pins = {n["id"]: n["sources"] for n in campaign["nodes"] if n.get("sources")}
for node, pin in pins.items():
inside(campaign_dir, str(pin["path"]), f"node {node}'s sources")
journal_path = campaign_dir / "run" / "journal.json"
journal = json.loads(journal_path.read_text(encoding="utf-8")) if journal_path.exists() else {}
run_dir = campaign_dir / "run"
numerical = [s for s in recipe.stages if s.kind != "prepare"]
# The start: the first numerical stage that starts from the sketch's conformer.
first = next(
(s for s in numerical if pins.get(f"{species_id}-{s.id}-run") is not None
and (campaign_dir / pins[f"{species_id}-{s.id}-run"]["path"] / "start").is_dir()),
None,
) # fmt: skip
if first is None:
raise RendersRefused(
[f"campaign {campaign_id} starts no stage of {species_id} from a conformer"]
)
pin = pins[f"{species_id}-{first.id}-run"]
sources = campaign_dir / pin["path"]
if source_digest(sources) != pin["sha256"]:
raise RendersRefused([f"{sources} no longer matches the campaign's pin"])
start = json.loads((sources / "start" / "geometry.json").read_text(encoding="utf-8"))
symbols = list(start["symbols"])
reference, _ = aligned_frame(np.asarray(start["positions_angstrom"], dtype=np.float64))
target.mkdir(parents=True, exist_ok=True)
rendered: list[tuple[str, Path, bool]] = []
stages: list[dict[str, Any]] = []
name = f"{species_id}-start.qvf"
(target / name).write_bytes(
_archive(
document, species_id, f"{species_id} start ({start.get('method')})", reference, None
)
)
rendered.append(("start", target / name, False))
stages.append(
{
"stage": "start", "kind": "prepare", "acceptance": "prepared in the sketch",
"method": start.get("method"), "conformer": start.get("conformer"),
"source": {"sources": pin["path"], "sources_sha256": pin["sha256"]},
"qvf": name, "rmsd_to_start_angstrom": 0.0,
}
) # fmt: skip
for stage in numerical:
check = f"{species_id}-{stage.id}-check"
accepted = _verified(run_dir, journal, check, "accepted")
row: dict[str, Any] = {"stage": stage.id, "kind": stage.kind}
if not accepted:
row |= {"acceptance": "not accepted", "rendered": False}
stages.append(row)
continue
verdict = json.loads((accepted["path"] / "verdict.json").read_text(encoding="utf-8"))
row |= {
"acceptance": "accepted" if verdict.get("accepted") else "rejected",
"artifact": {
"node": check,
"attempt": accepted["attempt"],
"manifest_sha256": accepted["manifest_sha256"],
"result_sha256": verdict.get("result_sha256"),
},
"review": verdict.get("review") or [],
}
for key in ("energy", "frequencies_cm", "program", "program_version"):
if verdict.get(key) is not None:
row[key] = verdict[key]
geometry_file = accepted["path"] / "geometry.json"
if not verdict.get("accepted") or not geometry_file.is_file():
# A Hessian or a single point shares the geometry it started from.
row["geometry"] = stage.geometry.stage
row["rendered"] = False
stages.append(row)
continue
geometry = json.loads(geometry_file.read_text(encoding="utf-8"))
if geometry.get("symbols") != symbols:
raise RendersRefused([f"{check}: the accepted geometry's atoms are not the start's"])
moved, rmsd = superpose(
reference, np.asarray(geometry["positions_angstrom"], dtype=np.float64)
)
trajectory = _trajectory(accepted["path"] / "result.qvf", symbols)
aligned = None
if trajectory is not None:
frames = [superpose(reference, frame) for frame in trajectory[0]]
aligned = (np.stack([f[0] for f in frames]), trajectory[1])
row["trajectory"] = {
"frames": len(frames),
"rmsd_to_start_angstrom": [round(f[1], 6) for f in frames],
}
name = f"{species_id}-{stage.id}.qvf"
title = f"{species_id} {stage.id}: accepted geometry, aligned to the start"
(target / name).write_bytes(_archive(document, species_id, title, moved, aligned))
rendered.append((stage.id, target / name, aligned is not None))
row |= {"qvf": name, "rmsd_to_start_angstrom": round(rmsd, 6), "rendered": True}
stages.append(row)
warnings: list[str] = []
command = None
if images:
from vibesketch.viewer import _command
command = _command()
if command is None:
warnings.append(
"vibe-view is not installed, so no images were drawn; the aligned archives are "
f"in {target}: vibe-view batch {target}/*.qvf"
)
by_stage = {row["stage"]: row for row in stages}
if command is not None:
width, height = size
for stage_id, archive, moving in rendered:
png = archive.with_suffix(".png")
capture = [
*command,
"capture",
str(archive),
"-o",
str(png),
"--size",
f"{width}x{height}",
]
if publication:
capture.append("--publication")
ok, detail = _vibe_view(capture, target)
by_stage[stage_id]["image"] = png.name if ok else None
if not ok:
warnings.append(f"{stage_id}: vibe-view capture failed: {detail}")
if moving:
gif = archive.with_name(archive.stem + "-trajectory.gif")
animate = [
*command, "animate", str(archive), "--kind", "trajectory", "--format", "gif",
"-o", str(gif), "--size", f"{width}x{height}",
] # fmt: skip
ok, detail = _vibe_view(animate, target)
by_stage[stage_id]["animation"] = gif.name if ok else None
if not ok:
warnings.append(f"{stage_id}: vibe-view animate failed: {detail}")
report = {
"schema": _SCHEMA,
"recipe": recipe_id,
"campaign": campaign_id,
"campaign_sha256": record.campaign_sha256,
"species": species_id,
"frame": (
"principal axes of the start geometry, centroid at the origin; every geometry and "
"trajectory frame superposed on it by a proper rotation (Kabsch)"
),
"images": {
"renderer": "vibe-view" if command is not None else None,
"size": f"{size[0]}x{size[1]}",
"publication": publication,
},
"stages": stages,
"warnings": warnings,
}
(target / "renders.json").write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
files = sorted(p.name for p in target.iterdir())
return StageRenders(target, report, files, warnings)
[docs]
def render_stages(
document: Document,
recipe_id: str,
campaign_id: str,
species_id: str,
directory: str | os.PathLike[str],
*,
size: tuple[int, int] = (900, 600),
publication: bool = False,
images: bool = True,
) -> StageRenders:
"""Renders of one species' start and accepted stages, in one orientation.
Writes into ``directory`` (new or empty) the aligned archives, the
record and, when ``images`` and vibe-view is installed, a PNG per
geometry (``size`` in pixels, vibe-view's publication preset with
``publication``) and a GIF per trajectory. Raises :class:`RendersRefused`
when the campaign, species or directory does not fit, or an artifact is
not the one vibe-queue verified. A campaign's files that are missing,
unreadable or point outside it are refused too, and a failure leaves the
directory as it was given (#69).
"""
from vibesketch.document.campaign import CampaignRefused, discard_output
target = Path(directory).expanduser().resolve()
existed = target.exists()
empty = not existed or (target.is_dir() and not any(target.iterdir()))
try:
return _render_stages(
document, recipe_id, campaign_id, species_id, target,
size=size, publication=publication, images=images,
) # fmt: skip
except RendersRefused:
if empty:
discard_output(target, existed)
raise
except CampaignRefused as error:
if empty:
discard_output(target, existed)
raise RendersRefused(error.reasons) from None
except (OSError, ValueError, KeyError, TypeError) as error:
if empty:
discard_output(target, existed)
raise RendersRefused([f"the campaign's files cannot be read: {error}"]) from None