Source code for vibesketch.inputs.vibeqc

"""The vibe-qc writer: a Python script that runs the job with ``vibeqc.run_job``.

The script follows vibe-view's vibe-qc templates (one ``Molecule``, one
``run_job`` call, QVF and Molden output) with one correction: charge and
multiplicity go to ``Molecule``, the only place vibe-qc reads them;
``run_job`` takes neither.

The translation is fixed by what vibe-qc's ``run_job`` accepts, read from its
source and refused where it ends (AGENTS.md, ground rule 5):

- **Methods.** HF becomes ``rhf`` or ``uhf``, a functional ``rks`` or
  ``uks`` with ``functional=``, never ``method="auto"``, which runs FCI or
  selected CI on molecules with up to eight electrons. The 3c composites,
  double hybrids, MP2 and coupled cluster are methods of their own.
  Functionals are the names vibe-qc's XC table resolves: ωB97X-D3 is not
  one, and is refused rather than swapped for ωB97X-D, a different
  functional.
- **Tasks.** ``optimize=True``; frequencies through ``hessian=True``, a
  finite-difference Hessian on HF and DFT surfaces only; full TDDFT through
  ``tddft=True, tddft_type="casida"``. vibe-qc has no NMR shieldings.
- **Keywords.** D3(BJ) or D4 on a functional; RIJCOSX as
  ``density_fit=True, cosx=True`` and RI-J as ``density_fit=True``, with
  vibe-qc choosing the fitting basis for the orbital basis unless a library
  fitting set is named; CPCM with vibe-qc's solvent presets (it has no SMD);
  threads, not processes; position constraints as ``geom_freeze``.

The same translation spells the ``job.spec`` of a QVF job container
(:meth:`VibeqcWriter.container_spec`), the file ``vibeqc run`` executes in
place (vibe-qc's ``qvf_job.py``, QVF spec § 5.9): the typed fields carry
method, functional, basis, charge, multiplicity and the runner's tasks
(``single_point``, ``optimize``, ``hessian``), and ``options`` carries the
remaining ``run_job`` keywords by name, which the runner refuses rather than
drops when it does not know one. TDDFT is a ``run_job`` keyword there, as
vibe-qc's own ``write_pending_qvf`` spells it.
"""

from __future__ import annotations

import json
import re
from collections.abc import Sequence
from typing import Any, ClassVar

from vibesketch.document.model import JobSpec, Species, Task
from vibesketch.elements import ATOMIC_NUMBER
from vibesketch.inputs.base import (
    InputFile,
    InputWriter,
    WriterRefusal,
    first_conformer,
    needs_conformer,
)
from vibesketch.units import ANGSTROM_TO_BOHR

__all__ = ["VibeqcWriter", "method_names"]

#: Functionals by the name a job gives (lower case) and the name vibe-qc's XC table resolves.
_FUNCTIONALS = {
    "lda": "lda", "svwn": "lda", "pbe": "pbe", "pbesol": "pbesol", "blyp": "blyp",
    "b3lyp": "b3lyp", "pbe0": "pbe0", "bhandhlyp": "bhandhlyp", "tpss": "tpss",
    "revtpss": "revtpss", "tpssh": "tpssh", "m06-l": "m06-l", "m06l": "m06-l",
    "m06-2x": "m06-2x", "m062x": "m06-2x", "scan": "scan", "r2scan": "r2scan",
    "r2scan0": "r2scan0", "mn15": "mn15", "wb97": "wb97", "wb97x": "wb97x",
    "wb97x-d": "wb97x-d", "wb97x-v": "wb97x-v", "wb97m-v": "wb97m-v", "hse06": "hse06",
    "cam-b3lyp": "cam-b3lyp", "lc-wpbe": "lc-wpbe",
}  # fmt: skip
#: Functionals with a dispersion correction of their own (VV10 or D2-like).
_DISPERSION_INCLUDED = {"wb97x-d", "wb97x-v", "wb97m-v"}
#: Methods run_job takes by name, by the name a job gives.
_METHODS = {
    "hf-3c": "hf-3c", "pbeh-3c": "pbeh-3c", "b97-3c": "b97-3c", "r2scan-3c": "r2scan-3c",
    "wb97x-3c": "wb97x-3c", "b2plyp": "b2plyp", "dsd-pbep86": "dsd-pbep86",
    "pwpb95": "pwpb95", "mp2": "mp2", "ccsd": "ccsd", "ccsd(t)": "ccsd(t)",
    "dlpno-ccsd": "dlpno-ccsd", "dlpno-ccsd(t)": "dlpno-ccsd(t)",
}  # fmt: skip
_COMPOSITES = {"hf-3c", "pbeh-3c", "b97-3c", "r2scan-3c", "wb97x-3c"}

#: Orbital basis sets vibe-qc ships (its basis_library), by its lower-case name.
_BASES = {
    "sto-3g", "sto-6g", "3-21g", "6-31g", "6-31g*", "6-31g**", "6-31+g**", "6-31++g**",
    "6-311g**", "6-311+g**", "6-311++g**", "6-311+g(2d,p)",
    "def2-sv(p)", "def2-svp", "def2-svpd", "def2-tzvp", "def2-tzvpd", "def2-tzvpp",
    "def2-tzvppd", "def2-qzvp", "def2-qzvpd", "def2-qzvpp", "def2-qzvppd",
    "def2-msvp", "def2-mtzvp", "def2-mtzvpp",
    "cc-pvdz", "cc-pvtz", "cc-pvqz", "cc-pv5z", "cc-pv6z",
    "aug-cc-pvdz", "aug-cc-pvtz", "aug-cc-pvqz", "aug-cc-pv5z",
    "cc-pcvdz", "cc-pcvtz", "cc-pcvqz", "aug-cc-pcvdz", "aug-cc-pcvtz", "aug-cc-pcvqz",
    *(f"pc-{n}" for n in range(5)), *(f"aug-pc-{n}" for n in range(5)),
    *(f"pcseg-{n}" for n in range(5)), *(f"aug-pcseg-{n}" for n in range(5)),
    "x2c-tzvpall", "lanl2dz", "dhf-svp", "dhf-tzvp", "dhf-tzvpp", "dhf-qzvp", "dhf-qzvpp",
}  # fmt: skip
#: Pople spellings vibe-qc maps to its files.
_POPLE_SPELLINGS = {"(d)": "*", "(d,p)": "**"}
#: vibe-qc's fitting sets: <orbital basis>-j, -jk, -jkfit, -c, -ri or -rifit.
_FITTING = re.compile(r"^[a-z0-9()+*-]+-(j|jk|jkfit|c|ri|rifit)$")

#: The container runner's tasks for vibe-sketch's (vibe-qc's qvf_job.py knows
#: single_point, optimize and hessian; TDDFT goes through run_job's keywords).
_CONTAINER_TASKS = {
    "sp": ["single_point"], "opt": ["optimize"], "freq": ["hessian"],
    "opt+freq": ["optimize", "hessian"], "tddft": ["single_point"],
}  # fmt: skip
#: run_job keywords a container states as typed job.spec fields or the runner sets itself.
_CONTAINER_TYPED = frozenset({"method", "functional", "basis", "optimize", "hessian"})

#: Solvents by the name a job gives and vibe-qc's CPCM preset.
_SOLVENTS = {
    "water": "water", "h2o": "water", "methanol": "methanol", "ethanol": "ethanol",
    "acetonitrile": "acetonitrile", "mecn": "acetonitrile", "acetone": "acetone",
    "dmso": "dmso", "dimethylsulfoxide": "dmso", "dmf": "dmf", "thf": "thf",
    "tetrahydrofuran": "thf", "dichloromethane": "dichloromethane",
    "ch2cl2": "dichloromethane", "dcm": "dichloromethane", "chloroform": "chloroform",
    "chcl3": "chloroform", "toluene": "toluene", "benzene": "benzene", "hexane": "n-hexane",
    "n-hexane": "n-hexane", "cyclohexane": "cyclohexane",
    "carbon tetrachloride": "carbon-tetrachloride", "ccl4": "carbon-tetrachloride",
    "isopropanol": "isopropanol", "n-butanol": "n-butanol",
}  # fmt: skip


def _basis(name: str) -> str | None:
    """vibe-qc's name for an orbital basis set it ships, or ``None``."""
    lowered = name.lower()
    for spelling, star in _POPLE_SPELLINGS.items():
        if lowered.startswith("6-3") and lowered.endswith(spelling):
            lowered = lowered[: -len(spelling)] + star
    return lowered if lowered in _BASES else None


def _approximation(name: str) -> str | None:
    """``RIJCOSX`` or ``RIJ`` for the approximations vibe-qc has, else ``None``."""
    upper = name.upper()
    if upper == "RIJCOSX":
        return "RIJCOSX"
    return "RIJ" if upper in ("RIJONX", "RI-J", "RIJ") else None


def _dispersion(name: str) -> str | None:
    """``d3bj`` or ``d4``, vibe-qc's spelling, else ``None``."""
    compact = name.lower().replace("(", "").replace(")", "")
    return compact if compact in ("d3bj", "d4") else None


def _method_refusals(route: _Route, spec: JobSpec) -> list[str]:
    """What this method cannot do with the rest of the specification."""
    reasons: list[str] = []
    if route.kind == "composite":
        if spec.basis is not None:
            reasons.append(f"{route.method} brings its own basis set; leave basis out")
    elif spec.basis is None:
        reasons.append(f"{spec.method} needs a basis set, for example def2-SVP")
    if spec.approximation is not None and route.kind not in ("hf", "dft"):
        reasons.append(f"vibe-sketch writes {spec.approximation} for HF and DFT only")
    if spec.dispersion is not None and _dispersion(spec.dispersion) is not None:
        if route.kind != "dft":
            reasons.append("vibe-qc applies a dispersion correction to a functional only")
        elif route.functional in _DISPERSION_INCLUDED:
            reasons.append(f"{spec.method} includes its own dispersion correction")
    if spec.task in ("freq", "opt+freq") and route.kind not in ("hf", "dft"):
        reasons.append(
            "vibe-qc computes frequencies on HF and DFT surfaces only; run them with HF or a "
            "functional"
        )
    if spec.task in ("opt", "opt+freq") and route.kind not in ("hf", "dft", "composite"):
        reasons.append(
            f"vibe-sketch writes vibe-qc optimisations with HF, DFT and the 3c composites, "
            f"not {spec.method}"
        )
    if spec.task == "tddft" and route.kind not in ("hf", "dft"):
        reasons.append(f"vibe-qc's TDDFT needs HF or a functional, not {spec.method}")
    return reasons


class _Route:
    """How one method is spelt in ``run_job``."""

    def __init__(self, method: str, functional: str | None, kind: str) -> None:
        self.method = method
        self.functional = functional
        #: ``hf``, ``dft``, ``composite``, ``double-hybrid`` or ``correlated``.
        self.kind = kind


def _route(name: str, multiplicity: int) -> _Route | None:
    lowered = name.lower()
    open_shell = multiplicity > 1
    if lowered in ("hf", "rhf", "uhf", "scf"):
        return _Route("uhf" if open_shell else "rhf", None, "hf")
    if lowered in _FUNCTIONALS:
        return _Route("uks" if open_shell else "rks", _FUNCTIONALS[lowered], "dft")
    if lowered in _METHODS:
        method = _METHODS[lowered]
        if method in _COMPOSITES:
            return _Route(method, None, "composite")
        if method in ("b2plyp", "dsd-pbep86", "pwpb95"):
            return _Route(method, None, "double-hybrid")
        return _Route(method, None, "correlated")
    return None


[docs] def method_names(name: str, multiplicity: int) -> tuple[str, str | None] | None: """The method and functional ``run_job`` takes for a job's method, or ``None``. A QVF ``job.spec`` names methods the same way (``rks`` with ``b3lyp``, ``uhf``, ``ccsd(t)``), so the QVF writer asks here. """ route = _route(name, multiplicity) return None if route is None else (route.method, route.functional)
[docs] class VibeqcWriter(InputWriter): name: ClassVar[str] = "vibeqc" description: ClassVar[str] = "vibe-qc run script (Python): HF, DFT, 3c, MP2, CC; SP to TDDFT" tasks: ClassVar[frozenset[Task]] = frozenset({"sp", "opt", "freq", "opt+freq", "tddft"}) fields: ClassVar[frozenset[str]] = frozenset( {"method", "basis", "aux_basis", "approximation", "dispersion", "solvent", "states", "constraints", "resources"} ) # fmt: skip def _refusals(self, species: Species, spec: JobSpec) -> list[str]: reasons = needs_conformer(self.name, species) if species.conformers and first_conformer(species).electron_count(species.charge) == 0: reasons.append(f"species {species.id} has no electrons: there is nothing to compute") route = _route(spec.method, species.multiplicity) if spec.method else None if spec.method is not None and route is None: known = sorted({*_FUNCTIONALS, *_METHODS, "hf"}) reasons.append( f"vibe-qc has no method or functional {spec.method!r} in vibe-sketch's table; " f"known: {', '.join(known)}" ) if route is not None: reasons.extend(_method_refusals(route, spec)) if spec.basis is not None and _basis(spec.basis) is None: reasons.append(f"vibe-qc ships no basis set {spec.basis!r} in vibe-sketch's table") if spec.aux_basis is not None and not _FITTING.match(spec.aux_basis.lower()): reasons.append( f"vibe-qc picks the fitting basis for the orbital basis; {spec.aux_basis!r} is " "not one of its fitting sets (such as def2-tzvp-jk), so leave aux_basis out" ) if spec.approximation is not None and _approximation(spec.approximation) is None: reasons.append(f"vibe-qc has no approximation {spec.approximation!r}") if spec.dispersion is not None and _dispersion(spec.dispersion) is None: reasons.append(f"vibe-qc has D3(BJ) and D4, not {spec.dispersion!r}") if spec.solvent is not None: if spec.solvent.model != "CPCM": reasons.append(f"vibe-qc has CPCM, not {spec.solvent.model}") if spec.solvent.name.lower() not in _SOLVENTS: reasons.append(f"vibe-qc has no CPCM preset for {spec.solvent.name!r}") for constraint in spec.constraints or []: if constraint.kind != "position": reasons.append( f"vibe-qc freezes atoms (geom_freeze) but holds no {constraint.kind}s fixed" ) if spec.resources is not None and spec.resources.maxcore_mb is not None: reasons.append( "vibe-qc takes a total memory budget, not memory per core; leave maxcore_mb out" ) return reasons
[docs] def file_stem(self, species: Species, spec: JobSpec) -> str: """``<species>_<method>_<basis>[_<options>]``, as the ORCA writer names its files.""" parts = [species.id, _slug(spec.method or "")] if spec.basis is not None: parts.append(_slug(spec.basis)) tags = [] if spec.dispersion is not None: tags.append(_dispersion(spec.dispersion) or spec.dispersion) if spec.approximation is not None: tags.append(spec.approximation) if spec.solvent is not None: tags.append(f"{spec.solvent.model}-{spec.solvent.name}") if spec.task != "sp": tags.extend(spec.task.split("+")) parts.extend(sorted(_slug(t) for t in tags)) return "_".join(parts)
[docs] def run_arguments(self, species: Species, spec: JobSpec) -> list[tuple[str, Any]]: """The ``run_job`` keywords for a job, in order, as Python values. Only called once there is nothing to refuse: the run script writes them as literals, the QVF job container as ``job.spec`` fields. """ route = _route(spec.method or "", species.multiplicity) assert route is not None # refusals ran first arguments: list[tuple[str, Any]] = [("method", route.method)] if route.functional is not None: arguments.append(("functional", route.functional)) if spec.basis is not None: arguments.append(("basis", _basis(spec.basis) or "")) if spec.dispersion is not None: arguments.append(("dispersion", _dispersion(spec.dispersion) or "")) if spec.approximation is not None: arguments.append(("density_fit", True)) if _approximation(spec.approximation) == "RIJCOSX": arguments.append(("cosx", True)) if spec.aux_basis is not None: arguments.append(("aux_basis", spec.aux_basis.lower())) if spec.solvent is not None: arguments.append(("solvent", _SOLVENTS[spec.solvent.name.lower()])) if spec.task in ("opt", "opt+freq"): arguments.append(("optimize", True)) if spec.task in ("freq", "opt+freq"): arguments.append(("hessian", True)) if spec.task == "tddft": arguments.extend( [("tddft", True), ("tddft_n_states", spec.states), ("tddft_type", "casida")] ) frozen = [c.atoms[0] for c in spec.constraints or [] if c.kind == "position"] if frozen: arguments.append(("geom_freeze", frozen)) if spec.resources is not None and spec.resources.nprocs > 1: arguments.append(("num_threads", spec.resources.nprocs)) return arguments
[docs] def container_spec(self, species: Species, spec: JobSpec) -> dict[str, Any]: """The ``job.spec`` fields of a QVF job container for this job (spec § 5.9). What vibe-qc's container runner reads: ``method``, ``functional`` (JSON null for a method that takes none, left out where the method brings its own), ``basis``, ``charge``, ``multiplicity``, its ``tasks``, and every other ``run_job`` keyword under ``options`` by name. The dispersion, solvation and resources are stated as typed fields too, for any reader. Raises :class:`WriterRefusal` with the same reasons the run script would. """ reasons = self.refusals(species, spec) if reasons: raise WriterRefusal(self.name, reasons) route = _route(spec.method or "", species.multiplicity) assert route is not None arguments = dict(self.run_arguments(species, spec)) fields: dict[str, Any] = { "job_type": "molecular", "method": route.method, "charge": species.charge, "multiplicity": species.multiplicity, "tasks": list(_CONTAINER_TASKS[spec.task]), } if route.functional is not None: fields["functional"] = route.functional elif route.kind in ("hf", "correlated"): fields["functional"] = None if "basis" in arguments: fields["basis"] = arguments["basis"] if "dispersion" in arguments: fields["dispersion"] = arguments["dispersion"] if spec.solvent is not None: fields["solvation"] = {"model": "cpcm", "solvent": arguments["solvent"]} if spec.resources is not None: fields["resources"] = {"n_cores": spec.resources.nprocs} options = {key: value for key, value in arguments.items() if key not in _CONTAINER_TYPED} if options: fields["options"] = options return fields
def _render(self, species: Species, spec: JobSpec) -> Sequence[InputFile]: from vibesketch import __version__ conformer = first_conformer(species) stem = self.file_stem(species, spec) what = species.name.value if species.name is not None else species.smiles or species.id arguments = [(key, _literal(value)) for key, value in self.run_arguments(species, spec)] arguments.extend( [("output", _text(stem)), ("output_qvf", "True"), ("write_molden_file", "True")] ) rows = [ f" ({ATOMIC_NUMBER[symbol]:>3}, {x:15.10f}, {y:15.10f}, {z:15.10f}), # {symbol}" for symbol, (x, y, z) in zip(conformer.symbols, conformer.xyz_angstrom, strict=True) ] lines = [ f"# {species.id}: {what} — {spec.task} with {spec.method}" + (f"/{spec.basis}" if spec.basis else ""), f"# Written by vibe-sketch {__version__}; geometry: {conformer.method}.", f"# Run: python {stem}.py", "from vibeqc import Atom, Molecule, run_job", "", f"ANGSTROM_TO_BOHR = {ANGSTROM_TO_BOHR!r} # CODATA 2018, as vibe-qc uses", "", "# Atomic number and position in angstrom, one atom per line.", "GEOMETRY = [", *rows, "]", "", "molecule = Molecule(", " [Atom(z, [x * ANGSTROM_TO_BOHR, y * ANGSTROM_TO_BOHR, w * ANGSTROM_TO_BOHR])", " for z, x, y, w in GEOMETRY],", f" charge={species.charge},", f" multiplicity={species.multiplicity},", ")", "", "run_job(", " molecule,", *(f" {key}={value}," for key, value in arguments), ")", ] return [InputFile(f"{stem}.py", "\n".join(lines) + "\n", "text/x-python")]
def _text(value: str) -> str: """A Python string literal in double quotes.""" return json.dumps(value) def _literal(value: Any) -> str: """A run script's literal for a keyword's value.""" return _text(value) if isinstance(value, str) else repr(value) def _slug(text: str) -> str: return re.sub(r"[^a-z0-9]+", "-", text.lower()).strip("-")