"""BIPOLE structure optimization — G3.
Atomic position relaxation using analytic BIPOLE gradients wrapped
in scipy L-BFGS-B. Lattice parameter optimization via energy-only
Nelder-Mead (stress tensor pending C++ implementation). Combined
cell + position relaxation alternates the two until convergence.
Usage:
from vibeqc.bipole_optimize import relax_atoms, relax_cell, relax_full
# Atomic relaxation only
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 Any, List, Optional, Sequence, Tuple, Union
import numpy as np
from scipy.optimize import minimize
from ._vibeqc_core import (
Atom,
BasisSet,
BlochKMesh,
Functional,
PeriodicKSOptions,
PeriodicRHFOptions,
PeriodicSystem,
monkhorst_pack,
)
from .bipole_gradient import (
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",
]
[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}")
return result.energy, result
def _compute_gradient(system, basis, result, method, lattice_opts):
"""Compute the BIPOLE gradient for any method."""
if method == "RHF":
return compute_bipole_gradient_rhf(
system,
basis,
result,
lattice_opts=lattice_opts,
)
elif method == "UHF":
return compute_bipole_gradient_uhf(
system,
basis,
result,
lattice_opts=lattice_opts,
)
elif method == "RKS":
return compute_bipole_gradient_rks(
system,
basis,
result,
lattice_opts=lattice_opts,
)
elif method == "UKS":
return compute_bipole_gradient_uks(
system,
basis,
result,
lattice_opts=lattice_opts,
)
else:
raise ValueError(f"Unknown method: {method}")
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: BlochKMesh,
method: str = "RHF",
*,
functional: Optional[str] = None,
max_iter: int = 30,
conv_tol_grad: float = 1e-4,
scf_conv_tol: float = 1e-7,
cutoff_bohr: float = 8.0,
freeze_indices: Optional[Sequence[int]] = None,
output_trajectory: Optional["str | Path"] = None,
**bipole_kwargs,
) -> OptimizeResult:
"""Relax atomic positions using analytic BIPOLE gradients + L-BFGS-B.
Parameters
----------
system : PeriodicSystem
Initial geometry (lattice fixed).
basis_name : str
Basis set name (rebuild per geometry step).
kmesh : BlochKMesh
k-point mesh.
method : str
"RHF", "UHF", "RKS", or "UKS".
functional : str, optional
XC functional for RKS/UKS.
max_iter : int
Maximum optimization steps.
conv_tol_grad : float
Gradient convergence tolerance (Ha/bohr).
scf_conv_tol : float
SCF energy convergence tolerance.
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`` § "Periodic reaction paths").
Default ``None`` ⇒ no trajectory output (no overhead).
Returns
-------
OptimizeResult
"""
method_upper = method.upper()
opts = (
PeriodicKSOptions() if method_upper in ("RKS", "UKS") else PeriodicRHFOptions()
)
opts.lattice_opts.cutoff_bohr = cutoff_bohr
opts.lattice_opts.nuclear_cutoff_bohr = cutoff_bohr
opts.max_iter = 50
opts.use_diis = True
opts.conv_tol_energy = scf_conv_tol
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))
def objective(x: np.ndarray) -> float:
sys = _flat_to_system(system, x)
basis = BasisSet(sys.unit_cell_molecule(), basis_name)
e, res = _run_scf(
sys, basis, kmesh, opts, method_upper, functional, **bipole_kwargs
)
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)
e, res = _run_scf(
sys, basis, kmesh, opts, method_upper, functional, **bipole_kwargs
)
# Convert Cartesian gradient to fractional gradient
lattice = np.asarray(sys.lattice, dtype=float)
grad_cart = _compute_gradient(sys, basis, res, method_upper, opts.lattice_opts)
# 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)
print(
f"\nAtomic relaxation: {res.nit} iters, "
f"E = {res.fun:.8f} Ha, "
f"|grad| = {float(np.linalg.norm(res.jac)):.4e}, "
f"converged={res.success}"
)
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 res.success 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, res.success)
def relax_cell(
system: PeriodicSystem,
basis_name: str,
kmesh: BlochKMesh,
method: str = "RHF",
*,
functional: Optional[str] = None,
max_iter: int = 20,
scf_conv_tol: float = 1e-7,
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.
"""
method_upper = method.upper()
opts = (
PeriodicKSOptions() if method_upper in ("RKS", "UKS") else PeriodicRHFOptions()
)
opts.lattice_opts.cutoff_bohr = cutoff_bohr
opts.lattice_opts.nuclear_cutoff_bohr = cutoff_bohr
opts.max_iter = 50
opts.use_diis = True
opts.conv_tol_energy = scf_conv_tol
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)
e, _ = _run_scf(
sys, basis, kmesh, 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: BlochKMesh,
method: str = "RHF",
*,
functional: Optional[str] = None,
max_iter: int = 20,
scf_conv_tol: float = 1e-7,
cutoff_bohr: float = 8.0,
**bipole_kwargs,
) -> OptimizeResult:
"""Relax lattice using force-based stress + L-BFGS-B.
Uses ``compute_stress_tensor`` for an analytic cell gradient,
replacing the energy-only Nelder-Mead approach.
"""
from .bipole_gradient import compute_stress_tensor
method_upper = method.upper()
opts = (
PeriodicKSOptions() if method_upper in ("RKS", "UKS") else PeriodicRHFOptions()
)
opts.lattice_opts.cutoff_bohr = cutoff_bohr
opts.lattice_opts.nuclear_cutoff_bohr = cutoff_bohr
opts.max_iter = 50
opts.use_diis = True
opts.conv_tol_energy = scf_conv_tol
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)
e, _ = _run_scf(
sys, basis, kmesh, opts, method_upper, functional, **bipole_kwargs
)
return e
def gradient(strain_pct: np.ndarray) -> np.ndarray:
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)
e, res = _run_scf(
sys, basis, kmesh, opts, method_upper, functional, **bipole_kwargs
)
grad_cart = _compute_gradient(sys, basis, res, method_upper, opts.lattice_opts)
stress = compute_stress_tensor(sys, grad_cart)
# dE/d(strain%) = V * stress * 0.01 (Voigt: xx,yy,zz,yz,xz,xy)
V = float(abs(np.linalg.det(new_lattice)))
dstress = (
np.array(
[
stress[0, 0],
stress[1, 1],
stress[2, 2],
stress[1, 2],
stress[0, 2],
stress[0, 1],
]
)
* V
* 0.01
)
return dstress
res = minimize(
objective,
np.zeros(6),
method="L-BFGS-B",
jac=gradient,
options={"maxiter": max_iter, "gtol": 1e-6},
)
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,
)
print(
f"\nCell relaxation (gradient): {res.nit} iters, "
f"E = {res.fun:.8f} Ha, converged={res.success}"
)
return OptimizeResult(sys_opt, res.fun, None, res.nit, res.success)
def relax_full(
system: PeriodicSystem,
basis_name: str,
kmesh: BlochKMesh,
method: str = "RHF",
*,
functional: Optional[str] = None,
max_outer: int = 5,
max_atom_iter: int = 20,
max_cell_iter: int = 10,
conv_tol_grad: float = 1e-4,
scf_conv_tol: float = 1e-7,
cutoff_bohr: float = 8.0,
**bipole_kwargs,
) -> OptimizeResult:
"""Full structure optimization: alternate atomic + cell relaxation.
Parameters
----------
max_outer : int
Maximum number of outer cell+atom cycles.
max_atom_iter, max_cell_iter
Max iterations per inner relaxation step.
"""
current = system
for outer in range(max_outer):
# Relax atoms at fixed lattice
atom_result = relax_atoms(
current,
basis_name,
kmesh,
method,
functional=functional,
max_iter=max_atom_iter,
conv_tol_grad=conv_tol_grad,
scf_conv_tol=scf_conv_tol,
cutoff_bohr=cutoff_bohr,
**bipole_kwargs,
)
current = atom_result.system
# Relax lattice at fixed (relaxed) atoms
cell_result = relax_cell(
current,
basis_name,
kmesh,
method,
functional=functional,
max_iter=max_cell_iter,
scf_conv_tol=scf_conv_tol,
cutoff_bohr=cutoff_bohr,
**bipole_kwargs,
)
current = cell_result.system
print(f" Outer cycle {outer + 1}/{max_outer}: E = {cell_result.energy:.8f} Ha")
return OptimizeResult(current, cell_result.energy, None, outer + 1, True)