Source code for vibeqc.solvers._common

"""Common interfaces and data structures for non-mean-field solvers.

All solvers in this package share a common protocol:
    result = solver.solve(hamiltonian, nelec, norb, options) -> SolverResult

A Hamiltonian carries one- and two-electron integrals in an orthonormal
(spatial-orbital) basis.  The orbital provider is responsible for
delivering the transformation from AO → MO; the solvers themselves are
orbital-basis-agnostic.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Optional, Protocol

import numpy as np

# ── Hamiltonian ───────────────────────────────────────────────────────────




@dataclass
class Hamiltonian:
    """One- and two-electron integrals in an orthonormal spatial-orbital basis.

    Attributes
    ----------
    h1e : (norb, norb) ndarray
        One-electron (core) Hamiltonian h_{pq}.
    h2e : (norb, norb, norb, norb) ndarray
        Two-electron integrals in **physicist's** notation:
        g_{pqrs} = (pr|qs) = ∫∫ φ_p*(r₁) φ_q*(r₂) r₁₂⁻¹ φ_r(r₁) φ_s(r₂).
        Stored as a 4-index tensor with anti-symmetry NOT folded in —
        the solvers apply Slater–Condon rules that need the full tensor.
    nuclear_repulsion : float
        Nuclear-nuclear repulsion energy E_nuc (Hartree).
    norb : int
        Number of spatial orbitals.
    nelec : int
        Number of electrons.
    ms2 : int
        2 × S_z = n_alpha − n_beta.  0 for closed-shell singlets.
    description : str
        Human-readable label for logging / diagnostics.
    """

    h1e: np.ndarray
    h2e: np.ndarray
    nuclear_repulsion: float = 0.0
    norb: int = 0
    nelec: int = 0
    ms2: int = 0
    description: str = ""

    def __post_init__(self) -> None:
        if self.norb == 0 and self.h1e is not None:
            self.norb = self.h1e.shape[0]



# ── Solver options ────────────────────────────────────────────────────────


@dataclass
class SolverOptions:
    """Base options shared by all non-mean-field solvers.

    Sub-classes add method-specific controls.
    """

    #: Target number of determinants / bond-dimension / constraint set size.
    #: Interpretation is method-specific.
    target_size: int = 100

    #: Energy convergence threshold (Hartree).
    conv_tol_energy: float = 1e-6

    #: Maximum number of macro-iterations.
    max_iter: int = 50

    #: Verbosity: 0 = silent, 1 = per-iteration, 2 = debug.
    verbose: int = 0

    #: Random seed for reproducibility.
    random_seed: int = 42


# ── Solver result ─────────────────────────────────────────────────────────




@dataclass
class SolverResult:
    """Common result container for all non-mean-field solvers.

    Fields marked ``method-specific`` are populated by individual backends
    and may be ``None``.
    """

    #: Total energy (Hartree).  Always includes nuclear repulsion.
    energy: float

    #: Method name, e.g. ``"selected_ci"``, ``"dmrg"``, ``"v2rdm"``.
    method: str

    #: Whether the solver considers itself converged.
    converged: bool = True

    #: Number of macro-iterations / sweeps taken.
    n_iter: int = 0

    #: Energy trace per macro-iteration (if tracked).
    energy_trace: list[float] = field(default_factory=list)

    # ── method-specific ──

    #: Wavefunction coefficients (determinant × configuration), CI only.
    ci_coeffs: Optional[np.ndarray] = None

    #: Determinant / configuration labels, CI only.
    ci_labels: Optional[list[tuple[int, ...]]] = None

    #: 1-RDM (norb, norb), if computed by the solver.
    rdm1: Optional[np.ndarray] = None

    #: 2-RDM (norb, norb, norb, norb), if computed by the solver.
    rdm2: Optional[np.ndarray] = None

    #: Truncation error / discarded weight, DMRG.
    truncation_error: Optional[float] = None

    #: Bond dimension, DMRG.
    bond_dim: Optional[int] = None

    #: N-representability residual norm, v2RDM.
    constraint_residual: Optional[float] = None

    #: Perturbative correction, Selected-CI.
    pt2_correction: Optional[float] = None

    #: Transcorrelated diagnostics dict.
    tc_diagnostics: Optional[dict] = None

    @property
    def energy_total(self) -> float:
        """Alias for ``energy``."""
        return self.energy



# ── Solver protocol ───────────────────────────────────────────────────────


class SolverProtocol(Protocol):
    """Structural protocol for a non-mean-field solver."""

    def solve(
        self,
        hamiltonian: Hamiltonian,
        options: SolverOptions | None = None,
    ) -> SolverResult: ...


# ── Utility ───────────────────────────────────────────────────────────────


def _physicist_to_chemist(h2e: np.ndarray) -> np.ndarray:
    """Convert physicist's ERI g_{pqrs} = (pr|qs) → chemist's (pq|rs).

    (pq|rs) = g_{prqs}
    """
    return h2e.transpose(0, 2, 1, 3)


def _chemist_to_physicist(eri_chem: np.ndarray) -> np.ndarray:
    """Convert chemist's (pq|rs) → physicist's g_{pqrs} = (pr|qs)."""
    return eri_chem.transpose(0, 2, 1, 3)


def _antisymmetrize_h2e(h2e_phys: np.ndarray) -> np.ndarray:
    """Build the antisymmetrized two-electron tensor <pq||rs>.

    <pq||rs> = g_{pqrs} - g_{pqsr}
    """
    return h2e_phys - h2e_phys.transpose(0, 1, 3, 2)