angular ket and radial ket cache matrix elements

Author: johannes-moegerle

src/rydstate/angular/angular_ket.py

@@ -22,6 +22,7 @@
     is_angular_operator_type,
     is_not_set,
     is_unknown,
+    lru_cache,
     minus_one_pow,
     try_trivial_spin_addition,
 )
@@ -600,6 +601,7 @@ def calc_reduced_overlap(self, other: AngularKetBase[Any]) -> float:
 
         raise NotImplementedError(f"This method is not yet implemented for {self!r} and {other!r}.")
 
+    @lru_cache(maxsize=1_000_000)
     def calc_reduced_matrix_element(  # noqa: C901, PLR0912
         self: Self, other: AngularKetBase[Any], operator: AngularOperatorType, kappa: int
     ) -> float:

src/rydstate/radial/radial_ket.py

@@ -8,7 +8,7 @@
 from mpmath import whitw
 from scipy.special import gamma
 
-from rydstate.angular.utils import is_unknown
+from rydstate.angular.utils import is_unknown, lru_cache
 from rydstate.metaclass_cache import CachedABCMeta
 from rydstate.radial.numerov import _run_numerov_integration_python, run_numerov_integration
 from rydstate.radial.radial_matrix_element import calc_radial_matrix_element_from_w_z
@@ -554,9 +554,12 @@ def calc_matrix_element(
             The radial matrix element in the desired unit.
 
         """
-        # Ensure wavefunctions are integrated before accessing the grid
-        radial_matrix_element_au = calc_radial_matrix_element_from_w_z(
-            self.z_list, self.w_list, other.z_list, other.w_list, k_radial, integration_method
+        ket1, ket2 = self, other
+        if hash(self) > hash(other):
+            # we sort the kets by their hash to increase the cache hit rate of calc_matrix_element_au_cached
+            ket1, ket2 = other, self
+        radial_matrix_element_au = calc_matrix_element_au_cached(
+            ket1, ket2, k_radial, integration_method=integration_method
         )
 
         if unit == "a.u.":
@@ -565,3 +568,16 @@ def calc_matrix_element(
         if unit is None:
             return radial_matrix_element
         return radial_matrix_element.to(unit).magnitude
+
+
+@lru_cache(maxsize=100_000)
+def calc_matrix_element_au_cached(
+    ket1: RadialKet,
+    ket2: RadialKet,
+    k_radial: int,
+    *,
+    integration_method: INTEGRATION_METHODS = "sum",
+) -> float:
+    return calc_radial_matrix_element_from_w_z(
+        ket1.z_list, ket1.w_list, ket2.z_list, ket2.w_list, k_radial, integration_method
+    )