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