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 collections.abc import Sequence
       7           1 : from typing import TYPE_CHECKING, Any, TypeGuard
       8             : 
       9           1 : import numpy as np
      10           1 : from typing_extensions import TypeAliasType
      11             : 
      12           1 : from pairinteraction import _backend
      13           1 : from pairinteraction.ket import KetPair, KetPairReal
      14           1 : from pairinteraction.state.state_atom import StateAtom
      15           1 : from pairinteraction.state.state_base import StateBase
      16             : 
      17             : if TYPE_CHECKING:
      18             :     from pairinteraction.ket.ket_atom import KetAtom
      19             : 
      20           1 : logger = logging.getLogger(__name__)
      21             : 
      22             : 
      23           1 : StateAtomTuple = TypeAliasType("StateAtomTuple", tuple[StateAtom, StateAtom] | Sequence[StateAtom])
      24             : 
      25             : 
      26           1 : def is_state_pair_like(obj: Any) -> TypeGuard[StatePairLike]:
      27           1 :     return isinstance(obj, StatePair) or is_state_atom_tuple(obj)
      28             : 
      29             : 
      30           1 : def is_state_atom_tuple(obj: Any) -> TypeGuard[tuple[StateAtom, StateAtom]]:
      31           1 :     return hasattr(obj, "__len__") and len(obj) == 2 and all(isinstance(x, StateAtom) for x in obj)
      32             : 
      33             : 
      34           1 : class StatePair(StateBase[KetPair]):
      35             :     """Pair state of two atoms.
      36             : 
      37             :     Currently StatePair objects don't offer any additional functionality.
      38             : 
      39             :     """
      40             : 
      41           1 :     _cpp: _backend.BasisPairComplex
      42           1 :     _ket_class = KetPair
      43             : 
      44           1 :     def __init__(self, ket: KetPair, basis: Any) -> None:
      45             :         """Initialize a state object representing a ket in a given basis.
      46             : 
      47             :         Args:
      48             :             ket: The ket to represent in the state.
      49             :             basis: The basis to which the state belongs.
      50             : 
      51             :         """
      52           0 :         raise NotImplementedError(
      53             :             "StatePair objects cannot be created directly. "
      54             :             "You can use `basis_pair.get_corresponding_state(ket)` or `basis_pair.get_state(i)` instead."
      55             :         )
      56             : 
      57           1 :     def get_label(
      58             :         self,
      59             :         stop_after_num_kets: int = 3,
      60             :         stop_after_accumulated_overlap: float = 0.95,
      61             :         considered_num_kets: int | None = None,
      62             :     ) -> str:
      63             :         """Label representing the state.
      64             : 
      65             :         Args:
      66             :             stop_after_num_kets: Maximum number of kets to include in the label.
      67             :             stop_after_accumulated_overlap: Stop including kets in the label,
      68             :                 if the accumulated overlap of the included kets exceeds this value.
      69             :             considered_num_kets: The number of kets to consider in the pair basis.
      70             :                 Default None uses a heuristic to determine a suitable number of kets.
      71             : 
      72             :         Returns:
      73             :             The label of the ket in the given format.
      74             : 
      75             :         """
      76           1 :         if isinstance(self._cpp, _backend.BasisPairComplex):
      77           1 :             from pairinteraction.basis import BasisAtom, BasisPair
      78           1 :             from pairinteraction.system import SystemAtom
      79             : 
      80           1 :             basis_atom_class = BasisAtom
      81           1 :             system_atom_class = SystemAtom
      82           1 :             basis_pair_class = BasisPair
      83             :         else:
      84           1 :             from pairinteraction.basis import BasisAtomReal, BasisPairReal
      85           1 :             from pairinteraction.system import SystemAtomReal
      86             : 
      87           1 :             basis_atom_class = BasisAtomReal
      88           1 :             system_atom_class = SystemAtomReal
      89           1 :             basis_pair_class = BasisPairReal
      90             : 
      91           1 :         basis_atoms_cpp = [self._cpp.get_basis1(), self._cpp.get_basis2()]
      92           1 :         basis_atoms = [
      93             :             basis_atom_class._from_cpp_object(basis_atom_cpp.canonicalized()) for basis_atom_cpp in basis_atoms_cpp
      94             :         ]
      95           1 :         system_atoms = [system_atom_class(basis_atom) for basis_atom in basis_atoms]
      96             : 
      97           1 :         coeffs = np.abs(self.get_coefficients())
      98           1 :         ket_pair = self.get_ket(int(np.argmax(coeffs)))
      99             : 
     100             :         # manually find the corresponding kets of the state atoms, to avoid warning messages
     101           1 :         ket_atom_tuple: list[KetAtom] = []
     102           1 :         for state_atom in ket_pair.state_atoms:
     103           1 :             overlaps = np.abs(state_atom.get_coefficients()) ** 2
     104           1 :             ket_atom_tuple.append(state_atom.get_ket(int(np.argmax(overlaps))))
     105             : 
     106             :         # heuristic to quickly find a basis, which includes the stop_after_num_kets most contributing kets
     107           1 :         is_converged = False
     108           1 :         considered_num_kets_list = [100, 1_000, 10_000] if considered_num_kets is None else [considered_num_kets]
     109           1 :         for number_of_kets in considered_num_kets_list:
     110           1 :             canonical_basis_pair = basis_pair_class.from_kets(
     111             :                 ket_atom_tuple, system_atoms, number_of_kets=number_of_kets, warn_number_of_kets=False
     112             :             )
     113             : 
     114           1 :             amplitudes = canonical_basis_pair.get_amplitudes(self)
     115           1 :             overlaps = np.abs(amplitudes) ** 2
     116             : 
     117           1 :             _stop_after_num_kets = min(stop_after_num_kets, len(overlaps))
     118           1 :             largest_inds = np.argpartition(overlaps, -_stop_after_num_kets)[-_stop_after_num_kets:]
     119           1 :             largest_inds = largest_inds[np.argsort(overlaps[largest_inds])[::-1]]
     120             : 
     121           1 :             acc_overlaps = np.cumsum(overlaps[largest_inds])
     122           1 :             max_ind_to_include = np.searchsorted(acc_overlaps, stop_after_accumulated_overlap * self.norm**2)
     123           1 :             largest_inds = largest_inds[: max_ind_to_include + 1]
     124             : 
     125             :             # if the overlap of the smallest ket we still include is larger than the remaining contributions,
     126             :             # we can be sure, that the label is accurate and the heuristic was successful.
     127           1 :             remaining_contributions = self.norm**2 - np.sum(overlaps)
     128           1 :             if overlaps[largest_inds[-1]] >= remaining_contributions:
     129           1 :                 is_converged = True
     130           1 :                 break
     131             : 
     132           1 :         label = ""
     133           1 :         accumulated_ov = 0.0
     134           1 :         for ind in largest_inds:
     135           1 :             coeff = np.real_if_close(amplitudes[ind])
     136           1 :             canonical_ket = canonical_basis_pair.get_ket(ind)
     137           1 :             label += f"{coeff:.2f} |{canonical_ket.get_label()}⟩"
     138           1 :             accumulated_ov += overlaps[ind]
     139           1 :             label += " + "
     140             : 
     141           1 :         if accumulated_ov <= self.norm**2 - 100 * np.finfo(float).eps:
     142           1 :             label += "..."
     143             :         else:
     144           1 :             label = label[:-3]  # Remove the last " + "
     145           1 :         label = label.replace("+ -", "- ")
     146             : 
     147           1 :         if not is_converged:
     148           0 :             logger.warning(
     149             :                 "The label '%s' may not be printing the largest contributions. "
     150             :                 "Consider calling get_label with a larger 'considered_num_kets'.",
     151             :                 label,
     152             :             )
     153             : 
     154           1 :         return label
     155             : 
     156           1 :     def get_amplitude(self, other: Any) -> Any:
     157           0 :         raise NotImplementedError("StatePair.get_amplitude not implemented yet")
     158             : 
     159           1 :     def get_overlap(self, other: Any) -> Any:
     160           0 :         raise NotImplementedError("StatePair.get_overlap not implemented yet")
     161             : 
     162           1 :     def get_matrix_element(self, other: Any, *args: Any, **kwargs: Any) -> Any:
     163           0 :         raise NotImplementedError("StatePair.get_matrix_element not implemented yet")
     164             : 
     165             : 
     166           1 : class StatePairReal(StatePair):
     167           1 :     _cpp: _backend.BasisPairReal  # type: ignore [assignment]
     168           1 :     _ket_class = KetPairReal
     169             : 
     170             : 
     171           1 : StatePairLike = TypeAliasType("StatePairLike", StatePair | StateAtomTuple)