Add caching

docs/examples/benchmark/benchmark_njit.ipynb

@@ -29,7 +29,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -45,13 +45,18 @@
     "    state = rydstate.RydbergStateSQDTAlkali(species, n, l=l, j=l + 0.5)\n",
     "    potential = state.potential\n",
     "    nu = state.nu\n",
+    "    RadialKet.clear_cached_instances()\n",
     "    radial = RadialKet(nu, potential, dz=1e-3, use_njit=use_njit)\n",
     "    radial.integrate_wavefunction()\n",
     "\n",
     "    results = []\n",
     "    for species, n, l, use_njit in test_cases:\n",
     "        # Setup the test function\n",
-    "        stmt = \"radial = RadialKet(nu, potential, dz=1e-3, use_njit=use_njit)\\nradial.integrate_wavefunction()\"\n",
+    "        stmt = (\n",
+    "            \"RadialKet.clear_cached_instances()\\n\"\n",
+    "            \"radial = RadialKet(nu, potential, dz=1e-3, use_njit=use_njit)\\n\"\n",
+    "            \"radial.integrate_wavefunction()\"\n",
+    "        )\n",
     "\n",
     "        # Time the integration multiple times and take average/std\n",
     "        globals_dict = {\n",
@@ -92,8 +97,8 @@
       "----------------------------------------------------------------------\n",
       " species   n   l   use_njit  time (ms)   std (ms)\n",
       "----------------------------------------------------------------------\n",
-      "       H 100  80       True      14.74       0.96\n",
-      "       H 100  80      False     140.28       4.93\n"
+      "       H 100  80       True      13.75       0.20\n",
+      "       H 100  80      False     136.34      10.99\n"
      ]
     }
    ],

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
 
@@ -26,6 +27,7 @@
     try_trivial_spin_addition,
 )
 from rydstate.angular.wigner_symbols import calc_wigner_3j, clebsch_gordan_6j, clebsch_gordan_9j
+from rydstate.metaclass_cache import CachedABCMeta
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
@@ -41,7 +43,7 @@
 T_Unknown = TypeVar("T_Unknown", AllKnown, Unknown, Any)
 
 
-class AngularKetBase(ABC, Generic[GenericT_Unknown]):
+class AngularKetBase(ABC, Generic[GenericT_Unknown], metaclass=CachedABCMeta):
     """Base class for a angular ket (i.e. a simple canonical spin ketstate)."""
 
     # We use __slots__ to prevent dynamic attributes and make the objects immutable after initialization
@@ -58,6 +60,8 @@ class AngularKetBase(ABC, Generic[GenericT_Unknown]):
         "quantum_numbers",
         "_allow_unknown",
         "_initialized",
+        "_reduced_matrix_element_cache",
+        "__weakref__",
     )
 
     quantum_number_names: ClassVar[tuple[AngularMomentumQuantumNumbers, ...]]
@@ -117,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
 
@@ -598,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.
 
@@ -610,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/metaclass_cache.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+import threading
+import weakref
+from abc import ABCMeta
+from inspect import Parameter, signature
+from typing import TYPE_CHECKING, TypeVar
+
+if TYPE_CHECKING:
+    from inspect import Signature
+
+CachedT = TypeVar("CachedT", bound="CachedABCMeta")
+
+
+class CachedABCMeta(ABCMeta):
+    """Metaclass that reuses live instances with equivalent constructor arguments."""
+
+    # Per-class storage: declared here so type checkers know every cached class carries them,
+    # but actually created per class in __init__ (see below).
+    _instances: weakref.WeakValueDictionary[tuple[tuple[str, object], ...], object]
+    _signature: Signature | None
+    _instances_lock: threading.RLock
+
+    def __init__(cls, name: str, bases: tuple[type, ...], namespace: dict[str, object], **kwargs: object) -> None:
+        super().__init__(name, bases, namespace, **kwargs)
+        # Each cached class gets its own cache, signature and lock.
+        cls._instances = weakref.WeakValueDictionary()
+        cls._signature = None
+        cls._instances_lock = threading.RLock()
+
+    def clear_cached_instances(cls) -> None:
+        """Clear currently cached instances of this class."""
+        with cls._instances_lock:
+            cls._instances.clear()
+
+    def __call__(cls: type[CachedT], *args: object, **kwargs: object) -> CachedT:  # type: ignore [misc]
+        with cls._instances_lock:
+            if cls._signature is None:
+                cls._signature = signature(cls.__init__)
+
+            bound_arguments = cls._signature.bind(None, *args, **kwargs)
+            bound_arguments.apply_defaults()
+            constructor_arguments_list: list[tuple[str, object]] = []
+            for name, value in list(bound_arguments.arguments.items())[1:]:
+                parameter_kind = cls._signature.parameters[name].kind
+                normalized_value = value
+                if parameter_kind is Parameter.VAR_POSITIONAL:
+                    normalized_value = tuple(value)
+                elif parameter_kind is Parameter.VAR_KEYWORD:
+                    normalized_value = tuple(sorted(value.items()))
+                constructor_arguments_list.append((name, normalized_value))
+            key = tuple(constructor_arguments_list)
+            try:
+                hash(key)
+            except TypeError as exc:
+                raise TypeError(
+                    f"Arguments to cached class {cls.__name__} must be hashable, but received {key!r}."
+                ) from exc
+
+            instance = cls._instances.get(key)
+            if instance is None:
+                instance = ABCMeta.__call__(cls, *args, **kwargs)
+                cls._instances[key] = instance
+            return instance  # type: ignore [return-value]

src/rydstate/radial/radial_ket.py

@@ -2,13 +2,15 @@
 
 import logging
 import math
+import weakref
 from typing import TYPE_CHECKING, Literal, overload
 
 import numpy as np
 from mpmath import whitw
 from scipy.special import gamma
 
 from rydstate.angular.utils import is_unknown
+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
 from rydstate.species.utils import calc_energy_from_nu
@@ -26,7 +28,7 @@
 WavefunctionSignConvention = Literal["positive_at_outer_bound", "n_l_1"] | None
 
 
-class RadialKet:
+class RadialKet(metaclass=CachedABCMeta):
     r"""Class representing a radial Rydberg state."""
 
     def __init__(
@@ -70,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:
@@ -514,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
@@ -553,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

src/rydstate/species/element_properties.py

@@ -1,18 +1,18 @@
 from __future__ import annotations
 
 from abc import ABC
-from functools import cache, cached_property
+from functools import cached_property
 from typing import TYPE_CHECKING, ClassVar, overload
 
+from rydstate.metaclass_cache import CachedABCMeta
 from rydstate.species.utils import get_all_subclasses
 from rydstate.units import rydberg_constant, ureg
 
 if TYPE_CHECKING:
-    from rydstate.species.utils import cache  # type: ignore [assignment]  # noqa: TC004
     from rydstate.units import PintFloat
 
 
-class ElementProperties(ABC):
+class ElementProperties(ABC, metaclass=CachedABCMeta):
     """Base class for all element properties classes.
 
     For the electronic ground state configurations and sorted shells,
@@ -111,7 +111,6 @@ def reduced_mass_au(self) -> float:
         return self.get_corrected_rydberg_constant("hartree") / rydberg_constant.to("hartree").m
 
 
-@cache
 def get_element_properties(species: str) -> ElementProperties:
     """Get an instance of the subclass of ElementProperties for the given species."""
     possible_subclasses = get_all_subclasses(ElementProperties, species)

src/rydstate/species/mqdt.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
 from abc import ABC
-from functools import cache, cached_property
+from functools import cached_property
 from typing import TYPE_CHECKING, Any, ClassVar, overload
 
+from rydstate.metaclass_cache import CachedABCMeta
 from rydstate.species.fmodel import FModelSQDT
 from rydstate.species.utils import get_all_subclasses
 from rydstate.units import ureg
@@ -12,11 +13,10 @@
     from rydstate.angular.angular_ket import AngularKetFJ
     from rydstate.angular.core_ket import CoreKet
     from rydstate.species.fmodel import FModel
-    from rydstate.species.utils import cache  # type: ignore [assignment]  # noqa: TC004
     from rydstate.units import PintFloat
 
 
-class MQDT(ABC):
+class MQDT(ABC, metaclass=CachedABCMeta):
     """Base class for all MQDT classes."""
 
     species: ClassVar[str]
@@ -98,7 +98,6 @@ def get_mqdt_models(self, outer_channel: AngularKetFJ[Any]) -> list[FModel]:
         return models
 
 
-@cache
 def get_mqdt(species: str, tag: str | None = None) -> MQDT:
     """Get an instance of the subclass of MQDT for the given species and tag."""
     subclasses = get_all_subclasses(MQDT, species, tag)

src/rydstate/species/potential.py

@@ -8,6 +8,7 @@
 import numpy as np
 
 from rydstate.angular.utils import is_unknown
+from rydstate.metaclass_cache import CachedABCMeta
 from rydstate.species.element_properties import get_element_properties
 from rydstate.species.utils import get_all_subclasses
 
@@ -21,7 +22,7 @@
 logger = logging.getLogger(__name__)
 
 
-class Potential(ABC):
+class Potential(ABC, metaclass=CachedABCMeta):
     """Base class for all potential classes."""
 
     species: ClassVar[str]

src/rydstate/species/sqdt.py

@@ -5,13 +5,14 @@
 import re
 from abc import ABC
 from fractions import Fraction
-from functools import cache, cached_property
+from functools import cached_property
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, ClassVar, overload
 
 import numpy as np
 
 from rydstate.angular.utils import check_spin_addition_rule, get_possible_quantum_number_values, is_unknown
+from rydstate.metaclass_cache import CachedABCMeta
 from rydstate.species.element_properties import get_element_properties
 from rydstate.species.utils import (
     calc_modified_ritz_formula,
@@ -24,17 +25,14 @@
 if TYPE_CHECKING:
     from rydstate.angular.angular_ket import AngularKetBase
     from rydstate.angular.utils import Unknown
-    from rydstate.species.utils import (  # type: ignore [assignment]
-        RydbergRitzParameters,
-        cache,  # noqa: TC004
-    )
+    from rydstate.species.utils import RydbergRitzParameters
     from rydstate.units import PintFloat
 
 
 logger = logging.getLogger(__name__)
 
 
-class SQDT(ABC):
+class SQDT(ABC, metaclass=CachedABCMeta):
     """Base class for all SQDT classes."""
 
     species: ClassVar[str]
@@ -255,7 +253,6 @@ def calc_nu(
         return n - delta_nlj
 
 
-@cache
 def get_sqdt(species: str, tag: str | None = None) -> SQDT:
     """Get an instance of the subclass of SQDT for the given species and tag."""
     subclasses = get_all_subclasses(SQDT, species, tag)

tests/test_generate_database.py

@@ -74,8 +74,6 @@ def test_generate_matrix_elements_table(species: str, conn: sqlite3.Connection)
         assert len(rows) > 2
 
     states = basis.states
-    for state in states:
-        state.radial.integrate_wavefunction()
     for row in rows_by_table["matrix_elements_d"]:
         state1, state2 = states[row[0]], states[row[1]]
         sign1 = (-1) ** (state1.radial.n_expected - state1.radial.l_r - 1)  # type: ignore [operator]

tests/test_hydrogen.py

@@ -33,7 +33,6 @@ def test_hydrogen_wavefunctions(species: str, n: int, l: int, run_backward: bool
 
     # Setup radial wavefunction and run the numerov integration
     radial = RadialKet(state.nu, state.potential, n_expected=n, run_backward=run_backward, sign_convention="n_l_1")
-    radial.integrate_wavefunction()
 
     # Get analytical solution from sympy
     if n <= 35: