"""Saved staged workflows on a document: the publication baseline, set, show, remove.
A recipe is additive to the document (schema v1 gains an optional
``recipes`` list) and program-neutral; what each stage would send to vibe-qc,
and what it refuses, comes from :mod:`vibesketch.inputs.recipe`. The
baseline is the workflow of the vibyverse manuscript (#42): the sketch's
MMFF94 preparation, a MACE-OFF23 optimisation, a PBE/def2-SVP optimisation
without dispersion, a same-surface finite-difference Hessian and a
DLPNO-CCSD(T)/cc-pVDZ single point, each numerical stage starting from the
accepted geometry of its predecessor.
"""
from __future__ import annotations
from collections.abc import Sequence
from typing import Any
from vibesketch.document.model import (
DftOptimization,
DftSurface,
DlpnoOptions,
Document,
GeometryRef,
HessianOptions,
MaceOptions,
Recipe,
Stage,
)
__all__ = [
"UnknownRecipe",
"baseline_recipe",
"recipe_report",
"remove_recipe",
"set_recipe",
]
[docs]
class UnknownRecipe(KeyError):
"""The document has no recipe with that id."""
def __str__(self) -> str:
return f"the document has no recipe {self.args[0]!r}"
[docs]
def baseline_recipe(
species: Sequence[str], *, identifier: str = "recipe1", name: str | None = None
) -> Recipe:
"""The publication baseline for ``species``: prepare, MACE, PBE, Hessian, DLPNO."""
surface = DftSurface()
return Recipe(
id=identifier,
name=name or "MACE → PBE/def2-SVP → Hessian → DLPNO-CCSD(T)",
species=list(species),
stages=[
Stage(
id="prepare",
kind="prepare",
label="ETKDG and MMFF94 preparation in vibe-sketch",
geometry=GeometryRef(conformer=0),
),
Stage(
id="mace",
kind="mace",
label="MACE-OFF23 medium optimisation",
parent="prepare",
geometry=GeometryRef(stage="prepare"),
mace=MaceOptions(),
),
Stage(
id="dft",
kind="dft",
label="PBE/def2-SVP optimisation, no dispersion",
parent="mace",
geometry=GeometryRef(stage="mace"),
dft=DftOptimization(surface=surface),
),
Stage(
id="hessian",
kind="hessian",
label="PBE/def2-SVP finite-difference Hessian",
parent="dft",
geometry=GeometryRef(stage="dft"),
hessian=HessianOptions(surface=surface),
),
Stage(
id="dlpno",
kind="dlpno",
label="DLPNO-CCSD(T)/cc-pVDZ single point, TightPNO, chemical frozen core",
parent="hessian",
geometry=GeometryRef(stage="dft"),
dlpno=DlpnoOptions(),
),
],
)
[docs]
def set_recipe(document: Document, recipe: Recipe) -> Document:
"""``document`` with ``recipe`` added, or replacing the recipe with its id.
The document's validation checks the recipe's species exist and its id
clashes with no species, reaction or job.
"""
recipes = [r for r in document.recipes if r.id != recipe.id]
position = next((i for i, r in enumerate(document.recipes) if r.id == recipe.id), None)
if position is None:
recipes.append(recipe)
else:
recipes.insert(position, recipe)
rebuilt = document.model_copy(update={"recipes": recipes})
return Document.model_validate(rebuilt.to_json_data())
[docs]
def remove_recipe(document: Document, identifier: str) -> Document:
"""``document`` without the recipe ``identifier``; raises :class:`UnknownRecipe`."""
if all(r.id != identifier for r in document.recipes):
raise UnknownRecipe(identifier)
recipes = [r for r in document.recipes if r.id != identifier]
rebuilt = document.model_copy(update={"recipes": recipes})
return Document.model_validate(rebuilt.to_json_data())
[docs]
def recipe_report(document: Document, identifier: str | None = None) -> list[dict[str, Any]]:
"""Each recipe (or the one named) with every stage resolved per species.
A stage's entry says what vibe-qc would run, what the operator supplies
at launch, and the refusals that stop it. Raises :class:`UnknownRecipe`.
"""
from vibesketch.inputs.recipe import resolve_recipe
chosen = document.recipes
if identifier is not None:
chosen = [r for r in document.recipes if r.id == identifier]
if not chosen:
raise UnknownRecipe(identifier)
species = {s.id: s for s in document.species}
report = []
for recipe in chosen:
resolved = resolve_recipe(recipe, species)
report.append(
{
"recipe": recipe.model_dump(mode="json", exclude_none=True),
"stages": {
sid: [stage.to_dict() for stage in stages] for sid, stages in resolved.items()
},
# In stage order, each stage once.
"refused": [
stage.id
for stage in recipe.stages
if any(
r.stage == stage.id and r.refused for rs in resolved.values() for r in rs
)
],
}
)
return report