improve

Author: johannes-moegerle

src/rydstate/radial/radial_ket.py

@@ -525,8 +525,11 @@ def calc_matrix_element(
             The radial matrix element in the desired unit.
 
         """
+        ket1, ket2 = self, other
+        if hash(self) > hash(other):
+            ket1, ket2 = other, self
         radial_matrix_element_au = calc_matrix_element_au_cached(
-            self, other, k_radial, integration_method=integration_method
+            ket1, ket2, k_radial, integration_method=integration_method
         )
 
         if unit == "a.u.":
@@ -537,7 +540,7 @@ def calc_matrix_element(
         return radial_matrix_element.to(unit).magnitude
 
 
-@lru_cache(maxsize=100_000)
+@lru_cache(maxsize=1000_000)
 def calc_matrix_element_au_cached(
     ket1: RadialKet,
     ket2: RadialKet,

src/rydstate/radial/radial_matrix_element.py

@@ -5,6 +5,7 @@
 
 import numpy as np
 import scipy.integrate
+from numba import njit
 
 if TYPE_CHECKING:
     from rydstate.units import NDArray
@@ -49,7 +50,9 @@ def calc_radial_matrix_element_from_w_z(
 
     """
     # Find overlapping grid range
-    z_min = max(z1[0], z2[0])
+    min_ind1 = first_nonzero(w1)
+    min_ind2 = first_nonzero(w2)
+    z_min = max(z1[min_ind1], z2[min_ind2])
     z_max = min(z1[-1], z2[-1])
     if z_max <= z_min:
         logger.debug("No overlapping grid points between states, returning 0")
@@ -61,7 +64,7 @@ def calc_radial_matrix_element_from_w_z(
         ind = int((z_min - z1[0]) / dz + 0.5)
         z1 = z1[ind:]
         w1 = w1[ind:]
-    elif z2[0] < z_min - dz / 2:
+    if z2[0] < z_min - dz / 2:
         ind = int((z_min - z2[0]) / dz + 0.5)
         z2 = z2[ind:]
         w2 = w2[ind:]
@@ -70,7 +73,7 @@ def calc_radial_matrix_element_from_w_z(
         ind = int((z1[-1] - z_max) / dz + 0.5)
         z1 = z1[:-ind]
         w1 = w1[:-ind]
-    elif z2[-1] > z_max + dz / 2:
+    if z2[-1] > z_max + dz / 2:
         ind = int((z2[-1] - z_max) / dz + 0.5)
         z2 = z2[:-ind]
         w2 = w2[:-ind]
@@ -120,3 +123,11 @@ def _integrate(integrand: NDArray, dz: float, method: INTEGRATION_METHODS) -> fl
         raise ValueError(f"Invalid integration method: {method}")
 
     return float(value)
+
+
+@njit(cache=True)
+def first_nonzero(a: np.ndarray) -> int:
+    for i in range(a.size):
+        if a[i] != 0:
+            return i
+    return -1