Add mqdt

pyproject.toml

@@ -44,6 +44,9 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
+mqdt = [
+  "juliacall >= 0.9.24",
+]
 tests = [
   "pytest >= 8.0",
   "nbmake >= 1.3",

src/rydstate/__init__.py

@@ -1,6 +1,7 @@
-from rydstate import angular, radial, species
-from rydstate.basis import BasisSQDTAlkali, BasisSQDTAlkalineFJ, BasisSQDTAlkalineJJ, BasisSQDTAlkalineLS
+from rydstate import angular, basis, radial, species
+from rydstate.basis import BasisMQDT, BasisSQDTAlkali, BasisSQDTAlkalineFJ, BasisSQDTAlkalineJJ, BasisSQDTAlkalineLS
 from rydstate.rydberg import (
+    RydbergStateMQDT,
     RydbergStateSQDT,
     RydbergStateSQDTAlkali,
     RydbergStateSQDTAlkalineFJ,
@@ -10,16 +11,19 @@
 from rydstate.units import ureg
 
 __all__ = [
+    "BasisMQDT",
     "BasisSQDTAlkali",
     "BasisSQDTAlkalineFJ",
     "BasisSQDTAlkalineJJ",
     "BasisSQDTAlkalineLS",
+    "RydbergStateMQDT",
     "RydbergStateSQDT",
     "RydbergStateSQDTAlkali",
     "RydbergStateSQDTAlkalineFJ",
     "RydbergStateSQDTAlkalineJJ",
     "RydbergStateSQDTAlkalineLS",
     "angular",
+    "basis",
     "radial",
     "species",
     "ureg",

src/rydstate/angular/__init__.py

@@ -1,7 +1,9 @@
 from rydstate.angular.angular_ket import AngularKetFJ, AngularKetJJ, AngularKetLS
+from rydstate.angular.angular_ket_dummy import AngularKetDummy
 from rydstate.angular.angular_state import AngularState
 
 __all__ = [
+    "AngularKetDummy",
     "AngularKetFJ",
     "AngularKetJJ",
     "AngularKetLS",

src/rydstate/angular/angular_ket.py

@@ -13,34 +13,24 @@
     is_angular_operator_type,
 )
 from rydstate.angular.utils import (
-    calc_wigner_3j,
+    InvalidQuantumNumbersError,
     check_spin_addition_rule,
-    clebsch_gordan_6j,
-    clebsch_gordan_9j,
     get_possible_quantum_number_values,
     minus_one_pow,
     try_trivial_spin_addition,
 )
+from rydstate.angular.wigner_symbols import calc_wigner_3j, clebsch_gordan_6j, clebsch_gordan_9j
 from rydstate.species import SpeciesObject
 
 if TYPE_CHECKING:
-    from typing_extensions import Self
+    from typing_extensions import Never, Self
 
     from rydstate.angular.angular_matrix_element import AngularMomentumQuantumNumbers, AngularOperatorType
     from rydstate.angular.angular_state import AngularState
+    from rydstate.angular.utils import CouplingScheme
 
 logger = logging.getLogger(__name__)
 
-CouplingScheme = Literal["LS", "JJ", "FJ"]
-
-
-class InvalidQuantumNumbersError(ValueError):
-    def __init__(self, ket: AngularKetBase, msg: str = "") -> None:
-        _msg = f"Invalid quantum numbers for {ket!r}"
-        if len(msg) > 0:
-            _msg += f"\n  {msg}"
-        super().__init__(_msg)
-
 
 class AngularKetBase(ABC):
     """Base class for a angular ket (i.e. a simple canonical spin ketstate)."""
@@ -224,6 +214,9 @@ def to_state(self, coupling_scheme: Literal["JJ"]) -> AngularState[AngularKetJJ]
     @overload
     def to_state(self, coupling_scheme: Literal["FJ"]) -> AngularState[AngularKetFJ]: ...
 
+    @overload
+    def to_state(self, coupling_scheme: Literal["Dummy"]) -> Never: ...
+
     @overload
     def to_state(self: Self) -> AngularState[Self]: ...
 
@@ -737,53 +730,3 @@ def sanity_check(self, msgs: list[str] | None = None) -> None:
             msgs.append(f"{self.f_c=}, {self.j_r=}, {self.f_tot=} don't satisfy spin addition rule.")
 
         super().sanity_check(msgs)
-
-
-def quantum_numbers_to_angular_ket(
-    species: str | SpeciesObject,
-    s_c: float | None = None,
-    l_c: int = 0,
-    j_c: float | None = None,
-    f_c: float | None = None,
-    s_r: float = 0.5,
-    l_r: int | None = None,
-    j_r: float | None = None,
-    s_tot: float | None = None,
-    l_tot: int | None = None,
-    j_tot: float | None = None,
-    f_tot: float | None = None,
-    m: float | None = None,
-) -> AngularKetBase:
-    r"""Return an AngularKet object in the corresponding coupling scheme from the given quantum numbers.
-
-    Args:
-        species: Atomic species.
-        s_c: Spin quantum number of the core electron (0 for Alkali, 0.5 for divalent atoms).
-        l_c: Orbital angular momentum quantum number of the core electron.
-        j_c: Total angular momentum quantum number of the core electron.
-        f_c: Total angular momentum quantum number of the core (core electron + nucleus).
-        s_r: Spin quantum number of the rydberg electron always 0.5)
-        l_r: Orbital angular momentum quantum number of the rydberg electron.
-        j_r: Total angular momentum quantum number of the rydberg electron.
-        s_tot: Total spin quantum number of all electrons.
-        l_tot: Total orbital angular momentum quantum number of all electrons.
-        j_tot: Total angular momentum quantum number of all electrons.
-        f_tot: Total angular momentum quantum number of the atom (rydberg electron + core)
-        m: Total magnetic quantum number.
-          Optional, only needed for concrete angular matrix elements.
-
-    """
-    if all(qn is None for qn in [j_c, f_c, j_r]):
-        return AngularKetLS(
-            s_c=s_c, l_c=l_c, s_r=s_r, l_r=l_r, s_tot=s_tot, l_tot=l_tot, j_tot=j_tot, f_tot=f_tot, m=m, species=species
-        )
-    if all(qn is None for qn in [s_tot, l_tot, f_c]):
-        return AngularKetJJ(
-            s_c=s_c, l_c=l_c, j_c=j_c, s_r=s_r, l_r=l_r, j_r=j_r, j_tot=j_tot, f_tot=f_tot, m=m, species=species
-        )
-    if all(qn is None for qn in [s_tot, l_tot, j_tot]):
-        return AngularKetFJ(
-            s_c=s_c, l_c=l_c, j_c=j_c, f_c=f_c, s_r=s_r, l_r=l_r, j_r=j_r, f_tot=f_tot, m=m, species=species
-        )
-
-    raise ValueError("Invalid combination of angular quantum numbers provided.")

src/rydstate/angular/angular_ket_dummy.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, ClassVar
+
+from rydstate.angular.angular_ket import AngularKetBase
+from rydstate.angular.angular_matrix_element import is_angular_operator_type
+from rydstate.angular.utils import InvalidQuantumNumbersError
+
+if TYPE_CHECKING:
+    from typing_extensions import Self
+
+    from rydstate.angular.angular_matrix_element import AngularOperatorType
+
+logger = logging.getLogger(__name__)
+
+
+class AngularKetDummy(AngularKetBase):
+    """Dummy spin ket for unknown quantum numbers."""
+
+    __slots__ = ("name",)
+    quantum_number_names: ClassVar = ("f_tot",)
+    coupled_quantum_numbers: ClassVar = {}
+    coupling_scheme = "Dummy"
+
+    name: str
+    """Name of the dummy ket."""
+
+    def __init__(
+        self,
+        name: str,
+        f_tot: float,
+        m: float | None = None,
+    ) -> None:
+        """Initialize the Spin ket."""
+        self.name = name
+
+        self.f_tot = f_tot
+        self.m = m
+
+        super()._post_init()
+
+    def sanity_check(self, msgs: list[str] | None = None) -> None:
+        """Check that the quantum numbers are valid."""
+        msgs = msgs if msgs is not None else []
+
+        if self.m is not None and not -self.f_tot <= self.m <= self.f_tot:
+            msgs.append(f"m must be between -f_tot and f_tot, but {self.f_tot=}, {self.m=}")
+
+        if msgs:
+            msg = "\n  ".join(msgs)
+            raise InvalidQuantumNumbersError(self, msg)
+
+    def __repr__(self) -> str:
+        args = f"{self.name}, f_tot={self.f_tot}"
+        if self.m is not None:
+            args += f", m={self.m}"
+        return f"{self.__class__.__name__}({args})"
+
+    def __str__(self) -> str:
+        return self.__repr__().replace("AngularKet", "")
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, AngularKetBase):
+            raise NotImplementedError(f"Cannot compare {self!r} with {other!r}.")
+        if not isinstance(other, AngularKetDummy):
+            return False
+        return self.name == other.name and self.f_tot == other.f_tot and self.m == other.m
+
+    def __hash__(self) -> int:
+        return hash((self.name, self.f_tot, self.m))
+
+    def calc_reduced_overlap(self, other: AngularKetBase) -> float:
+        return int(self == other)
+
+    def calc_reduced_matrix_element(
+        self: Self,
+        other: AngularKetBase,  # noqa: ARG002
+        operator: AngularOperatorType,
+        kappa: int,  # noqa: ARG002
+    ) -> float:
+        if not is_angular_operator_type(operator):
+            raise NotImplementedError(f"calc_reduced_matrix_element is not implemented for operator {operator}.")
+
+        # ignore contributions from dummy kets
+        return 0

src/rydstate/angular/angular_matrix_element.py

@@ -7,7 +7,8 @@
 import numpy as np
 from typing_extensions import TypeGuard
 
-from rydstate.angular.utils import calc_wigner_3j, calc_wigner_6j, minus_one_pow
+from rydstate.angular.utils import minus_one_pow
+from rydstate.angular.wigner_symbols import calc_wigner_3j, calc_wigner_6j
 
 if TYPE_CHECKING:
     from typing_extensions import ParamSpec

src/rydstate/angular/angular_state.py

@@ -12,15 +12,16 @@
     AngularKetJJ,
     AngularKetLS,
 )
+from rydstate.angular.angular_ket_dummy import AngularKetDummy
 from rydstate.angular.angular_matrix_element import is_angular_momentum_quantum_number
 
 if TYPE_CHECKING:
     from collections.abc import Iterator, Sequence
 
-    from typing_extensions import Self
+    from typing_extensions import Never, Self
 
-    from rydstate.angular.angular_ket import CouplingScheme
     from rydstate.angular.angular_matrix_element import AngularMomentumQuantumNumbers, AngularOperatorType
+    from rydstate.angular.utils import CouplingScheme
 
 logger = logging.getLogger(__name__)
 
@@ -32,21 +33,28 @@ class AngularState(Generic[_AngularKet]):
     def __init__(
         self, coefficients: Sequence[float], kets: Sequence[_AngularKet], *, warn_if_not_normalized: bool = True
     ) -> None:
-        self.coefficients = np.array(coefficients)
-        self.kets = kets
+        """Initialize an angular state as a linear combination of angular kets.
+
+        All kets must be of the same type (coupling scheme), and no duplicate kets are allowed.
+        Dummy kets (AngularKetDummy) are ignored in the state representation,
+        however adding them is recommended for normalization purposes.
+        """
+        self._coefficients = np.array(coefficients)
+        self._kets = kets
+        self._warn_if_not_normalized = warn_if_not_normalized
 
         if len(coefficients) != len(kets):
             raise ValueError("Length of coefficients and kets must be the same.")
         if len(kets) == 0:
             raise ValueError("At least one ket must be provided.")
-        if not all(type(ket) is type(kets[0]) for ket in kets):
+        if not all(type(ket) is type(self.kets[0]) for ket in self.kets):
             raise ValueError("All kets must be of the same type.")
-        if len(set(kets)) != len(kets):
-            raise ValueError("AngularState initialized with duplicate kets.")
+        if len(set(self.kets)) != len(self.kets):
+            raise ValueError("AngularState initialized with duplicate kets: %s", self.kets)
         if abs(self.norm - 1) > 1e-10 and warn_if_not_normalized:
             logger.warning("AngularState initialized with non-normalized coefficients: %s, %s", coefficients, kets)
         if self.norm > 1:
-            self.coefficients /= self.norm
+            self._coefficients /= self.norm
 
     def __iter__(self) -> Iterator[tuple[float, _AngularKet]]:
         return zip(self.coefficients, self.kets).__iter__()
@@ -59,6 +67,16 @@ def __str__(self) -> str:
         terms = [f"{coeff}*{ket!s}" for coeff, ket in self]
         return f"{', '.join(terms)}"
 
+    @property
+    def kets(self) -> list[_AngularKet]:
+        return [ket for ket in self._kets if not isinstance(ket, AngularKetDummy)]
+
+    @property
+    def coefficients(self) -> np.ndarray:
+        return np.array(
+            [coeff for coeff, ket in zip(self._coefficients, self._kets) if not isinstance(ket, AngularKetDummy)]
+        )
+
     @property
     def coupling_scheme(self) -> CouplingScheme:
         """Return the coupling scheme of the state."""
@@ -67,7 +85,7 @@ def coupling_scheme(self) -> CouplingScheme:
     @property
     def norm(self) -> float:
         """Return the norm of the state (should be 1)."""
-        return np.linalg.norm(self.coefficients)  # type: ignore [return-value]
+        return np.linalg.norm(self._coefficients)  # type: ignore [return-value]
 
     @overload
     def to(self, coupling_scheme: Literal["LS"]) -> AngularState[AngularKetLS]: ...
@@ -78,6 +96,9 @@ def to(self, coupling_scheme: Literal["JJ"]) -> AngularState[AngularKetJJ]: ...
     @overload
     def to(self, coupling_scheme: Literal["FJ"]) -> AngularState[AngularKetFJ]: ...
 
+    @overload
+    def to(self, coupling_scheme: Literal["Dummy"]) -> Never: ...
+
     def to(self, coupling_scheme: CouplingScheme) -> AngularState[Any]:
         """Convert to specified coupling scheme.
 
@@ -98,7 +119,8 @@ def to(self, coupling_scheme: CouplingScheme) -> AngularState[Any]:
                 else:
                     kets.append(scheme_ket)
                     coefficients.append(coeff * scheme_coeff)
-        return AngularState(coefficients, kets, warn_if_not_normalized=abs(self.norm - 1) < 1e-10)
+        warn_if_not_normalized = self._warn_if_not_normalized and (abs(self.norm - 1) < 1e-10)
+        return AngularState(coefficients, kets, warn_if_not_normalized=warn_if_not_normalized)
 
     def calc_exp_qn(self, q: AngularMomentumQuantumNumbers) -> float:
         """Calculate the expectation value of a quantum number q.