From 6208e72a6b43510a29acdd3a341260145aec4584 Mon Sep 17 00:00:00 2001 From: Matt Davis Date: Thu, 5 Feb 2015 11:35:44 -0800 Subject: [PATCH 1/2] Add ability to cache injectable results on inputs If you have a utility function registered as an injectable but it is not being automatically called you can now get caching of results based on the inputs so long as the inputs are all hashable. Use the memoize=True keyword along with autocall=False to get this behavior. Note that this only applies if you're using the utility via the simulation framework (via injection or get_injectable). --- urbansim/sim/simulation.py | 84 ++++++++++++++++++++++++--- urbansim/sim/tests/test_simulation.py | 64 ++++++++++++++++++++ 2 files changed, 139 insertions(+), 9 deletions(-) diff --git a/urbansim/sim/simulation.py b/urbansim/sim/simulation.py index da6c40bf..d0cfc139 100644 --- a/urbansim/sim/simulation.py +++ b/urbansim/sim/simulation.py @@ -6,6 +6,7 @@ import warnings from collections import Callable, namedtuple from contextlib import contextmanager +from functools import wraps import pandas as pd import tables @@ -27,6 +28,7 @@ _TABLE_CACHE = {} _COLUMN_CACHE = {} _INJECTABLE_CACHE = {} +_MEMOIZED = {} _CS_FOREVER = 'forever' _CS_ITER = 'iteration' @@ -48,6 +50,9 @@ def clear_sim(): _TABLE_CACHE.clear() _COLUMN_CACHE.clear() _INJECTABLE_CACHE.clear() + for m in _MEMOIZED.values(): + m.value.clear_cached() + _MEMOIZED.clear() logger.debug('simulation state cleared') @@ -66,12 +71,16 @@ def clear_cache(scope=None): _TABLE_CACHE.clear() _COLUMN_CACHE.clear() _INJECTABLE_CACHE.clear() + for m in _MEMOIZED.values(): + m.value.clear_cached() logger.debug('simulation cache cleared') else: for d in (_TABLE_CACHE, _COLUMN_CACHE, _INJECTABLE_CACHE): items = toolz.valfilter(lambda x: x.scope == scope, d) for k in items: del d[k] + for m in toolz.filter(lambda x: x.scope == scope, _MEMOIZED.values()): + m.value.clear_cached() logger.debug('cleared cached values with scope {!r}'.format(scope)) @@ -1022,8 +1031,52 @@ def _columns_for_table(table_name): if tname == table_name} +def _memoize_function(f, name, cache_scope=_CS_FOREVER): + """ + Wraps a function for memoization and ties it's cache into the + simulation cacheing system. + + Parameters + ---------- + f : function + name : str + Name of injectable. + cache_scope : {'step', 'iteration', 'forever'}, optional + Scope for which to cache data. Default is to cache forever + (or until manually cleared). 'iteration' caches data for each + complete iteration of the simulation, 'step' caches data for + a single step of the simulation. + + """ + cache = {} + + @wraps(f) + def wrapper(*args, **kwargs): + try: + cache_key = ( + args or None, frozenset(kwargs.items()) if kwargs else None) + in_cache = cache_key in cache + except TypeError: + raise TypeError( + 'function arguments must be hashable for memoization') + + if _CACHING and in_cache: + return cache[cache_key] + else: + result = f(*args, **kwargs) + cache[cache_key] = result + return result + + wrapper.cache = cache + wrapper.clear_cached = lambda: cache.clear() + _MEMOIZED[name] = CacheItem(name, wrapper, cache_scope) + + return wrapper + + def add_injectable( - name, value, autocall=True, cache=False, cache_scope=_CS_FOREVER): + name, value, autocall=True, cache=False, cache_scope=_CS_FOREVER, + memoize=False): """ Add a value that will be injected into other functions. @@ -1049,18 +1102,30 @@ def add_injectable( (or until manually cleared). 'iteration' caches data for each complete iteration of the simulation, 'step' caches data for a single step of the simulation. + memoize : bool, optional + If autocall is False it is still possible to cache function results + by setting this flag to True. Cached values are stored in a dictionary + keyed by argument values, so the argument values must be hashable. + Memoized functions have their caches cleared according to the same + rules as universal caching. + + """ + if isinstance(value, Callable): + if autocall: + value = _InjectableFuncWrapper( + name, value, cache=cache, cache_scope=cache_scope) + # clear any cached data from a previously registered value + value.clear_cached() + elif not autocall and memoize: + value = _memoize_function(value, name, cache_scope=cache_scope) - """ - if isinstance(value, Callable) and autocall: - value = _InjectableFuncWrapper( - name, value, cache=cache, cache_scope=cache_scope) - # clear any cached data from a previously registered value - value.clear_cached() logger.debug('registering injectable {!r}'.format(name)) _INJECTABLES[name] = value -def injectable(name=None, autocall=True, cache=False, cache_scope=_CS_FOREVER): +def injectable( + name=None, autocall=True, cache=False, cache_scope=_CS_FOREVER, + memoize=False): """ Decorates functions that will be injected into other functions. @@ -1078,7 +1143,8 @@ def decorator(func): else: n = func.__name__ add_injectable( - n, func, autocall=autocall, cache=cache, cache_scope=cache_scope) + n, func, autocall=autocall, cache=cache, cache_scope=cache_scope, + memoize=memoize) return func return decorator diff --git a/urbansim/sim/tests/test_simulation.py b/urbansim/sim/tests/test_simulation.py index 33ab619f..064a3b64 100644 --- a/urbansim/sim/tests/test_simulation.py +++ b/urbansim/sim/tests/test_simulation.py @@ -598,6 +598,54 @@ def inj(): assert i()() == 16 +def test_memoized_injectable(): + outside = 'x' + + @sim.injectable(autocall=False, memoize=True) + def x(s): + return outside + s + + assert 'x' in sim._MEMOIZED + + getx = lambda: sim.get_injectable('x') + + assert hasattr(getx(), 'cache') + assert hasattr(getx(), 'clear_cached') + + assert getx()('y') == 'xy' + outside = 'z' + assert getx()('y') == 'xy' + + getx().clear_cached() + + assert getx()('y') == 'zy' + + +def test_memoized_injectable_cache_off(): + outside = 'x' + + @sim.injectable(autocall=False, memoize=True) + def x(s): + return outside + s + + getx = lambda: sim.get_injectable('x')('y') + + sim.disable_cache() + + assert getx() == 'xy' + outside = 'z' + assert getx() == 'zy' + + sim.enable_cache() + outside = 'a' + + assert getx() == 'zy' + + sim.disable_cache() + + assert getx() == 'ay' + + def test_clear_cache_all(df): @sim.table(cache=True) def table(): @@ -611,18 +659,25 @@ def z(table): def x(): return 'x' + @sim.injectable(autocall=False, memoize=True) + def y(s): + return s + 'y' + sim.eval_variable('table.z') sim.eval_variable('x') + sim.get_injectable('y')('x') assert sim._TABLE_CACHE.keys() == ['table'] assert sim._COLUMN_CACHE.keys() == [('table', 'z')] assert sim._INJECTABLE_CACHE.keys() == ['x'] + assert sim._MEMOIZED['y'].value.cache == {(('x',), None): 'xy'} sim.clear_cache() assert sim._TABLE_CACHE == {} assert sim._COLUMN_CACHE == {} assert sim._INJECTABLE_CACHE == {} + assert sim._MEMOIZED['y'].value.cache == {} def test_clear_cache_scopes(df): @@ -638,30 +693,39 @@ def z(table): def x(): return 'x' + @sim.injectable(autocall=False, memoize=True, cache_scope='iteration') + def y(s): + return s + 'y' + sim.eval_variable('table.z') sim.eval_variable('x') + sim.get_injectable('y')('x') assert sim._TABLE_CACHE.keys() == ['table'] assert sim._COLUMN_CACHE.keys() == [('table', 'z')] assert sim._INJECTABLE_CACHE.keys() == ['x'] + assert sim._MEMOIZED['y'].value.cache == {(('x',), None): 'xy'} sim.clear_cache(scope='step') assert sim._TABLE_CACHE.keys() == ['table'] assert sim._COLUMN_CACHE.keys() == [('table', 'z')] assert sim._INJECTABLE_CACHE == {} + assert sim._MEMOIZED['y'].value.cache == {(('x',), None): 'xy'} sim.clear_cache(scope='iteration') assert sim._TABLE_CACHE.keys() == ['table'] assert sim._COLUMN_CACHE == {} assert sim._INJECTABLE_CACHE == {} + assert sim._MEMOIZED['y'].value.cache == {} sim.clear_cache(scope='forever') assert sim._TABLE_CACHE == {} assert sim._COLUMN_CACHE == {} assert sim._INJECTABLE_CACHE == {} + assert sim._MEMOIZED['y'].value.cache == {} def test_cache_scope(df): From 68e9ffb1f74f770275e77b8d3df83052babc8e94 Mon Sep 17 00:00:00 2001 From: Matt Davis Date: Thu, 5 Feb 2015 11:50:02 -0800 Subject: [PATCH 2/2] Update injectable docs with info on memoization [ci-skip] --- docs/sim/index.rst | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/docs/sim/index.rst b/docs/sim/index.rst index 89d241d9..09f90c48 100644 --- a/docs/sim/index.rst +++ b/docs/sim/index.rst @@ -340,9 +340,21 @@ Use the :py:func:`~urbansim.sim.simulation.add_injectable` function or the Be default injectable functions are evaluated before injection and the return value is passed into other functions. Use ``autocall=False`` to disable this -behavior and instead inject the wrapped function itself. -Like tables and columns, injectable functions can have their results -cached with ``cache=True``. +behavior and instead inject the function itself. +Like tables and columns, injectable functions that are automatically evaluated +can have their results cached with ``cache=True``. + +Functions that are not automatically evaluated can also have their results +cached using the ``memoize=True`` keyword along with ``autocall=False``. +A memoized injectable will cache results based on the function inputs, +so this only works if the function inputs are hashable +(usable as dictionary keys). +Memoized functions can have their caches cleared manually using their +``clear_cached`` function attribute. +The caches of memoized functions are also hooked into the global simulation +caching system, +so you can also manage their caches via the ``cache_scope`` keyword argument +and the :py:func:`~urbansim.sim.simulation.clear_cache` function. An example of the above injectables in IPython: