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 collections.abc import Sequence
       6           1 : from typing import TYPE_CHECKING, Any, Literal, TypeGuard, overload
       7             : 
       8           1 : import numpy as np
       9           1 : from typing_extensions import TypeAliasType
      10             : 
      11           1 : from pairinteraction.ket.ket_atom import KetAtom
      12           1 : from pairinteraction.ket.ket_base import KetBase
      13             : 
      14             : if TYPE_CHECKING:
      15             :     from pairinteraction import _backend
      16             :     from pairinteraction.state import StateAtom
      17             :     from pairinteraction.system.system_atom import SystemAtom
      18             :     from pairinteraction.units import PintFloat
      19             : 
      20             : 
      21           1 : KetAtomTuple = TypeAliasType("KetAtomTuple", tuple[KetAtom, KetAtom] | Sequence[KetAtom])
      22             : 
      23             : 
      24           1 : def is_ket_pair_like(obj: Any) -> TypeGuard[KetPairLike]:
      25           1 :     return isinstance(obj, KetPair) or is_ket_atom_tuple(obj)
      26             : 
      27             : 
      28           1 : def is_ket_atom_tuple(obj: Any) -> TypeGuard[KetAtomTuple]:
      29           1 :     return hasattr(obj, "__len__") and len(obj) == 2 and all(isinstance(x, KetAtom) for x in obj)
      30             : 
      31             : 
      32           1 : class KetPair(KetBase):
      33             :     """Ket for a pair state of two atoms.
      34             : 
      35             :     For pair systems, we choose KetPair object as the product states of the single-atom eigenstates.
      36             :     Thus, the Ket pair objects depend on the system and the applied fields.
      37             :     Therefore for different pair systems the KetPair objects are not necessarily orthogonal anymore.
      38             : 
      39             :     Currently one cannot create a KetPair object directly, but they are used in the background when creating a
      40             :     :class:`pairinteraction.BasisPair` object.
      41             : 
      42             :     """
      43             : 
      44           1 :     _cpp: _backend.KetPairComplex
      45             : 
      46           1 :     def __init__(self) -> None:
      47             :         """Creating a KetPair object directly is not possible."""  # noqa: D401
      48           0 :         raise NotImplementedError("KetPair objects cannot be created directly.")
      49             : 
      50           1 :     @property
      51           1 :     def m(self) -> float:
      52             :         """The magnetic quantum number m (int or half-int)."""
      53           0 :         return self._cpp.get_quantum_number_m()
      54             : 
      55           1 :     def get_label(
      56             :         self,
      57             :         fmt: Literal["raw", "ket", "bra", "detailed"] = "raw",
      58             :         *,
      59             :         stop_after_num_kets: int = 3,
      60             :         stop_after_accumulated_overlap: float = 0.95,
      61             :     ) -> str:
      62             :         """Label representing the ket pair.
      63             : 
      64             :         Args:
      65             :             fmt: The format of the label, i.e. whether to return the raw label, or the label in ket or bra notation.
      66             :             stop_after_num_kets: Maximum number of single atom kets to include in the label for each StateAtom.
      67             :             stop_after_accumulated_overlap: Stop including kets in the single atom label,
      68             :                 if the accumulated overlap of the included kets exceeds this value.
      69             : 
      70             :         Returns:
      71             :             A string representation of the ket pair.
      72             : 
      73             :         """
      74           1 :         if fmt == "detailed":
      75           0 :             atom_labels = [
      76             :                 atom.get_label(stop_after_num_kets, stop_after_accumulated_overlap) for atom in self.state_atoms
      77             :             ]
      78           0 :             return f"({atom_labels[0]}) ⊗ ({atom_labels[1]})"
      79           1 :         return super().get_label(fmt)
      80             : 
      81           1 :     def _get_raw_label(self) -> str:
      82           1 :         precision = 100 * np.finfo(float).eps
      83           1 :         labels = []
      84           1 :         for state_atom in self.state_atoms:
      85           1 :             ket_idx = state_atom.get_corresponding_ket_index()
      86           1 :             coefficient = state_atom.get_coefficients()[ket_idx]
      87           1 :             optional_tilde = "~" if abs(coefficient - 1.0) > precision else ""
      88           1 :             labels.append(optional_tilde + state_atom.get_ket(ket_idx).get_label("raw"))
      89           1 :         return "; ".join(labels)
      90             : 
      91           1 :     @property
      92           1 :     def state_atoms(self) -> tuple[StateAtom, StateAtom]:
      93             :         """Return the state atoms of the ket pair."""
      94           1 :         from pairinteraction.state import StateAtom, StateAtomReal
      95             : 
      96           1 :         _state_atom_class = StateAtomReal if isinstance(self, KetPairReal) else StateAtom
      97             : 
      98           1 :         state_atoms = []
      99           1 :         for atomic_state in self._cpp.get_atomic_states():
     100           1 :             state = _state_atom_class._from_cpp_object(atomic_state)
     101           1 :             state_atoms.append(state)
     102           1 :         return tuple(state_atoms)  # type: ignore [return-value]
     103             : 
     104             : 
     105           1 : class KetPairReal(KetPair):
     106           1 :     _cpp: _backend.KetPairReal  # type: ignore [assignment]
     107             : 
     108             : 
     109           1 : KetPairLike = TypeAliasType("KetPairLike", KetPair | KetAtomTuple)
     110             : 
     111             : 
     112           1 : def get_ketpairlike_m(ket: KetPair | KetAtomTuple) -> float:
     113           1 :     if is_ket_atom_tuple(ket):
     114           1 :         m1 = ket[0].m
     115           1 :         m2 = ket[1].m
     116           1 :         return m1 + m2
     117           0 :     if isinstance(ket, KetPair):
     118           0 :         return ket.m
     119           0 :     raise TypeError(f"Unknown type: {type(ket)=}")
     120             : 
     121             : 
     122             : @overload
     123             : def get_ketpairlike_energy(
     124             :     ket: KetPair | KetAtomTuple, system_atoms: Sequence[SystemAtom], unit: None
     125             : ) -> PintFloat: ...
     126             : 
     127             : 
     128             : @overload
     129             : def get_ketpairlike_energy(ket: KetPair | KetAtomTuple, system_atoms: Sequence[SystemAtom], unit: str) -> float: ...
     130             : 
     131             : 
     132           1 : def get_ketpairlike_energy(
     133             :     ket: KetPair | KetAtomTuple, system_atoms: Sequence[SystemAtom], unit: str | None
     134             : ) -> float | PintFloat:
     135           1 :     if is_ket_atom_tuple(ket):
     136           1 :         energy1 = system_atoms[0].get_corresponding_energy(ket[0], unit)
     137           1 :         energy2 = system_atoms[1].get_corresponding_energy(ket[1], unit)
     138           1 :         return energy1 + energy2
     139           0 :     if isinstance(ket, KetPair):
     140           0 :         return ket.get_energy(unit)
     141           0 :     raise TypeError(f"Unknown type: {type(ket)=}")