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 : import logging
       6           1 : from abc import ABC, abstractmethod
       7           1 : from typing import TYPE_CHECKING, Any, Generic, TypeAlias, TypeVar
       8             : 
       9           1 : import numpy as np
      10             : 
      11             : if TYPE_CHECKING:
      12             :     from typing_extensions import Self
      13             : 
      14             :     from pairinteraction import _backend
      15             :     from pairinteraction.ket import KetBase
      16             :     from pairinteraction.units import NDArray
      17             : 
      18           1 : logger = logging.getLogger(__name__)
      19             : 
      20           1 : KetType = TypeVar("KetType", bound="KetBase")
      21           1 : UnionCPPBasis: TypeAlias = "_backend.BasisAtomComplex | _backend.BasisPairComplex"
      22             : 
      23             : 
      24           1 : class StateBase(ABC, Generic[KetType]):
      25             :     """Base class for all State objects.
      26             : 
      27             :     The state objects are meant to represent a set of kets, that span a Hilbert space
      28             :     and store a coefficient vector, that describes the state in terms of the kets.
      29             : 
      30             :     Basically just a wrapper around the python Basis classes, but with more convenience functions specific to a state.
      31             : 
      32             :     Note, in the cpp code we dont have an explicit state class, but use a Basis object, which has only one state
      33             : 
      34             :     All state objects share a few common attributes and methods, that are defined in this base class, e.g.:
      35             :         - the number of kets,
      36             :         - the kets of the state,
      37             :         - the coefficient vector as 1d-array,
      38             :         - ...
      39             :     """
      40             : 
      41           1 :     _cpp: UnionCPPBasis
      42           1 :     _ket_class: type[KetType]  # should be ClassVar, but cannot be nested yet
      43             : 
      44           1 :     def __init__(self) -> None:
      45           1 :         self._kets_cache: dict[int, KetType] = {}
      46             : 
      47           1 :     @classmethod
      48           1 :     def _from_cpp_object(cls: type[Self], cpp_obj: UnionCPPBasis) -> Self:
      49           1 :         obj = cls.__new__(cls)
      50           1 :         obj._cpp = cpp_obj
      51           1 :         obj._kets_cache = {}
      52           1 :         return obj
      53             : 
      54           1 :     def __repr__(self) -> str:
      55           0 :         return f"{type(self).__name__}({self.get_label()})"
      56             : 
      57           1 :     def __str__(self) -> str:
      58           1 :         return self.get_label()
      59             : 
      60           1 :     def get_label(self, stop_after_num_kets: int = 3, stop_after_accumulated_overlap: float = 0.95) -> str:
      61             :         """Label representing the state.
      62             : 
      63             :         Args:
      64             :             stop_after_num_kets: Maximum number of kets to include in the label.
      65             :             stop_after_accumulated_overlap: Stop including kets in the label,
      66             :                 if the accumulated overlap of the included kets exceeds this value.
      67             : 
      68             :         Returns:
      69             :             The label of the ket in the given format.
      70             : 
      71             :         """
      72           1 :         coefficients = self.get_coefficients()
      73           1 :         sorted_inds = np.argsort(np.abs(coefficients))[::-1]
      74           1 :         norm_squared = self.norm**2
      75           1 :         label = ""
      76           1 :         overlap = 0
      77           1 :         for i, ind in enumerate(sorted_inds, 1):
      78           1 :             label += f"{np.real_if_close(coefficients[ind]):.2f} {self.get_ket(ind).get_label('ket')}"
      79           1 :             overlap += abs(coefficients[ind]) ** 2
      80           1 :             if overlap > (stop_after_accumulated_overlap * norm_squared) or i >= stop_after_num_kets:
      81           1 :                 break
      82           1 :             label += " + "
      83           1 :         if overlap <= norm_squared - 100 * np.finfo(float).eps:
      84           0 :             label += " + ... "
      85             : 
      86           1 :         return label.replace("+ -", "- ")
      87             : 
      88           1 :     @property
      89           1 :     def kets(self) -> list[KetType]:
      90             :         """Return a list containing the kets of the basis."""
      91           1 :         return [self.get_ket(i) for i in range(self.number_of_kets)]
      92             : 
      93           1 :     def get_ket(self, index: int) -> KetType:
      94             :         """Return the ket at the given index."""
      95           1 :         if index not in self._kets_cache:
      96           1 :             if index < 0 or index >= self.number_of_kets:
      97           0 :                 raise IndexError(f"Ket index {index} out of range (number of kets: {self.number_of_kets}).")
      98           1 :             ket_cpp = self._cpp.get_ket(index)
      99           1 :             self._kets_cache[index] = self._ket_class._from_cpp_object(ket_cpp)
     100           1 :         return self._kets_cache[index]
     101             : 
     102           1 :     @property
     103           1 :     def number_of_kets(self) -> int:
     104             :         """Return the number of kets in the basis."""
     105           1 :         return self._cpp.get_number_of_kets()
     106             : 
     107           1 :     @property
     108           1 :     def norm(self) -> np.floating:
     109             :         """Return the norm of the state."""
     110           1 :         return np.linalg.norm(self.get_coefficients())
     111             : 
     112           1 :     def get_coefficients(self) -> NDArray:
     113             :         """Return the coefficients of the state as a 1d-array.
     114             : 
     115             :         The coefficients are stored in a numpy.array with shape (number_of_kets,).
     116             : 
     117             :         The coefficients are normalized, i.e. the sum of the absolute values of the coefficients is equal to 1.
     118             : 
     119             :         """
     120           1 :         coefficients = self._cpp.get_coefficients().toarray().ravel()
     121           1 :         return np.real_if_close(coefficients)
     122             : 
     123           1 :     def get_corresponding_ket(self) -> KetType:
     124             :         """Return the ket with the maximal overlap with self."""
     125           1 :         return self.get_ket(self.get_corresponding_ket_index())
     126             : 
     127           1 :     def get_corresponding_ket_index(self) -> int:
     128             :         """Return the ket index with the maximal overlap with self."""
     129           1 :         overlaps = np.abs(self.get_coefficients()) ** 2
     130           1 :         err_msg = "ket for the state"
     131           1 :         return get_index_with_largest_overlap(overlaps, err_msg=err_msg)
     132             : 
     133             :     @abstractmethod
     134             :     def get_amplitude(self, other: Any) -> Any: ...
     135             : 
     136             :     @abstractmethod
     137             :     def get_overlap(self, other: Any) -> Any: ...
     138             : 
     139             :     @abstractmethod
     140             :     def get_matrix_element(self, other: Any, *args: Any, **kwargs: Any) -> Any: ...
     141             : 
     142             : 
     143           1 : def get_index_with_largest_overlap(overlaps: NDArray, err_msg: str) -> int:
     144           1 :     if len(overlaps) == 0:
     145           0 :         raise ValueError(f"Cannot find the corresponding {err_msg}: No overlaps, this should not happen.")
     146           1 :     if len(overlaps) == 1:
     147           1 :         largest_id = 0
     148           1 :         largest_overlap = overlaps[0]
     149           1 :         second_largest_overlap = 0
     150             :     else:
     151             :         # Find the indices of the two largest overlaps
     152             :         # this is more efficient than sorting the entire array
     153           1 :         ids = np.argpartition(overlaps, -2)[-2:]
     154           1 :         ids = ids[np.argsort(overlaps[ids])[::-1]]
     155           1 :         largest_id = int(ids[0])
     156           1 :         largest_overlap = overlaps[ids[0]]
     157           1 :         second_largest_overlap = overlaps[ids[1]]
     158             : 
     159           1 :     if largest_overlap == 0:
     160           0 :         raise ValueError(f"Cannot find the corresponding {err_msg}: All overlaps are 0.")
     161           1 :     if largest_overlap < 0.5 + 100 * np.finfo(float).eps:
     162           1 :         logger.warning(
     163             :             "Cannot find the uniquely corresponding %s: "
     164             :             "Largest overlap=%.3f <= 0.5, the second largest overlap is %.3f. "
     165             :             "Still returning the result with the largest overlap.",
     166             :             *(err_msg, largest_overlap, second_largest_overlap),
     167             :         )
     168             : 
     169           1 :     return largest_id