various small improvements

Author: johannes-moegerle

src/rydstate/angular/angular_ket.py

@@ -159,7 +159,7 @@ def __str__(self) -> str:
 
     def __eq__(self, other: object) -> bool:
         if not isinstance(other, AngularKetBase):
-            raise NotImplementedError(f"Cannot compare {self!r} with {other!r}.")
+            return NotImplemented
         if type(self) is not type(other):
             return False
         if self.m != other.m:

src/rydstate/basis/basis_sqdt.py

@@ -97,13 +97,14 @@ def __init__(
         _s_r = 0.5
         for _n in range(n[0], n[1] + 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]):
+                allowed = [self.species.is_allowed_shell(_n, _l_r, s) for s in [0, 1]]
+                if not all(allowed):
+                    if any(allowed):
+                        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),
+                        )
                     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)):
@@ -130,13 +131,14 @@ def __init__(
         _s_r = 0.5
         for _n in range(n[0], n[1] + 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]):
+                allowed = [self.species.is_allowed_shell(_n, _l_r, s) for s in [0, 1]]
+                if not all(allowed):
+                    if any(allowed):
+                        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),
+                        )
                     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):

src/rydstate/rydberg/rydberg_sqdt.py

@@ -413,6 +413,8 @@ def _get_transition_rates_au(
             \Gamma^{spontaneous}_{self \to other} = \frac{4}{3} \frac{\alpha}{c^2} \omega^3
                 |\langle self || r^k_radial \hat{O}_{k_angular} || other \rangle|^2
 
+        where :math:`\alpha = 1/c` in atomic units.
+
         and
 
         .. math::
@@ -443,27 +445,25 @@ def _get_transition_rates_au(
         )
         electric_dipole_moments_au = np.zeros(len(relevant_states))
         for q in [-1, 0, 1]:
-            # the different entries are only at most once nonzero -> we can just add the arrays
             el_di_m = np.array(
                 [s.calc_matrix_element(self, "electric_dipole", q, unit="a.u.") for s in relevant_states]
             )
-            electric_dipole_moments_au += el_di_m
+            electric_dipole_moments_au += np.abs(el_di_m) ** 2
 
         transition_rates_au = (
-            (4 / 3)
-            * np.abs(electric_dipole_moments_au) ** 2
-            * energy_differences_au**2
-            / ureg.Quantity(1, "speed_of_light").to_base_units().magnitude ** 3
+            (4 / 3) * electric_dipole_moments_au / ureg.Quantity(1, "speed_of_light").to_base_units().magnitude ** 3
         )
 
         if only_spontaneous:
-            transition_rates_au *= energy_differences_au
+            transition_rates_au *= energy_differences_au**3
         else:
             assert temperature_au is not None, "Temperature must be given for black body transitions."
             if temperature_au == 0:
                 transition_rates_au *= 0
             else:  # for numerical stability we use 1 / exprel(x) = x / (exp(x) - 1)
-                transition_rates_au *= temperature_au / exprel(energy_differences_au / temperature_au)
+                transition_rates_au *= (
+                    energy_differences_au**2 * temperature_au / exprel(energy_differences_au / temperature_au)
+                )
 
         mask = transition_rates_au != 0
         relevant_states_masked = [ket for ket, is_relevant in zip(relevant_states, mask, strict=True) if is_relevant]

src/rydstate/species/sqdt/species_object_sqdt.py

@@ -164,11 +164,11 @@ def get_ionization_energy(self, unit: None = None) -> PintFloat: ...
     @overload
     def get_ionization_energy(self, unit: str) -> float: ...
 
-    def get_ionization_energy(self, unit: str | None = "hartree") -> PintFloat | float:
+    def get_ionization_energy(self, unit: str | None = None) -> PintFloat | float:
         """Return the ionization energy in the desired unit.
 
         Args:
-            unit: Desired unit for the ionization energy. Default is atomic units "hartree".
+            unit: Desired unit for the ionization energy. Default is None (returns a Pint quantity).
 
         Returns:
             Ionization energy in the desired unit.

tests/test_lifetimes.py

@@ -1,9 +1,8 @@
 import numpy as np
 import pytest
-from rydstate import RydbergStateSQDTAlkali
-from rydstate.angular.angular_ket import AngularKetLS
-from rydstate.rydberg.rydberg_sqdt import RydbergStateSQDT
-from rydstate.species.sqdt.species_object_sqdt import SpeciesObjectSQDT
+from rydstate import RydbergStateSQDT, RydbergStateSQDTAlkali
+from rydstate.angular import AngularKetLS
+from rydstate.species import SpeciesObjectSQDT
 
 
 # Reference values from NIST Atomic Spectra Database (ASD), Einstein A coefficients: