"""Unified D3(BJ) dispersion correction entry point.
Two backends are available:
* **``"builtin"``** — the native C++ implementation in vibe-qc's core
(:func:`vibeqc._vibeqc_core.compute_d3bj`). Self-contained, no external
dependencies. Ships with a minimal C6 reference table (Phase D1a:
H, C, N, O, F with one CN reference each) that exercises the
framework but does not reproduce Grimme's reference energies
quantitatively. Use for qualitative checks or element coverage
beyond what your external install has.
* **``"dftd3"``** — a thin wrapper over `Grimme's reference Fortran
library <https://github.com/dftd3/simple-dftd3>`_, distributed as
the ``dftd3`` PyPI package (LGPL-3). Full element coverage up to
plutonium, the canonical CN-dependent c6ab table baked into the
shared library, and parameter sets for ~80 DFT functionals. This
is what PySCF, ORCA, and Psi4 use under the hood for dispersion.
vibe-qc treats ``dftd3`` as an **optional** dependency — it is
listed in ``pyproject.toml`` under the ``dispersion`` extra, so
``pip install -e '.[dispersion]'`` from the vibe-qc checkout (or a
plain ``pip install dftd3`` into the venv) pulls it in. If it is
not installed, ``backend="dftd3"`` raises ``ImportError`` with a
clear message.
* **``"auto"``** (default) — prefers ``dftd3`` when available, falls
back to ``"builtin"``. This is the right choice for most users:
accuracy when they have the external library, a reasonable answer
when they don't.
Until Phase D1b-2 ships the full CN-dependent c6ab grid inside the
vibe-qc C++ core, the ``"dftd3"`` backend is the way to get
quantitative D3-BJ energies.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Optional, Union
import numpy as np
from ._vibeqc_core import (
D3BJParams,
DispersionResult,
Molecule,
compute_d3bj as _compute_d3bj_builtin,
d3bj_params_for as _d3bj_params_for_builtin,
)
__all__ = [
"DispersionResult",
"D3BJParams",
"compute_d3bj",
"d3bj_params_for",
"dftd3_available",
]
def dftd3_available() -> bool:
"""Return True if the optional ``dftd3`` backend can be imported.
Used by ``backend="auto"`` dispatch and by tests that conditionally
exercise the external reference implementation.
"""
try:
import dftd3.interface # noqa: F401
return True
except ImportError:
return False
def _resolve_backend(backend: str) -> str:
if backend == "auto":
return "dftd3" if dftd3_available() else "builtin"
if backend in ("builtin", "dftd3"):
return backend
raise ValueError(
f"compute_d3bj: unknown backend {backend!r} "
f"(expected 'auto', 'builtin', or 'dftd3')"
)
def d3bj_params_for(
functional: str,
*,
backend: str = "auto",
) -> Optional[D3BJParams]:
"""Look up the D3(BJ) damping parameters for a DFT functional.
The ``"builtin"`` backend ships a starter set of ~10 functionals
(D1a); the ``"dftd3"`` backend covers ~80, drawn from Grimme's
canonical ``parameters.toml``. ``"auto"`` (default) tries ``dftd3``
first then falls back to builtin.
Returns ``None`` if the functional is unknown to the chosen
backend.
"""
chosen = _resolve_backend(backend)
if chosen == "dftd3":
try:
return _params_from_dftd3(functional)
except ImportError:
if backend == "auto":
return _d3bj_params_for_builtin(functional)
raise
return _d3bj_params_for_builtin(functional)
def _params_from_dftd3(functional: str) -> Optional[D3BJParams]:
"""Read the canonical D3-BJ damping parameters for a functional
from the ``dftd3`` package's ``parameters.toml``. Returns ``None``
if the functional has no 'bj' entry."""
import importlib.resources as _res
try:
# Python 3.11+: tomllib is stdlib
import tomllib
except ImportError: # pragma: no cover
import tomli as tomllib
# The parameters file ships inside the dftd3 package.
with _res.files("dftd3").joinpath("parameters.toml").open("rb") as fh:
params = tomllib.load(fh)
# dftd3's TOML keys strip hyphens and lowercase; align.
key = functional.lower().replace("-", "")
entry = params.get("parameter", {}).get(key, {}).get("d3", {}).get("bj")
if entry is None:
return None
# Canonical D3-BJ has s6 = 1.0 unless overridden (double hybrids).
return D3BJParams(
s6=float(entry.get("s6", 1.0)),
s8=float(entry["s8"]),
a1=float(entry["a1"]),
a2=float(entry["a2"]),
)
[docs]
def compute_d3bj(
mol: Molecule,
params: Union[D3BJParams, str],
*,
with_gradient: bool = False,
backend: str = "auto",
) -> DispersionResult:
"""Pairwise D3(BJ) dispersion energy (and optionally gradient).
Parameters
----------
mol
:class:`vibeqc.Molecule`.
params
Either a :class:`D3BJParams` with explicit (s6, s8, a1, a2)
values, or the lowercase name of a DFT functional
(e.g. ``"pbe"``, ``"b3lyp"``). A string is resolved via
:func:`d3bj_params_for` with the same ``backend``.
with_gradient
If True, the returned ``DispersionResult.gradient`` is a
``(n_atoms, 3)`` array of ∂E_disp / ∂r_A in Hartree / bohr.
Only the geometric (r-derivative) contribution is included in
D1a/D1b; the CN-derivative piece arrives with the full
CN-dependent C6 table in D1b-2.
backend
``"builtin"``, ``"dftd3"``, or ``"auto"`` (default). See the
module docstring.
"""
chosen = _resolve_backend(backend)
if isinstance(params, str):
p = d3bj_params_for(params, backend=chosen)
if p is None:
raise ValueError(
f"compute_d3bj: no D3-BJ parameters for functional "
f"{params!r} (backend={chosen!r})"
)
params = p
if chosen == "dftd3":
return _compute_d3bj_dftd3(mol, params, with_gradient=with_gradient)
return _compute_d3bj_builtin(mol, params, with_gradient)
def _compute_d3bj_dftd3(
mol: Molecule,
params: D3BJParams,
*,
with_gradient: bool,
) -> DispersionResult:
"""Bit-exact reference implementation via the ``dftd3`` package.
Grimme's Fortran library takes positions in bohr and atomic numbers
as integer arrays, matching vibe-qc's native conventions.
"""
try:
from dftd3.interface import DispersionModel, RationalDampingParam
except ImportError as e:
raise ImportError(
"compute_d3bj(backend='dftd3') requires the optional "
"dftd3 package. Install it with `pip install dftd3` into "
"your vibe-qc venv, or `pip install -e '.[dispersion]'` "
"from the repo checkout."
) from e
atoms = mol.atoms
n = len(atoms)
numbers = np.array([a.Z for a in atoms], dtype=np.int32)
positions = np.array([list(a.xyz) for a in atoms], dtype=float)
if n == 0:
return DispersionResult()
model = DispersionModel(numbers, positions)
bj = RationalDampingParam(
s6=params.s6, s8=params.s8, a1=params.a1, a2=params.a2,
)
res = model.get_dispersion(bj, grad=with_gradient)
# Our DispersionResult is a pybind11-bound C++ struct that pybind11
# lets us construct positionally from defaults and then set fields
# on — but both fields are readonly. Work around by building via
# D3BJParams-style positional construction would require bindings
# changes; instead we use an escape hatch: run the *builtin* call
# with zeroed params to get a fresh DispersionResult whose fields
# we then mutate... except they're def_readonly. So: declare the
# dftd3 backend's return as a small shim that quacks like
# DispersionResult for the public API.
out = _DispersionResultShim(
energy=float(res["energy"]),
gradient=(np.asarray(res["gradient"], dtype=float)
if with_gradient else np.zeros((0, 3))),
)
return out # type: ignore[return-value]
@dataclass
class _DispersionResultShim:
"""Python shim quacking like the pybind11 DispersionResult.
We can't construct DispersionResult from Python (its fields are
def_readonly), so the dftd3 backend returns this dataclass with
the same attribute names. Callers treat the return opaquely via
``.energy`` and ``.gradient`` so the shim is indistinguishable
from the native type in practice.
"""
energy: float
gradient: np.ndarray