Add optional cache for on-demand variables (#156)

* cache option for on-demand variables

* unrelated tweaks and fixes

* fix existing tests

* update on_demand docstrings

* black

* update release notes

* add test

Author: benbovy

doc/whats_new.rst

@@ -25,6 +25,8 @@ Enhancements
 - Added ``%create_setup`` IPython (Jupyter) magic command to auto-generate code
   cells with a new simulation setup from a given model (:issue:`152`). The
   command is available after executing ``%load_ext xsimlab.ipython``.
+- Added an optional cache for on-demand variables (:issue:`156`). The ``@compute``
+  decorator now has a ``cache`` option (deactivated by default).
 
 Bug fixes
 ~~~~~~~~~

xsimlab/model.py

@@ -494,6 +494,8 @@ def __init__(self, processes):
         self._index_vars = builder.get_variables(var_type=VarType.INDEX)
         self._index_vars_dict = None
 
+        self._od_vars = builder.get_variables(var_type=VarType.ON_DEMAND)
+
         builder.ensure_no_intent_conflict()
 
         self._input_vars = builder.get_input_variables()
@@ -787,6 +789,12 @@ def _merge_and_update_state(self, out_states):
         for p_obj in self._processes.values():
             p_obj.__xsimlab_state__ = self._state
 
+    def _clear_od_cache(self):
+        """Clear cached values of on-demand variables."""
+
+        for key in self._od_vars:
+            self._state.pop(key, None)
+
     def execute(
         self,
         stage,
@@ -856,21 +864,23 @@ def execute(
         stage = SimulationStage(stage)
         execute_args = (stage, runtime_context, hooks, validate)
 
+        self._clear_od_cache()
+
         self._call_hooks(hooks, runtime_context, stage, "model", "pre")
 
         if parallel:
             dsk = self._build_dask_graph(execute_args)
             out_states = dsk_get(dsk, "_gather", scheduler=scheduler)
 
             # TODO: without this -> flaky tests (don't know why)
-            # state is not well updated -> error when writing output vars in store
+            # state is not properly updated -> error when writing output vars in store
             if isinstance(scheduler, Client):
                 time.sleep(0.001)
 
             self._merge_and_update_state(out_states)
 
         else:
-            for p_name, p_obj in self._processes.items():
+            for p_obj in self._processes.values():
                 self._execute_process(p_obj, *execute_args)
 
         self._call_hooks(hooks, runtime_context, stage, "model", "post")

xsimlab/process.py

@@ -180,9 +180,14 @@ def get_from_state(self):
         return self.__xsimlab_state__[key]
 
     def get_on_demand(self):
-        p_name, v_name = self.__xsimlab_od_keys__[var_name]
-        p_obj = self.__xsimlab_model__._processes[p_name]
-        return getattr(p_obj, v_name)
+        key = self.__xsimlab_od_keys__[var_name]
+        try:
+            # get from cache
+            return self.__xsimlab_state__[key]
+        except KeyError:
+            p_name, v_name = key
+            p_obj = self.__xsimlab_model__._processes[p_name]
+            return getattr(p_obj, v_name)
 
     def put_in_state(self, value):
         key = self.__xsimlab_state_keys__[var_name]
@@ -235,16 +240,27 @@ def _make_property_on_demand(var):
     This property is a simple wrapper around the variable's compute method.
 
     """
-    if "compute" not in var.metadata:
+    if "compute_method" not in var.metadata:
         raise KeyError(
             "No compute method found for on_demand variable "
             f"'{var.name}'. A method decorated with '@{var.name}.compute' "
             "is required in the class definition."
         )
 
-    get_method = var.metadata["compute"]
+    var_name = var.name
+    compute_method = var.metadata["compute_method"]
+    compute_cache = var.metadata["compute_cache"]
+
+    def compute_and_cache(self):
+        value = compute_method(self)
+
+        if compute_cache:
+            key = self.__xsimlab_od_keys__[var_name]
+            self.__xsimlab_state__[key] = value
+
+        return value
 
-    return property(fget=get_method, doc=var_details(var))
+    return property(fget=compute_and_cache, doc=var_details(var))
 
 
 def _make_property_group(var):
@@ -315,6 +331,7 @@ def runtime(meth=None, args=None):
         - ``batch`` : current simulation number in the batch
         - ``sim_start`` : simulation start (date)time
         - ``sim_end`` : simulation end (date)time
+        - ``nsteps``: total number of simulation steps
         - ``step`` : current step number
         - ``step_start`` : current step start (date)time
         - ``step_end``: current step end (date)time

xsimlab/tests/fixture_process.py

@@ -119,6 +119,7 @@ def __init__(self):
         model,
         state,
         state_keys={"some_var": ("some_process", "some_var")},
+        od_keys={"some_od_var": ("some_process", "some_od_var")},
     )
     another_process = _init_process(
         AnotherProcess,
@@ -146,6 +147,7 @@ def __init__(self):
             "group_var": [("some_process", "some_var")],
         },
         od_keys={
+            "od_var": ("example_process", "od_var"),
             "in_foreign_od_var": ("some_process", "some_od_var"),
             "group_var": [("some_process", "some_od_var")],
         },

xsimlab/tests/test_model.py

@@ -319,3 +319,51 @@ def test_context_manager(self):
 
     def test_repr(self, simple_model, simple_model_repr):
         assert repr(simple_model) == simple_model_repr
+
+
+def test_on_demand_cache():
+    @xs.process
+    class P1:
+        var = xs.on_demand(dims="x")
+        cached_var = xs.on_demand(dims="x")
+
+        @var.compute
+        def _compute_var(self):
+            return np.random.rand(10)
+
+        @cached_var.compute(cache=True)
+        def _compute_cached_var(self):
+            return np.random.rand(10)
+
+    @xs.process
+    class P2:
+        var = xs.foreign(P1, "var")
+        cached_var = xs.foreign(P1, "cached_var")
+        view = xs.variable(dims="x", intent="out")
+        cached_view = xs.variable(dims="x", intent="out")
+
+        def run_step(self):
+            self.view = self.var
+            self.cached_view = self.cached_var
+
+    @xs.process
+    class P3:
+        p1_view = xs.foreign(P1, "var")
+        p1_cached_view = xs.foreign(P1, "cached_var")
+        p2_view = xs.foreign(P2, "view")
+        p2_cached_view = xs.foreign(P2, "cached_view")
+
+        def initialize(self):
+            self._p1_cached_view_init = self.p1_cached_view
+
+        def run_step(self):
+            # P1.var's compute method called twice
+            assert not np.all(self.p1_view == self.p2_view)
+            # P1.cached_var's compute method called once
+            assert self.p1_cached_view is self.p2_cached_view
+            # check cache cleared between simulation stages
+            assert not np.all(self.p1_cached_view == self._p1_cached_view_init)
+
+    model = xs.Model({"p1": P1, "p2": P2, "p3": P3})
+    model.execute("initialize", {})
+    model.execute("run_step", {})

xsimlab/variable.py

@@ -23,14 +23,22 @@ class VarIntent(Enum):
     INOUT = "inout"
 
 
-def compute(self, method):
+def compute(self, method=None, *, cache=False):
     """A decorator that, when applied to an on-demand variable, returns a
     value for that variable.
 
     """
-    self.metadata["compute"] = method
 
-    return method
+    def attach_to_metadata(method):
+        self.metadata["compute_method"] = method
+        self.metadata["compute_cache"] = cache
+
+        return method
+
+    if method is None:
+        return attach_to_metadata
+    else:
+        return attach_to_metadata(method)
 
 
 # monkey patch, waiting for cleaner solution:
@@ -284,13 +292,8 @@ def on_demand(
     Like other variables, such variable should be declared in a
     process class. Additionally, it requires its own method to compute
     its value, which must be defined in the same class and decorated
-    (e.g., using `@myvar.compute` if the name of the variable is
-    `myvar`).
-
-    An on-demand variable is always an output variable (i.e., intent='out').
-
-    Its computation usually involves other variables, although this is
-    not required.
+    (e.g., using ``@myvar.compute`` if the name of the variable is
+    ``myvar``).
 
     These variables may be useful, e.g., for model diagnostics.
 
@@ -318,6 +321,17 @@ def on_demand(
         and 'object_codec'. See :func:`zarr.creation.create` for details
         about these options. Other keys are ignored.
 
+    Notes
+    -----
+    An on-demand variable is always an output variable (i.e., intent='out').
+
+    Its computation usually involves other variables, although this is
+    not required.
+
+    It is possible to cache its value at each simulation stage, by applying
+    the compute decorator like this: ``@myvar.compute(cache=True)``. This is
+    useful if the variable is meant to be accessed many times in other processes.
+
     See Also
     --------
     :func:`variable`