Source code for hqs_spin_mapper.spin_mapper_protocols

r"""Module containing the protocols for transformable objects of the HQS Spin Mapper."""

# Copyright © 2021-2025 HQS Quantum Simulations GmbH. All Rights Reserved.
import sys
import textwrap
import numpy as np

from typing import Protocol, Optional, Any, Union, Tuple, Set, runtime_checkable

from dataclasses import dataclass

try:
    from lattice_functions.Fermions import (
        ExpressionSpinful,
        ExpressionSpinful_complex,
    )
except ImportError as e:
    if "mkl" in str(e):
        print(
            textwrap.dedent(
                """\
                ===== ERROR: Missing MKL =====

                The Intel Math Kernel Library (MKL) could not be found. Please ensure
                that the MKL is installed properly.

                When using HQStage, run

                   hqstage envs install-mkl

                to install the MKL. In other cases, please consult the manual."""
            ),
            file=sys.stderr,
        )
        sys.exit(1)
    else:
        raise

__all__ = [
    "InsufficientProtocolError",
    "Hamiltonian_Matrices",
    "Supports_SW_Transformation",
    "Supports_Spin_Optimization",
]




class InsufficientProtocolError(Exception):
    r"""Custom exception for the case that objects do not satisfy the required protocol."""



    def __init__(self, obj: Any, specified_protocol: Any) -> None:
        r"""Instantiate the error and print the error message.

        Args:
            obj (Any): Object that doesn't satisfy the specified protocol.
            specified_protocol (Any): The specified protocol that is required.
        """
        print(f"{type(obj)} does not satisfy the required {type(specified_protocol)} protocol.")






@dataclass
class Hamiltonian_Matrices:
    r"""Collection of the matrix representations of the system Hamiltonian.

    H0_uu (np.ndarray): Quadratic Hamiltonian, spin directions up-up
    H0_dd (np.ndarray): Quadratic Hamiltonian, spin directions down-down
    H0_uu (np.ndarray): Quadratic Hamiltonian, spin directions up-down
    HU (np.ndarray): Density-Density interaction matrix or interaction tensors
    Jz (np.ndarray): Sz-Sz interaction matrix
    Jp (np.ndarray): S+-S- interaction matrix
    Jc (np.ndarray): \vec{S}x\vec{S} interaction matrix
    HR (np.ndarray): Remainder bath interaction tensor
    """

    H0_uu: np.ndarray
    H0_dd: np.ndarray
    H0_ud: np.ndarray
    HU: np.ndarray
    Jz: Optional[np.ndarray] = None
    Jp: Optional[np.ndarray] = None
    Jc: Optional[np.ndarray] = None
    # Tensor containing the terms in the Hamiltonian that cannot be transformed
    HR: Optional[np.ndarray] = None





@runtime_checkable
class Supports_Spin_Optimization(Protocol):
    r"""Protocol for transformable systems with the prerequisites for spin-orbital optimization."""

    @property
    def tolerance(self) -> float:
        r"""Return the tolerated parity deviation :math:`\delta` in :math:`P=(-1 + \delta)` for an 
        orbital to be considered a spin-orbital."""

    @property
    def optimization_steps(self) -> int:
        r"""Return the number of optimization loops."""

    @property
    def immutable_indices(self) -> Set[int]:
        r"""Return the orbital indices that should remain unaltered by the optimization."""

    @property
    def rdm1(self) -> Union[np.ndarray, Tuple[np.ndarray, ...]]:
        r"""Return the matrices of the 1-RDM (:math:`\langle c^{\dagger}_{i\sigma} c_{j\sigma'}\rangle`)."""

    @rdm1.setter
    def rdm1(self, value: Union[np.ndarray, Tuple[np.ndarray, ...]]) -> None:
        r"""Set the one-particle reduced one-electron density-matrices.

        Args:
            value (np.ndarray): Updated 1-RDM
        """

    @property
    def rdm2(self) -> np.ndarray:
        r"""Return the tensor of the 2-RDM 
        (:math:`\langle c^{\dagger}_{i\uparrow} c^{\dagger}_{j\downarrow} c_{k\downarrow} c_{l\uparrow}\rangle`)."""

    @rdm2.setter
    def rdm2(self, value: np.ndarray) -> None:
        r"""Set the two-particle reduced density-matrix 2-RDM.

        Args:
            value (np.ndarray): Updated 2-RDM
        """





@runtime_checkable
class Supports_SW_Transformation(Protocol):
    r"""Protocol for systems that allow for a Schrieffer-Wolff transformation."""

    @property
    def hamiltonians(self) -> Hamiltonian_Matrices:
        r"""Return the matrices of the terms in the Hamiltonian describing the system."""

    @hamiltonians.setter
    def hamiltonians(self, value: Hamiltonian_Matrices) -> None:
        r"""Set the matrices of the terms in the Hamiltonian describing the system.

        Args:
            value (Hamiltonian_Matrices): New collection of Hamiltonian matrices
        """

    @property
    def rotation_matrix(self) -> np.ndarray:
        r"""Return the rotation matrix that transforms the system to the spin-optimized basis."""

    @rotation_matrix.setter
    def rotation_matrix(self, value: np.ndarray) -> None:
        r"""Set the rotation matrix that transforms the system to the spin-optimized basis.

        Args:
            value (np.ndarray): Updated rotation matrix
        """

    @property
    def spin_indices(self) -> Set[int]:
        r"""Return the set of indices of the spin-like orbitals."""

    @spin_indices.setter
    def spin_indices(self, value: Set[int]) -> None:
        r"""Set the set of indices of the spin-like orbitals.

        Args:
            value (Set[int]): Updated indices of the spin-like orbitals
        """

    @property
    def system_size(self) -> int:
        r"""Return the system size i.e. the number of orbitals."""

    @property
    def site_type(self) -> str:
        r"""Return the type of orbitals (e.g. ``spinful_fermions`` or ``spinless_fermions``)."""

    @property
    def prefactor_cutoff(self) -> float:
        r"""Return the cutoff for the prefactors of terms to be considered in the transformation."""

    @prefactor_cutoff.setter
    def prefactor_cutoff(self, value: float) -> None:
        r"""Set the prefactor cutoff for terms considered in the Schrieffer-Woff transformation.

        Args:
            value (float): Lower bound for the absolute value of prefactors of terms considered.
        """

    @property
    def generator(self) -> Union[ExpressionSpinful, ExpressionSpinful_complex]:
        r"""Return the generator :math:`S` of the Schrieffer-Wolff transformation."""

    @generator.setter
    def generator(self, value: Union[ExpressionSpinful, ExpressionSpinful_complex]) -> None:
        r"""Set the generator :math:`S` of the Schrieffer-Wolff transformation.

        Args:
            value (Union[ExpressionSpinful, ExpressionSpinful_complex]): 
            Generator of the Schrieffer-Wolff transformation
        """

    @property
    def transformed_hamiltonian(self) -> Union[ExpressionSpinful, ExpressionSpinful_complex]:
        r"""Return the Schrieffer-Wolff transformed Hamiltonian :math:`\tilde{H}`."""

    @transformed_hamiltonian.setter
    def transformed_hamiltonian(
        self, value: Union[ExpressionSpinful, ExpressionSpinful_complex]
    ) -> None:
        r"""Set the Schrieffer-Wolff transformed Hamiltonian :math:`\tilde{H}`.

        Args:
            value (Union[ExpressionSpinful, ExpressionSpinful_complex]): 
            Schrieffer-Wolff transformed Hamiltonian
        """