Add more coupling schemes

src/rydstate/__init__.py

@@ -1,18 +1,22 @@
 from rydstate import angular, radial, species
-from rydstate.basis import BasisSQDTAlkali, BasisSQDTAlkalineLS
+from rydstate.basis import BasisSQDTAlkali, BasisSQDTAlkalineFJ, BasisSQDTAlkalineJJ, BasisSQDTAlkalineLS
 from rydstate.rydberg import (
     RydbergStateSQDT,
     RydbergStateSQDTAlkali,
+    RydbergStateSQDTAlkalineFJ,
     RydbergStateSQDTAlkalineJJ,
     RydbergStateSQDTAlkalineLS,
 )
 from rydstate.units import ureg
 
 __all__ = [
     "BasisSQDTAlkali",
+    "BasisSQDTAlkalineFJ",
+    "BasisSQDTAlkalineJJ",
     "BasisSQDTAlkalineLS",
     "RydbergStateSQDT",
     "RydbergStateSQDTAlkali",
+    "RydbergStateSQDTAlkalineFJ",
     "RydbergStateSQDTAlkalineJJ",
     "RydbergStateSQDTAlkalineLS",
     "angular",

src/rydstate/basis/__init__.py

@@ -1,3 +1,3 @@
-from rydstate.basis.basis_sqdt import BasisSQDTAlkali, BasisSQDTAlkalineLS
+from rydstate.basis.basis_sqdt import BasisSQDTAlkali, BasisSQDTAlkalineFJ, BasisSQDTAlkalineJJ, BasisSQDTAlkalineLS
 
-__all__ = ["BasisSQDTAlkali", "BasisSQDTAlkalineLS"]
+__all__ = ["BasisSQDTAlkali", "BasisSQDTAlkalineFJ", "BasisSQDTAlkalineJJ", "BasisSQDTAlkalineLS"]

src/rydstate/basis/basis_sqdt.py

@@ -1,9 +1,18 @@
 from __future__ import annotations
 
+import logging
+
 import numpy as np
 
 from rydstate.basis.basis_base import BasisBase
-from rydstate.rydberg import RydbergStateSQDTAlkali, RydbergStateSQDTAlkalineLS
+from rydstate.rydberg import (
+    RydbergStateSQDTAlkali,
+    RydbergStateSQDTAlkalineFJ,
+    RydbergStateSQDTAlkalineJJ,
+    RydbergStateSQDTAlkalineLS,
+)
+
+logger = logging.getLogger(__name__)
 
 
 class BasisSQDTAlkali(BasisBase[RydbergStateSQDTAlkali]):
@@ -48,3 +57,63 @@ def __init__(self, species: str, n_min: int = 1, n_max: int | None = None) -> No
                                 species, n=n, l=l, s_tot=s_tot, j_tot=j_tot, f_tot=float(f_tot)
                             )
                             self.states.append(state)
+
+
+class BasisSQDTAlkalineJJ(BasisBase[RydbergStateSQDTAlkalineJJ]):
+    def __init__(self, species: str, n_min: int = 0, n_max: int | None = None) -> None:
+        super().__init__(species)
+
+        if n_max is None:
+            raise ValueError("n_max must be given")
+
+        i_c = self.species.i_c if self.species.i_c is not None else 0
+        j_c = 0.5
+        s_r = 0.5
+        self.states = []
+        for n in range(n_min, n_max + 1):
+            for l_r in range(n):
+                if self.species.is_allowed_shell(n, l_r, 0) != self.species.is_allowed_shell(n, l_r, 1):
+                    logger.warning(
+                        "For l=%d, n=%d one of the singlet/triplet states is not allowed. "
+                        "In JJ coupling the state does not exist, thus skipping this shell",
+                        *(l_r, n),
+                    )
+                if not all(self.species.is_allowed_shell(n, l_r, s_tot) for s_tot in [0, 1]):
+                    continue
+                for j_r in np.arange(abs(l_r - s_r), l_r + s_r + 1):
+                    for j_tot in range(int(abs(j_r - j_c)), int(j_r + j_c + 1)):
+                        for f_tot in np.arange(abs(j_tot - i_c), j_tot + i_c + 1):
+                            state = RydbergStateSQDTAlkalineJJ(
+                                species, n=n, l=l_r, j_r=float(j_r), j_tot=j_tot, f_tot=float(f_tot)
+                            )
+                            self.states.append(state)
+
+
+class BasisSQDTAlkalineFJ(BasisBase[RydbergStateSQDTAlkalineFJ]):
+    def __init__(self, species: str, n_min: int = 0, n_max: int | None = None) -> None:
+        super().__init__(species)
+
+        if n_max is None:
+            raise ValueError("n_max must be given")
+
+        i_c = self.species.i_c if self.species.i_c is not None else 0
+        j_c = 0.5
+        s_r = 0.5
+        self.states = []
+        for n in range(n_min, n_max + 1):
+            for l_r in range(n):
+                if self.species.is_allowed_shell(n, l_r, 0) != self.species.is_allowed_shell(n, l_r, 1):
+                    logger.warning(
+                        "For l=%d, n=%d one of the singlet/triplet states is not allowed. "
+                        "In FJ coupling the state does not exist, thus skipping this shell",
+                        *(l_r, n),
+                    )
+                if not all(self.species.is_allowed_shell(n, l_r, s_tot) for s_tot in [0, 1]):
+                    continue
+                for j_r in np.arange(abs(l_r - s_r), l_r + s_r + 1):
+                    for f_c in np.arange(abs(j_c - i_c), j_c + i_c + 1):
+                        for f_tot in np.arange(abs(f_c - j_r), f_c + j_r + 1):
+                            state = RydbergStateSQDTAlkalineFJ(
+                                species, n=n, l=l_r, j_r=float(j_r), f_c=float(f_c), f_tot=float(f_tot)
+                            )
+                            self.states.append(state)

src/rydstate/rydberg/__init__.py

@@ -1,13 +1,15 @@
 from rydstate.rydberg.rydberg_sqdt import (
     RydbergStateSQDT,
     RydbergStateSQDTAlkali,
+    RydbergStateSQDTAlkalineFJ,
     RydbergStateSQDTAlkalineJJ,
     RydbergStateSQDTAlkalineLS,
 )
 
 __all__ = [
     "RydbergStateSQDT",
     "RydbergStateSQDTAlkali",
+    "RydbergStateSQDTAlkalineFJ",
     "RydbergStateSQDTAlkalineJJ",
     "RydbergStateSQDTAlkalineLS",
 ]

src/rydstate/rydberg/rydberg_sqdt.py

@@ -15,7 +15,7 @@
 from rydstate.units import BaseQuantities, MatrixElementOperatorRanks, ureg
 
 if TYPE_CHECKING:
-    from rydstate.angular.angular_ket import AngularKetBase, AngularKetJJ, AngularKetLS
+    from rydstate.angular.angular_ket import AngularKetBase, AngularKetFJ, AngularKetJJ, AngularKetLS
     from rydstate.units import MatrixElementOperator, PintFloat
 
 
@@ -418,11 +418,74 @@ def nu(self) -> float:
         if self._nu is not None:
             return self._nu
         assert self.n is not None
-        nu_singlet = self.species.calc_nu(self.n, self.l, self.j_tot, s_tot=0)
-        nu_triplet = self.species.calc_nu(self.n, self.l, self.j_tot, s_tot=1)
-        if abs(nu_singlet - nu_triplet) > 1e-10:
+        nus = [self.species.calc_nu(self.n, self.l, self.j_tot, s_tot=s_tot) for s_tot in [0, 1]]
+
+        if any(abs(nu - nus[0]) > 1e-10 for nu in nus[1:]):
             raise ValueError(
                 "RydbergStateSQDTAlkalineJJ is intended for high-l states only, "
                 "where the quantum defects are the same for singlet and triplet states."
             )
-        return nu_singlet
+        return nus[0]
+
+
+class RydbergStateSQDTAlkalineFJ(RydbergStateSQDT):
+    """Create an Alkaline Rydberg state, including the radial and angular states."""
+
+    angular: AngularKetFJ
+
+    def __init__(
+        self,
+        species: str | SpeciesObject,
+        n: int,
+        l: int,
+        j_r: float,
+        f_c: float | None = None,
+        f_tot: float | None = None,
+        m: float | None = None,
+        nu: float | None = None,
+    ) -> None:
+        r"""Initialize the Rydberg state.
+
+        Args:
+            species: Atomic species.
+            n: Principal quantum number of the rydberg electron.
+            l: Orbital angular momentum quantum number of the rydberg electron.
+            j_r: Total angular momentum quantum number of the Rydberg electron.
+            f_c: Total angular momentum quantum number of the core (core electron + nucleus).
+            f_tot: Total angular momentum quantum number of the atom (rydberg electron + core)
+              Optional, only needed if the species supports hyperfine structure (i.e. species.i_c is not None or 0).
+            m: Total magnetic quantum number.
+              Optional, only needed for concrete angular matrix elements.
+            nu: Effective principal quantum number of the rydberg electron.
+              Optional, if not given it will be calculated from n, l.
+
+        """
+        super().__init__(species=species, n=n, nu=nu, l_r=l, j_r=j_r, f_c=f_c, f_tot=f_tot, m=m)
+
+        self.l = self.angular.l_r
+        self.j_r = self.angular.j_r
+        self.f_c = self.angular.f_c
+        self.f_tot = self.angular.f_tot
+        self.m = self.angular.m
+
+    def __repr__(self) -> str:
+        species, n, l, j_r, f_c, f_tot, m = self.species, self.n, self.l, self.j_r, self.f_c, self.f_tot, self.m
+        return f"{self.__class__.__name__}({species.name}, {n=}, {l=}, {j_r=}, {f_c=}, {f_tot=}, {m=})"
+
+    @cached_property
+    def nu(self) -> float:
+        if self._nu is not None:
+            return self._nu
+        assert self.n is not None
+        nus = [
+            self.species.calc_nu(self.n, self.l, float(j_tot), s_tot=s_tot)
+            for s_tot in [0, 1]
+            for j_tot in np.arange(abs(self.j_r - 1 / 2), self.j_r + 1 / 2 + 1)
+        ]
+
+        if any(abs(nu - nus[0]) > 1e-10 for nu in nus[1:]):
+            raise ValueError(
+                "RydbergStateSQDTAlkalineFJ is intended for high-l states only, "
+                "where the quantum defects are the same for singlet and triplet states."
+            )
+        return nus[0]

tests/test_basis.py

@@ -1,7 +1,11 @@
+from typing import TYPE_CHECKING, Any
+
 import numpy as np
 import pytest
-from rydstate import BasisSQDTAlkali
-from rydstate.basis.basis_sqdt import BasisSQDTAlkalineLS
+from rydstate import BasisSQDTAlkali, BasisSQDTAlkalineFJ, BasisSQDTAlkalineJJ, BasisSQDTAlkalineLS
+
+if TYPE_CHECKING:
+    from rydstate.basis.basis_base import BasisBase
 
 
 @pytest.mark.parametrize("species_name", ["Rb", "Na", "H"])
@@ -62,3 +66,13 @@ def test_alkaline_basis(species_name: str) -> None:
     me_matrix = basis.calc_reduced_matrix_elements(basis, "electric_dipole", unit="e a0")
     assert np.shape(me_matrix) == (len(basis.states), len(basis.states))
     assert np.count_nonzero(me_matrix) > 0
+
+    basis = BasisSQDTAlkalineLS(species_name, n_min=30, n_max=35)
+    basis.filter_states("l_r", (6, 10))
+    for basis_class in [BasisSQDTAlkalineJJ, BasisSQDTAlkalineFJ]:
+        basis2: BasisBase[Any] = basis_class(species_name, n_min=30, n_max=35)  # type: ignore [assignment]
+        basis2.filter_states("l_r", (6, 10))
+        assert len(basis2.states) == len(basis.states)
+        trafo = basis.calc_reduced_overlaps(basis2)
+        trafo_inv = basis2.calc_reduced_overlaps(basis)
+        assert np.allclose(trafo @ trafo_inv, np.eye(len(basis.states)), atol=1e-3)