feat: E_PT2 energy correction for HI-NQS (Issue #35 Tier 1 #2)

CHANGELOG.md

@@ -21,6 +21,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - 14 new CAS molecules in registry (26 total): N₂-CAS(10,12/15/17/20/26), Cr₂ + variants up to 72Q, Benzene CAS(6,15)
 - IBM `solve_fermion` auto-enabled when `qiskit_addon_sqd` is installed (α×β Cartesian product, dramatically better accuracy)
 - `_train_nqs_teacher` raises `ValueError` when `energy_weight > 0` without `hamiltonian`
+- `compute_e_pt2`: EN-PT2 energy correction (`E_PT2 = Σ |⟨x|H|Ψ₀⟩|² / (E₀ - H_xx)`)
+- `HINQSSQDConfig.compute_pt2_correction`: opt-in E_PT2 correction after final iteration
+- `use_ibm_solver` changed to tri-state (`None`=auto, `True`=force, `False`=disable)
 - `_compute_cas_integrals` helper with auto-CASCI fallback for large active spaces
 - `MolecularHamiltonian.build_sparse_hamiltonian()` for O(nnz) sparse H construction
 - Sparse eigenvalue dispatch in `gpu_solve_fermion` for basis > 8K configs

docs/api_reference.md

@@ -774,6 +774,10 @@ Score candidate configs by Epstein-Nesbet PT2 importance: `score(x) = |⟨x|H|Φ
 
 Keep only the highest-|c_i|² configs (ASCI pattern). Returns trimmed basis and coefficients.
 
+#### `compute_e_pt2(basis, coeffs, hamiltonian, e0) -> float`
+
+Compute EN-PT2 energy correction: `E_PT2 = Σ_{x∉V} |⟨x|H|Ψ₀⟩|² / (E₀ - H_xx)`. Requires full variational basis and matching eigenvector. Returns typically negative correction.
+
 #### `compute_temperature(iteration, max_iterations, t_init, t_final) -> float`
 
 Linear temperature annealing from `t_init` to `t_final` over iterations.

src/qvartools/methods/nqs/_pt2_helpers.py

@@ -3,13 +3,16 @@
 Standalone pure functions for:
 
 - Epstein-Nesbet PT2 scoring
+- EN-PT2 energy correction (E_PT2)
 - Coefficient-based basis eviction (ASCI pattern)
 - Linear temperature annealing
 
 Functions
 ---------
 compute_pt2_scores
     Score candidates by EN-PT2 importance.
+compute_e_pt2
+    Compute EN-PT2 energy correction.
 evict_by_coefficient
     Keep highest-|c_i|² configs.
 compute_temperature
@@ -114,6 +117,110 @@ def compute_pt2_scores(
     return scores
 
 
+def compute_e_pt2(
+    basis: torch.Tensor,
+    coeffs: np.ndarray,
+    hamiltonian: Any,
+    e0: float,
+) -> float:
+    r"""Compute Epstein-Nesbet second-order perturbation energy correction.
+
+    Sums over all determinants connected to the basis but NOT in the basis::
+
+        E_PT2 = Σ_{x ∉ V} |⟨x|H|Ψ₀⟩|² / (E₀ - H_xx)
+
+    where ``Ψ₀ = Σ_i c_i |x_i⟩`` and the sum runs over external
+    determinants reachable via single and double excitations.  For
+    real-valued Hamiltonians, ``|⟨x|H|Ψ₀⟩|² = ⟨x|H|Ψ₀⟩²``.
+
+    Parameters
+    ----------
+    basis : torch.Tensor
+        **Full** variational basis, shape ``(n_basis, n_qubits)``.
+        Must be the complete basis used for the diagonalisation that
+        produced ``coeffs`` and ``e0``.
+    coeffs : np.ndarray
+        Ground-state eigenvector from diagonalising H in ``basis``,
+        shape ``(n_basis,)``.
+    hamiltonian
+        Hamiltonian with ``get_connections``, ``diagonal_element``, and
+        ``diagonal_elements_batch``.
+    e0 : float
+        Variational ground-state energy from the same diagonalisation
+        as ``coeffs``.
+
+    Returns
+    -------
+    float
+        E_PT2 correction (typically negative).
+    """
+    if coeffs.shape[0] != basis.shape[0]:
+        raise ValueError(
+            f"coeffs length ({coeffs.shape[0]}) must match basis size "
+            f"({basis.shape[0]})"
+        )
+
+    basis_hash_list = config_integer_hash(basis)
+    basis_hash_set: set = set(basis_hash_list)
+
+    # Accumulate coupling ⟨x|H|Ψ₀⟩ for each external determinant x.
+    # Store config tensors temporarily for batched diagonal computation.
+    external_coupling: dict = {}  # hash -> coupling
+    external_config: dict = {}  # hash -> config tensor
+
+    for idx in range(basis.shape[0]):
+        c_i = float(coeffs[idx])
+        if abs(c_i) < 1e-14:
+            continue
+
+        connections, h_elements = hamiltonian.get_connections(basis[idx])
+        if connections is None or len(connections) == 0:
+            continue
+
+        conn_hashes = config_integer_hash(connections)
+        for j in range(len(connections)):
+            h_conn = conn_hashes[j]
+            if h_conn in basis_hash_set:
+                continue
+            contrib = float(h_elements[j]) * c_i
+            if h_conn in external_coupling:
+                external_coupling[h_conn] += contrib
+            else:
+                external_coupling[h_conn] = contrib
+                external_config[h_conn] = connections[j].detach().cpu().clone()
+
+    if not external_coupling:
+        return 0.0
+
+    # Compute diagonal elements in chunks to bound peak memory
+    ext_hashes = list(external_coupling.keys())
+    diag_batch_size = 4096
+
+    e_pt2 = 0.0
+    for start in range(0, len(ext_hashes), diag_batch_size):
+        chunk_hashes = ext_hashes[start : start + diag_batch_size]
+        ext_configs = torch.stack([external_config[h] for h in chunk_hashes])
+        h_diag = hamiltonian.diagonal_elements_batch(ext_configs)
+        h_diag_np = h_diag.detach().cpu().numpy().astype(np.float64)
+
+        for k, h_ext in enumerate(chunk_hashes):
+            h_xx = float(h_diag_np[k])
+            if not math.isfinite(h_xx):
+                continue
+            coupling = external_coupling[h_ext]
+            denom = e0 - h_xx
+            if abs(denom) < 1e-14:
+                sign = -1.0 if denom <= 0 else 1.0
+                denom = sign * 1e-14
+                logger.debug("Near-zero PT2 denominator clamped for hash %s", h_ext)
+            e_pt2 += coupling**2 / denom
+
+        del ext_configs, h_diag, h_diag_np
+
+    del external_config
+    return e_pt2
+
+
 def evict_by_coefficient(
     basis: torch.Tensor,
     coeffs: np.ndarray,

src/qvartools/methods/nqs/hi_nqs_sqd.py

@@ -35,6 +35,7 @@
 )
 from qvartools._utils.gpu.diagnostics import gpu_solve_fermion
 from qvartools.methods.nqs._pt2_helpers import (
+    compute_e_pt2,
     compute_pt2_scores,
     compute_temperature,
     evict_by_coefficient,
@@ -131,6 +132,12 @@ class HINQSSQDConfig:
         Weight for the REINFORCE energy loss term (default ``0.0``).
     entropy_weight : float
         Weight for the entropy regularisation term (default ``0.0``).
+    compute_pt2_correction : bool
+        Compute EN-PT2 energy correction after final iteration
+        (default ``False``).  When ``True``, ``metadata`` includes
+        ``e_pt2`` and ``corrected_energy = pt2_e0 + e_pt2``, where
+        ``pt2_e0`` is the full-basis diagonalisation energy used for the
+        PT2 correction (distinct from the main returned ``energy`` value).
     """
 
     n_iterations: int = 10
@@ -155,6 +162,7 @@ class HINQSSQDConfig:
     teacher_weight: float = 1.0
     energy_weight: float = 0.0
     entropy_weight: float = 0.0
+    compute_pt2_correction: bool = False
 
 
 # ---------------------------------------------------------------------------
@@ -641,6 +649,46 @@ def run_hi_nqs_sqd(
             best_energy,
         )
 
+    # --- Optional E_PT2 correction ---
+    metadata: dict[str, Any] = {
+        "energy_history": energy_history,
+        "n_iterations": len(energy_history),
+        "final_basis_size": int(cumulative_basis.shape[0]),
+    }
+
+    if cfg.compute_pt2_correction and cumulative_basis.shape[0] > 0:
+        t_pt2 = time.perf_counter()
+        # Full-basis diag to get consistent (E₀, Ψ₀) for E_PT2
+        n_full = cumulative_basis.shape[0]
+        if n_full <= 8_000:  # Align with CIPSI _SPARSE_DIAG_THRESHOLD
+            h_full = hamiltonian.matrix_elements_fast(cumulative_basis)
+            h_np = h_full.detach().cpu().numpy().astype(np.float64)
+            h_np = 0.5 * (h_np + h_np.T)
+            evals, evecs = np.linalg.eigh(h_np)
+            pt2_e0 = float(evals[0])
+            pt2_coeffs = evecs[:, 0]
+        else:
+            from scipy.sparse.linalg import eigsh as sp_eigsh
+
+            h_sp = hamiltonian.build_sparse_hamiltonian(cumulative_basis)
+            evals, evecs = sp_eigsh(h_sp.tocsr(), k=1, which="SA")
+            pt2_e0 = float(evals[0])
+            pt2_coeffs = evecs[:, 0]
+
+        e_pt2 = compute_e_pt2(cumulative_basis, pt2_coeffs, hamiltonian, pt2_e0)
+        pt2_wall = time.perf_counter() - t_pt2
+        metadata["e_pt2"] = e_pt2
+        metadata["pt2_e0"] = pt2_e0
+        metadata["corrected_energy"] = pt2_e0 + e_pt2
+        metadata["pt2_wall_time"] = pt2_wall
+        logger.info(
+            "  E_var=%.8f, E_PT2=%.6f Ha, corrected=%.8f Ha (%.1fs)",
+            pt2_e0,
+            e_pt2,
+            pt2_e0 + e_pt2,
+            pt2_wall,
+        )
+
     wall_time = time.perf_counter() - t_start
 
     return SolverResult(
@@ -649,9 +697,5 @@ def run_hi_nqs_sqd(
         wall_time=wall_time,
         method="HI+NQS+SQD",
         converged=converged,
-        metadata={
-            "energy_history": energy_history,
-            "n_iterations": len(energy_history),
-            "final_basis_size": int(cumulative_basis.shape[0]),
-        },
+        metadata=metadata,
     )

tests/test_methods/test_initial_basis.py

@@ -97,6 +97,7 @@ def _make_initial_basis(n_configs=3):
 class TestRunHiNqsSqdInitialBasis:
     """Test initial_basis kwarg for run_hi_nqs_sqd."""
 
+    @patch("qvartools.methods.nqs.hi_nqs_sqd._IBM_SQD_AVAILABLE", False)
     @patch("qvartools.methods.nqs.hi_nqs_sqd.gpu_solve_fermion")
     def test_accepts_initial_basis_none(
         self, mock_solver, mol_info, minimal_config_sqd
@@ -116,6 +117,7 @@ def test_accepts_initial_basis_none(
         assert isinstance(result, SolverResult)
         assert result.method == "HI+NQS+SQD"
 
+    @patch("qvartools.methods.nqs.hi_nqs_sqd._IBM_SQD_AVAILABLE", False)
     @patch("qvartools.methods.nqs.hi_nqs_sqd.gpu_solve_fermion")
     def test_accepts_initial_basis_tensor(
         self, mock_solver, mol_info, minimal_config_sqd
@@ -137,6 +139,7 @@ def test_accepts_initial_basis_tensor(
         # With warm-start, final basis should be non-empty
         assert result.metadata["final_basis_size"] > 0
 
+    @patch("qvartools.methods.nqs.hi_nqs_sqd._IBM_SQD_AVAILABLE", False)
     @patch("qvartools.methods.nqs.hi_nqs_sqd.gpu_solve_fermion")
     def test_initial_basis_deduplicates(
         self, mock_solver, mol_info, minimal_config_sqd