Source code for vibeqc.solvers._selected_ci

"""Selected Configuration Interaction solver.

Implements a CIPSI-style (Configuration Interaction using a Perturbative
Selection made Iteratively) selected-CI algorithm:

1. Start with a reference determinant (usually the HF determinant).
2. Generate all singles and doubles from the current space.
3. Estimate each candidate's first-order perturbative contribution:
       ΔE_D ≈ |⟨D|H|Ψ_0⟩|² / (E_0 − ⟨D|H|D⟩)
4. Add the top-N candidates to the variational space.
5. Diagonalize the Hamiltonian in the expanded space.
6. Repeat until convergence (energy change, determinant count, or PT2 estimate).

Optionally, compute a final PT2 correction using the Epstein-Nesbet denominator.

Supports both spin-restricted (closed-shell ``Det``) and spin-unrestricted
(``SpinDet`` = ``(alpha_occ, beta_occ)``) determinant spaces.

References
----------
* Huron, Malrieu & Rancurel, *J. Chem. Phys.* 58, 5745 (1973).
* Evangelisti, Daudey & Malrieu, *Chem. Phys.* 75, 91 (1983).
* Tubman et al., *JCTC* 17, 151 (2021) — ASCI / SHCI modern variants.
"""

from __future__ import annotations

import time
from dataclasses import dataclass, field
from typing import Optional, Union

import numpy as np
from scipy.linalg import eigh as scipy_eigh

from ._common import Hamiltonian, SolverOptions, SolverResult
from ._determinant import Det, SpinDet, reference_determinant
from ._slater_condon import (
    build_hamiltonian_matrix,
    build_hamiltonian_matrix_unrestricted,
    diagonal_matrix_element,
    diagonal_matrix_element_unrestricted,
    double_excitation_matrix_element,
    hamiltonian_matrix_element_unrestricted,
    single_excitation_matrix_element,
)


[docs] @dataclass class SelectedCIOptions(SolverOptions): """Options for the Selected-CI solver. Attributes ---------- target_size : int Maximum number of determinants in the variational space. max_iter : int Maximum selection + diagonalization cycles. conv_tol_energy : float Convergence threshold on total energy change (Hartree). pt2_threshold : float Perturbative threshold for selecting new determinants. Lower = more determinants selected per iteration. selection_growth_factor : float Max ratio by which the determinant space can grow each iteration. max_det_per_iter : int Hard cap on new determinants per iteration. do_pt2_correction : bool Compute final PT2 energy correction after convergence. spin_restricted : bool Use spin-restricted (closed-shell) determinant basis. When True, each Det represents a doubly-occupied spatial-orbital configuration (spin-summed, S_z = 0). When False, uses SpinDet = (alpha_occ, beta_occ) pairs — required for open-shell or broken-spin-symmetry systems (multiplicity > 1). use_davidson : bool Use Davidson iterative diagonalization for ndet > davidson_threshold. davidson_threshold : int Switch to Davidson when ndet exceeds this count. davidson_max_iter : int Maximum Davidson iterations. davidson_conv_tol : float Davidson residual convergence tolerance. """ pt2_threshold: float = 1e-6 selection_growth_factor: float = 2.0 max_det_per_iter: int = 5000 do_pt2_correction: bool = True spin_restricted: bool = True use_davidson: bool = True davidson_threshold: int = 50 davidson_max_iter: int = 100 davidson_conv_tol: float = 1e-8
@dataclass class _SelectionCandidate: """Internal: a candidate determinant with its PT2 weight estimate.""" det: Union[Det, SpinDet] weight: float coupling: float # |⟨D|H|Ψ_0⟩| class SelectedCISolver: """CIPSI-style selected configuration interaction solver. Supports both spin-restricted (``spin_restricted=True``) and spin-unrestricted (``spin_restricted=False``) determinant spaces. """ def __init__(self, options: Optional[SelectedCIOptions] = None): self.options = options or SelectedCIOptions() self._rng = np.random.default_rng(self.options.random_seed) def solve( self, hamiltonian: Hamiltonian, options: Optional[SelectedCIOptions] = None, ) -> SolverResult: opts = options or self.options h1e = hamiltonian.h1e h2e = hamiltonian.h2e enuc = hamiltonian.nuclear_repulsion norb = hamiltonian.norb nelec = hamiltonian.nelec if nelec % 2 != 0 and opts.spin_restricted: raise ValueError( "Spin-restricted Selected-CI requires an even number of electrons. " "Use spin_restricted=False for open-shell systems." ) # ── Initialize with reference determinant ───────────────────── if opts.spin_restricted: nocc = nelec // 2 ref_det = tuple(range(nocc)) det_space: list = [ref_det] else: nalpha = (nelec + hamiltonian.ms2) // 2 nbeta = (nelec - hamiltonian.ms2) // 2 ref_a, ref_b = reference_determinant(nalpha, nbeta, spin_restricted=False) det_space: list = [(ref_a, ref_b)] energy_trace: list[float] = [] prev_energy = float("inf") converged = False t_start = time.perf_counter() # Save last diagonalization results for PT2 and final result eigenvalues = None eigenvectors = None coeffs = None for iteration in range(opts.max_iter): # ── Build + diagonalize H in current space ──────────────── ndet = len(det_space) if ndet <= 0: break if opts.spin_restricted: H_mat = build_hamiltonian_matrix(det_space, h1e, h2e) else: H_mat = build_hamiltonian_matrix_unrestricted(det_space, h1e, h2e) # Use Davidson for larger spaces to avoid O(N^3) dense diagonalization if opts.use_davidson and ndet > opts.davidson_threshold: eigenvalues, eigenvectors = self._davidson( H_mat, det_space, h1e, h2e, opts ) else: eigenvalues, eigenvectors = scipy_eigh(H_mat) e_var = eigenvalues[0] + enuc coeffs = eigenvectors[:, 0] energy_trace.append(e_var) delta_e = abs(e_var - prev_energy) prev_energy = e_var if opts.verbose >= 1: print( f" Selected-CI iter {iteration + 1:3d}: " f"ndet={ndet:6d}, E_var={e_var:.10f}, ΔE={delta_e:.2e}" ) # ── Check convergence ───────────────────────────────────── if delta_e < opts.conv_tol_energy and iteration > 0: converged = True if opts.verbose: print(f" Selected-CI converged in {iteration + 1} iterations.") break if ndet >= opts.target_size: if opts.verbose: print(f" Selected-CI: reached target size {opts.target_size}.") converged = True break # ── Selection step ──────────────────────────────────────── candidates = self._select_candidates( det_space, coeffs, h1e, h2e, norb, opts, spin_restricted=opts.spin_restricted, e_var_0=eigenvalues[0], ) if not candidates: if opts.verbose: print(f" Selected-CI: no new candidates (full space reached?).") converged = True break # Add top candidates to space new_dets = [c.det for c in candidates] det_space_set = set(det_space) for d in new_dets: if d not in det_space_set: det_space.append(d) # Safety: don't exceed target size by too much if len(det_space) > int(opts.target_size * opts.selection_growth_factor): # Trim back to target_size based on coefficient magnitude coeff_abs = np.abs(coeffs) keep_idx = np.argsort(-coeff_abs)[: opts.target_size] det_space = [det_space[i] for i in keep_idx] # ── PT2 correction ──────────────────────────────────────────── pt2_corr = 0.0 if opts.do_pt2_correction and eigenvalues is not None and coeffs is not None: pt2_corr = self._compute_pt2( det_space, eigenvalues[0], coeffs, h1e, h2e, norb, opts, spin_restricted=opts.spin_restricted, ) dt = time.perf_counter() - t_start return SolverResult( energy=e_var + pt2_corr, method=f"selected_ci(ndet={len(det_space)})", converged=converged, n_iter=iteration + 1, energy_trace=energy_trace, ci_coeffs=coeffs, ci_labels=det_space, pt2_correction=pt2_corr, ) def _select_candidates( self, det_space: list, coeffs: np.ndarray, h1e: np.ndarray, h2e: np.ndarray, norb: int, opts: SelectedCIOptions, *, spin_restricted: bool = True, e_var_0: float = 0.0, ) -> list[_SelectionCandidate]: """Generate candidates via singles/doubles from top determinants. When ``spin_restricted=False``, ``det_space`` is ``list[SpinDet]`` and all five spin-sector excitation classes are explored (α singles, β singles, αα doubles, ββ doubles, αβ doubles). """ if spin_restricted: return self._select_candidates_restricted( det_space, coeffs, h1e, h2e, norb, opts, e_var_0=e_var_0 ) else: return self._select_candidates_unrestricted( det_space, coeffs, h1e, h2e, norb, opts, e_var_0=e_var_0 ) def _select_candidates_restricted( self, det_space: list[Det], coeffs: np.ndarray, h1e: np.ndarray, h2e: np.ndarray, norb: int, opts: SelectedCIOptions, *, e_var_0: float = 0.0, ) -> list[_SelectionCandidate]: """Spin-restricted candidate selection (closed-shell determinants).""" candidates: dict[Det, _SelectionCandidate] = {} det_set = set(det_space) # Focus on determinants with significant weight c_abs = np.abs(coeffs) c_threshold = 0.01 significant = np.where(c_abs > c_threshold)[0] if len(significant) == 0: significant = [np.argmax(c_abs)] for I in significant: occ_I = det_space[I] occ_set = set(occ_I) vir_set = set(range(norb)) - occ_set vir_list = sorted(vir_set) # Singles for i in occ_set: for a in vir_list: new_det = tuple(sorted((occ_set - {i}) | {a})) if new_det in candidates or new_det in det_set: continue H_Ia = single_excitation_matrix_element(occ_I, i, a, h1e, h2e) coupling = abs(H_Ia * coeffs[I]) diag_new = diagonal_matrix_element(new_det, h1e, h2e) denom = e_var_0 - diag_new if abs(denom) < 1e-12: continue weight = coupling * coupling / abs(denom) candidates[new_det] = _SelectionCandidate( det=new_det, weight=weight, coupling=coupling ) # Doubles occ_list = sorted(occ_set) for idx_i in range(len(occ_list)): i_d = occ_list[idx_i] for idx_j in range(idx_i + 1, len(occ_list)): j_d = occ_list[idx_j] for idx_a in range(len(vir_list)): a_d = vir_list[idx_a] for idx_b in range(idx_a + 1, len(vir_list)): b_d = vir_list[idx_b] new_det = tuple(sorted((occ_set - {i_d, j_d}) | {a_d, b_d})) if new_det in candidates or new_det in det_set: continue H_Iab = double_excitation_matrix_element( occ_I, i_d, j_d, a_d, b_d, h2e ) coupling = abs(H_Iab * coeffs[I]) diag_new = diagonal_matrix_element(new_det, h1e, h2e) denom = e_var_0 - diag_new if abs(denom) < 1e-12: continue weight = coupling * coupling / abs(denom) candidates[new_det] = _SelectionCandidate( det=new_det, weight=weight, coupling=coupling ) # Sort by weight descending, filter by threshold sorted_cands = sorted(candidates.values(), key=lambda c: c.weight, reverse=True) selected = [c for c in sorted_cands if c.weight > opts.pt2_threshold] max_new = max( 10, min( opts.max_det_per_iter, int(len(det_space) * (opts.selection_growth_factor - 1.0)), ), ) return selected[:max_new] def _select_candidates_unrestricted( self, det_space: list[SpinDet], coeffs: np.ndarray, h1e: np.ndarray, h2e: np.ndarray, norb: int, opts: SelectedCIOptions, *, e_var_0: float = 0.0, ) -> list[_SelectionCandidate]: """Spin-unrestricted candidate selection over all five excitation classes.""" candidates: dict[SpinDet, _SelectionCandidate] = {} det_set = set(det_space) # Focus on determinants with significant weight c_abs = np.abs(coeffs) c_threshold = 0.01 significant = np.where(c_abs > c_threshold)[0] if len(significant) == 0: significant = [np.argmax(c_abs)] for I in significant: aI, bI = det_space[I] occ_a = set(aI) occ_b = set(bI) vir_a = sorted(set(range(norb)) - occ_a) vir_b = sorted(set(range(norb)) - occ_b) # ── α single excitations ────────────────────────────────── for i in occ_a: for a in vir_a: new_a = tuple(sorted((occ_a - {i}) | {a})) new_det = (new_a, bI) if new_det in candidates or new_det in det_set: continue coupling = abs( hamiltonian_matrix_element_unrestricted( new_a, bI, aI, bI, h1e, h2e ) * coeffs[I] ) diag_new = diagonal_matrix_element_unrestricted(new_a, bI, h1e, h2e) denom = e_var_0 - diag_new if abs(denom) < 1e-12: continue weight = coupling * coupling / abs(denom) candidates[new_det] = _SelectionCandidate( det=new_det, weight=weight, coupling=coupling ) # ── β single excitations ────────────────────────────────── for i in occ_b: for a in vir_b: new_b = tuple(sorted((occ_b - {i}) | {a})) new_det = (aI, new_b) if new_det in candidates or new_det in det_set: continue coupling = abs( hamiltonian_matrix_element_unrestricted( aI, new_b, aI, bI, h1e, h2e ) * coeffs[I] ) diag_new = diagonal_matrix_element_unrestricted(aI, new_b, h1e, h2e) denom = e_var_0 - diag_new if abs(denom) < 1e-12: continue weight = coupling * coupling / abs(denom) candidates[new_det] = _SelectionCandidate( det=new_det, weight=weight, coupling=coupling ) # ── αα double excitations ───────────────────────────────── a_list = sorted(occ_a) for idx_i in range(len(a_list)): i = a_list[idx_i] for idx_j in range(idx_i + 1, len(a_list)): j = a_list[idx_j] for idx_a in range(len(vir_a)): a = vir_a[idx_a] for idx_b in range(idx_a + 1, len(vir_a)): b = vir_a[idx_b] new_a = tuple(sorted((occ_a - {i, j}) | {a, b})) new_det = (new_a, bI) if new_det in candidates or new_det in det_set: continue coupling = abs( hamiltonian_matrix_element_unrestricted( new_a, bI, aI, bI, h1e, h2e ) * coeffs[I] ) diag_new = diagonal_matrix_element_unrestricted( new_a, bI, h1e, h2e ) denom = e_var_0 - diag_new if abs(denom) < 1e-12: continue weight = coupling * coupling / abs(denom) candidates[new_det] = _SelectionCandidate( det=new_det, weight=weight, coupling=coupling ) # ── ββ double excitations ───────────────────────────────── b_list = sorted(occ_b) for idx_i in range(len(b_list)): i = b_list[idx_i] for idx_j in range(idx_i + 1, len(b_list)): j = b_list[idx_j] for idx_a in range(len(vir_b)): a = vir_b[idx_a] for idx_b in range(idx_a + 1, len(vir_b)): b = vir_b[idx_b] new_b = tuple(sorted((occ_b - {i, j}) | {a, b})) new_det = (aI, new_b) if new_det in candidates or new_det in det_set: continue coupling = abs( hamiltonian_matrix_element_unrestricted( aI, new_b, aI, bI, h1e, h2e ) * coeffs[I] ) diag_new = diagonal_matrix_element_unrestricted( aI, new_b, h1e, h2e ) denom = e_var_0 - diag_new if abs(denom) < 1e-12: continue weight = coupling * coupling / abs(denom) candidates[new_det] = _SelectionCandidate( det=new_det, weight=weight, coupling=coupling ) # ── αβ double excitations ───────────────────────────────── for i in occ_a: for j in occ_b: for a in vir_a: for b in vir_b: new_a = tuple(sorted((occ_a - {i}) | {a})) new_b = tuple(sorted((occ_b - {j}) | {b})) new_det = (new_a, new_b) if new_det in candidates or new_det in det_set: continue coupling = abs( hamiltonian_matrix_element_unrestricted( new_a, new_b, aI, bI, h1e, h2e ) * coeffs[I] ) diag_new = diagonal_matrix_element_unrestricted( new_a, new_b, h1e, h2e ) denom = e_var_0 - diag_new if abs(denom) < 1e-12: continue weight = coupling * coupling / abs(denom) candidates[new_det] = _SelectionCandidate( det=new_det, weight=weight, coupling=coupling ) # Sort by weight descending, filter by threshold sorted_cands = sorted(candidates.values(), key=lambda c: c.weight, reverse=True) selected = [c for c in sorted_cands if c.weight > opts.pt2_threshold] max_new = max( 10, min( opts.max_det_per_iter, int(len(det_space) * (opts.selection_growth_factor - 1.0)), ), ) return selected[:max_new] def _compute_pt2( self, det_space: list, e_var: float, coeffs: np.ndarray, h1e: np.ndarray, h2e: np.ndarray, norb: int, opts: SelectedCIOptions, *, spin_restricted: bool = True, ) -> float: """Epstein-Nesbet PT2 correction from all singles/doubles outside space.""" if spin_restricted: return self._compute_pt2_restricted( det_space, e_var, coeffs, h1e, h2e, norb ) else: return self._compute_pt2_unrestricted( det_space, e_var, coeffs, h1e, h2e, norb ) def _compute_pt2_restricted( self, det_space: list[Det], e_var: float, coeffs: np.ndarray, h1e: np.ndarray, h2e: np.ndarray, norb: int, ) -> float: """PT2 for spin-restricted (closed-shell) determinants.""" e2 = 0.0 det_set = set(det_space) for det_I, c_I in zip(det_space, coeffs): occ_set = set(det_I) vir_set = set(range(norb)) - occ_set vir_list = sorted(vir_set) # Singles for i in occ_set: for a in vir_list: new_det = tuple(sorted((occ_set - {i}) | {a})) if new_det in det_set: continue H_Ia = single_excitation_matrix_element(det_I, i, a, h1e, h2e) val = c_I * H_Ia diag_new = diagonal_matrix_element(new_det, h1e, h2e) denom = e_var - diag_new if abs(denom) > 1e-12: e2 += val * val / denom # Doubles occ_list = sorted(occ_set) for idx_i in range(len(occ_list)): i_d = occ_list[idx_i] for idx_j in range(idx_i + 1, len(occ_list)): j_d = occ_list[idx_j] for idx_a in range(len(vir_list)): a_d = vir_list[idx_a] for idx_b in range(idx_a + 1, len(vir_list)): b_d = vir_list[idx_b] new_det = tuple(sorted((occ_set - {i_d, j_d}) | {a_d, b_d})) if new_det in det_set: continue H_Iab = double_excitation_matrix_element( det_I, i_d, j_d, a_d, b_d, h2e ) val = c_I * H_Iab diag_new = diagonal_matrix_element(new_det, h1e, h2e) denom = e_var - diag_new if abs(denom) > 1e-12: e2 += val * val / denom return e2 def _compute_pt2_unrestricted( self, det_space: list[SpinDet], e_var: float, coeffs: np.ndarray, h1e: np.ndarray, h2e: np.ndarray, norb: int, ) -> float: """PT2 for spin-unrestricted determinants over all five excitation classes.""" e2 = 0.0 det_set = set(det_space) for (aI, bI), c_I in zip(det_space, coeffs): occ_a = set(aI) occ_b = set(bI) vir_a = sorted(set(range(norb)) - occ_a) vir_b = sorted(set(range(norb)) - occ_b) # ── α single excitations ────────────────────────────────── for i in occ_a: for a in vir_a: new_a = tuple(sorted((occ_a - {i}) | {a})) new_det = (new_a, bI) if new_det in det_set: continue H_Ia = hamiltonian_matrix_element_unrestricted( new_a, bI, aI, bI, h1e, h2e ) val = c_I * H_Ia diag_new = diagonal_matrix_element_unrestricted(new_a, bI, h1e, h2e) denom = e_var - diag_new if abs(denom) > 1e-12: e2 += val * val / denom # ── β single excitations ────────────────────────────────── for i in occ_b: for a in vir_b: new_b = tuple(sorted((occ_b - {i}) | {a})) new_det = (aI, new_b) if new_det in det_set: continue H_Ia = hamiltonian_matrix_element_unrestricted( aI, new_b, aI, bI, h1e, h2e ) val = c_I * H_Ia diag_new = diagonal_matrix_element_unrestricted(aI, new_b, h1e, h2e) denom = e_var - diag_new if abs(denom) > 1e-12: e2 += val * val / denom # ── αα double excitations ───────────────────────────────── a_list = sorted(occ_a) for idx_i in range(len(a_list)): i = a_list[idx_i] for idx_j in range(idx_i + 1, len(a_list)): j = a_list[idx_j] for idx_a in range(len(vir_a)): a = vir_a[idx_a] for idx_b in range(idx_a + 1, len(vir_a)): b = vir_a[idx_b] new_a = tuple(sorted((occ_a - {i, j}) | {a, b})) new_det = (new_a, bI) if new_det in det_set: continue H_Iab = hamiltonian_matrix_element_unrestricted( new_a, bI, aI, bI, h1e, h2e ) val = c_I * H_Iab diag_new = diagonal_matrix_element_unrestricted( new_a, bI, h1e, h2e ) denom = e_var - diag_new if abs(denom) > 1e-12: e2 += val * val / denom # ── ββ double excitations ───────────────────────────────── b_list = sorted(occ_b) for idx_i in range(len(b_list)): i = b_list[idx_i] for idx_j in range(idx_i + 1, len(b_list)): j = b_list[idx_j] for idx_a in range(len(vir_b)): a = vir_b[idx_a] for idx_b in range(idx_a + 1, len(vir_b)): b = vir_b[idx_b] new_b = tuple(sorted((occ_b - {i, j}) | {a, b})) new_det = (aI, new_b) if new_det in det_set: continue H_Iab = hamiltonian_matrix_element_unrestricted( aI, new_b, aI, bI, h1e, h2e ) val = c_I * H_Iab diag_new = diagonal_matrix_element_unrestricted( aI, new_b, h1e, h2e ) denom = e_var - diag_new if abs(denom) > 1e-12: e2 += val * val / denom # ── αβ double excitations ───────────────────────────────── for i in occ_a: for j in occ_b: for a in vir_a: for b in vir_b: new_a = tuple(sorted((occ_a - {i}) | {a})) new_b = tuple(sorted((occ_b - {j}) | {b})) new_det = (new_a, new_b) if new_det in det_set: continue H_Iab = hamiltonian_matrix_element_unrestricted( new_a, new_b, aI, bI, h1e, h2e ) val = c_I * H_Iab diag_new = diagonal_matrix_element_unrestricted( new_a, new_b, h1e, h2e ) denom = e_var - diag_new if abs(denom) > 1e-12: e2 += val * val / denom return e2 def _davidson( self, H_mat: np.ndarray, det_space: list, h1e: np.ndarray, h2e: np.ndarray, opts: SelectedCIOptions, ) -> tuple[np.ndarray, np.ndarray]: """Davidson iterative diagonalization using the pre-built H matrix for sigma-vector products: σ = H @ v. Returns (eigenvalues[0:1], eigenvectors[:, 0:1]) mimicking scipy_eigh.""" ndet = len(det_space) n_guess = min(4, ndet) max_iter = opts.davidson_max_iter tol = opts.davidson_conv_tol # Initial subspace: diagonal-preconditioned guess diag = np.diag(H_mat) idx = np.argsort(diag)[:n_guess] V = np.zeros((ndet, n_guess)) for k, i in enumerate(idx): V[i, k] = 1.0 # Orthonormalize initial guess V, _ = np.linalg.qr(V) for _ in range(max_iter): # Subspace Hamiltonian HV = H_mat @ V H_sub = V.T @ HV evals_sub, evecs_sub = np.linalg.eigh(H_sub) e_sub = evals_sub[0] c_sub = evecs_sub[:, 0] # Ritz vector and residual v = V @ c_sub sigma = HV @ c_sub r = sigma - e_sub * v r_norm = np.linalg.norm(r) if r_norm < tol: break # Precondition: (diag - e_sub)^{-1} denom = diag - e_sub denom[np.abs(denom) < 1e-12] = 1e-12 t = r / denom # Orthogonalize against current subspace t = t - V @ (V.T @ t) t_norm = np.linalg.norm(t) if t_norm < 1e-14: break t = t / t_norm # Expand subspace V = np.column_stack([V, t]) # Re-orthonormalize V, _ = np.linalg.qr(V) # Return in scipy_eigh-compatible format eigenvalues = np.array([e_sub]) eigenvectors = v.reshape(ndet, 1) return eigenvalues, eigenvectors
[docs] def solve_selected_ci( hamiltonian: Hamiltonian, options: Optional[SelectedCIOptions] = None, ) -> SolverResult: """One-shot Selected-CI solve. Parameters ---------- hamiltonian : Hamiltonian One- and two-electron integrals in an orthonormal spatial-orbital basis. options : SelectedCIOptions, optional Returns ------- SolverResult """ solver = SelectedCISolver(options) return solver.solve(hamiltonian)