Line data    Source code

       1             : # SPDX-FileCopyrightText: 2024 PairInteraction Developers
       2             : # SPDX-License-Identifier: LGPL-3.0-or-later
       3           1 : from __future__ import annotations
       4             : 
       5           1 : from abc import ABC, abstractmethod
       6           1 : from typing import TYPE_CHECKING, Any, Generic, TypeAlias, TypeVar
       7             : 
       8           1 : import numpy as np
       9             : 
      10           1 : from pairinteraction.state.state_base import get_index_with_largest_overlap
      11             : 
      12             : if TYPE_CHECKING:
      13             :     from scipy.sparse import csr_matrix
      14             :     from typing_extensions import Self
      15             : 
      16             :     from pairinteraction import _backend
      17             :     from pairinteraction.ket import KetBase
      18             :     from pairinteraction.state import StateBase
      19             : 
      20           1 : KetType = TypeVar("KetType", bound="KetBase")
      21           1 : StateType = TypeVar("StateType", bound="StateBase[Any]")
      22           1 : UnionCPPBasis: TypeAlias = "_backend.BasisAtomComplex | _backend.BasisPairComplex"
      23             : 
      24             : 
      25           1 : class BasisBase(ABC, Generic[KetType, StateType]):
      26             :     """Base class for all Basis objects.
      27             : 
      28             :     The basis objects are meant to represent a set of kets, that span a Hilbert space and store a coefficient matrix,
      29             :     that describe the basis states in terms of the kets.
      30             : 
      31             :     All basis objects share a few common attributes and methods, that are defined in this base class, e.g.:
      32             :         - the number of kets and states,
      33             :         - the kets of the basis,
      34             :         - the coefficients stored as scipy sparse matrix,
      35             :         - ...
      36             :     """
      37             : 
      38           1 :     _cpp: UnionCPPBasis
      39           1 :     _ket_class: type[KetType]  # should be ClassVar, but cannot be nested yet
      40           1 :     _state_class: type[StateType]  # should be ClassVar, but cannot be nested yet
      41             : 
      42           1 :     def _post_init(self) -> None:
      43           1 :         if self.number_of_kets == 0:
      44           0 :             raise ValueError("Cannot create a basis with zero kets.")
      45           1 :         self._kets_cache: dict[int, KetType] = {}
      46             : 
      47           1 :     @classmethod
      48           1 :     def _from_cpp_object(cls: type[Self], cpp_obj: UnionCPPBasis, *args: Any, **kwargs: Any) -> Self:
      49           1 :         assert len(kwargs) == len(args) == 0, "No additional arguments expected."
      50           1 :         obj = cls.__new__(cls)
      51           1 :         obj._cpp = cpp_obj
      52           1 :         obj._post_init()
      53           1 :         return obj
      54             : 
      55           1 :     def _from_cpp_object_additional_args(self) -> tuple[Any, ...]:
      56             :         """Extra positional args ``_from_cpp_object`` needs to rebuild an object of this type."""
      57           1 :         return ()
      58             : 
      59           1 :     def __repr__(self) -> str:
      60           1 :         args = f"{self.get_ket(0)} ... {self.get_ket(self.number_of_kets - 1)}"
      61           1 :         return f"{type(self).__name__}({args})"
      62             : 
      63           1 :     def __str__(self) -> str:
      64           1 :         return self.__repr__()
      65             : 
      66           1 :     @property
      67           1 :     def kets(self) -> list[KetType]:
      68             :         """Return a list containing the kets of the basis."""
      69           1 :         return [self.get_ket(i) for i in range(self.number_of_kets)]
      70             : 
      71           1 :     def get_ket(self, index: int) -> KetType:
      72             :         """Return the ket at the given index."""
      73           1 :         if index not in self._kets_cache:
      74           1 :             if index < 0 or index >= self.number_of_kets:
      75           0 :                 raise IndexError(f"Ket index {index} out of range (number of kets: {self.number_of_kets}).")
      76           1 :             ket_cpp = self._cpp.get_ket(index)
      77           1 :             self._kets_cache[index] = self._ket_class._from_cpp_object(ket_cpp)
      78           1 :         return self._kets_cache[index]
      79             : 
      80           1 :     @property  # not cached_property since State objects are mutable
      81           1 :     def states(self) -> list[StateType]:
      82             :         """Return a list containing the states of the basis."""
      83           0 :         return [self.get_state(i) for i in range(self.number_of_states)]
      84             : 
      85           1 :     def get_state(self, index: int) -> StateType:
      86             :         """Return the state at the given index."""
      87           1 :         if index < 0 or index >= self.number_of_states:
      88           0 :             raise IndexError(f"State index {index} out of range (number of states: {self.number_of_states}).")
      89           1 :         state_cpp = self._cpp.get_state(index)
      90           1 :         return self._state_class._from_cpp_object(state_cpp)
      91             : 
      92           1 :     @property
      93           1 :     def number_of_kets(self) -> int:
      94             :         """Return the number of kets in the basis."""
      95           1 :         return self._cpp.get_number_of_kets()
      96             : 
      97           1 :     @property
      98           1 :     def number_of_states(self) -> int:
      99             :         """Return the number of states in the basis."""
     100           1 :         return self._cpp.get_number_of_states()
     101             : 
     102           1 :     def get_coefficients(self) -> csr_matrix:
     103             :         """Return the coefficients of the basis as a sparse matrix.
     104             : 
     105             :         The coefficients are stored in a sparse matrix with shape (number_of_kets, number_of_states),
     106             :         where the first index correspond to the kets and the second index correspond to the states.
     107             :         For example `basis.get_coefficients()[i, j]` is the i-th coefficient
     108             :         (i.e. the coefficient corresponding to the i-th ket) of the j-th state.
     109             : 
     110             :         The coefficients are normalized, i.e. the sum of the absolute values of the coefficients
     111             :         in each row is equal to 1.
     112             : 
     113             :         """
     114           1 :         coefficients = self._cpp.get_coefficients()
     115           1 :         coefficients.data = np.real_if_close(coefficients.data)
     116           1 :         return coefficients
     117             : 
     118           1 :     def get_corresponding_ket(self: Self, state: StateType) -> KetType:
     119             :         """Return the ket of the basis with the maximal overlap with the given state."""
     120           1 :         return self.get_ket(self.get_corresponding_ket_index(state))
     121             : 
     122           1 :     def get_corresponding_ket_index(self, state: StateType) -> int:
     123             :         """Return the index of the ket of the basis with the maximal overlap with the given state."""
     124           1 :         canonical_basis = self.canonicalized()
     125           1 :         overlaps = canonical_basis.get_overlaps(state)
     126           1 :         err_msg = "ket for the given state in the basis"
     127           1 :         return get_index_with_largest_overlap(overlaps, err_msg=err_msg)
     128             : 
     129           1 :     def get_corresponding_state(self, ket: KetBase) -> StateType:
     130             :         """Return the state of the basis with the maximal overlap with the given ket."""
     131           1 :         return self.get_state(self.get_corresponding_state_index(ket))
     132             : 
     133           1 :     def get_corresponding_state_index(self, ket: KetBase) -> int:
     134             :         """Return the index of the state of the basis with the maximal overlap with the given ket."""
     135           1 :         overlaps = self.get_overlaps(ket)
     136           1 :         err_msg = "state for the given ket in the basis"
     137           1 :         return get_index_with_largest_overlap(overlaps, err_msg=err_msg)
     138             : 
     139           1 :     def canonicalized(self: Self) -> Self:
     140             :         """Return the canonical basis with identity coefficients."""
     141           1 :         return type(self)._from_cpp_object(self._cpp.canonicalized(), *self._from_cpp_object_additional_args())
     142             : 
     143           1 :     @property
     144           1 :     def is_canonical(self) -> bool:
     145             :         """Return True if the basis is in canonical order with identity coefficients."""
     146           0 :         return self._cpp.is_canonical()
     147             : 
     148           1 :     def merge(self: Self, other: Self) -> Self:
     149             :         """Return a canonical basis containing the kets from both bases.
     150             : 
     151             :         The bases must be compatible, i.e. they must share the same species, database, ...
     152             :         For BasisPair they also must share the same underlying atomic basis.
     153             :         Both bases must be canonical (i.e. have identity coefficients).
     154             :         """
     155           1 :         if type(self) is not type(other):
     156           0 :             raise TypeError(
     157             :                 f"Can only merge {type(self).__name__} with the same basis type, but got {type(other).__name__}."
     158             :             )
     159           1 :         return type(self)._from_cpp_object(self._cpp.merge(other._cpp), *self._from_cpp_object_additional_args())  # type: ignore [arg-type]
     160             : 
     161             :     @abstractmethod
     162             :     def get_amplitudes(self, other: Any) -> Any: ...
     163             : 
     164             :     @abstractmethod
     165             :     def get_overlaps(self, other: Any) -> Any: ...
     166             : 
     167             :     @abstractmethod
     168             :     def get_matrix_elements(self, other: Any, *args: Any, **kwargs: Any) -> Any: ...