"""BIPOLE structure optimization.
Atomic position relaxation wrapped in scipy L-BFGS-B. Lattice parameter
optimization via energy-only Nelder-Mead (``relax_cell``) or finite-
difference strain gradients (``relax_cell_gradient``). Combined cell +
position relaxation (``relax_full``) alternates the two until convergence.
Atomic forces and cell-strain gradients default to finite differences of the
BIPOLE energy
(``force_mode="fd"``). This is deliberate: the analytic BIPOLE gradient
(:mod:`vibeqc.bipole_gradient`) is still a research-preview surface:
RHF/UHF Γ and maintained KS Γ cases are pinned, but broader KS,
multi-k, finite-temperature, and meta-GGA certification remain open.
The FD path differentiates the real total energy and is correct by
construction (cost: ~6N SCFs per gradient eval). Pass
``force_mode="analytic"`` only for research on the analytic gradient
itself -- a warning is emitted and FD remains the production optimizer
default.
Usage:
from vibeqc.bipole_optimize import relax_atoms, relax_cell, relax_full
# Atomic relaxation only (FD forces by default)
result = relax_atoms(system, basis, kmesh, method="RHF")
# Lattice relaxation only (FD of energy)
system_opt = relax_cell(system, basis, kmesh, method="RHF")
# Full cell + atom relaxation
system_opt = relax_full(system, basis, kmesh, method="RHF")
"""
from __future__ import annotations
from typing import TYPE_CHECKING, Any, List, Optional, Sequence, Tuple, Union
import numpy as np
from scipy.optimize import minimize
if TYPE_CHECKING:
from .kpoints import KPoints
from ._vibeqc_core import (
Atom,
BasisSet,
BlochKMesh,
Functional,
LatticeSumOptions,
PeriodicKSOptions,
PeriodicRHFOptions,
PeriodicSystem,
bloch_kmesh_from_lists,
monkhorst_pack,
)
from .bipole_gradient import (
compute_bipole_gradient_fd,
compute_bipole_gradient_rhf,
compute_bipole_gradient_rks,
compute_bipole_gradient_uhf,
compute_bipole_gradient_uks,
)
from .pbc_bipole import run_pbc_bipole_rhf
from .pbc_bipole_rks import run_pbc_bipole_rks
from .pbc_bipole_uhf import run_pbc_bipole_uhf
from .pbc_bipole_uks import run_pbc_bipole_uks
__all__ = [
"relax_atoms",
"relax_cell",
"relax_cell_gradient",
"relax_full",
"OptimizeResult",
]
KMeshInput = Union[BlochKMesh, "KPoints"]
def _as_bloch_kmesh(kmesh: KMeshInput) -> BlochKMesh:
from .kpoints import as_bloch_kmesh
return as_bloch_kmesh(kmesh)
# Central-difference half-step (in % strain) for the cell-strain gradient.
# 0.05 % => ~5e-4 strain, well inside the SCF energy's smooth regime.
_CELL_FD_STRAIN_PCT = 0.05
_SCF_OPTION_FIELDS = (
"functional",
"grid",
"max_iter",
"conv_tol_energy",
"conv_tol_grad",
"damping",
"dynamic_damping",
"dynamic_damping_min",
"dynamic_damping_max",
"fock_mixing",
"use_diis",
"diis_start_iter",
"diis_subspace_size",
"scf_accelerator",
"ediis_diis_switch_threshold",
"initial_guess",
"level_shift",
"level_shift_warmup_cycles",
"smearing_temperature",
"quadratic_fallback_iter",
"quadratic_fallback_shift",
"quadratic_fallback_max_step",
"use_periodic_becke",
"becke_image_radius_bohr",
)
_LATTICE_OPTION_FIELDS = (
"cutoff_bohr",
"nuclear_cutoff_bohr",
"coulomb_method",
"screening_overlap_threshold",
"screening_exchange_threshold",
"schwarz_threshold",
"schwarz_threshold_forces",
)
def _central_fd_gradient(energy_fn, x0, h: float) -> np.ndarray:
"""Central finite-difference gradient ``d(energy_fn)/dx`` at ``x0``.
A descent-consistent gradient: it IS the finite difference of the exact
scalar ``energy_fn`` being minimised, so ``-g`` is always a descent
direction (for ``h`` small enough that the central difference resolves
the local slope). Used for the cell-strain gradient in
:func:`relax_cell_gradient` -- replacing the force virial, which is
inconsistent with the energy objective there."""
x0 = np.asarray(x0, dtype=float)
g = np.zeros(len(x0))
for i in range(len(x0)):
xp = x0.copy(); xp[i] += h
xm = x0.copy(); xm[i] -= h
g[i] = (energy_fn(xp) - energy_fn(xm)) / (2.0 * h)
return g
def _copy_lattice_options(src) -> LatticeSumOptions:
"""Return a fresh lattice-options object with the public fields copied."""
out = LatticeSumOptions()
for name in _LATTICE_OPTION_FIELDS:
if hasattr(src, name):
setattr(out, name, getattr(src, name))
return out
def _make_scf_options(
method_upper: str,
*,
cutoff_bohr: float,
scf_conv_tol: Optional[float],
scf_max_iter: Optional[int],
scf_options: Optional[Union[PeriodicRHFOptions, PeriodicKSOptions]] = None,
):
"""Build optimizer-local SCF options without mutating caller state."""
is_ks = method_upper in ("RKS", "UKS")
opts = PeriodicKSOptions() if is_ks else PeriodicRHFOptions()
if scf_options is None:
opts.max_iter = 50
opts.use_diis = True
opts.conv_tol_energy = 1e-7
else:
for name in _SCF_OPTION_FIELDS:
if hasattr(scf_options, name) and hasattr(opts, name):
setattr(opts, name, getattr(scf_options, name))
if hasattr(scf_options, "lattice_opts"):
opts.lattice_opts = _copy_lattice_options(scf_options.lattice_opts)
if scf_conv_tol is not None:
opts.conv_tol_energy = float(scf_conv_tol)
if scf_max_iter is not None:
opts.max_iter = int(scf_max_iter)
opts.lattice_opts.cutoff_bohr = float(cutoff_bohr)
opts.lattice_opts.nuclear_cutoff_bohr = float(cutoff_bohr)
return opts
def _kmesh_for_system(kmesh, reference_system: PeriodicSystem, target_system: PeriodicSystem):
"""Return ``kmesh`` represented on ``target_system``'s reciprocal lattice.
``BlochKMesh.kpoints`` are Cartesian vectors. Atomic relaxations can reuse
them because the lattice is fixed, but cell relaxations strain the lattice:
the same fractional Monkhorst-Pack or explicit k-list must be converted to
the target reciprocal basis before each strained SCF. Symmetry-reduced IBZ
meshes are expanded to their full MP mesh for strained cells so downstream
BIPOLE code sees uniform full-mesh weights instead of non-uniform IBZ
weights without matching target-cell symmetry metadata.
"""
if np.allclose(
np.asarray(reference_system.lattice, dtype=float),
np.asarray(target_system.lattice, dtype=float),
atol=1e-14,
rtol=0.0,
):
return kmesh
kpts = np.asarray(
getattr(kmesh, "kpoints", getattr(kmesh, "kpoints_cart", [])),
dtype=float,
).reshape(-1, 3)
if kpts.size == 0:
raise ValueError("_kmesh_for_system: kmesh has no k-points")
weights = np.asarray(getattr(kmesh, "weights", []), dtype=float).reshape(-1)
if weights.size == 0:
weights = np.full(kpts.shape[0], 1.0 / float(kpts.shape[0]))
if weights.shape != (kpts.shape[0],):
raise ValueError(
"_kmesh_for_system: kpoint/weight size mismatch "
f"({kpts.shape[0]} k-points, {weights.shape[0]} weights)"
)
mesh_raw = getattr(kmesh, "mesh", None)
shift_raw = getattr(kmesh, "is_shift", getattr(kmesh, "shift", None))
mesh = tuple(int(x) for x in mesh_raw) if mesh_raw is not None else ()
shift = tuple(int(x) for x in shift_raw) if shift_raw is not None else ()
ir_mapping = np.asarray(getattr(kmesh, "ir_mapping", []), dtype=int).reshape(-1)
if len(mesh) == 3:
shift = shift if len(shift) == 3 else (0, 0, 0)
full_n = int(mesh[0] * mesh[1] * mesh[2])
is_ibz = ir_mapping.size == full_n and kpts.shape[0] != full_n
is_full_mp = False
if kpts.shape[0] == full_n:
try:
ref_full = monkhorst_pack(
reference_system,
list(mesh),
list(shift),
False,
)
ref_kpts = np.asarray(ref_full.kpoints, dtype=float).reshape(-1, 3)
ref_weights = np.asarray(ref_full.weights, dtype=float).reshape(-1)
is_full_mp = (
ref_kpts.shape == kpts.shape
and np.allclose(ref_kpts, kpts, atol=1e-12)
and np.allclose(ref_weights, weights, atol=1e-12)
)
except Exception:
is_full_mp = False
if is_full_mp or is_ibz:
return monkhorst_pack(target_system, list(mesh), list(shift), False)
B_ref = np.asarray(reference_system.reciprocal_lattice(), dtype=float)
B_target = np.asarray(target_system.reciprocal_lattice(), dtype=float)
k_frac = np.linalg.solve(B_ref, kpts.T).T
k_target = (B_target @ k_frac.T).T
return bloch_kmesh_from_lists(
[np.asarray(k, dtype=float) for k in k_target],
[float(w) for w in weights],
)
class SCFNonConvergence(RuntimeError):
"""An SCF inside an optimization did not converge.
Raised by :func:`_run_scf` so optimizer objectives can distinguish
"this trial geometry's SCF failed" (recoverable: return a penalty
barrier so the line search backs off) from genuine programming
errors. A non-converged energy is never *used* as a real objective
value -- it is replaced by the barrier."""
[docs]
class OptimizeResult:
"""Container for optimization results."""
[docs]
def __init__(self, system, energy, gradient, n_iter, converged):
self.system = system
self.energy = energy
self.gradient = gradient
self.n_iter = n_iter
self.converged = converged
def _run_scf(system, basis, kmesh, opts, method, functional, **kwargs):
"""Run a single BIPOLE SCF and return (energy, result)."""
if method == "RHF":
result = run_pbc_bipole_rhf(
system,
basis,
kmesh,
opts,
progress=False,
**kwargs,
)
elif method == "UHF":
result = run_pbc_bipole_uhf(
system,
basis,
kmesh,
opts,
progress=False,
**kwargs,
)
elif method == "RKS":
result = run_pbc_bipole_rks(
system,
basis,
kmesh,
opts,
functional=functional,
progress=False,
**kwargs,
)
elif method == "UKS":
result = run_pbc_bipole_uks(
system,
basis,
kmesh,
opts,
functional=functional,
progress=False,
**kwargs,
)
else:
raise ValueError(f"Unknown method: {method}")
if not bool(getattr(result, "converged", True)):
n_iter = getattr(result, "n_iter", "unknown")
energy = getattr(result, "energy", None)
energy_txt = "unknown" if energy is None else f"{float(energy):.12g}"
raise SCFNonConvergence(
"bipole_optimize: "
f"{method} SCF did not converge during optimization "
f"(n_iter={n_iter}, energy={energy_txt}). Refusing to optimize "
"against a non-converged energy; increase scf_max_iter, loosen "
"scf_conv_tol, or adjust the SCF accelerator/smearing options."
)
return result.energy, result
def _compute_gradient(
system, basis, result, method, lattice_opts,
*, kmesh=None, dft_plus_u=None,
):
"""Compute the BIPOLE gradient for any method.
When ``dft_plus_u`` is set, the multi-k +U Pulay overlap-derivative
contribution is added; ``kmesh`` must be the same
:class:`BlochKMesh` the SCF was run on.
"""
kwargs = {"lattice_opts": lattice_opts}
if dft_plus_u:
kwargs["kmesh"] = kmesh
kwargs["dft_plus_u"] = dft_plus_u
if method == "RHF":
return compute_bipole_gradient_rhf(system, basis, result, **kwargs)
elif method == "UHF":
return compute_bipole_gradient_uhf(system, basis, result, **kwargs)
elif method == "RKS":
return compute_bipole_gradient_rks(system, basis, result, **kwargs)
elif method == "UKS":
return compute_bipole_gradient_uks(system, basis, result, **kwargs)
else:
raise ValueError(f"Unknown method: {method}")
def _compute_forces(
system, basis, result, method, lattice_opts,
*, kmesh, basis_name, opts, functional, bipole_kwargs,
force_mode, fd_step_bohr,
):
"""Cartesian BIPOLE gradient ``(n_atoms, 3)`` in Ha/bohr.
``force_mode="fd"`` (default) uses the exact finite-difference
gradient :func:`compute_bipole_gradient_fd` (differentiates the real
total energy -- correct by construction; +U is handled automatically).
``force_mode="analytic"`` uses the research-preview analytic gradient
surface (emits a warning; see :mod:`vibeqc.bipole_gradient` for the
currently maintained cases).
"""
if force_mode == "fd":
return np.asarray(
compute_bipole_gradient_fd(
system, basis_name, kmesh, opts,
method=method, functional=functional,
step_bohr=fd_step_bohr, **bipole_kwargs,
)
)
if force_mode == "analytic":
return _compute_gradient(
system, basis, result, method, lattice_opts,
kmesh=kmesh, dft_plus_u=bipole_kwargs.get("dft_plus_u"),
)
raise ValueError(
f"force_mode must be 'fd' or 'analytic'; got {force_mode!r}"
)
def _atoms_to_flat(system: PeriodicSystem) -> np.ndarray:
"""Flatten atomic positions to a 1D array (fractional coordinates)."""
lattice = np.asarray(system.lattice, dtype=float)
inv_lat = np.linalg.inv(lattice)
frac = []
for atom in system.unit_cell:
frac.extend(inv_lat @ np.asarray(atom.xyz, dtype=float))
return np.array(frac, dtype=float)
def _flat_to_system(
system_template: PeriodicSystem,
x: np.ndarray,
) -> PeriodicSystem:
"""Rebuild system from flat fractional coordinates."""
lattice = np.asarray(system_template.lattice, dtype=float)
n_atoms = len(system_template.unit_cell)
new_atoms = []
for i in range(n_atoms):
frac = x[3 * i : 3 * i + 3]
cart = lattice @ frac
new_atoms.append(
Atom(int(system_template.unit_cell[i].Z), list(cart)),
)
return PeriodicSystem(
system_template.dim,
lattice,
new_atoms,
charge=system_template.charge,
multiplicity=system_template.multiplicity,
)
def relax_atoms(
system: PeriodicSystem,
basis_name: str,
kmesh: KMeshInput,
method: str = "RHF",
*,
functional: Optional[str] = None,
force_mode: str = "fd",
fd_step_bohr: float = 1e-3,
max_iter: int = 30,
conv_tol_grad: float = 1e-4,
scf_conv_tol: Optional[float] = None,
scf_max_iter: Optional[int] = None,
scf_options: Optional[Union[PeriodicRHFOptions, PeriodicKSOptions]] = None,
cutoff_bohr: float = 8.0,
freeze_indices: Optional[Sequence[int]] = None,
output_trajectory: Optional["str | Path"] = None,
**bipole_kwargs,
) -> OptimizeResult:
"""Relax atomic positions with L-BFGS-B and BIPOLE forces.
Forces default to the exact finite-difference gradient
(``force_mode="fd"``); see the module docstring for why the analytic
gradient is not used by default.
Parameters
----------
system : PeriodicSystem
Initial geometry (lattice fixed).
basis_name : str
Basis set name (rebuild per geometry step).
kmesh : BlochKMesh or KPoints
k-point mesh.
method : str
"RHF", "UHF", "RKS", or "UKS".
functional : str, optional
XC functional for RKS/UKS.
force_mode : str
``"fd"`` (default) for the exact finite-difference gradient, or
``"analytic"`` for the research-preview analytic gradient (emits a
warning; maintained cases are documented in
:mod:`vibeqc.bipole_gradient`).
fd_step_bohr : float
Central-difference half-step (bohr) used when ``force_mode="fd"``.
max_iter : int
Maximum optimization steps.
conv_tol_grad : float
Gradient convergence tolerance (Ha/bohr).
scf_conv_tol : float
Optional SCF energy convergence tolerance override. Defaults to the
supplied ``scf_options`` value, or ``1e-7`` when no options are supplied.
scf_max_iter : int, optional
Optional SCF iteration limit override. Defaults to the supplied
``scf_options`` value, or ``50`` when no options are supplied.
scf_options : PeriodicRHFOptions / PeriodicKSOptions, optional
SCF controls to copy into every geometry point (smearing, FMIXING,
level shift, accelerator, grid, etc.). The object is copied before the
optimizer sets its lattice cutoffs, so caller state is not mutated.
cutoff_bohr : float
Lattice cutoff for integrals.
freeze_indices : sequence of int, optional
Atom indices (into ``system.unit_cell``) to hold fixed during
the relaxation. The standard surface-catalysis pattern: pass
``SlabInfo.bottom_layer_indices(n)`` from :func:`vibeqc.build.slab`
to freeze the bottom N layers of a slab. The SCF + gradient
still see every atom; the optimizer simply zeros the gradient
components on the frozen atoms so their positions never move.
output_trajectory : str or Path, optional
Path stem (``.qvf`` appended automatically). When set, the
relaxation collects a per-step (geometry, energy) frame via
scipy's ``callback=`` and writes a vibe-view-renderable QVF
archive on exit. Each frame is one accepted L-BFGS-B step;
the initial geometry is frame 0 and the converged geometry is
the last frame. Periodic systems ship as QVF v2 with the
per-frame lattice + dim attached, so vibe-view renders the
cell + wraps atoms across periodic boundaries (see
``docs/user_guide/vibe_view.md`` Sec. "Periodic reaction paths").
Default ``None`` => no trajectory output (no overhead).
Returns
-------
OptimizeResult
"""
kmesh = _as_bloch_kmesh(kmesh)
method_upper = method.upper()
opts = _make_scf_options(
method_upper,
cutoff_bohr=cutoff_bohr,
scf_conv_tol=scf_conv_tol,
scf_max_iter=scf_max_iter,
scf_options=scf_options,
)
history: List[Tuple[float, float]] = []
n_atoms_total = len(system.unit_cell)
if freeze_indices is None:
frozen_set: set[int] = set()
else:
frozen_set = {int(i) for i in freeze_indices}
bad = [i for i in frozen_set
if i < 0 or i >= n_atoms_total]
if bad:
raise ValueError(
f"relax_atoms: freeze_indices {bad} out of range "
f"[0, {n_atoms_total - 1}]"
)
# Bound matrix: shape (3 * n_atoms, ) with (fixed,fixed) bounds for frozen
# entries (in fractional coords) so L-BFGS-B literally cannot move
# them. Free atoms use (-inf, inf).
bounds: Optional[list[tuple[Optional[float], Optional[float]]]] = None
if frozen_set:
bounds = []
x0_full = _atoms_to_flat(system)
for atom_i in range(n_atoms_total):
if atom_i in frozen_set:
for k in range(3):
fixed = float(x0_full[3 * atom_i + k])
bounds.append((fixed, fixed))
else:
for _ in range(3):
bounds.append((None, None))
# L-BFGS-B optimizes in *fractional* coordinates, so its first unit
# trial step in a large box moves atoms by ~one lattice vector -- a
# geometry whose SCF legitimately may not converge. A line-search
# probe failing is recoverable: return a finite penalty barrier so
# the line search backs off (the non-converged energy itself is
# never used). Only a failure at the *initial* geometry -- where
# there is nothing to back off to -- aborts the optimization.
_PENALTY_HA = 1.0e6
_had_good_eval: dict[str, bool] = {"ok": False}
def objective(x: np.ndarray) -> float:
sys = _flat_to_system(system, x)
basis = BasisSet(sys.unit_cell_molecule(), basis_name)
try:
e, res = _run_scf(
sys, basis, kmesh, opts, method_upper, functional, **bipole_kwargs
)
except SCFNonConvergence:
if not _had_good_eval["ok"]:
raise
history.append((_PENALTY_HA, 0.0))
return _PENALTY_HA
_had_good_eval["ok"] = True
history.append((e, 0.0))
return e
def gradient(x: np.ndarray) -> np.ndarray:
sys = _flat_to_system(system, x)
basis = BasisSet(sys.unit_cell_molecule(), basis_name)
try:
e, res = _run_scf(
sys, basis, kmesh, opts, method_upper, functional, **bipole_kwargs
)
except SCFNonConvergence:
if not _had_good_eval["ok"]:
raise
# Penalty region: zero slope + the barrier value recorded by
# objective() force the line search to shrink the step.
history.append((_PENALTY_HA, 0.0))
return np.zeros(3 * len(system.unit_cell))
_had_good_eval["ok"] = True
# Convert Cartesian gradient to fractional gradient. Forces use
# the FD path by default (analytic is a research preview). FD
# differentiates the real energy, so dft_plus_u is handled
# automatically; the analytic branch threads it explicitly.
lattice = np.asarray(sys.lattice, dtype=float)
grad_cart = _compute_forces(
sys, basis, res, method_upper, opts.lattice_opts,
kmesh=kmesh, basis_name=basis_name, opts=opts,
functional=functional, bipole_kwargs=bipole_kwargs,
force_mode=force_mode, fd_step_bohr=fd_step_bohr,
)
# Zero the Cartesian gradient on frozen atoms so the reported
# |grad| reflects only the free degrees of freedom (the bounds
# already keep frozen positions in place; this just keeps the
# convergence metric honest).
if frozen_set:
for a in frozen_set:
grad_cart[a, :] = 0.0
# dE/d(frac) = lattice^T . dE/d(cart)
grad_frac = np.zeros_like(grad_cart)
for a in range(len(sys.unit_cell)):
grad_frac[a, :] = lattice.T @ grad_cart[a, :]
history.append((e, float(np.linalg.norm(grad_cart))))
return grad_frac.ravel()
# Optional per-step trajectory capture. scipy's L-BFGS-B fires
# callback(x) after each accepted step -- exactly one frame per
# step, with the initial geometry recorded up front.
trajectory_frames: list[PeriodicSystem] = []
trajectory_energies: list[float] = []
if output_trajectory is not None:
trajectory_frames.append(system)
# Energy of the initial frame; cheap because the first
# objective() call evaluates exactly this geometry.
last_step_energy: dict[str, float] = {}
def _wrapped_objective(x: np.ndarray) -> float:
e = objective(x)
last_step_energy["value"] = e
return e
def _capture(x: np.ndarray) -> None:
if output_trajectory is None:
return
trajectory_frames.append(_flat_to_system(system, x))
trajectory_energies.append(
last_step_energy.get("value", float("nan"))
)
x0 = _atoms_to_flat(system)
res = minimize(
_wrapped_objective if output_trajectory is not None else objective,
x0,
method="L-BFGS-B",
jac=gradient,
bounds=bounds,
options={"maxiter": max_iter, "gtol": conv_tol_grad},
callback=_capture if output_trajectory is not None else None,
)
sys_opt = _flat_to_system(system, res.x)
# Independent gradient gate: scipy reports res.success on EITHER the
# gtol OR the (default) ftol criterion, so a success flag alone can
# claim convergence at a non-stationary geometry. Require the actual
# max-component (fractional) gradient to meet conv_tol_grad
# (2026-05-31 audit, F1; same gate as optimize_molecule).
from .molecular_optimize import _gradient_converged
converged, grad_max = _gradient_converged(
bool(res.success), res.jac, conv_tol_grad
)
print(
f"\nAtomic relaxation: {res.nit} iters, "
f"E = {res.fun:.8f} Ha, "
f"max|grad| = {grad_max:.4e}, "
f"converged={converged}"
)
if output_trajectory is not None:
# The initial frame's energy isn't captured by callback (which
# fires after the first step). Patch it in by evaluating the
# first-frame energy from the first objective call -- the
# history list captures that.
if history and len(trajectory_energies) < len(trajectory_frames):
trajectory_energies.insert(0, history[0][0])
# If the optimizer terminated cleanly, the last accepted x is
# res.x -- append it if callback didn't catch it (e.g.
# converged-on-step-0).
if len(trajectory_frames) == 1:
trajectory_frames.append(sys_opt)
trajectory_energies.append(float(res.fun))
from .output.formats.qvf import write_reaction_path_qvf
n = len(trajectory_frames)
waypoints: list[dict[str, Any]] = [
{
"frame_index": 0,
"label": "start",
"kind": "reactant",
"energy_eh": float(trajectory_energies[0]),
},
{
"frame_index": n - 1,
"label": "converged" if converged else "stopped",
"kind": "product",
"energy_eh": float(trajectory_energies[-1]),
},
]
# Reaction coordinate = step index normalised to [0, 1].
rc = [i / max(1, n - 1) for i in range(n)]
write_reaction_path_qvf(
output_trajectory,
frames=trajectory_frames,
energies=trajectory_energies,
waypoints=waypoints,
reaction_coordinate=rc,
method=method_upper,
basis=basis_name,
functional=functional,
)
return OptimizeResult(sys_opt, res.fun, res.jac, res.nit, converged)
def relax_cell(
system: PeriodicSystem,
basis_name: str,
kmesh: KMeshInput,
method: str = "RHF",
*,
functional: Optional[str] = None,
max_iter: int = 20,
scf_conv_tol: Optional[float] = None,
scf_max_iter: Optional[int] = None,
scf_options: Optional[Union[PeriodicRHFOptions, PeriodicKSOptions]] = None,
cutoff_bohr: float = 8.0,
**bipole_kwargs,
) -> OptimizeResult:
"""Relax lattice parameters via energy-only Nelder-Mead.
Optimizes the 6 independent lattice strain components.
Atomic positions are NOT relaxed -- use relax_full for both.
Parameters
----------
system, basis_name, kmesh, method, functional
As in relax_atoms.
max_iter : int
Maximum Nelder-Mead iterations.
scf_conv_tol, scf_max_iter, scf_options
As in :func:`relax_atoms`.
"""
kmesh = _as_bloch_kmesh(kmesh)
method_upper = method.upper()
opts = _make_scf_options(
method_upper,
cutoff_bohr=cutoff_bohr,
scf_conv_tol=scf_conv_tol,
scf_max_iter=scf_max_iter,
scf_options=scf_options,
)
ref_lattice = np.asarray(system.lattice, dtype=float)
atoms = list(system.unit_cell)
def objective(strain: np.ndarray) -> float:
# strain is [e_xx, e_yy, e_zz, e_yz, e_xz, e_xy] in Voigt notation
e_xx, e_yy, e_zz, e_yz, e_xz, e_xy = strain
strain_matrix = (
np.array(
[
[e_xx, e_xy, e_xz],
[e_xy, e_yy, e_yz],
[e_xz, e_yz, e_zz],
]
)
* 0.01
) # scale to percent
new_lattice = ref_lattice @ (np.eye(3) + strain_matrix)
sys = PeriodicSystem(
system.dim,
new_lattice,
atoms,
charge=system.charge,
multiplicity=system.multiplicity,
)
basis = BasisSet(sys.unit_cell_molecule(), basis_name)
km = _kmesh_for_system(kmesh, system, sys)
e, _ = _run_scf(
sys, basis, km, opts, method_upper, functional, **bipole_kwargs
)
return e
res = minimize(
objective,
np.zeros(6),
method="Nelder-Mead",
options={"maxiter": max_iter, "xatol": 0.01, "fatol": 1e-6},
)
# Build optimized system
e_xx, e_yy, e_zz, e_yz, e_xz, e_xy = res.x
strain_matrix = (
np.array(
[
[e_xx, e_xy, e_xz],
[e_xy, e_yy, e_yz],
[e_xz, e_yz, e_zz],
]
)
* 0.01
)
opt_lattice = ref_lattice @ (np.eye(3) + strain_matrix)
sys_opt = PeriodicSystem(
system.dim,
opt_lattice,
atoms,
charge=system.charge,
multiplicity=system.multiplicity,
)
print(
f"\nCell relaxation: {res.nit} iters, "
f"E = {res.fun:.8f} Ha, converged={res.success}"
)
print(f" Lattice vectors (bohr):")
for i in range(3):
print(
f" a{i + 1} = [{opt_lattice[i, 0]:10.6f} {opt_lattice[i, 1]:10.6f} {opt_lattice[i, 2]:10.6f}]"
)
return OptimizeResult(sys_opt, res.fun, None, res.nit, res.success)
def relax_cell_gradient(
system: PeriodicSystem,
basis_name: str,
kmesh: KMeshInput,
method: str = "RHF",
*,
functional: Optional[str] = None,
force_mode: str = "fd",
fd_step_bohr: float = 1e-3,
max_iter: int = 20,
scf_conv_tol: Optional[float] = None,
scf_max_iter: Optional[int] = None,
scf_options: Optional[Union[PeriodicRHFOptions, PeriodicKSOptions]] = None,
cutoff_bohr: float = 8.0,
**bipole_kwargs,
) -> OptimizeResult:
"""Relax lattice (6 Voigt strains) with L-BFGS-B on the exact energy.
The L-BFGS-B gradient is the **central finite difference of the SCF
energy** w.r.t. each strain component -- consistent by construction with
the energy being minimised. (Before 2026-06-05 this used the force
virial ``compute_stress_tensor``, which is wrong for this objective --
it assumes the atoms scale with strain and misses the explicit
lattice/Ewald + Gaussian Pulay stress, coming out opposite in sign on
H₂/STO-3G. Use :func:`relax_cell` for a gradient-free Nelder-Mead
alternative.) ``force_mode`` / ``fd_step_bohr`` are accepted for API
compatibility but unused (the cell gradient is energy-FD, not the atomic
force).
``scf_options`` is copied into every strained geometry so smearing,
FMIXING, level-shift, accelerator, and grid controls match the energy
route being optimized.
"""
kmesh = _as_bloch_kmesh(kmesh)
method_upper = method.upper()
opts = _make_scf_options(
method_upper,
cutoff_bohr=cutoff_bohr,
scf_conv_tol=scf_conv_tol,
scf_max_iter=scf_max_iter,
scf_options=scf_options,
)
ref_lattice = np.asarray(system.lattice, dtype=float)
atoms = list(system.unit_cell)
V0 = float(abs(np.linalg.det(ref_lattice)))
def objective(strain_pct: np.ndarray) -> float:
strain = strain_pct * 0.01
S = np.array(
[
[strain[0], strain[5], strain[4]],
[strain[5], strain[1], strain[3]],
[strain[4], strain[3], strain[2]],
]
)
new_lattice = ref_lattice @ (np.eye(3) + S)
sys = PeriodicSystem(
system.dim,
new_lattice,
atoms,
charge=system.charge,
multiplicity=system.multiplicity,
)
basis = BasisSet(sys.unit_cell_molecule(), basis_name)
km = _kmesh_for_system(kmesh, system, sys)
e, _ = _run_scf(
sys, basis, km, opts, method_upper, functional, **bipole_kwargs
)
return e
def gradient(strain_pct: np.ndarray) -> np.ndarray:
# Exact dE/dstrain via a central finite difference of the energy
# actually being minimised -- the production-correct cell gradient.
#
# NOTE: the force virial ``compute_stress_tensor`` (used here before
# 2026-06-05) is WRONG for this objective: it (a) assumes the atoms
# scale with the strain, but ``objective`` holds the Cartesian atom
# positions fixed, and (b) misses the explicit lattice / Ewald strain
# dependence + the Gaussian-basis Pulay stress. On H₂/STO-3G it came
# out *opposite in sign* to dE/dstrain, so L-BFGS-B walked the cell
# the wrong way. FD of the objective is consistent by construction.
return _central_fd_gradient(objective, strain_pct, _CELL_FD_STRAIN_PCT)
cell_gtol = 1e-6
res = minimize(
objective,
np.zeros(6),
method="L-BFGS-B",
jac=gradient,
options={"maxiter": max_iter, "gtol": cell_gtol},
)
strain = res.x * 0.01
S = np.array(
[
[strain[0], strain[5], strain[4]],
[strain[5], strain[1], strain[3]],
[strain[4], strain[3], strain[2]],
]
)
opt_lattice = ref_lattice @ (np.eye(3) + S)
sys_opt = PeriodicSystem(
system.dim,
opt_lattice,
atoms,
charge=system.charge,
multiplicity=system.multiplicity,
)
# Gradient-gated convergence (F1): res.success can be set by scipy's
# ftol criterion at a non-stationary cell, so require the actual
# max-component strain gradient to meet the gtol tolerance.
from .molecular_optimize import _gradient_converged
converged, grad_max = _gradient_converged(bool(res.success), res.jac, cell_gtol)
print(
f"\nCell relaxation (gradient): {res.nit} iters, "
f"E = {res.fun:.8f} Ha, max|grad| = {grad_max:.4e}, converged={converged}"
)
return OptimizeResult(sys_opt, res.fun, None, res.nit, converged)
def relax_full(
system: PeriodicSystem,
basis_name: str,
kmesh: KMeshInput,
method: str = "RHF",
*,
functional: Optional[str] = None,
force_mode: str = "fd",
fd_step_bohr: float = 1e-3,
max_outer: int = 5,
max_atom_iter: int = 20,
max_cell_iter: int = 10,
conv_tol_grad: float = 1e-4,
scf_conv_tol: Optional[float] = None,
scf_max_iter: Optional[int] = None,
scf_options: Optional[Union[PeriodicRHFOptions, PeriodicKSOptions]] = None,
cutoff_bohr: float = 8.0,
**bipole_kwargs,
) -> OptimizeResult:
"""Full structure optimization: alternate atomic + cell relaxation.
Atomic forces default to the exact finite-difference path
(``force_mode="fd"``); see the module docstring.
Parameters
----------
max_outer : int
Maximum number of outer cell+atom cycles.
max_atom_iter, max_cell_iter
Max iterations per inner relaxation step.
scf_conv_tol, scf_max_iter, scf_options
As in :func:`relax_atoms`.
"""
current = system
current_kmesh = _as_bloch_kmesh(kmesh)
for outer in range(max_outer):
# Relax atoms at fixed lattice
atom_result = relax_atoms(
current,
basis_name,
current_kmesh,
method,
functional=functional,
force_mode=force_mode,
fd_step_bohr=fd_step_bohr,
max_iter=max_atom_iter,
conv_tol_grad=conv_tol_grad,
scf_conv_tol=scf_conv_tol,
scf_max_iter=scf_max_iter,
scf_options=scf_options,
cutoff_bohr=cutoff_bohr,
**bipole_kwargs,
)
current = atom_result.system
# Relax lattice at fixed (relaxed) atoms
before_cell = current
cell_result = relax_cell(
current,
basis_name,
current_kmesh,
method,
functional=functional,
max_iter=max_cell_iter,
scf_conv_tol=scf_conv_tol,
scf_max_iter=scf_max_iter,
scf_options=scf_options,
cutoff_bohr=cutoff_bohr,
**bipole_kwargs,
)
current = cell_result.system
current_kmesh = _kmesh_for_system(current_kmesh, before_cell, current)
print(f" Outer cycle {outer + 1}/{max_outer}: E = {cell_result.energy:.8f} Ha")
# Outer convergence: both inner relaxations of THIS cycle converged
# (atoms already stationary AND cell already stationary) => the coupled
# structure is at a stationary point; stop early.
if atom_result.converged and cell_result.converged:
return OptimizeResult(
current, cell_result.energy, None, outer + 1, True)
# Exhausted max_outer cycles without both inner relaxations converging in
# the same cycle -- report the achieved (not assumed) convergence state.
return OptimizeResult(
current, cell_result.energy, None, max_outer,
bool(atom_result.converged and cell_result.converged))