sharedmethodcache

Added a per-instance method cache decorator, which can be shared among multiple methods with the same signature.

sharedmethodcache allows memoization similarily to functools.cache, but for instance methods of classes. The cache dictionary is added as a member to the method's owning instance. Further, multiple methods with the same signature (up to kwargs) may use the same cache dict by specifying the same cache_id. This makes sense for methods that produce the same results on identical inputs, but by different computational paths.

This decorator is currently employed in lbmpy!113 (merged) by abstract_equilibrium.py, but surely, more use cases will follow.

Example:

class Fib:
    def __init__(self):
        self.fib_rec_called = 0
        self.fib_iter_called = 0

    @sharedmethodcache("fib_cache")
    def fib_rec(self, n):
        self.fib_rec_called += 1
        return 1 if n <= 1 else self.fib_rec(n-1) + self.fib_rec(n-2)

    @sharedmethodcache("fib_cache")
    def fib_iter(self, n):
        self.fib_iter_called += 1
        f1, f2 = 0, 1
        for i in range(n):
            f2 = f1 + f2
            f1 = f2 - f1
        return f2


>>> fib = Fib()
>>> fib.fib_rec(13)
377
>>> fib.fib_cache
{(1,): 1,
 (0,): 1,
 (2,): 2,
 (3,): 3,
 (4,): 5,
 (5,): 8,
 (6,): 13,
 (7,): 21,
 (8,): 34,
 (9,): 55,
 (10,): 89,
 (11,): 144,
 (12,): 233,
 (13,): 377}
>>> fib.fib_rec_called
14
>>> fib.fib_iter(11)
144
>>> fib.fib_iter_called
0

pystencils/cache.py

@@ -32,6 +32,35 @@ def memorycache_if_hashable(maxsize=128, typed=False):
 
     return wrapper
 
+
+def sharedmethodcache(cache_id: str):
+    """Decorator for memoization of instance methods, allowing multiple methods to use the same cache.
+
+    This decorator caches results of instance methods per instantiated object of the surrounding class.
+    It allows multiple methods to use the same cache, by passing them the same `cache_id` string.
+    Cached values are stored in a dictionary, which is added as a member `self.<cache_id>` to the 
+    `self` object instance. Make sure that this doesn't cause any naming conflicts with other members!
+    Of course, for this to be useful, said methods must have the same signature (up to additional kwargs)
+    and must return the same result when called with the same arguments."""
+    def _decorator(user_method):
+        def _decorated_func(self, *args, **kwargs):
+            objdict = self.__dict__
+            cache = objdict.setdefault(cache_id, dict())
+
+            key = args
+            for item in kwargs.items():
+                key += item
+
+            if key not in cache:
+                result = user_method(self, *args, **kwargs)
+                cache[key] = result
+                return result
+            else:
+                return cache[key]
+        return _decorated_func
+    return _decorator
+
+
 # Disable memory cache:
 # disk_cache = lambda o: o
 # disk_cache_no_fallback = lambda o: o