Include parent classes in process name/class reverse lookup (#45)

* include parent classes in process name/class reverse lookup

* add tests

* DOC: add subsection in create_model section

* avoid conda solver conflicts for doc environment

* add what's new entry

Author: benbovy

doc/create_model.rst

@@ -291,3 +291,52 @@ In this latter case, users will have to provide initial values of
    possible to modify these instances by adding, updating or removing
    processes. Both methods ``.update_processes()`` and
    ``.drop_processes()`` always return new instances of ``Model``.
+
+Customize existing processes
+----------------------------
+
+Sometimes we only want to update an existing model with very minor
+changes.
+
+As an example, let's update ``model2`` by using a fixed grid (i.e.,
+with hard-coded values for grid spacing and length). One way to
+achieve this is to create a small new process class that sets
+the values of ``spacing`` and ``length``:
+
+.. literalinclude:: scripts/advection_model.py
+   :lines: 151-158
+
+However, one drawback of this "additive" approach is that the number
+of processes in a model might become unnecessarily high:
+
+.. literalinclude:: scripts/advection_model.py
+   :lines: 161-161
+
+Alternatively, it is possible to write a process class that inherits
+from ``UniformGrid1D``, in which we can re-declare variables *and/or*
+re-define "runtime" methods:
+
+.. literalinclude:: scripts/advection_model.py
+   :lines: 164-172
+
+We can here directly update the model and replace the original process
+``UniformGrid1D`` by the inherited class ``FixedGrid``. Foreign
+variables that refer to ``UniformGrid1D`` will still correctly point
+to the ``grid`` process in the updated model:
+
+.. literalinclude:: scripts/advection_model.py
+   :lines: 175-175
+
+.. warning::
+
+   This feature is experimental! It may be removed in a next version of
+   xarray-simlab.
+
+   In particular, linking foreign variables in a model is ambiguous
+   when both conditions below are met:
+
+   - two different processes in a model inherit from a common class
+     (except ``object``)
+
+   - a third process in the same model has a foreign variable that
+     links to that common class

doc/environment.yml

@@ -5,8 +5,8 @@ channels:
 dependencies:
   - attrs=19.1.0
   - python=3.7
-  - numpy=1.16.0
-  - pandas=0.23.3
+  - numpy=1.17.2
+  - pandas=0.25.1
   - xarray=0.13.0
   - ipython=7.8.0
   - matplotlib=3.0.2

doc/scripts/advection_model.py

@@ -146,3 +146,30 @@ def initialize(self):
 
 
 model4 = model2.drop_processes('init')
+
+
+@xs.process
+class FixedGridParams(object):
+    spacing = xs.foreign(UniformGrid1D, 'spacing', intent='out')
+    length = xs.foreign(UniformGrid1D, 'length', intent='out')
+
+    def initialize(self):
+        self.spacing = 0.01
+        self.length = 1.
+
+
+model5 = model2.update_processes({'fixed_grid_params': FixedGridParams})
+
+
+@xs.process
+class FixedGrid(UniformGrid1D):
+    spacing = xs.variable(description='uniform spacing', intent='out')
+    length = xs.variable(description='total length', intent='out')
+
+    def initialize(self):
+        self.spacing = 0.01
+        self.length = 1.
+        super(FixedGrid, self).initialize()
+
+
+model6 = model2.update_processes({'grid': FixedGrid})

doc/whats_new.rst

@@ -6,6 +6,16 @@ Release Notes
 v0.3.0 (Unreleased)
 -------------------
 
+Breaking changes
+~~~~~~~~~~~~~~~~
+
+- It is now possible to use class inheritance to customize a process
+  without re-writing the class from scratch and without breaking the
+  links between (foreign) variables when replacing the process in a
+  model (:issue:`45`). Although it should work just fine in most
+  cases, there are potential caveats. This should be considered as an
+  experimental, possibly breaking change.
+
 Enhancements
 ~~~~~~~~~~~~
 

xsimlab/model.py

@@ -43,7 +43,7 @@ def __init__(self, processes_cls):
         self._processes_cls = processes_cls
         self._processes_obj = {k: cls() for k, cls in processes_cls.items()}
 
-        self._reverse_lookup = {cls: k for k, cls in processes_cls.items()}
+        self._reverse_lookup = self._get_reverse_lookup(processes_cls)
 
         self._input_vars = None
 
@@ -53,6 +53,24 @@ def __init__(self, processes_cls):
         # a cache for group keys
         self._group_keys = {}
 
+    def _get_reverse_lookup(self, processes_cls):
+        """Return a dictionary with process classes as keys and process names
+        as values.
+
+        Additionally, the returned dictionary maps all parent classes
+        to one (str) or several (list) process names.
+
+        """
+        reverse_lookup = defaultdict(list)
+
+        for p_name, p_cls in processes_cls.items():
+            # exclude `object` base class from lookup
+            for cls in p_cls.mro()[:-1]:
+                reverse_lookup[cls].append(p_name)
+
+        return {k: v[0] if len(v) == 1 else v
+                for k, v in reverse_lookup.items()}
+
     def bind_processes(self, model_obj):
         for p_name, p_obj in self._processes_obj.items():
             p_obj.__xsimlab_model__ = model_obj
@@ -92,6 +110,17 @@ def _get_var_key(self, p_name, var):
                     .format(target_p_cls.__name__, var.name, p_name)
                 )
 
+            elif isinstance(target_p_name, list):
+                raise ValueError(
+                    "Process class {!r} required by foreign variable '{}.{}' "
+                    "is used (possibly via one its child classes) by multiple "
+                    "processes: {}"
+                    .format(
+                        target_p_cls.__name__, p_name, var.name,
+                        ', '.join(['{!r}'.format(n) for n in target_p_name])
+                    )
+                )
+
             store_key, od_key = self._get_var_key(target_p_name, target_var)
 
         elif var_type == VarType.GROUP:

xsimlab/tests/test_model.py

@@ -109,6 +109,22 @@ def test_get_stage_processes(self, model):
         expected = [model['roll'], model['profile']]
         assert model._p_run_step == expected
 
+    def test_process_inheritance(self, model):
+        @xs.process
+        class InheritedProfile(Profile):
+            pass
+
+        new_model = model.update_processes(
+            {'profile': InheritedProfile})
+
+        assert type(new_model['profile']) is InheritedProfile
+        assert isinstance(new_model['profile'], Profile)
+
+        with pytest.raises(ValueError) as excinfo:
+            invalid_model = model.update_processes(
+                {'profile2': InheritedProfile})
+        assert "multiple processes" in str(excinfo.value)
+
 
 class TestModel(object):