"""Reaction-path inputs from a sketched reaction: NEB-TS for ORCA, CI-NEB for vibe-qc (#29).
Draw a reaction, then write the inputs for a nudged elastic band between its
reactants and its products. Three things must hold, and each is checked:
* **The atoms are paired.** :mod:`vibesketch.chem.mapping` pairs every
reactant atom with a product atom by the fewest bonds formed and broken.
The pairing is shown before anything is written
(:func:`reaction_mapping`), and the user can correct it by pinning atoms
(``m1:3=m3:5``: reactant species m1's atom 3 becomes product species
m3's atom 5). An ambiguous pairing is refused, never guessed.
* **The ends match.** :mod:`vibesketch.chem.endpoints` builds the reactant
and product complexes with their atoms in one order, the product turned
onto the reactant so the band does not rotate.
* **The writer can say it.** The ORCA writer reuses the keyword model of
:mod:`vibesketch.inputs.orca` and adds ``NEB-TS`` (climbing-image NEB, then
a transition-state optimisation from the climbing image), ``NEB-CI`` or
``NEB`` with a ``%neb`` block; the product geometry goes in its own XYZ
file. vibe-qc's ``run_neb`` has the climbing image but no
transition-state refinement, so its writer writes ``neb-ci`` and ``neb``
and refuses ``neb-ts`` with that reason. Both optimise the two ends
first, as the built complexes are no minima.
Reagents drawn at the arrow are not part of the path. A species takes part
once per side; ``2 A -> B`` is refused as unbalanced, since the drawing's
species list does not say how many molecules react.
"""
from __future__ import annotations
import json
import re
from collections.abc import Sequence
from dataclasses import dataclass, field
from typing import Any, Literal
from pydantic import BaseModel, ConfigDict, Field
from vibesketch.document import model
from vibesketch.document.model import Document, JobSpec, Species
from vibesketch.inputs.base import InputBundle, InputFile, WriterRefusal
from vibesketch.inputs.registry import UnknownReference
__all__ = [
"NEB_PROGRAMS",
"NebSpec",
"ReactionPath",
"mapping_text",
"neb_inputs",
"parse_pin",
"reaction_mapping",
"reaction_path",
]
#: The programs this module writes reaction-path inputs for.
NEB_PROGRAMS = ("orca", "vibeqc")
_PIN = re.compile(r"^\s*([A-Za-z0-9_.-]+):(\d+)\s*=\s*([A-Za-z0-9_.-]+):(\d+)\s*$")
[docs]
class NebSpec(BaseModel):
"""What to compute along the path: the level of theory and the band.
``kind`` is ``neb-ts`` (climbing-image NEB, then a transition-state
optimisation; ORCA), ``neb-ci`` (climbing-image NEB) or ``neb`` (a
plain band). ``images`` counts the images between the two ends.
"""
model_config = ConfigDict(extra="forbid")
kind: Literal["neb-ts", "neb-ci", "neb"] = "neb-ts"
method: str = Field(min_length=1)
basis: str | None = None
aux_basis: str | None = None
approximation: str | None = None
dispersion: str | None = None
# Through the module: the documentation would find two Solvents otherwise.
solvent: model.Solvent | None = None
resources: model.Resources | None = None
images: int = Field(default=8, ge=3, le=64)
#: Optimise both ends before the band; the built complexes are no minima.
preopt_ends: bool = True
[docs]
def job(self) -> JobSpec:
"""The level of theory as a job specification (an optimisation needs the same gradients)."""
return JobSpec(
task="opt",
method=self.method,
basis=self.basis,
aux_basis=self.aux_basis,
approximation=self.approximation,
dispersion=self.dispersion,
solvent=self.solvent,
resources=self.resources,
)
[docs]
def parse_pin(text: str) -> tuple[str, int, str, int]:
"""``m1:3=m3:5`` as (reactant species, atom, product species, atom), atoms 1-based."""
match = _PIN.match(text)
if match is None:
raise ValueError(
f"a pinned pair is written reactant:atom=product:atom, for example m1:3=m3:5; "
f"not {text!r}"
)
return match[1], int(match[2]), match[3], int(match[4])
[docs]
@dataclass(frozen=True)
class ReactionPath:
"""A mapped reaction and, when the mapping is clear, its two ends."""
reaction: str
mapping: Any | None # vibesketch.chem.mapping.AtomMapping, when the atoms could be paired
endpoints: Any | None # vibesketch.chem.endpoints.Endpoints, when the mapping is clear
charge: int
multiplicity: int
#: Why no inputs can be written for this reaction, if any reason.
refusals: tuple[str, ...] = ()
#: Where each species' heavy atoms are drawn on the canvas, for this reaction:
#: (species, atom index in the species) -> KET coordinates.
canvas: dict[tuple[str, int], tuple[float, float]] = field(default_factory=dict)
[docs]
def report(self) -> dict[str, Any]:
"""The mapping as the user reviews it before anything is written."""
mapping = self.mapping
if mapping is None:
return {
"reaction": self.reaction,
"charge": self.charge,
"multiplicity": self.multiplicity,
"atoms": [],
"changes": [],
"ambiguous": False,
"alternatives": [],
"refusals": list(self.refusals),
}
reactants, products = mapping.reactants, mapping.products
heavy = [i for i, atom in enumerate(reactants.mol.GetAtoms()) if atom.GetAtomicNum() > 1]
pairs = self.endpoints.pairs if self.endpoints is not None else mapping.pairs
centre = mapping.centre(pairs)
def where(side: Any, atom: int) -> list[float] | None:
point = self.canvas.get(side.owners[atom])
return None if point is None else [point[0], point[1]]
def entry(i: int) -> dict[str, Any]:
(r_species, r_atom), (p_species, p_atom) = (
reactants.owners[i],
products.owners[pairs[i]],
)
return {
"index": i + 1,
"reactant": reactants.label(i),
"product": products.label(pairs[i]),
"reactant_species": r_species,
"reactant_atom": r_atom + 1,
"product_species": p_species,
"product_atom": p_atom + 1,
"reactant_at": where(reactants, i),
"product_at": where(products, pairs[i]),
"centre": i in centre,
}
return {
"reaction": self.reaction,
"reactants": [{"species": s, "smiles": smiles} for s, smiles, _ in reactants.species],
"products": [{"species": s, "smiles": smiles} for s, smiles, _ in products.species],
"charge": self.charge,
"multiplicity": self.multiplicity,
"atoms": [entry(i) for i in heavy],
"changes": mapping.changes(),
"cost": mapping.cost,
"ambiguous": mapping.ambiguous,
"alternatives": [mapping.changes(alternative) for alternative in mapping.alternatives],
"choices": _choices(mapping) if mapping.ambiguous else [],
"pinned": [f"{reactants.label(a)}={products.label(b)}" for a, b in mapping.pinned],
"rmsd_angstrom": None if self.endpoints is None else round(self.endpoints.rmsd, 3),
"warnings": [] if self.endpoints is None else list(self.endpoints.warnings),
"refusals": list(self.refusals),
}
def _choices(mapping: Any) -> list[dict[str, Any]]:
"""Each equally cheap pairing with the pins that select it.
A pairing's pins are its pairs for the heavy atoms that some other
cheapest pairing pairs differently: pinning them leaves no doubt
between these pairings.
"""
reactants, products = mapping.reactants, mapping.products
pairings = [mapping.pairs, *mapping.alternatives]
heavy = [i for i, atom in enumerate(reactants.mol.GetAtoms()) if atom.GetAtomicNum() > 1]
found = []
for pairing in pairings:
differing = [
i
for i in heavy
if any(other[i] != pairing[i] for other in pairings if other is not pairing)
]
pins = []
for i in differing:
(r_species, r_atom), (p_species, p_atom) = (
reactants.owners[i],
products.owners[pairing[i]],
)
pins.append(f"{r_species}:{r_atom + 1}={p_species}:{p_atom + 1}")
found.append({"changes": mapping.changes(pairing), "pins": pins})
return found
def _species_of(document: Document, identifiers: Sequence[str]) -> list[Species]:
found = {s.id: s for s in document.species}
return [found[i] for i in identifiers]
def _spin(side: list[Species]) -> int | None:
"""The complex's multiplicity, when the species' own say it; ``None`` when ambiguous."""
open_shell = [s.multiplicity for s in side if s.multiplicity > 1]
if not open_shell:
return 1
if len(open_shell) == 1:
return open_shell[0]
return None
[docs]
def reaction_path(document: Document, reaction: str, pins: Sequence[str] = ()) -> ReactionPath:
"""Map ``reaction`` of ``document`` and build its two ends.
Raises :class:`UnknownReference` for a reaction or pinned species the
document lacks and ``ValueError`` for a malformed pin; everything else
that stops the path is in :attr:`ReactionPath.refusals`.
"""
from vibesketch.chem.endpoints import EndpointsFailed, build_endpoints
from vibesketch.chem.mapping import MappingFailed, map_reaction, side
chosen = next((r for r in document.reactions if r.id == reaction), None)
if chosen is None:
raise UnknownReference("reaction", reaction)
refusals: list[str] = []
if chosen.arrow not in ("forward", "equilibrium"):
refusals.append(
f"reaction {reaction} is drawn with a {chosen.arrow} arrow, which is no reaction "
"step a path can follow"
)
reactants, products = (
_species_of(document, chosen.reactants),
_species_of(document, chosen.products),
)
missing = [s.id for s in reactants + products if not s.smiles]
if missing:
refusals.append(f"species {', '.join(missing)} have no structure to map")
charge = sum(s.charge for s in reactants)
if charge != sum(s.charge for s in products):
refusals.append(
f"the charges differ: {charge} for the reactants, "
f"{sum(s.charge for s in products)} for the products"
)
spin_r, spin_p = _spin(reactants), _spin(products)
if spin_r is None or spin_p is None:
refusals.append(
"more than one open-shell species on a side: the complex's spin state is not "
"determined by the drawing"
)
elif spin_r != spin_p:
refusals.append(
f"the multiplicities differ: {spin_r} for the reactants, {spin_p} for the products; "
"a band needs one spin state"
)
if refusals:
return ReactionPath(reaction, None, None, charge, spin_r or 1, tuple(refusals))
left = side([(s.id, s.smiles or "") for s in reactants])
right = side([(s.id, s.smiles or "") for s in products])
pairs = [_pin_indices(left, right, pin) for pin in pins]
try:
mapping = map_reaction(left, right, pairs)
except MappingFailed as failure:
return ReactionPath(reaction, None, None, charge, spin_r or 1, (str(failure),))
canvas = _canvas_locations(document, reaction)
if mapping.ambiguous:
reason = (
f"the atom mapping is ambiguous: {len(mapping.alternatives) + 1} pairings break "
"and form as few bonds; pin atoms (reactant:atom=product:atom) to choose one"
)
return ReactionPath(reaction, mapping, None, charge, spin_r or 1, (reason,), canvas)
try:
endpoints = build_endpoints(mapping)
except EndpointsFailed as failure:
return ReactionPath(reaction, mapping, None, charge, spin_r or 1, (str(failure),), canvas)
return ReactionPath(reaction, mapping, endpoints, charge, spin_r or 1, (), canvas)
def _canvas_locations(
document: Document, reaction: str
) -> dict[tuple[str, int], tuple[float, float]]:
"""Where the reaction's species are drawn: (species, atom) -> KET coordinates.
The reaction's own molecules are found as derive finds them (the n-th
arrow with molecules on both sides is the n-th reaction), and each
species' atoms are matched onto its drawing. Empty when the canvas no
longer matches the reaction.
"""
from rdkit import rdBase
from vibesketch.chem.species import identity, mol_from_smiles
from vibesketch.document.derive import _roles
from vibesketch.document.ket import read_scene
ket: dict[str, Any] = dict(document.canvas.data)
position = next((i for i, r in enumerate(document.reactions) if r.id == reaction), None)
if not ket or position is None:
return {}
chosen = document.reactions[position]
by_key = {s.inchikey or s.smiles or s.id: s for s in document.species}
scene = read_scene(ket)
species_of: dict[int, str] = {}
for index, fragment in enumerate(scene.fragments):
if fragment.mol is not None:
found = by_key.get(identity(fragment.mol).key)
if found is not None:
species_of[index] = found.id
steps = [
(tails, heads)
for _, tails, _, heads, _ in _roles(scene)
if any(i in species_of for i in tails) and any(i in species_of for i in heads)
]
if position >= len(steps):
return {}
tails, heads = steps[position]
if {species_of[i] for i in tails if i in species_of} != set(chosen.reactants) or {
species_of[i] for i in heads if i in species_of
} != set(chosen.products):
return {}
smiles = {s.id: s.smiles for s in document.species}
found_at: dict[tuple[str, int], tuple[float, float]] = {}
blocker = rdBase.BlockLogs()
try:
for side in (tails, heads):
for index in side:
species = species_of.get(index)
if species is None or not smiles.get(species):
continue
if any(key[0] == species for key in found_at):
continue # the species' first drawing on this side
fragment = scene.fragments[index]
match = fragment.mol.GetSubstructMatch(mol_from_smiles(smiles[species] or ""))
atoms = ket[fragment.node]["atoms"]
for atom, matched in enumerate(match):
x, y = atoms[fragment.atoms[matched]]["location"][:2]
found_at[(species, atom)] = (float(x), float(y))
finally:
del blocker
return found_at
def _pin_indices(left: Any, right: Any, pin: str) -> tuple[int, int]:
r_species, r_atom, p_species, p_atom = parse_pin(pin)
def index(side: Any, species: str, atom: int, kind: str) -> int:
for identifier, _, first in side.species:
if identifier == species:
count = sum(1 for owner, _ in side.owners if owner == species)
if not 1 <= atom <= count:
raise ValueError(f"{kind} {species} has atoms 1 to {count}, not {atom}")
return int(first) + atom - 1
raise UnknownReference(f"{kind} species", species)
return index(left, r_species, r_atom, "reactant"), index(right, p_species, p_atom, "product")
[docs]
def mapping_text(report: dict[str, Any]) -> str:
"""The mapping report as lines a person reviews: pairs, bond changes, refusals."""
sides = " -> ".join(
" + ".join(entry["species"] for entry in report.get(key, []))
for key in ("reactants", "products")
)
lines = [
f"{report['reaction']}: {sides or 'not mapped'} "
f"(charge {report['charge']}, multiplicity {report['multiplicity']})"
]
if report["atoms"]:
lines.append(" atoms, reactant -> product (hydrogens follow their atoms):")
lines += [f" {atom['reactant']} -> {atom['product']}" for atom in report["atoms"]]
if report["changes"]:
lines.append(" bonds:")
lines += [f" {change}" for change in report["changes"]]
for number, alternative in enumerate(report["alternatives"], start=2):
lines.append(f" pairing {number}, as cheap:")
lines += [f" {change}" for change in alternative]
if report.get("pinned"):
lines.append(" pinned: " + ", ".join(report["pinned"]))
if report.get("rmsd_angstrom") is not None:
lines.append(f" ends built; they differ by {report['rmsd_angstrom']} Å RMS")
lines += [f" warning: {warning}" for warning in report.get("warnings", [])]
lines += [f" refused: {reason}" for reason in report["refusals"]]
return "\n".join(lines) + "\n"
[docs]
def reaction_mapping(document: Document, reaction: str, pins: Sequence[str] = ()) -> dict[str, Any]:
"""The atom mapping of ``reaction`` as the user reviews it, with any refusal."""
return reaction_path(document, reaction, pins).report()
# ── Writers ─────────────────────────────────────────────────────────────
def _complex(path: ReactionPath, spec: NebSpec) -> Species:
"""The reactant complex as one species, for the writers' keyword models."""
from vibesketch.document.model import Conformer
assert path.endpoints is not None and path.mapping is not None
names = [s for s, _, _ in path.mapping.reactants.species]
smiles = ".".join(smiles for _, smiles, _ in path.mapping.reactants.species)
return Species.model_validate(
{
"id": path.reaction,
"smiles": smiles,
"charge": path.charge,
"multiplicity": path.multiplicity,
"name": {"value": " + ".join(names), "source": "vibe-sketch", "confidence": "high"},
"conformers": [
Conformer(
method=path.endpoints.method,
symbols=list(path.endpoints.symbols),
xyz_angstrom=[list(p) for p in path.endpoints.reactant],
).model_dump()
],
}
)
def _xyz(symbols: Sequence[str], points: Sequence[Sequence[float]], comment: str) -> str:
lines = [str(len(symbols)), comment]
lines += [
f"{s:<2}{x:14.10f}{y:15.10f}{z:15.10f}"
for s, (x, y, z) in zip(symbols, points, strict=True)
]
return "\n".join(lines) + "\n"
def _slug(text: str) -> str:
return re.sub(r"[^a-z0-9]+", "-", text.lower()).strip("-")
def _stem(path: ReactionPath, spec: NebSpec) -> str:
parts = [path.reaction, _slug(spec.method)]
if spec.basis:
parts.append(_slug(spec.basis))
parts.append(spec.kind)
return "_".join(parts)
def _orca(path: ReactionPath, spec: NebSpec) -> tuple[list[str], list[InputFile]]:
from vibesketch import __version__
from vibesketch.inputs.orca import OrcaWriter, _blocks
writer = OrcaWriter()
species = _complex(path, spec)
job = spec.job()
reasons = [r.replace("optimisations", "reaction paths") for r in writer.refusals(species, job)]
if reasons:
return reasons, []
assert path.endpoints is not None and path.mapping is not None
stem = _stem(path, spec)
keyword = {"neb-ts": "NEB-TS", "neb-ci": "NEB-CI", "neb": "NEB"}[spec.kind]
# The level of theory as a single point's line, the band's keyword after it.
line = writer.keyword_line(species, job.model_copy(update={"task": "sp"})) + f" {keyword}"
reaction = (
" + ".join(s for s, _, _ in path.mapping.reactants.species)
+ " -> "
+ " + ".join(s for s, _, _ in path.mapping.products.species)
)
product = f"{stem}_product.xyz"
lines = [
line,
f"# {path.reaction}: {reaction} — {keyword}, {spec.images} images, "
f"{spec.method}" + (f"/{spec.basis}" if spec.basis else ""),
f"# vibe-sketch {__version__}; ends built by {path.endpoints.method}, mapped by the fewest "
"bonds changed; both ends list their atoms in one order",
*(f"# {change}" for change in path.mapping.changes()),
*_blocks(job),
"%neb",
f' NEB_End_XYZFile "{product}"',
f" Nimages {spec.images}",
f" PreOpt_Ends {'true' if spec.preopt_ends else 'false'}",
"end",
f"* xyz {path.charge} {path.multiplicity}",
*(
f"{s:<2}{x:14.10f}{y:15.10f}{z:15.10f}"
for s, (x, y, z) in zip(path.endpoints.symbols, path.endpoints.reactant, strict=True)
),
"*",
]
files = [
InputFile(f"{stem}.inp", "\n".join(lines) + "\n", "text/plain"),
InputFile(
product,
_xyz(path.endpoints.symbols, path.endpoints.product, f"{path.reaction} product end"),
"chemical/x-xyz",
),
]
return [], files
def _vibeqc(path: ReactionPath, spec: NebSpec) -> tuple[list[str], list[InputFile]]:
from vibesketch import __version__
from vibesketch.elements import ATOMIC_NUMBER
from vibesketch.inputs.vibeqc import _basis, _route
from vibesketch.units import ANGSTROM_TO_BOHR
reasons: list[str] = []
if spec.kind == "neb-ts":
reasons.append(
"vibe-qc's run_neb ends at the climbing image and has no transition-state "
"optimisation to follow it; ask for kind 'neb-ci' (or write ORCA's NEB-TS)"
)
route = _route(spec.method, path.multiplicity)
if route is None or route.kind not in ("hf", "dft"):
reasons.append(
"vibe-sketch writes vibe-qc reaction paths with HF or a functional, "
f"not {spec.method!r}"
)
if spec.basis is None:
reasons.append(f"{spec.method} needs a basis set, for example def2-SVP")
elif _basis(spec.basis) is None:
reasons.append(f"vibe-qc ships no basis set {spec.basis!r} in vibe-sketch's table")
for name, label in (
("aux_basis", "auxiliary basis set"),
("approximation", "integral approximation"),
("dispersion", "dispersion correction"),
("solvent", "solvation model"),
):
if getattr(spec, name) is not None:
reasons.append(f"vibe-sketch does not write a {label} into vibe-qc's run_neb yet")
if spec.resources is not None and spec.resources.maxcore_mb is not None:
reasons.append("vibe-qc takes no memory per core; leave maxcore_mb out")
if spec.resources is not None and spec.resources.nprocs > 1:
reasons.append(
"vibe-qc's run_neb cannot evaluate images in parallel processes yet (it cannot "
"pickle a Molecule; mpei/vibe-qc#564); leave nprocs at 1"
)
if reasons:
return reasons, []
assert route is not None and spec.basis is not None
assert path.endpoints is not None and path.mapping is not None
stem = _stem(path, spec)
basis = _basis(spec.basis)
symbols = path.endpoints.symbols
def rows(points: Sequence[Sequence[float]]) -> list[str]:
return [
f" ({ATOMIC_NUMBER[s]:>3}, {x:15.10f}, {y:15.10f}, {z:15.10f}), # {s}"
for s, (x, y, z) in zip(symbols, points, strict=True)
]
functional = f", functional={json.dumps(route.functional)}" if route.functional else ""
arguments = [
f" basis={json.dumps(basis)},",
f" n_images={spec.images},",
f" method={json.dumps(route.method)},",
*([f" functional={json.dumps(route.functional)},"] if route.functional else []),
f" climbing_image={spec.kind == 'neb-ci'},",
' interpolation="idpp",',
f" output={json.dumps(stem)},",
" n_jobs=1, # one image at a time: parallel workers cannot take a Molecule yet",
]
reaction = (
" + ".join(s for s, _, _ in path.mapping.reactants.species)
+ " -> "
+ " + ".join(s for s, _, _ in path.mapping.products.species)
)
preopt = [
"",
"# The ends were built, not optimised: optimise each first.",
"def optimised(molecule, name):",
f" result = optimize_molecule(molecule, {json.dumps(basis)}, "
f"method={json.dumps(route.method)}{functional})",
" if not result.converged:",
' raise SystemExit(f"the {name} end did not converge: {result.message}")',
" return result.system",
"",
'reactant = optimised(reactant, "reactant")',
'product = optimised(product, "product")',
]
lines = [
f"# {path.reaction}: {reaction} — "
+ ("climbing-image NEB" if spec.kind == "neb-ci" else "NEB")
+ f", {spec.images} images, {spec.method}/{spec.basis}",
f"# Written by vibe-sketch {__version__}; ends built by {path.endpoints.method}, mapped "
"by the fewest bonds changed; both ends list their atoms in one order.",
*(f"# {change}" for change in path.mapping.changes()),
f"# Run: python {stem}.py",
"from vibeqc import Atom, Molecule, optimize_molecule, run_neb",
"",
f"ANGSTROM_TO_BOHR = {ANGSTROM_TO_BOHR!r} # CODATA 2018, as vibe-qc uses",
"",
"# Atomic number and position in angstrom, one atom per line, the same order in both.",
"REACTANT = [",
*rows(path.endpoints.reactant),
"]",
"PRODUCT = [",
*rows(path.endpoints.product),
"]",
"",
"",
"def molecule(rows):",
" return Molecule(",
" [Atom(z, [x * ANGSTROM_TO_BOHR, y * ANGSTROM_TO_BOHR, w * ANGSTROM_TO_BOHR])",
" for z, x, y, w in rows],",
f" charge={path.charge},",
f" multiplicity={path.multiplicity},",
" )",
"",
"",
"reactant, product = molecule(REACTANT), molecule(PRODUCT)",
*(preopt if spec.preopt_ends else []),
"",
"result = run_neb(",
" reactant,",
" product,",
*arguments,
")",
"if result is None: # VIBEQC_DRY_RUN: the planned job is recorded, nothing computed",
" raise SystemExit(0)",
f"result.write_qvf({json.dumps(stem)})",
'print("highest image:", result.transition_state_index)',
]
return [], [InputFile(f"{stem}.py", "\n".join(lines) + "\n", "text/x-python")]