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, Literal, TypeAlias, overload
       7             : 
       8           1 : from pairinteraction.units import QuantityScalar
       9             : 
      10             : if TYPE_CHECKING:
      11             :     from typing_extensions import Self
      12             : 
      13             :     from pairinteraction import _backend
      14             :     from pairinteraction.units import PintFloat
      15             : 
      16           1 : UnionCPPKet: TypeAlias = "_backend.KetAtom | _backend.KetPairComplex"
      17             : 
      18             : 
      19           1 : class KetBase(ABC):
      20             :     """Base class for all Ket objects.
      21             : 
      22             :     The ket objects are meant to represent mathematically the canonical basis states, with respect to which
      23             :     the coefficient matrix of the basis objects are defined.
      24             :     For single atoms we simply choose the atomic states defined by their quantum numbers,
      25             :     therefore all KetAtom objects are orthogonal to each other.
      26             :     For pair systems, we choose the product states of the single-atom eigenstates, which depends on the system
      27             :     and the applied fields. Thus for different pair systems the KetPair objects are not necessarily orthogonal anymore.
      28             : 
      29             :     All ket objects share a few common attributes and methods, that are defined in this base class.
      30             :     E.g. each ket has a total momentum quantum number f, a magnetic quantum number m, a parity, an energy,
      31             :     as well as a label that represents the ket.
      32             :     """
      33             : 
      34           1 :     _cpp: UnionCPPKet
      35             : 
      36           1 :     @classmethod
      37           1 :     def _from_cpp_object(cls: type[Self], cpp_obj: UnionCPPKet) -> Self:
      38           1 :         obj = cls.__new__(cls)
      39           1 :         obj._cpp = cpp_obj
      40           1 :         return obj
      41             : 
      42           1 :     def __repr__(self) -> str:
      43           0 :         return f"{type(self).__name__}({self.get_label('raw')})"
      44             : 
      45           1 :     def __str__(self) -> str:
      46           1 :         return self.get_label("ket")
      47             : 
      48           1 :     def __hash__(self) -> int:
      49           1 :         return self._cpp.__hash__()
      50             : 
      51           1 :     def __eq__(self, other: object) -> bool:
      52           1 :         if not isinstance(other, KetBase):
      53           0 :             return NotImplemented
      54           1 :         if type(self._cpp) is not type(other._cpp):
      55           0 :             return False
      56           1 :         return self._cpp == other._cpp  # type: ignore [operator]
      57             : 
      58           1 :     def get_label(self, fmt: Literal["raw", "ket", "bra"] = "raw") -> str:
      59             :         """Label representing the ket.
      60             : 
      61             :         Args:
      62             :             fmt: The format of the label, i.e. whether to return the raw label, or the label in ket or bra notation.
      63             : 
      64             :         Returns:
      65             :             The label of the ket in the given format.
      66             : 
      67             :         """
      68           1 :         raw = self._get_raw_label()
      69           1 :         if fmt == "raw":
      70           1 :             return raw
      71           1 :         if fmt == "ket":
      72           1 :             return f"|{raw}⟩"
      73           1 :         if fmt == "bra":
      74           1 :             return f"⟨{raw}|"
      75           0 :         raise ValueError(f"Unknown fmt {fmt}")
      76             : 
      77           1 :     @abstractmethod
      78           1 :     def _get_raw_label(self) -> str:
      79             :         """Return the raw label of the ket."""
      80             : 
      81             :     @overload
      82             :     def get_energy(self, unit: None = None) -> PintFloat: ...
      83             : 
      84             :     @overload
      85             :     def get_energy(self, unit: str) -> float: ...
      86             : 
      87           1 :     def get_energy(self, unit: str | None = None) -> float | PintFloat:
      88             :         """Get the energy of the ket in the given unit.
      89             : 
      90             :         Args:
      91             :             unit: The unit to which to convert the energy to.
      92             :                 Default None will return a `pint.Quantity`.
      93             : 
      94             :         Returns:
      95             :             The energy as float if a unit was given, otherwise a `pint.Quantity`.
      96             : 
      97             :         """
      98           1 :         energy_au = self._cpp.get_energy()
      99           1 :         return QuantityScalar.convert_au_to_user(energy_au, "energy", unit)