"""Generate checked teaching sketches, exchange files and calculation inputs. Run: python named_reactions.py NEW_DIRECTORY Requires vibe-sketch with RDKit and the optional QVF integration. No quantum calculation or queue submission is performed. """ from __future__ import annotations import argparse import ast import json from pathlib import Path from rdkit import Chem from rdkit.Chem import rdMolDescriptors from vibesketch.document import Canvas, Document, Job, JobSpec, Name, new_document, save from vibesketch.document.clean import clean from vibesketch.document.derive import derive from vibesketch.document.embed import embed from vibesketch.document.formats import ExportRefused, export, import_into from vibesketch.document.ket import write_scene from vibesketch.inputs import WriterRefusal, write_input, write_inputs REACTIONS = { "diels-alder": "C=CC=C.C=C>>C1=CCCCC1", "fischer": "CC(=O)O.CCO>>CC(=O)OCC.O", "suzuki": "Brc1ccccc1.OB(O)c1ccccc1>>c1ccc(-c2ccccc2)cc1", } def document(smiles: str) -> Document: result = import_into(None, smiles, "smiles") if result.warnings: raise RuntimeError(result.warnings) return result.document def formulae(doc: Document) -> list[str]: return [rdMolDescriptors.CalcMolFormula(Chem.MolFromSmiles(s.smiles)) for s in doc.species] def placed(smiles: str, x: float, y: float) -> dict: node = write_scene([Chem.MolFromSmiles(smiles)])["mol0"] points = [a["location"] for a in node["atoms"]] cx = (min(p[0] for p in points) + max(p[0] for p in points)) / 2 cy = (min(p[1] for p in points) + max(p[1] for p in points)) / 2 for atom in node["atoms"]: atom["location"] = [atom["location"][0] - cx + x, atom["location"][1] - cy + y, 0] return node def arrow(tail: tuple[float, float], head: tuple[float, float]) -> dict: return { "type": "arrow", "data": { "mode": "open-angle", "pos": [ {"x": tail[0], "y": tail[1], "z": 0}, {"x": head[0], "y": head[1], "z": 0}, ], }, } def sequence() -> Document: """Two reactions sharing cyclohexene: cycloaddition, then hydrogenation.""" ket = {"root": {"nodes": [], "connections": [], "templates": []}} for index, (smiles, x) in enumerate( [ ("C=CC=C", 0), ("C=C", 4), ("C1=CCCCC1", 10), ("[H][H]", 14), ("C1CCCCC1", 20), ] ): ket[f"mol{index}"] = placed(smiles, x, 0) ket["root"]["nodes"].append({"$ref": f"mol{index}"}) ket["root"]["nodes"].extend( [ {"type": "plus", "location": [2.3, 0, 0]}, arrow((6, 0), (8, 0)), {"type": "plus", "location": [12.4, 0, 0]}, arrow((16, 0), (18, 0)), ] ) doc = derive(new_document().model_copy(update={"canvas": Canvas(format="ket", data=ket)})) assert len(doc.document.reactions) == 2, doc.warnings return doc.document def mechanism() -> Document: """Concerted Diels–Alder electron-pair movements, drawn as KET annotations.""" def atom(x, y): return {"label": "C", "location": [x, y, 0]} ket = { "root": {"nodes": [{"$ref": f"mol{i}"} for i in range(3)]}, "mol0": { "type": "molecule", "atoms": [atom(0, 1), atom(-0.866, 0.5), atom(-0.866, -0.5), atom(0, -1)], "bonds": [ {"type": 2, "atoms": [0, 1]}, {"type": 1, "atoms": [1, 2]}, {"type": 2, "atoms": [2, 3]}, ], }, "mol1": { "type": "molecule", "atoms": [atom(1.6, -0.5), atom(1.6, 0.5)], "bonds": [{"type": 2, "atoms": [0, 1]}], }, "mol2": placed("C1=CCCCC1", 6.3, 0), } ket["root"]["nodes"].append(arrow((3.2, 0), (4.8, 0))) for tail, head, height in [ ((-0.433, 0.75), (0.8, 0.75), 0.4), ((1.6, 0), (0.8, -0.75), 0.4), ((-0.433, -0.75), (-0.866, 0), 0.3), ]: curved = arrow(tail, head) curved["data"].update(mode="elliptical-arc-arrow-filled-triangle", height=height) ket["root"]["nodes"].append(curved) result = derive(new_document().model_copy(update={"canvas": Canvas(format="ket", data=ket)})) assert len(result.document.reactions) == 1 assert formulae(result.document) == ["C4H6", "C2H4", "C6H10"] return result.document def bundle_files(output: Path, doc: Document, program: str, spec: JobSpec) -> None: output.mkdir(parents=True, exist_ok=False) (output / "spec.json").write_text(spec.model_dump_json(indent=2), encoding="utf-8") bundle = write_input(doc, species=doc.species[0].id, program=program, spec=spec) for file in bundle.files: if file.name.endswith(".py"): ast.parse(file.content) if file.name.endswith(".json"): json.loads(file.content) (output / file.name).write_text(file.content, encoding="utf-8") def generate(output: Path) -> None: output.mkdir(parents=True, exist_ok=False) documents = {name: document(smiles) for name, smiles in REACTIONS.items()} assert formulae(documents["diels-alder"]) == ["C4H6", "C2H4", "C6H10"] documents.update( { "two-step": sequence(), "diels-alder-mechanism": mechanism(), "stereochemistry": document("C[C@H](O)C(=O)O.F/C=C/F.[13CH4].[NH4+].[CH3]"), "lactic-acid": document("C[C@H](O)C(=O)O"), "methyl-radical": document("[CH3]"), } ) assert documents["methyl-radical"].species[0].multiplicity == 2 for name, doc in documents.items(): save(doc, output / f"{name}.vsk") (output / f"{name}.svg").write_bytes(export(doc, "svg").content) # Preserve drawing identities while standardizing coordinates and making 3D. cleaned, cleanup = clean(documents["lactic-acid"]) prepared, report = embed(cleaned.document, count=3) assert not report.failed, report.failed assert prepared.species[0].smiles == documents["lactic-acid"].species[0].smiles assert 1 <= len(prepared.species[0].conformers) <= 3 cached, cache_report = embed(prepared) assert cache_report.kept == [prepared.species[0].id] assert cached.species[0].conformers == prepared.species[0].conformers prepared, redo_report = embed(prepared, [prepared.species[0].id], count=3, replace=True) assert not redo_report.failed prepared.species[0].label = "1" prepared.species[0].name = Name(value="lactic acid", source="user", confidence="high") save(prepared, output / "lactic-acid-3d.vsk") (output / "lactic-acid.png").write_bytes(export(prepared, "png").content) exchange = output / "formats" exchange.mkdir() for fmt in ("vsk", "ket", "mol", "sdf", "rxn", "smiles", "inchi", "qvf"): payload = export(documents["diels-alder"], fmt) (exchange / f"diels-alder.{payload.extension}").write_bytes(payload.content) restored = import_into(None, payload.content, fmt).document assert sorted(formulae(restored)) == ["C2H4", "C4H6", "C6H10"] cdxml = """ """ (exchange / "ethanol.cdxml").write_text(cdxml, encoding="utf-8") assert formulae(import_into(None, cdxml, "cdxml").document) == ["C2H6O"] # A deliberately large *synthetic structure*, not a named reaction example. # Supply coordinates directly so this demonstrates V3000 without an expensive layout. atoms = [{"label": "C", "location": [i * 0.87, (i % 2) * 0.5, 0]} for i in range(1000)] ket = { "root": {"nodes": [{"$ref": "mol0"}]}, "mol0": { "type": "molecule", "atoms": atoms, "bonds": [{"type": 1, "atoms": [i, i + 1]} for i in range(999)], }, } large = import_into(None, json.dumps(ket), "ket").document molfile = export(large, "mol").content assert b"V3000" in molfile (exchange / "large-v3000.mol").write_bytes(molfile) assert len(import_into(None, molfile, "mol").document.species) == 1 # Cover every task family; these are inputs, not results or method recommendations. for program, tasks in { "vibeqc": ["sp", "opt", "freq", "opt+freq", "tddft"], "orca": ["sp", "opt", "freq", "opt+freq", "tddft", "nmr"], "qcschema": ["sp", "freq"], }.items(): for task in tasks: spec = JobSpec.model_validate( { "task": task, "method": "HF" if program == "qcschema" else "B3LYP", "basis": "def2-SVP", **({"states": 3} if task == "tddft" else {}), } ) bundle_files(output / "inputs" / program / task, prepared, program, spec) for program in ("xyz", "qcschema-molecule"): bundle_files(output / "inputs" / program, prepared, program, JobSpec(task="geometry")) for program in ("vibeqc", "orca"): spec = JobSpec.model_validate( { "task": "opt", "method": "B3LYP", "basis": "def2-TZVP", "dispersion": "D3BJ", "approximation": "RIJCOSX", "aux_basis": "def2-tzvp-jk" if program == "vibeqc" else "def2/J", "solvent": {"model": "CPCM" if program == "vibeqc" else "SMD", "name": "water"}, "resources": {"nprocs": 2, **({"maxcore_mb": 1000} if program == "orca" else {})}, "constraints": [{"kind": "position", "atoms": [0]}], } ) bundle_files(output / "inputs" / program / "options", prepared, program, spec) coordinate_spec = JobSpec.model_validate( { "task": "opt", "method": "B3LYP", "basis": "def2-SVP", "constraints": [ {"kind": "distance", "atoms": [0, 1]}, {"kind": "angle", "atoms": [0, 1, 3]}, {"kind": "dihedral", "atoms": [0, 1, 3, 4]}, ], } ) bundle_files( output / "inputs" / "orca" / "internal-constraints", prepared, "orca", coordinate_spec ) # An existing Job can be selected by id, just as `input --job j1` does. prepared.jobs = [ Job( id="j1", species=prepared.species[0].id, program="vibeqc", spec=JobSpec(task="sp", method="B3LYP", basis="def2-SVP"), ) ] assert write_input(prepared, job="j1").files save(prepared, output / "lactic-acid-job.vsk") xyz_file = next((output / "inputs" / "xyz").glob("*.xyz")) imported = import_into(None, xyz_file.read_bytes(), "xyz").document assert imported.species[0].conformers refusals = {} for program, spec in [ ("vibeqc", JobSpec(task="nmr", method="B3LYP", basis="def2-SVP")), ("qcschema", JobSpec(task="opt", method="HF", basis="def2-SVP")), ("xyz", JobSpec(task="sp", method="HF", basis="def2-SVP")), ]: try: write_input(prepared, species=prepared.species[0].id, program=program, spec=spec) except WriterRefusal as exc: refusals[program] = exc.reasons else: raise AssertionError(f"Expected {program} to refuse {spec.task}") try: export(documents["two-step"], "rxn") except ExportRefused as exc: refusals["multi-step RXN"] = str(exc) else: raise AssertionError("RXN must refuse two reactions") batch_doc, report = embed(documents["diels-alder"]) assert not report.failed batch = write_inputs( batch_doc, program="vibeqc", spec=JobSpec(task="sp", method="B3LYP", basis="def2-SVP") ) assert len(batch.bundles) == 3 and not batch.refused save(batch_doc, output / "diels-alder-3d.vsk") (output / "refusals.json").write_text(json.dumps(refusals, indent=2), encoding="utf-8") print(f"Created {output}; all chemistry, round-trip, cache and writer checks passed.") print(f"Embedding warnings: {report.warnings}; cleanup warnings: {cleanup.left}") print("No calculation was run or submitted.") if __name__ == "__main__": parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("directory", type=Path) generate(parser.parse_args().directory)