add BasisSQDTAlkalineJJ and BasisSQDTAlkalineFJ

Author: johannes-moegerle

src/rydstate/__init__.py

@@ -1,5 +1,5 @@
 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,
@@ -11,6 +11,8 @@
 
 __all__ = [
     "BasisSQDTAlkali",
+    "BasisSQDTAlkalineFJ",
+    "BasisSQDTAlkalineJJ",
     "BasisSQDTAlkalineLS",
     "RydbergStateSQDT",
     "RydbergStateSQDTAlkali",

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)