Various small fixes

src/rydstate/radial/radial_matrix_element.py

@@ -101,7 +101,8 @@ def _sanity_check_integration(z1: NDArray, z2: NDArray) -> None:
     assert len(z1) == len(z2), f"Length mismatch: {len(z1)=} != {len(z2)=}"
     assert z1[0] - z2[0] < tol, f"First point mismatch: {z1[0]=} != {z2[0]=}"
     assert z1[1] - z2[1] < tol, f"Second point mismatch: {z1[1]=} != {z2[1]=}"
-    assert z1[2] - z2[2] < tol, f"Third point mismatch: {z1[2]=} != {z2[2]=}"
+    if len(z1) > 2:
+        assert z1[2] - z2[2] < tol, f"Third point mismatch: {z1[2]=} != {z2[2]=}"
     assert z1[-1] - z2[-1] < tol, f"Last point mismatch: {z1[-1]=} != {z2[-1]=}"
 
 

src/rydstate/radial/radial_state.py

@@ -197,11 +197,28 @@ def create_wavefunction(
         self._wavefunction.apply_sign_convention(sign_convention)
         self._grid = self._wavefunction.grid
 
+    def calc_overlap(self, other: RadialState, *, integration_method: INTEGRATION_METHODS = "sum") -> float:
+        r"""Calculate the overlap <self|other> of two radial states.
+
+        Args:
+            other: Other radial state
+            integration_method: Integration method to use
+
+        Returns:
+            The overlap inegral between self and other.
+
+        """
+        return self.calc_matrix_element(other, k_radial=0, unit="a.u.", integration_method=integration_method)
+
     @overload
-    def calc_matrix_element(self, other: RadialState, k_radial: int) -> PintFloat: ...
+    def calc_matrix_element(
+        self, other: RadialState, k_radial: int, *, integration_method: INTEGRATION_METHODS = "sum"
+    ) -> PintFloat: ...
 
     @overload
-    def calc_matrix_element(self, other: RadialState, k_radial: int, unit: str) -> float: ...
+    def calc_matrix_element(
+        self, other: RadialState, k_radial: int, unit: str, *, integration_method: INTEGRATION_METHODS = "sum"
+    ) -> float: ...
 
     def calc_matrix_element(
         self,
@@ -217,8 +234,8 @@ def calc_matrix_element(
 
         .. math::
             \int_{0}^{\infty} dr r^2 r^k_{radial} R_1(r) R_2(r)
-            = a_0^k_{radial} \int_{0}^{\infty} dx x^k_{radial} \tilde{u}_1(x) \tilde{u}_2(x)
-            = a_0^k_{radial} \int_{0}^{\infty} dz 2 z^{2 + 2k_{radial}} w_1(z) w_2(z)
+            = a_0^{k_{radial}} \int_{0}^{\infty} dx x^k_{radial} \tilde{u}_1(x) \tilde{u}_2(x)
+            = a_0^{k_{radial}} \int_{0}^{\infty} dz 2 z^{2 + 2k_{radial}} w_1(z) w_2(z)
 
         where R_1 and R_2 are the radial wavefunctions of self and other,
         and w(z) = z^{-1/2} \tilde{u}(z^2) = (r/_a_0)^{1/4} \sqrt{a_0} r R(r).

src/rydstate/radial/wavefunction.py

@@ -322,7 +322,8 @@ def sanity_check(self, z_stop: float, run_backward: bool) -> bool:  # noqa: C901
             if not run_backward and z_stop < grid.z_list[-1] + grid.dz / 2:
                 warning_msgs.append(f"The integration did not stop before z_stop, z={grid.z_list[-1]}")
         elif state.l_r == 0 and run_backward:
-            if grid.z_list[0] > 0.035:  # z_list[0] should run almost to zero for l=0
+            z_tol = 0.025 if species.number_valence_electrons == 1 else 0.05
+            if grid.z_list[0] > z_tol:  # z_list[0] should run almost to zero for l=0
                 warning_msgs.append(f"The integration for l=0 did stop at {grid.z_list[0]} (should be close to zero).")
 
         if warning_msgs:

src/rydstate/rydberg_state.py

@@ -68,6 +68,12 @@ def get_energy(self, unit: str | None = None) -> PintFloat | float:
             return energy
         return energy.to(unit, "spectroscopy").magnitude
 
+    def calc_reduced_overlap(self, other: RydbergStateBase) -> float:
+        """Calculate the reduced overlap <self|other> (ignoring the magnetic quantum number m)."""
+        radial_overlap = self.radial.calc_overlap(other.radial)
+        angular_overlap = self.angular.calc_reduced_overlap(other.angular)
+        return radial_overlap * angular_overlap
+
     @overload
     def calc_reduced_matrix_element(
         self, other: Self, operator: MatrixElementOperator, unit: None = None