Allows setting multiple groups for a variable (#71)

* rename var metadata key 'group' -> 'groups'

* some code style changes

* update variable API

* fix process tests

* fix failing tests and warnings

* update release notes

* update doc

* add multiple groups test

doc/examples/landscape-evolution-model.ipynb

@@ -2091,7 +2091,7 @@
     "    active_nodes = xs.foreign(ClosedBoundaryFaces, 'active_nodes')\n",
     "    x = xs.foreign(Grid2D, 'x')\n",
     "    \n",
-    "    uplift = xs.variable(intent='out', group='uplift')\n",
+    "    uplift = xs.variable(intent='out', groups='uplift')\n",
     "\n",
     "    def initialize(self):\n",
     "        mask = self.active_nodes\n",

doc/framework.rst

@@ -132,12 +132,12 @@ Group variables
 In some cases, using group variables may provide an elegant
 alternative to hard-coded links between processes.
 
-The membership of variables to a group is defined via their ``group``
-attribute. If you want to use in a separate process all the variables
-of a group, instead of explicitly declaring foreign variables you can
-declare a :func:`~xsimlab.group` variable. The latter behaves like an
-iterable of foreign variables pointing to each of the variables that
-are members of the group, across the model.
+The membership of variables to one or several groups is defined via their
+``groups`` attribute. If you want to use in a separate process all the variables
+of a group, instead of explicitly declaring foreign variables you can declare a
+:func:`~xsimlab.group` variable. The latter behaves like an iterable of foreign
+variables pointing to each of the variables that are members of the group,
+across the model.
 
 Note that group variables only support ``intent='in'``, i.e, group
 variables should only be used to get the values of multiple foreign

doc/scripts/advection_model.py

@@ -72,7 +72,7 @@ class AdvectionLax:
     v = xs.variable(dims=[(), 'x'], description='velocity')
     grid_spacing = xs.foreign(UniformGrid1D, 'spacing')
     u = xs.foreign(ProfileU, 'u')
-    u_advected = xs.variable(dims='x', intent='out', group='u_vars')
+    u_advected = xs.variable(dims='x', intent='out', groups='u_vars')
 
     @xs.runtime(args='step_delta')
     def run_step(self, dt):
@@ -115,7 +115,7 @@ class SourcePoint:
     loc = xs.variable(description='source location')
     flux = xs.variable(description='source flux')
     x = xs.foreign(UniformGrid1D, 'x')
-    u_source = xs.variable(dims='x', intent='out', group='u_vars')
+    u_source = xs.variable(dims='x', intent='out', groups='u_vars')
 
     @property
     def nearest_node(self):

doc/whats_new.rst

@@ -11,6 +11,18 @@ Breaking changes
 
 - Python 3.6 is now the oldest supported version (:issue:`70`).
 
+Depreciations
+~~~~~~~~~~~~~
+
+- Using the ``group`` parameter in ``xsimlab.variable`` and
+  ``xsimlab.on_demand`` is depreciated; use ``groups`` instead.
+
+Enhancements
+~~~~~~~~~~~~
+
+- It is now possible to assign multiple groups to a single variable
+  (:issue:`71`).
+
 Bug fixes
 ~~~~~~~~~
 
@@ -33,8 +45,8 @@ Breaking changes
   and ``Model.finalize`` have been removed in favor of
   ``Model.execute`` (:issue:`59`).
 
-Deprecations
-~~~~~~~~~~~~
+Depreciations
+~~~~~~~~~~~~~
 
 - ``run_step`` methods defined in process classes won't accept anymore
   current step duration as a positional argument by default. Use the

xsimlab/model.py

@@ -1,10 +1,9 @@
 from collections import OrderedDict, defaultdict
-from inspect import isclass
 
 from .variable import VarIntent, VarType
 from .process import (filter_variables, get_process_cls,
                       get_target_variable, SimulationStage)
-from .utils import AttrMapping, ContextMixin, has_method, variables_dict
+from .utils import AttrMapping, ContextMixin, variables_dict
 from .formatting import repr_model
 
 
@@ -39,6 +38,7 @@ class _ModelBuilder:
       step of a simulation
 
     """
+
     def __init__(self, processes_cls):
         self._processes_cls = processes_cls
         self._processes_obj = {k: cls() for k, cls in processes_cls.items()}
@@ -140,11 +140,8 @@ def _get_group_var_keys(self, group):
         store_keys = []
         od_keys = []
 
-        filter_group = lambda var: (var.metadata.get('group') == group and
-                                    var.metadata['var_type'] != VarType.GROUP)
-
         for p_name, p_obj in self._processes_obj.items():
-            for var in filter_variables(p_obj, func=filter_group).values():
+            for var in filter_variables(p_obj, group=group).values():
                 store_key, od_key = self._get_var_key(p_name, var)
 
                 if store_key is not None:
@@ -179,10 +176,9 @@ def ensure_no_intent_conflict(self):
         intent='out' targets the same variable.
 
         """
-        filter_out = lambda var: (
-            var.metadata['intent'] == VarIntent.OUT and
-            var.metadata['var_type'] != VarType.ON_DEMAND
-        )
+        def filter_out(var):
+            return (var.metadata['intent'] == VarIntent.OUT and
+                    var.metadata['var_type'] != VarType.ON_DEMAND)
 
         targets = defaultdict(list)
 
@@ -227,11 +223,12 @@ def get_input_variables(self):
           model inputs.
 
         """
-        filter_in = lambda var: (
-            var.metadata['var_type'] != VarType.GROUP and
-            var.metadata['intent'] != VarIntent.OUT
-        )
-        filter_out = lambda var: var.metadata['intent'] == VarIntent.OUT
+        def filter_in(var):
+            return (var.metadata['var_type'] != VarType.GROUP and
+                    var.metadata['intent'] != VarIntent.OUT)
+
+        def filter_out(var):
+            return var.metadata['intent'] == VarIntent.OUT
 
         in_keys = []
         out_keys = []
@@ -381,6 +378,7 @@ class Model(AttrMapping, ContextMixin):
     value before running the model.
 
     """
+
     def __init__(self, processes):
         """
         Parameters

xsimlab/process.py

@@ -92,7 +92,7 @@ def filter_variables(process, var_type=None, intent=None, group=None,
 
     if group is not None:
         vars = {k: v for k, v in vars.items()
-                if v.metadata.get('group') == group}
+                if group in v.metadata.get('groups', [])}
 
     if func is not None:
         vars = {k: v for k, v in vars.items() if func(v)}

xsimlab/tests/fixture_model.py

@@ -47,7 +47,7 @@ class Roll:
     shift = xs.variable(description=('shift profile by a nb. of points'),
                         attrs={'units': 'unitless'})
     u = xs.foreign(Profile, 'u')
-    u_diff = xs.variable(dims='x', group='diff', intent='out')
+    u_diff = xs.variable(dims='x', groups='diff', intent='out')
 
     def run_step(self):
         self.u_diff = np.roll(self.u, self.shift) - self.u
@@ -57,7 +57,7 @@ def run_step(self):
 class Add:
     offset = xs.variable(description=('offset * dt added every time step '
                                       'to profile u'))
-    u_diff = xs.variable(group='diff', intent='out')
+    u_diff = xs.variable(groups='diff', intent='out')
 
     @xs.runtime(args='step_delta')
     def run_step(self, dt):
@@ -67,7 +67,7 @@ def run_step(self, dt):
 @xs.process
 class AddOnDemand:
     offset = xs.variable(description='offset added to profile u')
-    u_diff = xs.on_demand(group='diff')
+    u_diff = xs.on_demand(groups='diff')
 
     @u_diff.compute
     def _compute_u_diff(self):

xsimlab/tests/fixture_process.py

@@ -10,8 +10,8 @@
 @xs.process
 class SomeProcess:
     """Just used for foreign variables in ExampleProcess."""
-    some_var = xs.variable(group='some_group', intent='out')
-    some_od_var = xs.on_demand(group='some_group')
+    some_var = xs.variable(groups='some_group', intent='out')
+    some_od_var = xs.on_demand(groups='some_group')
 
     @some_od_var.compute
     def compute_some_od_var(self):
@@ -29,7 +29,7 @@ class AnotherProcess:
 class ExampleProcess:
     """A process with complete interface for testing."""
     in_var = xs.variable(dims=['x', ('x', 'y')], description='input variable')
-    out_var = xs.variable(group='example_group', intent='out')
+    out_var = xs.variable(groups='example_group', intent='out')
     inout_var = xs.variable(intent='inout')
     od_var = xs.on_demand()
 
@@ -80,7 +80,7 @@ def in_var_details():
     - type : variable
     - intent : in
     - dims : (('x',), ('x', 'y'))
-    - group : None
+    - groups : ()
     - attrs : {}
     """)
 

xsimlab/tests/test_model.py

@@ -48,6 +48,21 @@ def test_set_process_keys(self, model, p_name,
         assert actual_store_keys == expected_store_keys
         assert actual_od_keys == expected_od_keys
 
+    def test_multiple_groups(self):
+        @xs.process
+        class A:
+            v = xs.variable(groups=['g1', 'g2'])
+
+        @xs.process
+        class B:
+            g1 = xs.group('g1')
+            g2 = xs.group('g2')
+
+        m = xs.Model({'a': A, 'b': B})
+
+        assert m.b.__xsimlab_store_keys__['g1'] == [('a', 'v')]
+        assert m.b.__xsimlab_store_keys__['g2'] == [('a', 'v')]
+
     def test_get_all_variables(self, model):
         assert all([len(t) == 2 for t in model.all_vars])
         assert all([p_name in model for p_name, _ in model.all_vars])

xsimlab/tests/test_variable.py

@@ -1,7 +1,7 @@
 import pytest
 
 from xsimlab.tests.fixture_process import ExampleProcess
-from xsimlab.variable import _as_dim_tuple, foreign
+from xsimlab.variable import _as_dim_tuple, _as_group_tuple, foreign
 
 
 @pytest.mark.parametrize("dims,expected", [
@@ -27,6 +27,24 @@ def test_as_dim_tuple_invalid():
     assert "('x',), ('y',) and ('x', 'y'), ('y', 'x')" in str(excinfo.value)
 
 
+@pytest.mark.parametrize("groups,group,expected", [
+    (None, None, ()),
+    ('group1', None, ('group1',)),
+    (['group1', 'group2'], None, ('group1', 'group2')),
+    ('group1', 'group2', ('group1', 'group2')),
+    ('group1', 'group1', ('group1',))
+])
+def test_as_group_tuple(groups, group, expected):
+    if group is not None:
+        with pytest.warns(FutureWarning):
+            actual = _as_group_tuple(groups, group)
+
+    else:
+        actual = _as_group_tuple(groups, group)
+
+    assert actual == expected
+
+
 def test_foreign():
     with pytest.raises(ValueError) as excinfo:
         foreign(ExampleProcess, 'some_var', intent='inout')

xsimlab/variable.py

@@ -1,5 +1,6 @@
 from enum import Enum
 import itertools
+import warnings
 
 import attr
 from attr._make import _CountingAttr
@@ -53,7 +54,9 @@ def _as_dim_tuple(dims):
         dims = [dims]
 
     # check ndim uniqueness could be simpler but provides detailed error msg
-    fget_ndim = lambda dims: len(dims)
+    def fget_ndim(dims):
+        return len(dims)
+
     dims_sorted = sorted(dims, key=fget_ndim)
     ndim_groups = [list(g)
                    for _, g in itertools.groupby(dims_sorted, fget_ndim)]
@@ -70,8 +73,29 @@ def _as_dim_tuple(dims):
     return tuple(dims)
 
 
-def variable(dims=(), intent='in', group=None, default=attr.NOTHING,
-             validator=None, description='', attrs=None):
+def _as_group_tuple(groups, group):
+    if groups is None:
+        groups = []
+    elif isinstance(groups, str):
+        groups = [groups]
+    else:
+        groups = list(groups)
+
+    if group is not None:
+        warnings.warn(
+            "Setting variable group using `group` is depreciated; "
+            "use `groups`.",
+            FutureWarning,
+            stacklevel=2
+        )
+        if group not in groups:
+            groups.append(group)
+
+    return tuple(groups)
+
+
+def variable(dims=(), intent='in', group=None, groups=None,
+             default=attr.NOTHING, validator=None, description='', attrs=None):
     """Create a variable.
 
     Variables store useful metadata such as dimension labels, a short
@@ -100,7 +124,9 @@ def variable(dims=(), intent='in', group=None, default=attr.NOTHING,
         process may update the value of the variable).
         (default: input).
     group : str, optional
-        Variable group.
+        Variable group (depreciated, use ``groups`` instead).
+    groups : str or list, optional
+        Variable group(s).
     default : any, optional
         Single default value for the variable, ignored when ``intent='out'``
         (default: NOTHING). A default value may also be set using a decorator.
@@ -128,7 +154,7 @@ def variable(dims=(), intent='in', group=None, default=attr.NOTHING,
     metadata = {'var_type': VarType.VARIABLE,
                 'dims': _as_dim_tuple(dims),
                 'intent': VarIntent(intent),
-                'group': group,
+                'groups': _as_group_tuple(groups, group),
                 'attrs': attrs or {},
                 'description': description}
 
@@ -143,7 +169,7 @@ def variable(dims=(), intent='in', group=None, default=attr.NOTHING,
                        init=_init, repr=_repr, kw_only=True)
 
 
-def on_demand(dims=(), group=None, description='', attrs=None):
+def on_demand(dims=(), group=None, groups=None, description='', attrs=None):
     """Create a variable that is computed on demand.
 
     Instead of being computed systematically at every step of a simulation
@@ -173,6 +199,9 @@ def on_demand(dims=(), group=None, description='', attrs=None):
         the variable accepts different numbers of dimensions.
         This should not include a time dimension, which may always be added.
     group : str, optional
+        Variable group (depreciated, use ``groups`` instead).
+    groups : str or list, optional
+        Variable group(s).
     description : str, optional
         Short description of the variable.
     attrs : dict, optional
@@ -187,7 +216,7 @@ def on_demand(dims=(), group=None, description='', attrs=None):
     metadata = {'var_type': VarType.ON_DEMAND,
                 'dims': _as_dim_tuple(dims),
                 'intent': VarIntent.OUT,
-                'group': group,
+                'groups': _as_group_tuple(groups, group),
                 'attrs': attrs or {},
                 'description': description}
 
@@ -258,7 +287,7 @@ def group(name):
 
     Parameters
     ----------
-    group : str
+    name : str
         Name of the group.
 
     See Also