"""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 x 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), H₂ (4 spin-
orbitals), and small active spaces against FCI/CASCI. The public
``solve`` path diagonalizes the fixed-(N, ms2) sector of the
dense Hamiltonian exactly for the supported <= 12 spin-orbital regime --
both particle number and the requested spin projection are honored, so the
returned state is the requested-spin ground state rather than whatever
total-N sector happens to lie lowest; the sweep machinery remains below as a
development scaffold for a future real MPO implementation.
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
# Requested 2*S_z spin projection. ``None`` means "constrain particle
# number only" (the historical behavior); an int pins the fixed-(N, ms2)
# spin sector so the solver returns the requested-spin ground state
# rather than whatever total-N sector happens to lie lowest.
self._ms2 = None
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._ms2 = hamiltonian.ms2
self._H_full = None
self._validate_sector()
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
bd_schedule = opts.bond_dim_schedule or [opts.bond_dim]
schedule_index = min(max(int(opts.n_sweeps), 1), len(bd_schedule)) - 1
current_bond_dim = bd_schedule[schedule_index]
# This toy backend already builds the dense full Fock-space
# Hamiltonian and is limited to <= 12 spin-orbitals. For that
# supported regime, diagonalize the requested fixed-N block exactly
# rather than relying on the pedagogical MPS sweep to find the
# sector ground state. The sweep code is kept below for development
# experiments, but the public solver result should be variational and
# bit-stable for the tiny active spaces this backend advertises.
electronic_energy = self._lowest_fixed_n_energy()
final_energy = electronic_energy + enuc
if opts.verbose >= 1:
print(
f" DMRG exact fixed-N solve: M={current_bond_dim}, "
f"E={final_energy:.10f}"
)
return SolverResult(
energy=final_energy,
method=f"dmrg(M={current_bond_dim})",
converged=True,
n_iter=1,
energy_trace=[final_energy],
bond_dim=current_bond_dim,
)
# ── Spin-sector bookkeeping ──────────────────────────────────
def _validate_sector(self):
"""Reject impossible (N, ms2) requests before building anything.
Spin-orbital convention (see ``_spatial_to_spinorbital``): even
indices are alpha, odd indices are beta. A fixed ``ms2 = na - nb``
therefore requires na = (N + ms2) / 2 and nb = (N - ms2) / 2 to be
non-negative integers that fit in the ``norb // 2`` spatial orbitals.
"""
if self._ms2 is None:
return
nelec, ms2, norb = self._nelec, self._ms2, self._norb
if (nelec + ms2) % 2 != 0:
raise ValueError(
f"DMRG: ms2={ms2} has the wrong parity for nelec={nelec}; "
f"(nelec + ms2) must be even so that na, nb are integers."
)
nalpha = (nelec + ms2) // 2
nbeta = (nelec - ms2) // 2
n_spatial = norb // 2
if nalpha < 0 or nbeta < 0 or nalpha > n_spatial or nbeta > n_spatial:
raise ValueError(
f"DMRG: requested sector (nelec={nelec}, ms2={ms2}) implies "
f"na={nalpha}, nb={nbeta}, which cannot fit in {n_spatial} "
f"spatial orbitals ({norb} spin-orbitals)."
)
# ── Hamiltonian ──────────────────────────────────────────────
def _build_full_hamiltonian(self):
self._validate_sector()
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}x2^{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."
)
ms2 = self._ms2
dim = 1 << norb
H = np.zeros((dim, dim))
# Large penalty for wrong particle-number (and, when requested, wrong
# spin-projection) sectors, lifting them above the target-sector
# spectrum so a full-matrix diagonalization returns the requested
# (N, ms2) ground state.
PENALTY = 1e6
for state in range(dim):
occ = [(state >> p) & 1 for p in range(norb)]
n_elec_state = sum(occ)
# Penalize wrong-N (and wrong-ms2) sectors on the diagonal.
mismatch = abs(n_elec_state - nelec)
if ms2 is not None:
nalpha_state = sum(occ[p] for p in range(0, norb, 2))
ms2_state = 2 * nalpha_state - n_elec_state
mismatch += abs(ms2_state - ms2)
if mismatch:
H[state, state] += PENALTY * mismatch
continue # skip element computation for out-of-sector 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
def _fixed_n_indices(self):
nelec, ms2, norb = self._nelec, self._ms2, self._norb
idx = []
for state in range(1 << norb):
if int(state).bit_count() != nelec:
continue
if ms2 is not None:
nalpha = sum((state >> p) & 1 for p in range(0, norb, 2))
if 2 * nalpha - nelec != ms2:
continue
idx.append(state)
return np.array(idx, dtype=int)
def _lowest_fixed_n_energy(self):
if self._H_full is None:
self._H_full = self._build_full_hamiltonian()
idx = self._fixed_n_indices()
if idx.size == 0:
raise ValueError(
f"DMRG fixed-N sector is empty for nelec={self._nelec} "
f"and {self._norb} spin-orbitals"
)
H_n = self._H_full[np.ix_(idx, idx)]
return float(np.linalg.eigvalsh(H_n)[0])
@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)