move stuff from angular ket to angular utils

Author: johannes-moegerle

src/rydstate/angular/angular_ket.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import contextlib
 import logging
 from abc import ABC
 from typing import TYPE_CHECKING, Any, ClassVar, Literal, overload
@@ -14,6 +13,7 @@
     is_angular_operator_type,
 )
 from rydstate.angular.utils import (
+    InvalidQuantumNumbersError,
     check_spin_addition_rule,
     get_possible_quantum_number_values,
     minus_one_pow,
@@ -27,19 +27,10 @@
 
     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)."""
@@ -736,55 +727,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.
-
-    """
-    with contextlib.suppress(InvalidQuantumNumbersError, ValueError):
-        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
-        )
-
-    with contextlib.suppress(InvalidQuantumNumbersError, ValueError):
-        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
-        )
-
-    with contextlib.suppress(InvalidQuantumNumbersError, ValueError):
-        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_state.py

@@ -19,8 +19,8 @@
 
     from typing_extensions import 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__)
 

src/rydstate/angular/utils.py

@@ -1,7 +1,24 @@
 from __future__ import annotations
 
+import contextlib
+from typing import TYPE_CHECKING, Literal
+
 import numpy as np
 
+if TYPE_CHECKING:
+    from rydstate.angular.angular_ket import AngularKetBase
+    from rydstate.species.species_object import SpeciesObject
+
+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)
+
 
 def minus_one_pow(n: float) -> int:
     """Calculate (-1)^n for an integer n and raise an error if n is not an integer."""
@@ -42,3 +59,57 @@ def get_possible_quantum_number_values(s_1: float, s_2: float, s_tot: float | No
     if s_tot is not None:
         return [s_tot]
     return [float(s) for s in np.arange(abs(s_1 - s_2), s_1 + s_2 + 1, 1)]
+
+
+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.
+
+    """
+    from rydstate.angular.angular_ket import AngularKetFJ, AngularKetJJ, AngularKetLS  # noqa: PLC0415
+
+    with contextlib.suppress(InvalidQuantumNumbersError, ValueError):
+        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
+        )
+
+    with contextlib.suppress(InvalidQuantumNumbersError, ValueError):
+        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
+        )
+
+    with contextlib.suppress(InvalidQuantumNumbersError, ValueError):
+        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/rydberg/rydberg_sqdt.py

@@ -7,7 +7,7 @@
 
 import numpy as np
 
-from rydstate.angular.angular_ket import quantum_numbers_to_angular_ket
+from rydstate.angular.utils import quantum_numbers_to_angular_ket
 from rydstate.radial import RadialKet
 from rydstate.rydberg.rydberg_base import RydbergStateBase
 from rydstate.species import SpeciesObject

tests/test_angular_matrix_elements.py

@@ -8,8 +8,9 @@
 from rydstate.angular.angular_matrix_element import AngularMomentumQuantumNumbers
 
 if TYPE_CHECKING:
-    from rydstate.angular.angular_ket import AngularKetBase, CouplingScheme
+    from rydstate.angular.angular_ket import AngularKetBase
     from rydstate.angular.angular_matrix_element import AngularOperatorType
+    from rydstate.angular.utils import CouplingScheme
 
 TEST_KET_PAIRS = [
     (