angular ket and radial ket cache matrix elements

Author: johannes-moegerle

src/rydstate/angular/angular_ket.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+import weakref
 from abc import ABC
 from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, TypeVar, overload
 
@@ -59,6 +60,7 @@ class AngularKetBase(ABC, Generic[GenericT_Unknown], metaclass=CachedABCMeta):
         "quantum_numbers",
         "_allow_unknown",
         "_initialized",
+        "_reduced_matrix_element_cache",
         "__weakref__",
     )
 
@@ -119,6 +121,10 @@ def __init__(  # noqa: C901, PLR0912
         Atomic species, e.g. 'Rb87', will not be used for calculation,
         only for convenience to infer the core electron spin and nuclear spin quantum numbers.
         """
+        self._reduced_matrix_element_cache: weakref.WeakKeyDictionary[
+            AngularKetBase[Any], dict[tuple[AngularOperatorType, int], float]
+        ] = weakref.WeakKeyDictionary()
+
         if species is not None:
             from rydstate.species.element_properties import get_element_properties  # noqa: PLC0415
 
@@ -600,8 +606,8 @@ def calc_reduced_overlap(self, other: AngularKetBase[Any]) -> float:
 
         raise NotImplementedError(f"This method is not yet implemented for {self!r} and {other!r}.")
 
-    def calc_reduced_matrix_element(  # noqa: C901, PLR0912
-        self: Self, other: AngularKetBase[Any], operator: AngularOperatorType, kappa: int
+    def calc_reduced_matrix_element(
+        self, other: AngularKetBase[Any], operator: AngularOperatorType, kappa: int
     ) -> float:
         r"""Calculate the reduced angular matrix element.
 
@@ -612,6 +618,15 @@ def calc_reduced_matrix_element(  # noqa: C901, PLR0912
             \left\langle self || \hat{O}^{(\kappa)} || other \right\rangle
 
         """
+        cache = self._reduced_matrix_element_cache.setdefault(other, {})
+        cache_key = (operator, kappa)
+        if cache_key not in cache:
+            cache[cache_key] = self._calc_reduced_matrix_element(other, operator, kappa)
+        return cache[cache_key]
+
+    def _calc_reduced_matrix_element(  # noqa: C901, PLR0912
+        self: Self, other: AngularKetBase[Any], operator: AngularOperatorType, kappa: int
+    ) -> float:
         if not is_angular_operator_type(operator):
             raise NotImplementedError(f"calc_reduced_matrix_element is not implemented for operator {operator}.")
 

src/rydstate/radial/radial_ket.py

@@ -2,6 +2,7 @@
 
 import logging
 import math
+import weakref
 from typing import TYPE_CHECKING, Literal, overload
 
 import numpy as np
@@ -71,6 +72,10 @@ def __init__(
                 The "n_l_1" convention requires ``n_expected`` to be set.
 
         """
+        self._matrix_element_cache: weakref.WeakKeyDictionary[RadialKet, dict[tuple[int, str], float]] = (
+            weakref.WeakKeyDictionary()
+        )
+
         self.potential = potential
 
         if not nu > 0:
@@ -515,7 +520,7 @@ def calc_overlap(self, other: RadialKet, *, integration_method: INTEGRATION_METH
 
     @overload
     def calc_matrix_element(
-        self, other: RadialKet, k_radial: int, *, integration_method: INTEGRATION_METHODS = "sum"
+        self, other: RadialKet, k_radial: int, *, unit: None = None, integration_method: INTEGRATION_METHODS = "sum"
     ) -> PintFloat: ...
 
     @overload
@@ -554,10 +559,17 @@ 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
-        )
+        if other not in self._matrix_element_cache and self in other._matrix_element_cache:
+            return other.calc_matrix_element(self, k_radial=k_radial, unit=unit, integration_method=integration_method)
+
+        cache = self._matrix_element_cache.setdefault(other, {})
+        cache_key = (k_radial, integration_method)
+        if cache_key not in cache:
+            cache[cache_key] = calc_radial_matrix_element_from_w_z(
+                self.z_list, self.w_list, other.z_list, other.w_list, k_radial, integration_method
+            )
+
+        radial_matrix_element_au = cache[cache_key]
 
         if unit == "a.u.":
             return radial_matrix_element_au