speed up check is not set

Author: johannes-moegerle

src/rydstate/angular/angular_ket.py

@@ -17,6 +17,7 @@
     NotSet,
     check_spin_addition_rule,
     get_possible_quantum_number_values,
+    is_not_set,
     minus_one_pow,
     try_trivial_spin_addition,
 )
@@ -115,7 +116,7 @@ def __init__(
         self.l_r = int(l_r)
 
         # f_tot will be set in the subclasses
-        self.m = NotSet if isinstance(m, NotSet) else float(m)
+        self.m = NotSet if is_not_set(m) else float(m)
 
     def _post_init(self) -> None:
         self.quantum_numbers = tuple(getattr(self, qn) for qn in self.quantum_number_names)
@@ -133,7 +134,7 @@ def sanity_check(self, msgs: list[str] | None = None) -> None:
         if self.s_r != 0.5:
             msgs.append(f"Rydberg electron spin s_r must be 1/2, but {self.s_r=}")
 
-        if not isinstance(self.m, NotSet) and not -self.f_tot <= self.m <= self.f_tot:
+        if not is_not_set(self.m) 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:
@@ -150,7 +151,7 @@ def __setattr__(self, key: str, value: object) -> None:
 
     def __repr__(self) -> str:
         args = ", ".join(f"{qn}={val}" for qn, val in zip(self.quantum_number_names, self.quantum_numbers, strict=True))
-        if not isinstance(self.m, NotSet):
+        if not is_not_set(self.m):
             args += f", m={self.m}"
         return f"{self.__class__.__name__}({args})"
 
@@ -467,16 +468,16 @@ def calc_matrix_element(self, other: AngularKetBase, operator: AngularOperatorTy
             The dimensionless angular matrix element.
 
         """
-        if isinstance(self.m, NotSet) or isinstance(other.m, NotSet):
-            raise RuntimeError("m must be set to calculate the matrix element.")  # noqa: TRY004
+        if is_not_set(self.m) or is_not_set(other.m):
+            raise RuntimeError("m must be set to calculate the matrix element.")
 
         prefactor = self._calc_wigner_eckart_prefactor(other, kappa, q)
         reduced_matrix_element = self.calc_reduced_matrix_element(other, operator, kappa)
         return prefactor * reduced_matrix_element
 
     def _calc_wigner_eckart_prefactor(self, other: AngularKetBase, kappa: int, q: int) -> float:
-        if isinstance(self.m, NotSet) or isinstance(other.m, NotSet):
-            raise RuntimeError("m must be set to calculate the Wigner-Eckart prefactor.")  # noqa: TRY004
+        if is_not_set(self.m) or is_not_set(other.m):
+            raise RuntimeError("m must be set to calculate the Wigner-Eckart prefactor.")
         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:

src/rydstate/angular/angular_state.py

@@ -13,7 +13,7 @@
     AngularKetLS,
 )
 from rydstate.angular.angular_matrix_element import is_angular_momentum_quantum_number
-from rydstate.angular.utils import NotSet
+from rydstate.angular.utils import is_not_set
 
 if TYPE_CHECKING:
     from collections.abc import Iterator, Sequence
@@ -210,8 +210,8 @@ def calc_matrix_element(
                     "Different m values are not supported yet for AngularState.calc_matrix_element."
                 )
 
-        if isinstance(self.kets[0].m, NotSet) or isinstance(other.kets[0].m, NotSet):
-            raise RuntimeError("m must be set for all kets to calculate the matrix element.")  # noqa: TRY004
+        if is_not_set(self.kets[0].m) or is_not_set(other.kets[0].m):
+            raise RuntimeError("m must be set for all kets to calculate the matrix element.")
 
         prefactor = self.kets[0]._calc_wigner_eckart_prefactor(other.kets[0], kappa, q)  # noqa: SLF001
         reduced_matrix_element = self.calc_reduced_matrix_element(other, operator, kappa)

src/rydstate/angular/utils.py

@@ -2,11 +2,13 @@
 
 import contextlib
 import typing as t
-from typing import TYPE_CHECKING, Literal
+from typing import TYPE_CHECKING, Any, Literal
 
 import numpy as np
 
 if TYPE_CHECKING:
+    from typing_extensions import TypeIs
+
     from rydstate.angular.angular_ket import AngularKetBase
     from rydstate.species.species_object import SpeciesObject
 
@@ -26,6 +28,11 @@ class NotSet(t.Protocol):
     def __not_set() -> None: ...
 
 
+def is_not_set(obj: Any) -> TypeIs[NotSet]:  # noqa: ANN401
+    """Check if the obj is the NotSet singleton."""
+    return obj is NotSet
+
+
 class InvalidQuantumNumbersError(ValueError):
     def __init__(self, ket: AngularKetBase, msg: str = "") -> None:
         _msg = f"Invalid quantum numbers for {ket!r}"

src/rydstate/basis/basis_sqdt.py

@@ -5,7 +5,8 @@
 
 import numpy as np
 
-from rydstate.angular.utils import NotSet
+from rydstate.angular import NotSet
+from rydstate.angular.utils import is_not_set
 from rydstate.basis.basis_base import BasisBase
 from rydstate.rydberg import (
     RydbergStateSQDT,
@@ -23,7 +24,7 @@
 
 class BasisSQDT(BasisBase[_RydbergStateSQDT]):
     def _get_m_range(self, m: tuple[float, float] | None | NotSet, f_tot: float | np.floating) -> list[NotSet | float]:
-        if isinstance(m, NotSet):
+        if is_not_set(m):
             return [NotSet]
         if m is None:
             m = (-np.inf, np.inf)

src/rydstate/rydberg/rydberg_sqdt.py

@@ -9,7 +9,7 @@
 from scipy.special import exprel
 
 from rydstate.angular import NotSet
-from rydstate.angular.utils import quantum_numbers_to_angular_ket
+from rydstate.angular.utils import is_not_set, quantum_numbers_to_angular_ket
 from rydstate.radial import RadialKet
 from rydstate.rydberg.rydberg_base import RydbergStateBase
 from rydstate.species import SpeciesObjectSQDT
@@ -428,8 +428,8 @@ def _get_transition_rates_au(
         basis_class = BasisSQDTAlkali if self.species.number_valence_electrons == 1 else BasisSQDTAlkalineLS
 
         m = self.angular.m
-        if isinstance(m, NotSet):
-            raise RuntimeError("m quantum number must be defined to calculate transition rates.")  # noqa: TRY004
+        if is_not_set(m):
+            raise RuntimeError("m quantum number must be defined to calculate transition rates.")
 
         basis = basis_class(self.species, n=(1, int(self.nu + 35)), m=(m - 1, m + 1))
         basis.filter_states("l_r", (self.angular.l_r - 1, self.angular.l_r + 1))