"""Native molecular geometry optimization — no ASE required.
Atomic position relaxation using analytic SCF gradients wrapped in
scipy L-BFGS-B. Supports RHF, UHF, RKS, UKS, and wavefunction methods
(selected_ci, dmrg, v2rdm, transcorrelated_ci — the latter four fall
back to central finite differences on the energy).
Dispersion corrections (D3-BJ) and implicit solvation (CPCM/COSMO)
are passed through transparently so the optimizer sees the total
energy + gradient.
Usage::
from vibeqc.molecular_optimize import optimize_molecule
result = optimize_molecule(
mol, basis_name="def2-svp", method="rks", functional="PBE",
)
# result.system — optimized Molecule (bohr)
# result.energy — final total energy (Ha)
# result.trajectory_frames — per-step geometries
# result.trajectory_energies — per-step energies
Integration with ``run_job`` / QVF writing is automatic: when
``optimize=True`` the trajectory data collected here is passed
through to ``write_qvf`` for vibe-view's animation player.
"""
from __future__ import annotations
from typing import Any, Optional
import numpy as np
from ._vibeqc_core import (
Atom,
BasisSet,
GradientOptions,
GridOptions,
Molecule,
RHFOptions,
RKSOptions,
UHFOptions,
UKSOptions,
run_rhf,
run_rks,
run_uhf,
run_uks,
)
__all__ = [
"MolecularOptimizeResult",
"optimize_molecule",
]
[docs]
class MolecularOptimizeResult:
"""Container for molecular geometry optimization results."""
def __init__(
self,
system: Molecule,
energy: float,
gradient: np.ndarray,
n_iter: int,
converged: bool,
trajectory_frames: Optional[list[Molecule]] = None,
trajectory_energies: Optional[list[float]] = None,
):
self.system = system
self.energy = energy
self.gradient = gradient
self.n_iter = n_iter
self.converged = converged
self.trajectory_frames = trajectory_frames or []
self.trajectory_energies = trajectory_energies or []
def __repr__(self) -> str:
return (
f"MolecularOptimizeResult("
f"energy={self.energy:.8f}, "
f"|grad|={float(np.linalg.norm(self.gradient)):.4e}, "
f"n_iter={self.n_iter}, "
f"converged={self.converged})"
)
# ---- SCF dispatch ---------------------------------------------------------
def _run_molecular_scf(
molecule: Molecule,
basis: BasisSet,
method: str,
*,
functional: Optional[str] = None,
rhf_options: Optional[RHFOptions] = None,
uhf_options: Optional[UHFOptions] = None,
rks_options: Optional[RKSOptions] = None,
uks_options: Optional[UKSOptions] = None,
solvent: Any = None,
progress: bool = False,
) -> tuple[float, Any]:
"""Run a single SCF and return (energy, result)."""
method_lower = method.lower()
if method_lower == "rhf":
opts = rhf_options or RHFOptions()
r = run_rhf(molecule, basis, opts)
elif method_lower == "uhf":
opts = uhf_options or UHFOptions()
r = run_uhf(molecule, basis, opts)
elif method_lower == "rks":
opts = rks_options or RKSOptions()
if functional and getattr(opts, "functional", None) in (None, "", "lda"):
opts.functional = functional
r = run_rks(molecule, basis, opts)
elif method_lower == "uks":
opts = uks_options or UKSOptions()
if functional and getattr(opts, "functional", None) in (None, "", "lda"):
opts.functional = functional
r = run_uks(molecule, basis, opts)
else:
raise ValueError(
f"Unknown method {method!r} for molecular optimization. "
f"Use 'rhf', 'uhf', 'rks', or 'uks'."
)
if solvent is not None:
from .solvation import run_cpcm_scf
from .solvation.driver import _solvent_aware_scf_result
sol = run_cpcm_scf(
molecule,
basis,
method=method_lower,
solvent=solvent,
options=opts,
)
return sol.energy, _solvent_aware_scf_result(sol)
return r.energy, r
def _compute_molecular_gradient(
molecule: Molecule,
basis: BasisSet,
scf_result: Any,
method: str,
*,
gradient_options: Optional[GradientOptions] = None,
grid_options: Optional[GridOptions] = None,
dispersion_params: Any = None,
) -> np.ndarray:
"""Compute the analytic nuclear gradient (Ha/bohr, n_atoms × 3).
When ``dispersion_params`` is provided, the D3-BJ gradient is
folded in. Returns the energy gradient ∇E (not the force).
"""
from ._vibeqc_core import (
compute_gradient as _grad_rhf,
)
from ._vibeqc_core import (
compute_gradient_rks as _grad_rks,
)
from ._vibeqc_core import (
compute_gradient_uhf as _grad_uhf,
)
from ._vibeqc_core import (
compute_gradient_uks as _grad_uks,
)
gopt = gradient_options or GradientOptions()
method_lower = method.lower()
if method_lower == "rhf":
grad = _grad_rhf(molecule, basis, scf_result, gopt)
elif method_lower == "uhf":
grad = _grad_uhf(molecule, basis, scf_result, gopt)
elif method_lower == "rks":
ggrid = grid_options or GridOptions()
grad = _grad_rks(molecule, basis, scf_result, ggrid, gopt)
elif method_lower == "uks":
ggrid = grid_options or GridOptions()
grad = _grad_uks(molecule, basis, scf_result, ggrid, gopt)
else:
raise ValueError(f"No analytic gradient for method {method!r}.")
grad = np.asarray(grad, dtype=float)
# Fold in dispersion gradient if requested.
if dispersion_params is not None:
from .dispersion import compute_d3bj
disp = compute_d3bj(molecule, dispersion_params, with_gradient=True)
grad = grad + np.asarray(disp.gradient, dtype=float)
return grad
# ---- Cartesian ↔ flat encoding -------------------------------------------
def _positions_to_flat(molecule: Molecule) -> np.ndarray:
"""Flatten Cartesian atom positions to a 1D array (bohr)."""
flat: list[float] = []
for atom in molecule.atoms:
flat.extend(atom.xyz)
return np.array(flat, dtype=float)
def _flat_to_molecule(
template: Molecule,
x: np.ndarray,
) -> Molecule:
"""Rebuild a Molecule from flat Cartesian coordinates (bohr)."""
n_atoms = len(list(template.atoms))
new_atoms: list[Atom] = []
for i in range(n_atoms):
xyz = [float(x[3 * i + c]) for c in range(3)]
new_atoms.append(Atom(int(template.atoms[i].Z), xyz))
return Molecule(new_atoms, template.charge, template.multiplicity)
# ---- FD fallback for methods without analytic gradients ------------------
def _gradient_via_central_difference(
molecule: Molecule,
basis_name: str,
method: str,
*,
functional: Optional[str] = None,
rhf_options: Any = None,
uhf_options: Any = None,
rks_options: Any = None,
uks_options: Any = None,
solvent: Any = None,
dispersion_params: Any = None,
step_bohr: float = 0.005,
) -> np.ndarray:
"""Central-difference energy gradient for wavefunction methods.
Two-point central difference on each Cartesian degree of freedom.
Returns ∇E (not forces), shape (n_atoms, 3), in Ha/bohr.
"""
n_atoms = len(list(molecule.atoms))
grad = np.zeros((n_atoms, 3), dtype=float)
for i in range(n_atoms):
for c in range(3):
pos = np.asarray([list(a.xyz) for a in molecule.atoms], dtype=float)
pos_plus = pos.copy()
pos_plus[i, c] += step_bohr
mol_plus = Molecule(
[Atom(int(a.Z), list(p)) for a, p in zip(molecule.atoms, pos_plus)],
molecule.charge,
molecule.multiplicity,
)
basis_plus = BasisSet(mol_plus, basis_name)
e_plus = _evaluate_energy(
mol_plus,
basis_plus,
method,
functional=functional,
rhf_options=rhf_options,
uhf_options=uhf_options,
rks_options=rks_options,
uks_options=uks_options,
solvent=solvent,
dispersion_params=dispersion_params,
)
pos_minus = pos.copy()
pos_minus[i, c] -= step_bohr
mol_minus = Molecule(
[Atom(int(a.Z), list(p)) for a, p in zip(molecule.atoms, pos_minus)],
molecule.charge,
molecule.multiplicity,
)
basis_minus = BasisSet(mol_minus, basis_name)
e_minus = _evaluate_energy(
mol_minus,
basis_minus,
method,
functional=functional,
rhf_options=rhf_options,
uhf_options=uhf_options,
rks_options=rks_options,
uks_options=uks_options,
solvent=solvent,
dispersion_params=dispersion_params,
)
grad[i, c] = (e_plus - e_minus) / (2.0 * step_bohr)
return grad
def _evaluate_energy(
molecule: Molecule,
basis: BasisSet,
method: str,
*,
functional: Optional[str] = None,
rhf_options: Any = None,
uhf_options: Any = None,
rks_options: Any = None,
uks_options: Any = None,
solvent: Any = None,
dispersion_params: Any = None,
) -> float:
"""Evaluate the total energy at a given geometry (Ha)."""
from .runner import _run_single_point
result = _run_single_point(
method,
molecule,
basis,
functional=functional,
rhf_options=rhf_options,
uhf_options=uhf_options,
rks_options=rks_options,
uks_options=uks_options,
solvent=solvent,
)
e = float(getattr(result, "energy", 0.0))
if dispersion_params is not None:
from .dispersion import compute_d3bj
disp = compute_d3bj(molecule, dispersion_params)
e += float(disp.energy)
return e
# ---- Public API -----------------------------------------------------------
[docs]
def optimize_molecule(
molecule: Molecule,
basis_name: str,
*,
method: str = "rhf",
functional: Optional[str] = None,
rhf_options: Optional[RHFOptions] = None,
uhf_options: Optional[UHFOptions] = None,
rks_options: Optional[RKSOptions] = None,
uks_options: Optional[UKSOptions] = None,
max_iter: int = 100,
conv_tol_grad: float = 4.5e-4,
conv_tol_energy: float = 1e-6,
gradient_options: Optional[GradientOptions] = None,
grid_options: Optional[GridOptions] = None,
dispersion_params: Any = None,
solvent: Any = None,
record_trajectory: bool = True,
progress: bool = False,
fd_step_bohr: float = 0.005,
) -> MolecularOptimizeResult:
"""Relax molecular geometry using analytic gradients + L-BFGS-B.
Parameters
----------
molecule
Starting geometry (Cartesian coordinates in bohr).
basis_name
Basis-set name (rebuilt at each geometry step).
method
``"rhf"``, ``"uhf"``, ``"rks"``, ``"uks"``, or a wavefunction
method (``"selected_ci"``, ``"dmrg"``, ``"v2rdm"``,
``"transcorrelated_ci"``). Wavefunction methods fall back to
central finite differences on the energy.
functional
XC functional string for ``"rks"`` / ``"uks"`` (e.g. ``"PBE"``).
rhf_options / uhf_options / rks_options / uks_options
Per-method SCF options. If ``None``, defaults are used.
max_iter
Maximum L-BFGS-B iterations.
conv_tol_grad
Gradient convergence tolerance (Ha/bohr). Default 4.5e-4
corresponds to ~0.01 eV/Å — tight enough for routine use.
conv_tol_energy
Energy convergence tolerance (Ha). Controls the scipy
``ftol`` parameter.
gradient_options
Options for the analytic gradient kernels (density fitting,
COSX, etc.).
grid_options
DFT integration grid options (RKS / UKS only).
dispersion_params
A :class:`D3BJParams` instance — if provided, the D3-BJ
energy and gradient are folded into the objective.
solvent
A :class:`SolventModel` or preset string / dict for CPCM
implicit solvation (v0.9.0).
record_trajectory
If True (default), collect per-step geometries and energies
for downstream visualisation (QVF animation player).
progress
If True, print per-step energy and gradient norms to stdout.
fd_step_bohr
Finite-difference step size for wavefunction-method gradients
(bohr). Default 0.005 (≈ 0.0026 Å).
Returns
-------
MolecularOptimizeResult
"""
from scipy.optimize import minimize
method_lower = method.lower()
_mean_field = {"rhf", "uhf", "rks", "uks"}
_has_analytic_gradient = method_lower in _mean_field
trajectory_frames: list[Molecule] = []
trajectory_energies: list[float] = []
_x0 = _positions_to_flat(molecule)
# Pre-construct a scipy gradient closure. The "force" minimizers
# expect ∂E/∂x (not -∂E/∂x), so we pass the gradient as-is.
if _has_analytic_gradient:
def _grad_fn(x: np.ndarray) -> np.ndarray:
mol = _flat_to_molecule(molecule, x)
basis = BasisSet(mol, basis_name)
e, res = _run_molecular_scf(
mol,
basis,
method_lower,
functional=functional,
rhf_options=rhf_options,
uhf_options=uhf_options,
rks_options=rks_options,
uks_options=uks_options,
solvent=solvent,
)
grad = _compute_molecular_gradient(
mol,
basis,
res,
method_lower,
gradient_options=gradient_options,
grid_options=grid_options,
dispersion_params=dispersion_params,
)
return grad.ravel()
def _energy_fn(x: np.ndarray) -> float:
mol = _flat_to_molecule(molecule, x)
basis = BasisSet(mol, basis_name)
e, _res = _run_molecular_scf(
mol,
basis,
method_lower,
functional=functional,
rhf_options=rhf_options,
uhf_options=uhf_options,
rks_options=rks_options,
uks_options=uks_options,
solvent=solvent,
)
if dispersion_params is not None:
from .dispersion import compute_d3bj
disp = compute_d3bj(mol, dispersion_params)
e += float(disp.energy)
return e
else:
# Wavefunction methods — FD on energy.
def _grad_fn(x: np.ndarray) -> np.ndarray:
mol = _flat_to_molecule(molecule, x)
return _gradient_via_central_difference(
mol,
basis_name,
method_lower,
functional=functional,
rhf_options=rhf_options,
uhf_options=uhf_options,
rks_options=rks_options,
uks_options=uks_options,
solvent=solvent,
dispersion_params=dispersion_params,
step_bohr=fd_step_bohr,
).ravel()
def _energy_fn(x: np.ndarray) -> float:
mol = _flat_to_molecule(molecule, x)
basis = BasisSet(mol, basis_name)
return _evaluate_energy(
mol,
basis,
method_lower,
functional=functional,
rhf_options=rhf_options,
uhf_options=uhf_options,
rks_options=rks_options,
uks_options=uks_options,
solvent=solvent,
dispersion_params=dispersion_params,
)
# Combined objective: scipy calls `fun` first, then `jac` at the
# same x. We evaluate energy once in `fun` and stash it so `jac`
# can reuse the SCF result in the analytic-gradient path. For FD
# methods the caching is in the gradient evaluation itself.
if _has_analytic_gradient:
_cache: dict[str, Any] = {
"result": None,
"mol": None,
"basis": None,
"energy": float("nan"),
}
def _fun_cached(x: np.ndarray) -> float:
mol = _flat_to_molecule(molecule, x)
basis = BasisSet(mol, basis_name)
e, res = _run_molecular_scf(
mol,
basis,
method_lower,
functional=functional,
rhf_options=rhf_options,
uhf_options=uhf_options,
rks_options=rks_options,
uks_options=uks_options,
solvent=solvent,
)
_cache["result"] = res
_cache["mol"] = mol
_cache["basis"] = basis
if dispersion_params is not None:
from .dispersion import compute_d3bj
disp = compute_d3bj(mol, dispersion_params)
e += float(disp.energy)
_cache["energy"] = e
return e
def _jac_cached(x: np.ndarray) -> np.ndarray:
# Reuse the cached SCF result to avoid double-running.
if _cache.get("result") is not None and _cache.get("mol") is not None:
grad = _compute_molecular_gradient(
_cache["mol"],
_cache["basis"],
_cache["result"],
method_lower,
gradient_options=gradient_options,
grid_options=grid_options,
dispersion_params=dispersion_params,
)
_cache["result"] = None # clear for next iteration
return grad.ravel()
# Fallback: re-evaluate (shouldn't normally happen).
return _grad_fn(x)
_objective = _fun_cached
_jacobian = _jac_cached
else:
_objective = _energy_fn
_jacobian = _grad_fn
# Callback to collect trajectory.
if record_trajectory:
def _callback(xk: np.ndarray) -> None:
mol_frame = _flat_to_molecule(molecule, xk)
trajectory_frames.append(mol_frame)
# scipy guarantee: fun(xk) was called just before the
# callback. Use the cached energy to avoid a duplicate
# SCF evaluation.
if _has_analytic_gradient:
e_frame = _cache.get("energy", float("nan"))
else:
e_frame = _energy_fn(xk)
trajectory_energies.append(e_frame)
if progress:
print(f" step {len(trajectory_frames):3d} E = {e_frame:14.8f} Ha")
else:
_callback = None # type: ignore[assignment]
# ---- run the scipy optimizer ------------------------------------------
if progress:
print(
f"\n Geometry optimization — {method.upper()}"
+ (f"/{functional}" if functional else "")
+ f" basis={basis_name}"
)
print(
f" n_atoms={len(list(molecule.atoms))}, "
f"max_iter={max_iter}, "
f"gtol={conv_tol_grad:.1e} Ha/bohr\n"
)
# Feed the energy through the objective so the cache is primed.
e_start = _objective(_x0)
if record_trajectory:
trajectory_frames.append(molecule)
trajectory_energies.append(e_start)
res = minimize(
_objective,
_x0,
method="L-BFGS-B",
jac=_jacobian if _has_analytic_gradient else _grad_fn,
callback=_callback,
options={
"maxiter": max_iter,
"gtol": conv_tol_grad,
"ftol": conv_tol_energy,
},
)
mol_opt = _flat_to_molecule(molecule, res.x)
# Final gradient norm for reporting.
grad_final = _grad_fn(res.x) if res.success else res.jac
if grad_final is not None:
grad_norm = float(np.linalg.norm(grad_final))
else:
grad_norm = float("nan")
if progress:
print(
f"\n Geometry optimization: {res.nit} iters, "
f"E = {res.fun:.8f} Ha, "
f"|grad| = {grad_norm:.4e}, "
f"converged={res.success}"
)
return MolecularOptimizeResult(
system=mol_opt,
energy=float(res.fun),
gradient=grad_final if grad_final is not None else np.array([]),
n_iter=int(res.nit),
converged=bool(res.success),
trajectory_frames=trajectory_frames if record_trajectory else None,
trajectory_energies=trajectory_energies if record_trajectory else None,
)