"""DMRG solver for ab initio Hamiltonians — EXPERIMENTAL / TOY.
Two-site DMRG using a dense full-Fock-space Hamiltonian for exact
environment contraction. This approach builds a 2^n × 2^n matrix
and is therefore limited to ≤ 6 spatial orbitals (12 spin-orbitals).
For larger systems the memory cost is prohibitive and the solver
will raise an explicit error.
For production DMRG, a proper complementary-operator MPO construction
(Keller et al., JCP 143, 244118, 2015) is planned for v0.10.0.
Currently verified: exact for He (2 spin-orbitals) and H₂ (4 spin-
orbitals) against FCI. Approximate for larger systems with known
environment-contraction limitations.
References: White (1992), Schollwöck (2011).
"""
from __future__ import annotations
import warnings
from dataclasses import dataclass, field
from typing import Optional
import numpy as np
from ._common import Hamiltonian, SolverOptions, SolverResult
[docs]
@dataclass
class DMRGOptions(SolverOptions):
"""Options for the DMRG solver."""
bond_dim: int = 64
bond_dim_schedule: list[int] = field(default_factory=lambda: [4, 8, 16, 32, 64])
n_sweeps: int = 10
n_lanczos_iter: int = 20
truncation_tol: float = 1e-8
orbital_order: Optional[list[int]] = None
noise: float = 0.0
class DMRGSolver:
"""Two-site DMRG solver."""
def __init__(self, options=None):
self.options = options or DMRGOptions()
self._rng = np.random.default_rng(self.options.random_seed)
self._h1e = None
self._h2e = None
self._norb = 0
self._nelec = 0
self._H_full = None
def solve(self, hamiltonian, options=None):
opts = options or self.options
enuc = hamiltonian.nuclear_repulsion
h1e, h2e, norb = self._spatial_to_spinorbital(
hamiltonian.h1e, hamiltonian.h2e, hamiltonian.norb
)
nelec = hamiltonian.nelec
self._h1e = h1e
self._h2e = h2e
self._norb = norb
self._nelec = nelec
self._H_full = None
if norb > 12:
warnings.warn(f"DMRG with {norb} spin-orbitals may be slow.")
order = opts.orbital_order
if order is not None:
h1e = h1e[np.ix_(order, order)]
h2e = h2e[np.ix_(order, order, order, order)]
self._h1e = h1e
self._h2e = h2e
mpo = [np.zeros((1, 1, 2, 2)) for _ in range(norb)]
mps = self._init_mps(norb, nelec, opts.bond_dim_schedule[0])
bd_schedule = opts.bond_dim_schedule or [opts.bond_dim]
sweep_energies = []
current_bond_dim = bd_schedule[0]
for sweep in range(opts.n_sweeps):
if sweep < len(bd_schedule):
current_bond_dim = bd_schedule[sweep]
se = self._run_sweep(mps, mpo, nelec, current_bond_dim, opts, "right")
se = self._run_sweep(mps, mpo, nelec, current_bond_dim, opts, "left")
sweep_energies.append(se)
if opts.verbose >= 1:
print(
f" DMRG sweep {sweep + 1:3d}: M={current_bond_dim}, E={se + enuc:.10f}"
)
if (
sweep > 0
and abs(sweep_energies[-1] - sweep_energies[-2]) < opts.conv_tol_energy
):
if opts.verbose:
print(f" DMRG converged after {sweep + 1} sweeps.")
break
final_energy = sweep_energies[-1] + enuc if sweep_energies else 0.0
return SolverResult(
energy=final_energy,
method=f"dmrg(M={current_bond_dim})",
converged=len(sweep_energies) > 1
and abs(sweep_energies[-1] - sweep_energies[-2]) < opts.conv_tol_energy,
n_iter=len(sweep_energies),
bond_dim=current_bond_dim,
)
# ── Hamiltonian ──────────────────────────────────────────────
def _build_full_hamiltonian(self):
h1e, h2e, norb = self._h1e, self._h2e, self._norb
nelec = self._nelec
# Hard guard: 2^14 = 16384 → 2 GB matrix. Limit to 12 spin-orbitals.
if norb > 12:
raise ValueError(
f"DMRG _build_full_hamiltonian: {norb} spin-orbitals "
f"requires a 2^{norb}×2^{norb} dense matrix "
f"({(1 << norb) ** 2 * 8 / (1024**3):.1f} GB). "
f"Maximum is 12 spin-orbitals (6 spatial). "
f"For larger systems, use Selected-CI or FCI instead."
)
dim = 1 << norb
H = np.zeros((dim, dim))
# Large penalty for wrong particle-number sectors
PENALTY = 1e6
for state in range(dim):
occ = [(state >> p) & 1 for p in range(norb)]
n_elec_state = sum(occ)
# Penalize wrong-N sectors on diagonal
if n_elec_state != nelec:
H[state, state] += PENALTY * abs(n_elec_state - nelec)
continue # skip matrix-element computation for wrong-N states
for p in range(norb):
if occ[p]:
H[state, state] += h1e[p, p]
for q in range(norb):
if occ[q]:
H[state, state] += 0.5 * (h2e[p, q, p, q] - h2e[p, q, q, p])
for p in range(norb):
if not occ[p]:
continue
for q in range(norb):
if occ[q]:
continue
ns = state ^ (1 << p) ^ (1 << q)
if ns <= state:
continue
# Only connect to same-N states
n_elec_ns = n_elec_state # single excitation preserves N
lo, hi = (p, q) if p < q else (q, p)
sign = 1
for r in range(lo + 1, hi):
if occ[r]:
sign = -sign
val = sign * h1e[p, q]
for r in range(norb):
if occ[r]:
val += sign * (h2e[p, r, q, r] - h2e[p, r, r, q])
H[state, ns] = val
H[ns, state] = val
for p in range(norb):
if not occ[p]:
continue
for q in range(p + 1, norb):
if not occ[q]:
continue
for r in range(norb):
if occ[r]:
continue
for s in range(r + 1, norb):
if occ[s]:
continue
ns = state ^ (1 << p) ^ (1 << q) ^ (1 << r) ^ (1 << s)
if ns <= state:
continue
# Double excitation also preserves N
sign = self._excitation_phase(list(occ), p, q, r, s)
val = sign * (h2e[p, q, r, s] - h2e[p, q, s, r])
H[state, ns] = val
H[ns, state] = val
return H
@staticmethod
def _excitation_phase(occ, p, q, r, s):
tmp, sign = occ[:], 1
if tmp[q] == 0:
return 0
for i in range(q + 1, len(tmp)):
if tmp[i]:
sign = -sign
tmp[q] = 0
if tmp[p] == 0:
return 0
for i in range(p + 1, len(tmp)):
if tmp[i]:
sign = -sign
tmp[p] = 0
if tmp[r]:
return 0
for i in range(r + 1, len(tmp)):
if tmp[i]:
sign = -sign
tmp[r] = 1
if tmp[s]:
return 0
for i in range(s + 1, len(tmp)):
if tmp[i]:
sign = -sign
return sign
# ── MPS ──────────────────────────────────────────────────────
def _init_mps(self, norb, nelec, bond_dim):
mps = []
left_dim = 1
for site in range(norb):
right_dim = min(bond_dim, 2 ** min(site + 1, norb - site - 1))
mps.append(self._rng.normal(0, 0.1, (left_dim, right_dim, 2)))
left_dim = right_dim
return self._canonicalize(mps, "right")
def _run_sweep(self, mps, mpo, nelec, bond_dim, opts, direction):
if self._H_full is None:
self._H_full = self._build_full_hamiltonian()
return self._sweep(mps, bond_dim, self._H_full, direction)
# ── Environment contraction ──────────────────────────────────
def _contract_left(self, mps, up_to):
if up_to < 0:
return np.ones((1, 1, 1))
result = mps[0].copy()
for i in range(1, up_to + 1):
result = np.tensordot(result, mps[i], axes=([1], [0]))
result = result.reshape(1, mps[i].shape[1], -1)
return result
def _contract_right(self, mps, from_site):
norb = len(mps)
if from_site >= norb:
return np.ones((1, 1))
result = mps[-1].copy()[:, 0, :]
for i in range(norb - 2, from_site - 1, -1):
result = np.einsum("ijk,jl->ikl", mps[i], result)
result = result.reshape(mps[i].shape[0], -1)
return result
def _build_heff(self, k, L_mat, R_mat, H_full):
norb = self._norb
dimL, dimR = 1 << k, 1 << (norb - k - 2)
bL, bR = L_mat.shape[0], R_mat.shape[0]
V = np.zeros((bL * 4 * bR, 1 << norb))
for iL in range(bL):
for p in range(4):
for iR in range(bR):
row = iL * (4 * bR) + p * bR + iR
for a in range(dimL):
lv = L_mat[iL, a]
if abs(lv) < 1e-15:
continue
for b in range(dimR):
rv = R_mat[iR, b]
if abs(rv) < 1e-15:
continue
V[row, a + p * dimL + b * (dimL * 4)] = lv * rv
return V @ H_full @ V.T
# ── Sweep core ──────────────────────────────────────────────
def _sweep(self, mps, bond_dim, H_full, direction):
norb = len(mps)
energy = 0.0
if direction == "right":
L_block = np.ones((1, 1, 1))
for k in range(norb - 1):
R_mat = self._contract_right(mps, k + 2)
L_mat = L_block[0, :, :]
H_eff = self._build_heff(k, L_mat, R_mat, H_full)
evals, evecs = np.linalg.eigh(H_eff)
energy = float(evals[0])
psi_gs = evecs[:, 0]
bL, bR = L_mat.shape[0], R_mat.shape[0]
two_site = psi_gs.reshape(bL, 2, 2, bR)
mat = two_site.transpose(0, 1, 3, 2).reshape(bL * 2, 2 * bR)
U, s, Vh = np.linalg.svd(mat, full_matrices=False)
keep = min(bond_dim, len(s))
# Left site: purely unitary (left-canonical) — no singular values
mps[k] = U[:, :keep].reshape(bL, 2, keep).transpose(0, 2, 1)
# Right site: absorbs all singular values
sVh = s[:keep, None] * Vh[:keep, :]
mps[k + 1] = sVh.reshape(keep, 2, bR).transpose(0, 2, 1)
L_new = np.tensordot(L_block, mps[k], axes=([1], [0]))
L_block = L_new.reshape(1, keep, -1)
else:
R_block = np.ones((1, 1))
for k in range(norb - 2, -1, -1):
L_mat = self._contract_left(mps, k - 1)[0, :, :]
H_eff = self._build_heff(k, L_mat, R_block, H_full)
evals, evecs = np.linalg.eigh(H_eff)
energy = float(evals[0])
psi_gs = evecs[:, 0]
bL, bR = L_mat.shape[0], R_block.shape[0]
two_site = psi_gs.reshape(bL, 2, 2, bR)
mat = two_site.transpose(0, 1, 3, 2).reshape(bL * 2, 2 * bR)
U, s, Vh = np.linalg.svd(mat, full_matrices=False)
keep = min(bond_dim, len(s))
# Left site: absorbs all singular values
Us = U[:, :keep] * s[:keep]
mps[k] = Us.reshape(bL, 2, keep).transpose(0, 2, 1)
# Right site: purely unitary (right-canonical) — no singular values
mps[k + 1] = Vh[:keep, :].reshape(keep, 2, bR).transpose(0, 2, 1)
R_new = np.tensordot(mps[k + 1], R_block, axes=([1], [0]))
R_block = R_new.reshape(keep, -1)
return energy
# ── Spin-orbital conversion ─────────────────────────────────
@staticmethod
def _spatial_to_spinorbital(h1e, h2e, norb_spatial):
n = 2 * norb_spatial
h1s = np.zeros((n, n))
for p in range(norb_spatial):
for q in range(norb_spatial):
h1s[2 * p, 2 * q] = h1s[2 * p + 1, 2 * q + 1] = h1e[p, q]
h2s = np.zeros((n, n, n, n))
for p in range(norb_spatial):
for q in range(norb_spatial):
for r in range(norb_spatial):
for s in range(norb_spatial):
v = h2e[p, q, r, s]
h2s[2 * p, 2 * q, 2 * r, 2 * s] = v
h2s[2 * p + 1, 2 * q + 1, 2 * r + 1, 2 * s + 1] = v
h2s[2 * p, 2 * q + 1, 2 * r, 2 * s + 1] = v
h2s[2 * p + 1, 2 * q, 2 * r + 1, 2 * s] = v
return h1s, h2s, n
# ── Canonicalization ────────────────────────────────────────
def _canonicalize(self, mps, form):
norb = len(mps)
if form == "right":
for i in range(norb - 1, 0, -1):
t = mps[i]
mat = t.transpose(0, 2, 1).reshape(t.shape[0], -1)
U, s, Vh = np.linalg.svd(mat, full_matrices=False)
mps[i] = Vh.reshape(len(s), t.shape[1], t.shape[2])
t_prev = np.tensordot(mps[i - 1], U * s, axes=([1], [0]))
mps[i - 1] = t_prev.transpose(0, 2, 1)
else:
for i in range(norb - 1):
t = mps[i]
mat = t.transpose(0, 2, 1).reshape(t.shape[0], -1)
U, s, Vh = np.linalg.svd(mat, full_matrices=False)
mps[i] = U.reshape(t.shape[0], len(s), t.shape[2])
t_next = np.tensordot(s[:, None] * Vh, mps[i + 1], axes=([1], [0]))
mps[i + 1] = t_next.transpose(0, 2, 1).reshape(
len(s), t_next.shape[2], t_next.shape[1]
)
return mps
[docs]
def solve_dmrg(hamiltonian, options=None):
return DMRGSolver(options).solve(hamiltonian)