"""Variational two-electron reduced density matrix (v2RDM) solver.
Minimises the electronic energy directly as a functional of the 2-RDM
under necessary N-representability constraints.
The energy functional is:
E[¹D, ²D] = Σ_{ij} h_{ij} ¹D_{ij}
+ ¼ Σ_{ijkl} g_{ijkl} ²D_{ijkl}
+ E_nuc
where ¹D is the 1-RDM and ²D is the 2-RDM, contracted by
¹D_{ij} = (1/(N−1)) Σ_k ²D_{ikjk}
Constraints enforced via augmented Lagrangian:
* ²D ≥ 0 (positive semidefinite 2-RDM)
* Tr(¹D) = N (fixed particle number)
* Hermiticity + antisymmetry of ²D
The algorithm is a boundary-point method: gradient step → antisymmetrise
→ PSD projection → update Lagrange multiplier for trace.
References
----------
* Mazziotti, *Phys. Rev. Lett.* 93, 213001 (2004).
* Mazziotti, *Acc. Chem. Res.* 39, 207 (2006).
* DePrince, *J. Chem. Phys.* 145, 164109 (2016).
"""
from __future__ import annotations
import time
import warnings
from dataclasses import dataclass
from typing import Optional
import numpy as np
from ._common import Hamiltonian, SolverOptions, SolverResult
[docs]
@dataclass
class V2RDMOptions(SolverOptions):
"""Options for the v2RDM solver.
Only the 'p' constraint (²D ≥ 0) is currently enforced.
Q and G constraint builders exist but feedback is not yet
implemented — a NotImplementedError is raised if 'q' or 'g'
appear in the constraints string.
"""
constraints: str = "p"
mu: float = 10.0
mu_factor: float = 1.2
mu_max: float = 1e8
outer_max_iter: int = 500
conv_tol_primal: float = 1e-6
conv_tol_dual: float = 1e-6
class V2RDMSolver:
"""Variational 2-RDM solver via augmented-Lagrangian / boundary-point."""
def __init__(self, options: Optional[V2RDMOptions] = None):
self.options = options or V2RDMOptions()
def solve(
self,
hamiltonian: Hamiltonian,
options: Optional[V2RDMOptions] = None,
) -> SolverResult:
opts = options or self.options
h1e = hamiltonian.h1e
h2e = hamiltonian.h2e
enuc = hamiltonian.nuclear_repulsion
norb = hamiltonian.norb
nelec = hamiltonian.nelec
if norb > 12:
warnings.warn(
f"v2RDM with {norb} orbitals uses O(n^6) storage. "
f"Consider active-space reduction."
)
# Initialise from HF 2-RDM
rdm2 = self._init_rdm2(norb, nelec)
rdm1 = self._contract_rdm1(rdm2, nelec)
energy_trace: list[float] = []
mu = opts.mu
lam = 0.0 # Lagrange multiplier for Tr(¹D) = N
t_start = time.perf_counter()
converged = False
for iteration in range(opts.outer_max_iter):
# ── Energy gradient w.r.t. ²D ────────────────────────────
grad_e = 0.5 * h2e.copy()
if nelec > 1:
for i in range(norb):
for k in range(norb):
grad_e[i, :, k, :] += h1e[i, k] / (nelec - 1) * np.eye(norb)
# ── Gradient of trace penalty: λ + μ·(Tr(¹D) − N) ───────
trace_defect = np.trace(rdm1) - nelec
penalty_coeff = lam + mu * trace_defect
grad_penalty = np.zeros_like(rdm2)
if nelec > 1:
for i in range(norb):
for k in range(norb):
grad_penalty[i, :, k, :] += (
penalty_coeff / (nelec - 1)
) * np.eye(norb)
grad_total = grad_e + grad_penalty
# ── Step size: 1/μ ───────────────────────────────────────
step = 1.0 / mu
rdm2_trial = rdm2 - step * grad_total
rdm2_trial = self._enforce_antisymmetry(rdm2_trial)
# ── PSD projection ───────────────────────────────────────
rdm2_proj = self._project_psd(rdm2_trial, norb)
rdm2_proj = self._enforce_antisymmetry(rdm2_proj)
# ── Update ───────────────────────────────────────────────
change = np.linalg.norm(rdm2_proj - rdm2)
rdm2 = rdm2_proj
rdm1 = self._contract_rdm1(rdm2, nelec)
# ── Q and G PSD projections (feed back via residual / μ) ─
_constr = opts.constraints.lower()
if "q" in _constr or "g" in _constr:
raise NotImplementedError(
"v2RDM Q/G constraint feedback is not yet implemented. "
"Only 'p' (²D ≥ 0) is currently enforced. "
"Remove 'q'/'g' from constraints or set constraints='p'."
)
# Update Lagrange multiplier
trace_defect = np.trace(rdm1) - nelec
lam = lam + mu * trace_defect
# Energy
energy = self._compute_energy(rdm2, rdm1, h1e, h2e, enuc)
energy_trace.append(energy)
# Residual (all requested constraints)
psd_res = self._compute_residual(rdm2, rdm1, norb, opts.constraints)
if opts.verbose >= 1 and (iteration < 10 or iteration % 20 == 0):
print(
f" v2RDM iter {iteration + 1:4d}: "
f"E={energy:.10f}, Tr(D)={np.trace(rdm1):.6f}, "
f"psd_res={psd_res:.2e}, change={change:.2e}"
)
if (
abs(trace_defect) < opts.conv_tol_primal
and psd_res < opts.conv_tol_primal
and change < opts.conv_tol_primal
):
converged = True
if opts.verbose:
print(f" v2RDM converged after {iteration + 1} iterations.")
break
# Increase penalty slowly
if iteration > 0 and iteration % 20 == 0:
mu = min(mu * opts.mu_factor, opts.mu_max)
dt = time.perf_counter() - t_start
final_energy = energy_trace[-1] if energy_trace else 0.0
return SolverResult(
energy=final_energy,
method=f"v2rdm({opts.constraints})",
converged=converged,
n_iter=iteration + 1,
energy_trace=energy_trace,
rdm1=rdm1,
rdm2=rdm2,
constraint_residual=abs(np.trace(rdm1) - nelec) if energy_trace else None,
)
# ── Initialisation ────────────────────────────────────────────────
def _init_rdm2(self, norb: int, nelec: int) -> np.ndarray:
"""Build the HF 2-RDM from the aufbau occupation."""
rdm1_hf = np.zeros((norb, norb))
nocc = nelec // 2
for i in range(nocc):
rdm1_hf[i, i] = 2.0
rdm2 = np.zeros((norb, norb, norb, norb))
for i in range(norb):
for j in range(norb):
for k in range(norb):
for l_ in range(norb):
rdm2[i, j, k, l_] = (
rdm1_hf[i, k] * rdm1_hf[j, l_]
- 0.5 * rdm1_hf[i, l_] * rdm1_hf[j, k]
)
return rdm2
def _contract_rdm1(self, rdm2: np.ndarray, nelec: int) -> np.ndarray:
"""¹D_{ij} = 1/(N−1) Σ_k ²D_{ikjk}."""
norb = rdm2.shape[0]
rdm1 = np.zeros((norb, norb))
if nelec > 1:
for i in range(norb):
for j in range(norb):
s = 0.0
for k in range(norb):
s += rdm2[i, k, j, k]
rdm1[i, j] = s / (nelec - 1)
return rdm1
# ── Energy ────────────────────────────────────────────────────────
def _compute_energy(
self,
rdm2: np.ndarray,
rdm1: np.ndarray,
h1e: np.ndarray,
h2e: np.ndarray,
enuc: float,
) -> float:
"""E = Σ h_{ij} ¹D_{ij} + ¼ Σ g_{ijkl} ²D_{ijkl} + E_nuc."""
e1 = np.sum(h1e * rdm1)
e2 = 0.5 * np.sum(h2e * rdm2)
return e1 + e2 + enuc
# ── Projections ───────────────────────────────────────────────────
def _project_psd(self, rdm2: np.ndarray, norb: int) -> np.ndarray:
"""Project onto the PSD cone by eigenvalue thresholding.
Reshape ²D → matrix M_{(ij),(kl)}, symmetrise, set negative
eigenvalues to 0, reshape back.
"""
mat = rdm2.reshape(norb * norb, norb * norb)
mat = 0.5 * (mat + mat.T)
evals, evecs = np.linalg.eigh(mat)
evals[evals < 0] = 0.0
mat_pos = evecs @ np.diag(evals) @ evecs.T
return mat_pos.reshape(norb, norb, norb, norb)
def _enforce_antisymmetry(self, rdm2: np.ndarray) -> np.ndarray:
"""Enforce ²D_{ijkl} = −²D_{jikl} = −²D_{ijlk} = ²D_{jilk}."""
norb = rdm2.shape[0]
result = np.zeros_like(rdm2)
for i in range(norb):
for j in range(norb):
for k in range(norb):
for l_ in range(norb):
result[i, j, k, l_] = 0.25 * (
rdm2[i, j, k, l_]
- rdm2[j, i, k, l_]
- rdm2[i, j, l_, k]
+ rdm2[j, i, l_, k]
)
return result
# ── Q and G matrix builders ───────────────────────────────────────
def _build_q_matrix(
self, rdm2: np.ndarray, rdm1: np.ndarray, norb: int
) -> np.ndarray:
"""Build the 2-hole RDM in the (ij,kl) supermatrix basis.
²Q_{ijkl} = δ_{ik}δ_{jl} − δ_{il}δ_{jk}
− δ_{ik}¹D_{jl} + δ_{il}¹D_{jk}
− δ_{jl}¹D_{ik} + δ_{jk}¹D_{il}
+ ²D_{ijkl}
"""
Q = np.zeros((norb, norb, norb, norb))
for i in range(norb):
for j in range(norb):
for k in range(norb):
for l_ in range(norb):
val = 0.0
if i == k and j == l_:
val += 1.0
if i == l_ and j == k:
val -= 1.0
if i == k:
val -= rdm1[j, l_]
if i == l_:
val += rdm1[j, k]
if j == l_:
val -= rdm1[i, k]
if j == k:
val += rdm1[i, l_]
val += rdm2[i, j, k, l_]
Q[i, j, k, l_] = val
return Q
def _build_g_matrix(
self, rdm2: np.ndarray, rdm1: np.ndarray, norb: int
) -> np.ndarray:
"""Build the particle-hole RDM in the (ij,kl) supermatrix basis.
²G_{ijkl} = δ_{jl}¹D_{ik} − ²D_{ilkj}
"""
G = np.zeros((norb, norb, norb, norb))
for i in range(norb):
for j in range(norb):
for k in range(norb):
for l_ in range(norb):
val = 0.0
if j == l_:
val += rdm1[i, k]
val -= rdm2[i, l_, k, j]
G[i, j, k, l_] = val
return G
def _compute_residual(
self,
rdm2: np.ndarray,
rdm1: np.ndarray,
norb: int,
constraints: str = "pqg",
) -> float:
"""Compute combined PSD constraint residual (sum of negative
eigenvalues across all requested constraints).
"""
res = 0.0
c_lower = constraints.lower()
# ²D ≥ 0 (2-particle RDM)
mat_d = rdm2.reshape(norb * norb, norb * norb)
mat_d = 0.5 * (mat_d + mat_d.T)
evals_d = np.linalg.eigvalsh(mat_d)
res += float(np.sum(np.abs(evals_d[evals_d < 0])))
# ²Q ≥ 0 (2-hole RDM)
if "q" in c_lower:
Q = self._build_q_matrix(rdm2, rdm1, norb)
mat_q = Q.reshape(norb * norb, norb * norb)
mat_q = 0.5 * (mat_q + mat_q.T)
evals_q = np.linalg.eigvalsh(mat_q)
res += float(np.sum(np.abs(evals_q[evals_q < 0])))
# ²G ≥ 0 (particle-hole RDM)
if "g" in c_lower:
G = self._build_g_matrix(rdm2, rdm1, norb)
mat_g = G.reshape(norb * norb, norb * norb)
mat_g = 0.5 * (mat_g + mat_g.T)
evals_g = np.linalg.eigvalsh(mat_g)
res += float(np.sum(np.abs(evals_g[evals_g < 0])))
return res
[docs]
def solve_v2rdm(
hamiltonian: Hamiltonian,
options: Optional[V2RDMOptions] = None,
) -> SolverResult:
"""One-shot v2RDM solve."""
solver = V2RDMSolver(options)
return solver.solve(hamiltonian)