Source code for vibeqc.solvers._hamiltonian

"""Hamiltonian construction from AO integrals.

Provides a clean interface for building :class:`Hamiltonian` objects
from vibe-qc's native C++ integral bindings, with optional orbital
transformation.
"""

from __future__ import annotations

from typing import Optional

import numpy as np

from .._vibeqc_core import (
    BasisSet,
    Molecule,
    compute_eri,
    compute_kinetic,
    compute_nuclear,
    compute_overlap,
)
from ._common import Hamiltonian, _chemist_to_physicist


def build_hamiltonian_ao(
    molecule: Molecule,
    basis: BasisSet,
    *,
    description: str = "",
) -> Hamiltonian:
    """Build a :class:`Hamiltonian` in the **AO basis**.

    The returned ``h1e`` and ``h2e`` are in the original (non-orthogonal)
    atomic-orbital basis.  Most solvers require an orthonormal basis;
    transform with :func:`transform_hamiltonian_to_mo` after orthogonalising.

    Parameters
    ----------
    molecule : Molecule
        Molecular geometry, charge, multiplicity.
    basis : BasisSet
        AO basis set.
    description : str
        Label for logging.

    Returns
    -------
    Hamiltonian
        With ``h1e``, ``h2e``, ``nuclear_repulsion``, ``nelec``, ``ms2``,
        and ``norb`` populated.
    """
    S = np.asarray(compute_overlap(basis), order="C")
    T = np.asarray(compute_kinetic(basis), order="C")
    V = np.asarray(compute_nuclear(basis, molecule), order="C")
    h1e_ao = T + V  # (norb, norb)

    # Chemist's notation (μν|λσ) from libint → physicist's g_{μνλσ}
    eri_chem_ao = np.asarray(compute_eri(basis), order="C")
    h2e_ao = _chemist_to_physicist(eri_chem_ao)  # g_{μνλσ} = (μλ|νσ)

    norb = h1e_ao.shape[0]
    nelec = molecule.n_electrons()
    multiplicity = molecule.multiplicity
    ms2 = multiplicity - 1  # 2*S_z

    return Hamiltonian(
        h1e=h1e_ao,
        h2e=h2e_ao,
        nuclear_repulsion=molecule.nuclear_repulsion(),
        norb=norb,
        nelec=nelec,
        ms2=ms2,
        description=description or f"{molecule.n_electrons()}e, {norb} orbitals",
    )




def build_hamiltonian_mo(
    molecule: Molecule,
    basis: BasisSet,
    mo_coeffs: np.ndarray,
    *,
    description: str = "",
) -> Hamiltonian:
    """Build a :class:`Hamiltonian` in the **MO basis**.

    Transforms the one- and two-electron integrals from the AO basis
    into the molecular-orbital basis given by ``mo_coeffs``.

    Parameters
    ----------
    molecule : Molecule
    basis : BasisSet
    mo_coeffs : (n_ao, n_mo) ndarray
        AO → MO coefficient matrix (columns are MOs).
    description : str

    Returns
    -------
    Hamiltonian
        With integrals in the MO basis. ``norb`` = n_mo.
    """
    ham_ao = build_hamiltonian_ao(molecule, basis, description=description)
    return transform_hamiltonian(ham_ao, mo_coeffs)





def transform_hamiltonian(
    ham: Hamiltonian,
    C: np.ndarray,
) -> Hamiltonian:
    """Transform a Hamiltonian into a new orthonormal basis given by ``C``.

    New one-electron matrix:  h̃_{pq} = Σ_{μν} C_{μp} h_{μν} C_{νq}
    New two-electron tensor:  g̃_{pqrs} = Σ_{μνλσ} C_{μp} C_{νq} g_{μνλσ} C_{λr} C_{σs}

    Parameters
    ----------
    ham : Hamiltonian
        Source Hamiltonian.
    C : (n_ao, n_mo) ndarray
        Transformation matrix.  Columns are the new basis vectors in the
        old (AO) basis.
    """
    h1e_new = np.einsum("up,uv,vq->pq", C, ham.h1e, C, optimize=True)
    n_mo = C.shape[1]
    h2e_new = np.einsum("ap,bq,cr,ds,abcd->pqrs", C, C, C, C, ham.h2e, optimize=True)

    return Hamiltonian(
        h1e=h1e_new,
        h2e=h2e_new,
        nuclear_repulsion=ham.nuclear_repulsion,
        norb=n_mo,
        nelec=ham.nelec,
        ms2=ham.ms2,
        description=ham.description,
    )





def get_hf_orbital_provider(
    molecule: Molecule,
    basis: BasisSet,
    *,
    method: str = "rhf",
) -> np.ndarray:
    """Run vibe-qc's RHF and return the MO coefficient matrix.

    This is a **convenience** function — the solver layer does not
    require HF orbitals.  Any source of orthonormal orbitals works.

    Parameters
    ----------
    molecule : Molecule
    basis : BasisSet
    method : str
        ``"rhf"`` (default) or ``"uhf"``.

    Returns
    -------
    mo_coeffs : (n_ao, n_ao) ndarray
        Columns are canonical HF molecular orbitals.
    """
    from .._vibeqc_core import RHFOptions, UHFOptions, run_rhf, run_uhf

    if method == "rhf":
        opts = RHFOptions()
        opts.max_iter = 200
        opts.conv_tol_energy = 1e-12
        opts.conv_tol_grad = 1e-10
        result = run_rhf(molecule, basis, opts)
        return np.asarray(result.mo_coeffs, order="C")
    elif method == "uhf":
        opts = UHFOptions()
        opts.max_iter = 200
        opts.conv_tol_energy = 1e-12
        opts.conv_tol_grad = 1e-10
        result = run_uhf(molecule, basis, opts)
        return np.asarray(result.mo_coeffs_alpha, order="C")
    else:
        raise ValueError(f"Unknown HF method: {method!r}")



def canonical_orthogonalize(
    S: np.ndarray,
    threshold: float = 1e-8,
) -> np.ndarray:
    """Return the canonical (Löwdin-symmetric) orthogonalisation matrix X.

    X = U s^{-1/2} U^T  where S = U s U^T is the overlap eigen-decomposition.
    This produces orthonormal orbitals closest to the original AO basis.

    Parameters
    ----------
    S : (n, n) ndarray
        Overlap matrix.
    threshold : float
        Eigenvalue threshold for linear-dependence removal.

    Returns
    -------
    X : (n, n_active) ndarray
        Orthogonalisation matrix: C_orth = X^T · C_ao.
    """
    evals, evecs = np.linalg.eigh(S)
    mask = evals > threshold
    s_inv_sqrt = np.diag(1.0 / np.sqrt(evals[mask]))
    return evecs[:, mask] @ s_inv_sqrt