impurity_opt_tools.py

from typing import Union
from numpy import ndarray as NDarray
from scipy.sparse import lil_matrix, csc_matrix
import numpy as np
from helper_functions import (
    get_overlap_matrix,
    get_reduced_hamiltonian,
    optimize_subspace,
    get_covariance,
    get_overlap_matrix_covariance,
    get_reduced_hamiltonian_covariance,
    get_ps_basis,
    slater_det,
    load_bin_basis,
)
from scipy.sparse.linalg import eigsh
from typing import Literal
from copy import deepcopy
from math import comb
import time


class ParameterMismatchError(Exception):
    """Custom exception for invalid input parameters to ImpurityHamiltonian class."""

    ...


class ImpurityHamiltonian:

    def __init__(self, parameters: dict):

        self._validate_input(parameters)

        self.U = parameters["U"]
        self.V_bath = parameters["V_bath"]
        self.e_bath = parameters["e_bath"]
        self.N_bath = parameters["N_bath"]
        self.N_imp = parameters["N_imp"]
        self.e_imp = parameters["e_imp"]
        self.mu = parameters["mu"]
        self.V_imp = parameters.get("V_imp")
        self.weight = parameters.get("weight")

        self.N = (self.N_bath + 1) * self.N_imp
        if self.N < 5:
            self.sparse_flag = False
        else:
            self.sparse_flag = True

    def _validate_input(self, parameters: dict):
        if not isinstance(parameters["U"], int):
            if not isinstance(parameters["U"], float):
                raise ParameterMismatchError("U must be an integer or float.")

        if not isinstance(parameters["V_bath"], Union[list, NDarray]):
            raise ParameterMismatchError("V_bath must be a list or an array.")

        if not isinstance(parameters["e_bath"], Union[list, NDarray]):
            raise ParameterMismatchError("e_bath must be a list or an array.")

        if parameters.get("weight") is not None:
            if not isinstance(parameters["weight"], float):
                if not isinstance(parameters["weight"], int):
                    raise ParameterMismatchError("weight must be an integer or float.")

        if not isinstance(parameters["N_bath"], int):
            raise ParameterMismatchError("N_bath must be an integer.")

        if not isinstance(parameters["N_imp"], int):
            raise ParameterMismatchError("N_imp must be an integer.")

        if not isinstance(parameters["mu"], float):
            if not isinstance(parameters["mu"], int):
                raise ParameterMismatchError("mu must be an integer or float.")

        if len(parameters["V_bath"]) != parameters["N_bath"] * parameters["N_imp"]:
            raise ParameterMismatchError(
                "V_bath must have the same length as N_bath * N_imp."
            )

        if len(parameters["e_bath"]) != parameters["N_bath"] * parameters["N_imp"]:
            raise ParameterMismatchError(
                "e_bath must have the same length as N_bath * N_imp."
            )

        if parameters["N_imp"] > 1:
            if parameters.get("V_imp") is not None:
                if not isinstance(parameters["V_imp"], Union[list, NDarray]):
                    raise ParameterMismatchError("V_imp must be a list or an array.")
                if len(parameters["V_imp"]) != comb(parameters["N_imp"], 2):
                    raise ParameterMismatchError(
                        "V_imp must have a value for each combination of two impurities."
                    )
                if not isinstance(parameters["e_imp"], Union[list, NDarray]):
                    raise ParameterMismatchError(
                        "e_imp must be a list or an array when there is more than one impurity site."
                    )
                if len(parameters["e_imp"]) != parameters["N_imp"]:
                    raise ParameterMismatchError(
                        "e_imp must have the same length as N_imp."
                    )
            else:
                if parameters["N_imp"] > 1:
                    raise ParameterMismatchError("V_imp must be provided if N_imp > 1.")
                # else:
                #     if not isinstance(parameters["e_imp"], float):
                #         if not isinstance(parameters["e_imp"], int):
                #             raise ParameterMismatchError(
                #                 "e_imp must be an integer or float."
                #             )

        if parameters.get("return format") is not None:
            if parameters["return format"] not in ["of", "pauli", "sparse", "dense"]:
                raise ParameterMismatchError(
                    "return format must be either 'of' (for openfermion operator), 'pauli' for coefficients and Pauli strings, and 'sparse' or 'dense' for a matrix."
                )

    def get_parameters(self):
        return {
            "U": self.U,
            "V_bath": self.V_bath,
            "e_bath": self.e_bath,
            "weight": self.weight,
            "N": self.N,
            "N_bath": self.N_bath,
            "N_imp": self.N_imp,
            "e_imp": self.e_imp,
            "V_imp": self.V_imp,
            "mu": self.mu,
        }

    def set_parameters(self, parameters: dict):
        self._validate_input(parameters)

        self.U = parameters["U"]
        self.V_bath = parameters["V_bath"]
        self.e_bath = parameters["e_bath"]
        self.N_bath = parameters["N_bath"]
        self.N_imp = parameters["N_imp"]
        self.e_imp = parameters["e_imp"]
        self.mu = parameters["mu"]
        self.V_imp = parameters.get("V_imp")
        self.weight = parameters.get("weight")

        self.N = (self.N_bath + 1) * self.N_imp
        if self.N < 5:
            self.sparse_flag = False
        else:
            self.sparse_flag = True

    def build_model(self, model_type: str = "exact", **kwargs):
        if model_type == "exact":
            return self._build_exact_model(**kwargs)
        elif model_type == "particle selected":
            return self._build_particle_selected_model(**kwargs)
        elif model_type == "single orbital":
            if self.weight is None:
                raise ParameterMismatchError(
                    "Weight must be provided for single orbital model."
                )
            return self._build_single_orbital_model()
        else:
            raise ValueError(
                "Invalid model_type. Must be 'exact', 'particle selected', or 'single orbital'."
            )

    def _build_exact_model(self, **kwargs):  # -> (
        #     NDarray
        #     | csc_matrix
        #     | tuple[list | NDarray, list | NDarray]
        #     # | openfermion.ops.operators.qubit_operator.QubitOperator
        # ):

        verbose = kwargs.get("verbose", False)
        params = self.get_parameters()
        N_ib = params["N_bath"] + 1
        N_imp = params["N_imp"]
        N_bath = params["N_bath"]
        N = N_ib * N_imp
        U = params["U"]
        mu = params["mu"]
        e_imp = params["e_imp"]
        e_bath = params["e_bath"]
        V_imp = params["V_imp"]
        V_bath = params["V_bath"]

        bitstrings = load_bin_basis(2 * N)
        bitarrays = np.array([[int(bit) for bit in bits] for bits in bitstrings])
        n_basis = len(bitstrings)

        if self.sparse_flag:
            H = lil_matrix((n_basis, n_basis), dtype=complex)
        else:
            H = np.zeros((n_basis, n_basis), dtype=complex)
        times = []
        for i, el1 in enumerate(bitarrays):
            start = time.time()
            if verbose:
                if i != 0:
                    if i % 1000 == 0:
                        time_avg = np.average(times)
                        times = []
                    if i < 1000:
                        time_avg = np.average(times)
                    print(
                        f"Processing basis state {i+1}/{n_basis}... Estimated time to completion: {(time_avg * n_basis * (1 - ((i+1)/n_basis)))/60:.2f} minutes",
                        end="\r",
                    )
            for j, el2 in enumerate(bitarrays):
                if np.array_equal(el1, el2):
                    # Diagonal terms
                    for k in range(N_imp):
                        imp_i = N_ib * k
                        up = el1[imp_i]
                        down = el1[imp_i + N]
                        if up:
                            H[i, j] += e_imp[k] - mu
                        if down:
                            H[i, j] += e_imp[k] - mu
                        if up and down:
                            H[i, j] += U

                        # Inter-impurity interactions
                        for l in range(k + 1, N_imp):
                            imp_j = N_ib * l
                            up_j = el1[imp_j]
                            down_j = el1[imp_j + N]
                            H[i, j] += U * (
                                up * up_j + up * down_j + down * up_j + down * down_j
                            )

                        # Bath terms
                        for b in range(N_bath):
                            bath_i = imp_i + b + 1
                            bath_idx = k * N_bath + b
                            if el1[bath_i]:
                                H[i, j] += e_bath[bath_idx]
                            if el1[bath_i + N]:
                                H[i, j] += e_bath[bath_idx]
                else:
                    # Off-diagonal (hopping) terms
                    diff = el1 != el2
                    hop_locations = np.where(diff)[0]
                    if len(hop_locations) == 2:
                        a, b = hop_locations
                        # Hopping between impurities
                        if a % N_ib == 0 and b % N_ib == 0 and a // N == b // N:
                            k = (a // N_ib) % N_imp
                            l = (b // N_ib) % N_imp
                            V_imp_i = int(l - (k + 1) + k / 2 * (2 * N_imp - 1 - k))
                            # Fermionic sign
                            count = np.sum(el1[a + 1 : b] | el2[a + 1 : b])
                            H[i, j] += (-1) ** count * V_imp[V_imp_i]
                        # Hopping between impurity and bath
                        elif a % N_ib == 0 and b // N_ib == a // N_ib:
                            k = (a // N_ib) % N_imp
                            bath_b = b % N_ib - 1
                            bath_idx = k * N_bath + bath_b
                            count = np.sum(el1[a + 1 : b] | el2[a + 1 : b])
                            H[i, j] += (-1) ** count * V_bath[bath_idx]
            end = time.time()
            times.append(end - start)

        return H

    def _build_particle_selected_model(self, **kwargs) -> NDarray | csc_matrix:
        if kwargs is not None:
            if kwargs.get("particle_selection") is not None:
                particle_selection = kwargs["particle_selection"]
            else:
                print("Particle selection not provided, defaulting to half-filling.")
                particle_selection = (self.N // 2, self.N // 2)
        else:
            print("Particle selection not provided, defaulting to half-filling.")
            particle_selection = (self.N // 2, self.N // 2)

        verbose = kwargs.get("verbose", False)
        spin_protected = kwargs.get("spin_protected", True)
        if spin_protected:
            particle_selected_basis = get_ps_basis(particle_selection, 2 * self.N)
        else:
            particle_selected_basis = get_ps_basis(sum(particle_selection), 2 * self.N)
        params = self.get_parameters()
        N_ib = params["N_bath"] + 1
        N_imp = params["N_imp"]
        N_bath = params["N_bath"]
        N = N_ib * N_imp
        n_basis = len(particle_selected_basis)
        U = params["U"]
        mu = params["mu"]
        e_imp = params["e_imp"]
        e_bath = params["e_bath"]
        V_imp = params["V_imp"]
        V_bath = params["V_bath"]

        if self.sparse_flag:
            H = lil_matrix((n_basis, n_basis), dtype=complex)
        else:
            H = np.zeros((n_basis, n_basis), dtype=complex)

        # Precompute bitstrings
        if verbose:
            print("Precomputing bitstrings...")
        bitstrings = [bin(b)[2:].zfill(2 * N) for b in particle_selected_basis]
        bitarrays = np.array([[int(bit) for bit in bits] for bits in bitstrings])

        times = []
        for i, el1 in enumerate(bitarrays):
            start = time.time()
            if verbose:
                if i != 0:
                    if i % 1000 == 0:
                        time_avg = np.average(times)
                        times = []
                    if i < 1000:
                        time_avg = np.average(times)
                    print(
                        f"Processing basis state {i+1}/{n_basis}... Estimated time to completion: {(time_avg * n_basis * (1 - ((i+1)/n_basis)))/60:.2f} minutes",
                        end="\r",
                    )
            for j, el2 in enumerate(bitarrays):
                if np.array_equal(el1, el2):
                    # Diagonal terms
                    for k in range(N_imp):
                        imp_i = N_ib * k
                        up = el1[imp_i]
                        down = el1[imp_i + N]
                        if up:
                            H[i, j] += e_imp[k] - mu
                        if down:
                            H[i, j] += e_imp[k] - mu
                        if up and down:
                            H[i, j] += U

                        # Inter-impurity interactions
                        for l in range(k + 1, N_imp):
                            imp_j = N_ib * l
                            up_j = el1[imp_j]
                            down_j = el1[imp_j + N]
                            H[i, j] += U * (
                                up * up_j + up * down_j + down * up_j + down * down_j
                            )

                        # Bath terms
                        for b in range(N_bath):
                            bath_i = imp_i + b + 1
                            bath_idx = k * N_bath + b
                            if el1[bath_i]:
                                H[i, j] += e_bath[bath_idx]
                            if el1[bath_i + N]:
                                H[i, j] += e_bath[bath_idx]
                else:
                    # Off-diagonal (hopping) terms
                    diff = el1 != el2
                    hop_locations = np.where(diff)[0]
                    if len(hop_locations) == 2:
                        a, b = hop_locations
                        # Hopping between impurities
                        if a % N_ib == 0 and b % N_ib == 0 and a // N == b // N:
                            k = (a // N_ib) % N_imp
                            l = (b // N_ib) % N_imp
                            V_imp_i = int(l - (k + 1) + k / 2 * (2 * N_imp - 1 - k))
                            # Fermionic sign
                            count = np.sum(el1[a + 1 : b] | el2[a + 1 : b])
                            H[i, j] += (-1) ** count * V_imp[V_imp_i]
                        # Hopping between impurity and bath
                        elif a % N_ib == 0 and b // N_ib == a // N_ib:
                            k = (a // N_ib) % N_imp
                            bath_b = b % N_ib - 1
                            bath_idx = k * N_bath + bath_b
                            count = np.sum(el1[a + 1 : b] | el2[a + 1 : b])
                            H[i, j] += (-1) ** count * V_bath[bath_idx]

            end = time.time()
            times.append(end - start)
        if self.sparse_flag:
            H = H.tocsc()
        elif isinstance(H, np.ndarray):
            pass  # already correct type
        else:
            H = np.array(H)
        return H

    def _build_single_orbital_model(self) -> NDarray:
        single_particle_hamiltonian = np.zeros((2 * self.N, 2 * self.N), dtype=complex)
        params = self.get_parameters()
        N_imp = params["N_imp"]
        N_bath = params["N_bath"]
        N_ib = N_bath + 1
        N = N_imp * N_ib
        V_imp_i = 0
        e_imp = params["e_imp"]
        e_bath = params["e_bath"]
        V_bath = params["V_bath"]
        mu = params["mu"]
        V_imp = params.get("V_imp", np.array([]))
        for i in range(N_imp):
            imp_i = i * N_ib
            end = (i + 1) * N_ib
            single_particle_hamiltonian[imp_i, imp_i] = e_imp[i] - mu
            single_particle_hamiltonian[imp_i + 1 : end, imp_i + 1 : end] = np.diag(
                e_bath[N_bath * i : N_bath * (i + 1)]
            )
            single_particle_hamiltonian[imp_i + 1 : end, imp_i] = V_bath[
                N_bath * i : N_bath * (i + 1)
            ]
            single_particle_hamiltonian[imp_i, imp_i + 1 : end] = V_bath[
                N_bath * i : N_bath * (i + 1)
            ]
            for j in range(i + 1, N_imp):
                imp_j = j * N_ib
                single_particle_hamiltonian[imp_i, imp_j] = V_imp[V_imp_i]
                single_particle_hamiltonian[imp_j, imp_i] = V_imp[V_imp_i]
                V_imp_i += 1

        single_particle_hamiltonian[N:, N:] += single_particle_hamiltonian[:N, :N]

        return single_particle_hamiltonian


class GaussianSubspace:

    def __init__(self, parameters: dict):
        """
        Initializes a Gaussian Subspace for the given parameters

        Parameters
        ----------
        parameters : `dict`
            System parameters
        """

        self.parameters = parameters
        self.parameters["V_bath"] = np.array(parameters["V_bath"], dtype=float)
        self.parameters["e_bath"] = np.array(parameters["e_bath"], dtype=float)
        self.parameters["e_imp"] = np.array(parameters["e_imp"], dtype=float)
        if "V_imp" in parameters:
            self.parameters["V_imp"] = np.array(parameters["V_imp"], dtype=float)
        self.subspace_vectors = None
        self.subspace_covariances = None

    def build_subspace_full_Hilbert(
        self,
        interaction_weights: Union[list, NDarray],
        hopping_weights: Union[list, NDarray],
        return_vectors: bool = False,
        verbose: bool = False,
        particle_selection: tuple[int, int] | Literal[None] = None,
        return_parameters=False,
        **kwargs,
    ):
        """
        Builds a Gaussian subspace in the full Hilbert space

        Parameters
        ----------
        interaction_weights : `Union[list, NDarray]`
            weights to offset the interaction system parameters by
        hopping_weights : `Union[list, NDarray]`
            weights to offset the hopping system parameters by
        return_vectors : `bool`
            should the subspace of vectors be returned?
        num_excited : `int`
            the number of excited states
        verbose : `bool`
            should there be verbose output?
        particle_selection : `tuple[int, int] | Literal[None] = None`
            the particle selection to use

        Returns
        -------
        subspace_vectors : `ndarray`
            the subspace of vectors, only returns if return_vectors is
            set to True
        """
        self.interaction_weights = interaction_weights
        self.hopping_weights = hopping_weights
        self.particle_selection = particle_selection
        hamiltonians = []
        modified_parameters = deepcopy(self.parameters)
        modified_parameters["U"] = 0
        parameters_set = []
        for h_weight in self.hopping_weights:
            for i_weight in self.interaction_weights:
                if self.parameters["N_imp"] > 1:
                    modified_parameters["e_imp"] = (
                        np.array(self.parameters["e_imp"])
                        + self.parameters["U"] * i_weight
                    )
                else:
                    modified_parameters["e_imp"] = (
                        self.parameters["e_imp"] + self.parameters["U"] * i_weight
                    )

                modified_parameters["V_bath"] = (
                    np.array(self.parameters["V_bath"]) * h_weight
                )
                model = ImpurityHamiltonian(modified_parameters)
                H = model.build_model("single orbital")
                hamiltonians.append(H)
                parameters_set.append(deepcopy(modified_parameters))

        N = 2 * model.N
        self.N = model.N
        if self.N < 5:
            self.sparse_flag = False
        else:
            self.sparse_flag = True

        self.spin_protected = kwargs.get("spin_protected", True)

        # Resolve particle sector once — it is constant across all Hamiltonians
        # whenever spin_protected=True or particle_selection is explicit.
        if particle_selection is None:
            self.particle_selection = (N // 4, N // 4)
            if self.spin_protected:
                self.s = sum(self.particle_selection)
            # else: self.s determined per-H from eigenvalues (see loop below)
        elif type(particle_selection) == int:
            self.s = particle_selection
        elif type(particle_selection) == tuple:
            self.s = sum(particle_selection)

        # Pre-compute the basis when it doesn't vary across Hamiltonians
        if self.spin_protected:
            self.basis = get_ps_basis(self.particle_selection, N)
        elif particle_selection is not None:
            self.basis = get_ps_basis(self.s, N)
        # else: non-spin-protected + None → basis recomputed per-H below

        return_full_vectors = kwargs.get("return_full_vectors", True)
        subspace_vectors = []

        for i, H in enumerate(hamiltonians):

            if verbose:
                print(f"Diagonalizing subspace Hamiltonian {i+1}...", end="\r")

            # Only recompute s/basis when filling depends on the specific Hamiltonian
            if not self.spin_protected and particle_selection is None:
                single_particle_energies, _ = np.linalg.eigh(H)
                self.s = len([e for e in single_particle_energies if e < 0])
                self.basis = get_ps_basis(self.s, N)

            if return_full_vectors:
                vector, nonint_energy = slater_det(
                    H, N, list(range(self.s)), self.basis
                )
            else:
                vector, nonint_energy = slater_det(
                    H, N, list(range(self.s)), self.basis, return_full_vector=False
                )
            subspace_vectors.append(vector)

        self.subspace_vectors = np.array(subspace_vectors)

        if return_vectors and return_parameters:
            return self.subspace_vectors, parameters_set
        if return_vectors:
            return self.subspace_vectors
        if return_parameters:
            return parameters_set

    def build_subspace_covariance(
        self,
        interaction_weights: np.ndarray | list,
        hopping_weights: np.ndarray | list,
        return_covariances: bool = False,
        particle_selection: tuple[int, int] | Literal[None] = None,
        return_parameters: bool = False,
    ) -> tuple[np.ndarray, np.ndarray] | np.ndarray:  # type: ignore
        """
        Builds a Gaussian subspace of covariance matrices

        Parameters
        ----------
        interaction_weights : `np.ndarray | list`
            weights to offset the interaction system parameters by
        hopping_weights : `np.ndarray | list`
            weights to offset the hopping system parameters by
        return_covariances : `bool`
            should the subspace of covariances be returned?
        particle_selection : `tuple[int, int] | Literal[none] = none`
            the particle selection to use

        Returns
        -------
        subspace_covariances : `ndarray`
            the subspace of covariances, only returns if return_covariances is
            set to True
        """
        N_states = len(interaction_weights) * len(hopping_weights)

        self.subspace_covariances = np.zeros(N_states, dtype=np.ndarray)
        self.subspace_parameters = np.zeros(N_states, dtype=dict)
        self.interaction_weights = np.array(interaction_weights, dtype=float)
        self.hopping_weights = np.array(hopping_weights, dtype=float)

        e_imp = np.array(self.parameters["e_imp"], dtype=float)
        V_bath = np.array(self.parameters["V_bath"], dtype=float)

        N_imp = self.parameters["N_imp"]
        N_bath = self.parameters["N_bath"]
        N = N_imp + N_imp * N_bath

        if type(particle_selection) is not tuple:
            self.particle_selection = (N // 2, N // 2)
        else:
            self.particle_selection = particle_selection

        mod_parameters = deepcopy(self.parameters)
        mod_parameters["U"] = 0
        for hi, hw in enumerate(hopping_weights):
            for ii, iw in enumerate(interaction_weights):
                new_e_imp = e_imp + self.parameters["U"] * iw
                new_V_bath = V_bath * hw

                mod_parameters["e_imp"] = new_e_imp
                mod_parameters["V_bath"] = new_V_bath
                h = ImpurityHamiltonian(mod_parameters).build_model("single orbital")

                self.subspace_covariances[hi * len(interaction_weights) + ii] = (
                    get_covariance(h, sum(self.particle_selection))
                )
                self.subspace_parameters[hi * len(interaction_weights) + ii] = {
                    "U": 0,
                    "mu": self.parameters["mu"],
                    "V_bath": np.array(new_V_bath),
                    "e_bath": np.array(self.parameters["e_bath"]),
                    "V_imp": np.array(self.parameters.get("V_imp", np.array([]))),
                    "e_imp": np.array(new_e_imp),
                    "N_bath": self.parameters["N_bath"],
                    "N_imp": self.parameters["N_imp"],
                    "weight": 0,
                }

        if return_covariances and return_parameters:
            return self.subspace_covariances, self.subspace_parameters
        elif return_covariances:
            return self.subspace_covariances
        elif return_parameters:
            return self.subspace_parameters

    def subspace_diagonalization(
        self,
        subspace=None,
        method="Full Hilbert",
        **kwargs,
    ):
        if method == "Full Hilbert":
            return self.subspace_diagonalization_full_Hilbert(subspace, **kwargs)
        elif method == "Covariance":
            return self.subspace_diagonalization_covariance(subspace, **kwargs)
        else:
            raise ValueError(
                f"{method} is an incorrect subspace diagonalization method. Choose one of the following:\n+ Full Hilbert\n+ Covariance"
            )

    def subspace_diagonalization_full_Hilbert(
        self, subspace_vectors=None, reference_Hamiltonian=None, **kwargs
    ):

        if subspace_vectors is None:
            if self.subspace_vectors is None:
                raise ValueError(
                    "Subspace vectors must be built first or provided as an arguement. Call build_subspace_full_Hilbert() method or provide the subspace manually."
                )
        else:
            self.subspace_vectors = subspace_vectors

        if reference_Hamiltonian is None:
            actual_parameters = self.parameters
            if self.particle_selection is not None:
                reference_Hamiltonian = ImpurityHamiltonian(
                    actual_parameters
                ).build_model()
                ps_Hamiltonian = ImpurityHamiltonian(actual_parameters).build_model(
                    "particle selected",
                    **{"particle_selection": self.particle_selection},
                )
                if self.sparse_flag:
                    energies, selected_evecs = eigsh(ps_Hamiltonian, k=25, which="SA")  # type: ignore
                    selected_evecs_full = lil_matrix(
                        (2 ** (2 * self.N), selected_evecs.shape[1]), dtype=complex
                    )
                    selected_evecs_full[self.basis, :] = selected_evecs
                    vectors = selected_evecs_full.tocsc()
                else:
                    energies, selected_evecs = np.linalg.eigh(ps_Hamiltonian)  # type: ignore
                    selected_evecs_full = np.zeros(
                        (2 ** (2 * self.N), selected_evecs.shape[1]), dtype=complex
                    )
                    selected_evecs_full[self.basis, :] = selected_evecs
                    vectors = selected_evecs_full

            else:
                reference_Hamiltonian = ImpurityHamiltonian(
                    actual_parameters
                ).build_model("exact")
                if self.sparse_flag:
                    energies, vectors = eigsh(reference_Hamiltonian, k=25, which="SA")
                else:
                    energies, vectors = np.linalg.eigh(reference_Hamiltonian)  # type: ignore
        else:
            if self.sparse_flag:
                energies, vectors = eigsh(reference_Hamiltonian, k=25, which="SA")
            else:
                energies, vectors = np.linalg.eigh(reference_Hamiltonian)  # type: ignore

        reduced_hamiltonian = get_reduced_hamiltonian(
            reference_Hamiltonian, self.subspace_vectors
        )
        overlap = get_overlap_matrix(self.subspace_vectors)

        exact_GS_energy = energies[0]
        if isinstance(vectors, np.ndarray):
            exact_GS_vector = vectors[:, 0]    # 1D dense — matches PS-sector subspace vectors
        else:
            exact_GS_vector = vectors[:, [0]]  # sparse (n, 1) — full Hilbert space path

        if kwargs.get("return_H_and_S", False):
            (
                optimal_EC_energy,
                optimal_EC_vector,
                all_subspace_en_candidates,
                chosen,
                H,
                S,
            ) = optimize_subspace(
                self.subspace_vectors,
                reduced_hamiltonian,
                overlap,
                reference_Hamiltonian,
                exact_GS_energy,
                exact_GS_vector,
                convergence_threshold=kwargs.get("convergence_threshold", 1e-6),
                discard_threshold=kwargs.get("discard_threshold", 1e-3),
                optimize_amplitudes=kwargs.get("optimize_amplitudes", False),
                return_H_and_S=kwargs.get("return_H_and_S", False),
                verbose=kwargs.get("verbose", False),
            )

            return optimal_EC_energy, optimal_EC_vector, chosen, H, S
        else:
            (
                optimal_EC_energy,
                optimal_EC_vector,
                all_subspace_en_candidates,
                chosen,
            ) = optimize_subspace(
                self.subspace_vectors,
                reduced_hamiltonian,
                overlap,
                reference_Hamiltonian,
                exact_GS_energy,
                exact_GS_vector,
                convergence_threshold=kwargs.get("convergence_threshold", 1e-6),
                discard_threshold=kwargs.get("discard_threshold", 1e-3),
                optimize_amplitudes=kwargs.get("optimize_amplitudes", False),
                return_H_and_S=False,
                verbose=kwargs.get("verbose", False),
            )
            return optimal_EC_energy, optimal_EC_vector, chosen

    def subspace_diagonalization_covariance(
        self,
        subspace_covariances=None,
        **kwargs,
    ):
        if subspace_covariances is None:
            if self.subspace_covariances is None:
                raise ValueError("Must provide subspace.")
        else:
            self.subspace_covariances = subspace_covariances

        overlap_matrix = get_overlap_matrix_covariance(self.subspace_covariances)
        reduced_hamiltonian = get_reduced_hamiltonian_covariance(
            self.subspace_covariances, overlap_matrix, self.parameters
        )

        return optimize_subspace(
            self.subspace_covariances,
            reduced_hamiltonian,
            overlap_matrix,
            convergence_threshold=kwargs.get("convergence_threshold", 1e-6),
            discard_threshold=kwargs.get(
                "discard_threshold", kwargs.get("discard_threshold", 1e-3)
            ),
            verbose=kwargs.get("verbose", False),
            method="Covariance",
            max_condition_number=kwargs.get("max_condition_number", 1e9),
            return_H_and_S=kwargs.get("return_H_and_S", False),
            save_file=kwargs.get("log_file", None),
        )