improve angular utils and angular init

Author: johannes-moegerle

src/ryd_numerov/angular/__init__.py

@@ -1,30 +1,9 @@
-from ryd_numerov.angular.angular_ket import AngularKetBase, AngularKetFJ, AngularKetJJ, AngularKetLS
-from ryd_numerov.angular.angular_matrix_element import (
-    calc_prefactor_of_operator_in_coupled_scheme,
-    calc_reduced_spherical_matrix_element,
-    calc_reduced_spin_matrix_element,
-)
+from ryd_numerov.angular.angular_ket import AngularKetFJ, AngularKetJJ, AngularKetLS
 from ryd_numerov.angular.angular_state import AngularState
-from ryd_numerov.angular.utils import (
-    calc_wigner_3j,
-    calc_wigner_6j,
-    calc_wigner_9j,
-    clebsch_gordan_6j,
-    clebsch_gordan_9j,
-)
 
 __all__ = [
-    "AngularKetBase",
     "AngularKetFJ",
     "AngularKetJJ",
     "AngularKetLS",
     "AngularState",
-    "calc_prefactor_of_operator_in_coupled_scheme",
-    "calc_reduced_spherical_matrix_element",
-    "calc_reduced_spin_matrix_element",
-    "calc_wigner_3j",
-    "calc_wigner_6j",
-    "calc_wigner_9j",
-    "clebsch_gordan_6j",
-    "clebsch_gordan_9j",
 ]

src/ryd_numerov/angular/angular_ket.py

@@ -18,6 +18,7 @@
     clebsch_gordan_6j,
     clebsch_gordan_9j,
     get_possible_quantum_number_list,
+    minus_one_pow,
     try_trivial_spin_addition,
 )
 from ryd_numerov.elements import BaseElement
@@ -449,7 +450,7 @@ def calc_matrix_element(self, other: AngularKetBase, operator: AngularOperatorTy
 
     def _calc_wigner_eckart_prefactor(self, other: AngularKetBase, kappa: int, q: int) -> float:
         assert self.m is not None and other.m is not None, "m must be set to calculate the Wigner-Eckart prefactor."  # noqa: PT018
-        return (-1) ** (self.f_tot - self.m) * calc_wigner_3j(self.f_tot, kappa, other.f_tot, -self.m, q, other.m)  # type: ignore [return-value]
+        return minus_one_pow(self.f_tot - self.m) * calc_wigner_3j(self.f_tot, kappa, other.f_tot, -self.m, q, other.m)
 
     def _kronecker_delta_non_involved_spins(self, other: AngularKetBase, qn: AngularMomentumQuantumNumbers) -> int:
         """Calculate the Kronecker delta for non involved angular momentum quantum numbers.

src/ryd_numerov/angular/angular_matrix_element.py

@@ -76,7 +76,8 @@ def calc_reduced_spherical_matrix_element(l_r_final: int, l_r_initial: int, kapp
 def calc_reduced_spin_matrix_element(s_final: float, s_initial: float) -> float:
     r"""Calculate the reduced spin matrix element (s_final || \hat{s} || s_initial).
 
-    The spin operator \hat{s} must be the operator corresponding to the quantum number s_final and s_initial.
+    The spin operator \hat{s} can be any of the AngularMomentumQuantumNumbers,
+    but must be corresponding to the given quantum number s_final and s_initial.
 
     We follow the convention of equation (5.4.3) from Edmonds 1985 "Angular Momentum in Quantum Mechanics".
     The matrix elements of the spin operators are given by:

src/ryd_numerov/angular/utils.py

@@ -21,11 +21,16 @@
     def lru_cache(maxsize: int) -> Callable[[Callable[P, R]], Callable[P, R]]: ...  # type: ignore [no-redef]
 
 
+# global variables to possibly improve the performance of wigner j calculations
+# in the public release we will always use CHECK_ARGS = True and USE_SYMMETRIES = False to reduce potential of bugs
+CHECK_ARGS = True
 USE_SYMMETRIES = False
 
 
 def sympify_args(func: Callable[P, R]) -> Callable[P, R]:
     """Check that quantum numbers are valid and convert to sympy.Integer (and half-integer)."""
+    if not CHECK_ARGS:
+        return func
 
     def check_arg(arg: float) -> Integer:
         if arg.is_integer():
@@ -43,8 +48,86 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
     return wrapper
 
 
+@lru_cache(maxsize=10_000)
+@sympify_args
 def calc_wigner_3j(j1: float, j2: float, j3: float, m1: float, m2: float, m3: float) -> float:
-    """Calculate the Wigner 3j symbol using symmetries and lru_cache to improve performance."""
+    """Calculate the Wigner 3j symbol using lru_cache to improve performance."""
+    return float(sympy_wigner_3j(j1, j2, j3, m1, m2, m3).evalf())
+
+
+@lru_cache(maxsize=100_000)
+@sympify_args
+def calc_wigner_6j(j1: float, j2: float, j3: float, j4: float, j5: float, j6: float) -> float:
+    """Calculate the Wigner 6j symbol using lru_cache to improve performance."""
+    return float(sympy_wigner_6j(j1, j2, j3, j4, j5, j6).evalf())
+
+
+@lru_cache(maxsize=10_000)
+@sympify_args
+def calc_wigner_9j(
+    j1: float, j2: float, j3: float, j4: float, j5: float, j6: float, j7: float, j8: float, j9: float
+) -> float:
+    """Calculate the Wigner 9j symbol using lru_cache to improve performance."""
+    return float(sympy_wigner_9j(j1, j2, j3, j4, j5, j6, j7, j8, j9).evalf())
+
+
+def clebsch_gordan_6j(j1: float, j2: float, j3: float, j12: float, j23: float, j_tot: float) -> float:
+    """Calculate the overlap between <((j1,j2)j12,j3)j_tot|(j1,(j2,j3)j23)j_tot>.
+
+    We follow the convention of equation (6.1.5) from Edmonds 1985 "Angular Momentum in Quantum Mechanics".
+
+    See Also:
+        - https://en.wikipedia.org/wiki/Racah_W-coefficient
+        - https://en.wikipedia.org/wiki/6-j_symbol
+
+    Args:
+        j1: Spin quantum number 1.
+        j2: Spin quantum number 2.
+        j3: Spin quantum number 3.
+        j12: Total spin quantum number of j1 + j2.
+        j23: Total spin quantum number of j2 + j3.
+        j_tot: Total spin quantum number of j1 + j2 + j3.
+
+    Returns:
+        The Clebsch-Gordan coefficient <((j1,j2)j12,j3)j_tot|(j1,(j2,j3)j23)j_tot>.
+
+    """
+    prefactor = minus_one_pow(j1 + j2 + j3 + j_tot) * math.sqrt((2 * j12 + 1) * (2 * j23 + 1))
+    wigner_6j = calc_wigner_6j(j1, j2, j12, j3, j_tot, j23)
+    return prefactor * wigner_6j
+
+
+def clebsch_gordan_9j(
+    j1: float, j2: float, j12: float, j3: float, j4: float, j34: float, j13: float, j24: float, j_tot: float
+) -> float:
+    """Calculate the overlap between <((j1,j2)j12,(j3,j4)j34))j_tot|((j1,j3)j13,(j2,j4)j24))j_tot>.
+
+    We follow the convention of equation (6.4.2) from Edmonds 1985 "Angular Momentum in Quantum Mechanics".
+
+    See Also:
+        - https://en.wikipedia.org/wiki/9-j_symbol
+
+    Args:
+        j1: Spin quantum number 1.
+        j2: Spin quantum number 2.
+        j12: Total spin quantum number of j1 + j2.
+        j3: Spin quantum number 1.
+        j4: Spin quantum number 2.
+        j34: Total spin quantum number of j1 + j2.
+        j13: Total spin quantum number of j1 + j3.
+        j24: Total spin quantum number of j2 + j4.
+        j_tot: Total spin quantum number of j1 + j2 + j3 + j4.
+
+    Returns:
+        The Clebsch-Gordan coefficient <((j1,j2)j12,(j3,j4)j34))j_tot|((j1,j3)j13,(j2,j4)j24))j_tot>.
+
+    """
+    prefactor = math.sqrt((2 * j12 + 1) * (2 * j34 + 1) * (2 * j13 + 1) * (2 * j24 + 1))
+    return prefactor * calc_wigner_9j(j1, j2, j12, j3, j4, j34, j13, j24, j_tot)
+
+
+def calc_wigner_3j_with_symmetries(j1: float, j2: float, j3: float, m1: float, m2: float, m3: float) -> float:
+    """Calculate the Wigner 3j symbol using symmetries to reduce the number of symbols, that are not cached."""
     symmetry_factor: float = 1
 
     # even permutation -> sort smallest j to be j1
@@ -65,17 +148,11 @@ def calc_wigner_3j(j1: float, j2: float, j3: float, m1: float, m2: float, m3: fl
 
     # TODO Regge symmetries
 
-    return symmetry_factor * _calc_wigner_3j(j1, j2, j3, m1, m2, m3)
+    return symmetry_factor * calc_wigner_3j(j1, j2, j3, m1, m2, m3)
 
 
-@lru_cache(maxsize=10_000)
-@sympify_args
-def _calc_wigner_3j(j1: float, j2: float, j3: float, m1: float, m2: float, m3: float) -> float:
-    return float(sympy_wigner_3j(j1, j2, j3, m1, m2, m3).evalf())
-
-
-def calc_wigner_6j(j1: float, j2: float, j3: float, j4: float, j5: float, j6: float) -> float:
-    """Calculate the Wigner 6j symbol using symmetries and lru_cache to improve performance."""
+def calc_wigner_6j_with_symmetries(j1: float, j2: float, j3: float, j4: float, j5: float, j6: float) -> float:
+    """Calculate the Wigner 6j symbol using symmetries to reduce the number of symbols, that are not cached."""
     # interchange upper and lower for 2 columns -> make j1 < j4 and j2 < j5
     if j4 < j1:
         j1, j2, j3, j4, j5, j6 = j4, j2, j6, j1, j5, j3  # noqa: PLW0127
@@ -91,19 +168,13 @@ def calc_wigner_6j(j1: float, j2: float, j3: float, j4: float, j5: float, j6: fl
     if j3 < j2:
         j1, j2, j3, j4, j5, j6 = j1, j3, j2, j4, j6, j5  # noqa: PLW0127
 
-    return _calc_wigner_6j(j1, j2, j3, j4, j5, j6)
-
+    return calc_wigner_6j(j1, j2, j3, j4, j5, j6)
 
-@lru_cache(maxsize=100_000)
-@sympify_args
-def _calc_wigner_6j(j1: float, j2: float, j3: float, j4: float, j5: float, j6: float) -> float:
-    return float(sympy_wigner_6j(j1, j2, j3, j4, j5, j6).evalf())
 
-
-def calc_wigner_9j(
+def calc_wigner_9j_with_symmetries(
     j1: float, j2: float, j3: float, j4: float, j5: float, j6: float, j7: float, j8: float, j9: float
 ) -> float:
-    """Calculate the Wigner 9j symbol using symmetries and lru_cache to improve performance."""
+    """Calculate the Wigner 9j symbol using symmetries to reduce the number of symbols, that are not cached."""
     symmetry_factor: float = 1
     js = [j1, j2, j3, j4, j5, j6, j7, j8, j9]
 
@@ -132,78 +203,22 @@ def calc_wigner_9j(
     if js[3] < js[1]:
         js = [js[0], js[3], js[6], js[1], js[4], js[7], js[2], js[5], js[8]]
 
-    return symmetry_factor * _calc_wigner_9j(*js)
-
-
-@lru_cache(maxsize=10_000)
-@sympify_args
-def _calc_wigner_9j(
-    j1: float, j2: float, j3: float, j4: float, j5: float, j6: float, j7: float, j8: float, j9: float
-) -> float:
-    return float(sympy_wigner_9j(j1, j2, j3, j4, j5, j6, j7, j8, j9).evalf())
-
-
-def clebsch_gordan_6j(j1: float, j2: float, j3: float, j12: float, j23: float, j_tot: float) -> float:
-    """Calculate the overlap between <((j1,j2)j12,j3)j_tot|(j1,(j2,j3)j23)j_tot>.
-
-    We follow the convention of equation (6.1.5) from Edmonds 1985 "Angular Momentum in Quantum Mechanics".
-
-    See Also:
-        - https://en.wikipedia.org/wiki/Racah_W-coefficient
-        - https://en.wikipedia.org/wiki/6-j_symbol
-
-    Args:
-        j1: Spin quantum number 1.
-        j2: Spin quantum number 2.
-        j3: Spin quantum number 3.
-        j12: Total spin quantum number of j1 + j2.
-        j23: Total spin quantum number of j2 + j3.
-        j_tot: Total spin quantum number of j1 + j2 + j3.
-
-    Returns:
-        The Clebsch-Gordan coefficient <((j1,j2)j12,j3)j_tot|(j1,(j2,j3)j23)j_tot>.
+    return symmetry_factor * calc_wigner_9j(*js)
 
-    """
-    prefactor = minus_one_pow(j1 + j2 + j3 + j_tot) * math.sqrt((2 * j12 + 1) * (2 * j23 + 1))
-    wigner_6j = calc_wigner_6j(j1, j2, j12, j3, j_tot, j23)
-    return prefactor * wigner_6j
 
-
-def clebsch_gordan_9j(
-    j1: float, j2: float, j12: float, j3: float, j4: float, j34: float, j13: float, j24: float, j_tot: float
-) -> float:
-    """Calculate the overlap between <((j1,j2)j12,(j3,j4)j34))j_tot|((j1,j3)j13,(j2,j4)j24))j_tot>.
-
-    We follow the convention of equation (6.4.2) from Edmonds 1985 "Angular Momentum in Quantum Mechanics".
-
-    See Also:
-        - https://en.wikipedia.org/wiki/9-j_symbol
-
-    Args:
-        j1: Spin quantum number 1.
-        j2: Spin quantum number 2.
-        j12: Total spin quantum number of j1 + j2.
-        j3: Spin quantum number 1.
-        j4: Spin quantum number 2.
-        j34: Total spin quantum number of j1 + j2.
-        j13: Total spin quantum number of j1 + j3.
-        j24: Total spin quantum number of j2 + j4.
-        j_tot: Total spin quantum number of j1 + j2 + j3 + j4.
-
-    Returns:
-        The Clebsch-Gordan coefficient <((j1,j2)j12,(j3,j4)j34))j_tot|((j1,j3)j13,(j2,j4)j24))j_tot>.
-
-    """
-    prefactor = math.sqrt((2 * j12 + 1) * (2 * j34 + 1) * (2 * j13 + 1) * (2 * j24 + 1))
-    return prefactor * calc_wigner_9j(j1, j2, j12, j3, j4, j34, j13, j24, j_tot)
+if USE_SYMMETRIES:
+    calc_wigner_3j = calc_wigner_3j_with_symmetries  # type: ignore [assignment]
+    calc_wigner_6j = calc_wigner_6j_with_symmetries  # type: ignore [assignment]
+    calc_wigner_9j = calc_wigner_9j_with_symmetries  # type: ignore [assignment]
 
 
 def minus_one_pow(n: float) -> int:
+    """Calculate (-1)^n for an integer n and raise an error if n is not an integer."""
     if n % 2 == 0:
         return 1
     if n % 2 == 1:
         return -1
-    raise ValueError(f"Invalid input {n}.")
+    raise ValueError(f"minus_one_pow: Invalid input {n=} is not an integer.")
 
 
 def try_trivial_spin_addition(s_1: float, s_2: float, s_tot: float | None, name: str) -> float:
@@ -235,9 +250,3 @@ def get_possible_quantum_number_list(s_1: float, s_2: float, s_tot: float | None
     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)]
-
-
-if not USE_SYMMETRIES:
-    calc_wigner_3j = _calc_wigner_3j
-    calc_wigner_6j = _calc_wigner_6j
-    calc_wigner_9j = _calc_wigner_9j

src/ryd_numerov/rydberg_state.py

@@ -17,7 +17,7 @@
 if TYPE_CHECKING:
     from typing_extensions import Self
 
-    from ryd_numerov.angular import AngularKetBase
+    from ryd_numerov.angular.angular_ket import AngularKetBase
     from ryd_numerov.units import PintFloat
 
 

tests/test_angular_matrix_elements.py

@@ -8,8 +8,7 @@
 from ryd_numerov.angular.angular_matrix_element import AngularMomentumQuantumNumbers
 
 if TYPE_CHECKING:
-    from ryd_numerov.angular import AngularKetBase
-    from ryd_numerov.angular.angular_ket import CouplingScheme
+    from ryd_numerov.angular.angular_ket import AngularKetBase, CouplingScheme
     from ryd_numerov.angular.angular_matrix_element import AngularOperatorType
 
 TEST_KET_PAIRS = [