Allow specification of dims instead of shape

RELEASE-NOTES.md

@@ -17,6 +17,7 @@
 - `pm.LKJCholeskyCov` now automatically computes and returns the unpacked Cholesky decomposition, the correlations and the standard deviations of the covariance matrix (see [#3881](https://github.com/pymc-devs/pymc3/pull/3881)).
 - `pm.Data` container can now be used for index variables, i.e with integer data and not only floats (issue [#3813](https://github.com/pymc-devs/pymc3/issues/3813), fixed by [#3925](https://github.com/pymc-devs/pymc3/pull/3925)).
 - `pm.Data` container can now be used as input for other random variables (issue [#3842](https://github.com/pymc-devs/pymc3/issues/3842), fixed by [#3925](https://github.com/pymc-devs/pymc3/pull/3925)).
+- Allow users to specify coordinates and dimension names instead of numerical shapes. This makes interoperability with ArviZ easier. ([see #3551](https://github.com/pymc-devs/pymc3/pull/3551))
 - Plots and Stats API sections now link to ArviZ documentation [#3927](https://github.com/pymc-devs/pymc3/pull/3927)
 - Add `SamplerReport` with properties `n_draws`, `t_sampling` and `n_tune` to SMC. `n_tune` is always 0 [#3931](https://github.com/pymc-devs/pymc3/issues/3931).
 
@@ -48,6 +49,7 @@
 ### New features
 - Implemented robust u turn check in NUTS (similar to stan-dev/stan#2800). See PR [#3605]
 - Add capabilities to do inference on parameters in a differential equation with `DifferentialEquation`. See [#3590](https://github.com/pymc-devs/pymc3/pull/3590) and [#3634](https://github.com/pymc-devs/pymc3/pull/3634).
+>>>>>>> master
 - Distinguish between `Data` and `Deterministic` variables when graphing models with graphviz. PR [#3491](https://github.com/pymc-devs/pymc3/pull/3491).
 - Sequential Monte Carlo - Approximate Bayesian Computation step method is now available. The implementation is in an experimental stage and will be further improved.
 - Added `Matern12` covariance function for Gaussian processes. This is the Matern kernel with nu=1/2.

pymc3/data.py

@@ -22,13 +22,15 @@
 import pymc3 as pm
 import theano.tensor as tt
 import theano
+import pandas as pd
 
 __all__ = [
     'get_data',
     'GeneratorAdapter',
     'Minibatch',
     'align_minibatches',
     'Data',
+    'TidyData',
 ]
 
 
@@ -479,7 +481,7 @@ class Data:
     https://docs.pymc.io/notebooks/data_container.html
     """
 
-    def __new__(self, name, value):
+    def __new__(self, name, value, *, dims=None, export_dims=False):
         if isinstance(value, list):
             value = np.array(value)
 
@@ -497,10 +499,130 @@ def __new__(self, name, value):
         # transforms it to something digestible for pymc3
         shared_object = theano.shared(pm.model.pandas_to_array(value), name)
 
+        if isinstance(dims, str):
+            dims = (dims,)
+        if dims is not None and len(dims) != shared_object.ndim:
+            raise ValueError('Length of `dims` must match the dimensionality '
+                             'of the dataset.')
+
+        coords = {}
+        if isinstance(value, (pd.Series, pd.DataFrame)):
+            name = None
+            if dims is not None:
+                name = dims[0]
+            if (name is None 
+                    and value.index.name is not None
+                    and value.index.name.isidentifier()):
+                name = value.index.name
+            if name is not None:
+                coords[name] = value.index
+        if isinstance(value, pd.DataFrame):
+            name = None
+            if dims is not None:
+                name = dims[1]
+            if (name is None
+                    and value.columns.name is not None
+                    and value.columns.name.isidentifier()):
+                name = value.columns.name
+            if name is not None:
+                coords[name] = value.columns
+        if isinstance(value, np.ndarray) and dims is not None:
+            if len(dims) != value.ndim:
+                raise ValueError('Invalid data shape %s. The rank of the dataset '
+                                 'must match the length of `dims`.'
+                                 % value.shape)
+            for size, dim in zip(value.shape, dims):
+                coord = model.coords.get(dim, None)
+                if coord is None:
+                    coords[dim] = pd.RangeIndex(size, name=dim)
+
+        if export_dims:
+            model.add_coords(coords)
+
         # To draw the node for this variable in the graphviz Digraph we need
         # its shape.
         shared_object.dshape = tuple(shared_object.shape.eval())
+        if dims is not None:
+            shape_dims = model.shape_from_dims(dims)
+            if shared_object.dshape != shape_dims:
+                raise ValueError('Invalid shape. It is %s but the dimensions '
+                                 'suggest %s.'
+                                 % (shared_object.dshape, shape_dims))
+
+        model.add_random_variable(shared_object, dims=dims)
+        return shared_object
 
-        model.add_random_variable(shared_object)
 
-        return shared_object
+class _IndexAccessor:
+    def __init__(self, data):
+        self._data = data
+
+    def __getitem__(self, key):
+        category = self._data._col_as_category(key)
+        vals = self._data.data.reset_index().loc[:, key]
+        return pd.Categorical(vals, dtype=category).codes
+
+
+class TidyData:
+    def __init__(self, data, copy_data=True, import_dims=None, model=None):
+        self.data = data
+        self._shared_vars = {}
+        self._category_cols = {}
+        self._index_dict = _IndexAccessor(self)
+
+        if import_dims is not None:
+            model = pm.model.modelcontext(model)
+            coords = self._extract_coords(import_dims)
+            model.add_coords(coords)
+
+    @property
+    def idxs(self):
+        return self._index_dict
+
+    def _col_as_category(self, key):
+        if key in self._category_cols:
+            return self._category_cols[key]
+        data = self.data.reset_index()
+        values = data.loc[:, key]
+        if values.dtype.name != 'category':
+            values = values.astype('category')
+        self._category_cols[key] = values.dtype
+        return values.dtype
+
+    def __getitem__(self, key):
+        if key not in self.data.columns:
+            raise KeyError('Unknown column %s' % key)
+        if key in self._shared_vars:
+            return self._shared_vars[key]
+
+        shared_var = theano.shared(self.data.loc[:, key].values)
+        self._shared_vars[key] = shared_var
+        return shared_var
+
+    def _extract_coords(self, dims):
+        data = self.data
+        dims = set(dims)
+        coords = {}
+
+        if data.index.name is not None and data.index.name in dims:
+            dims.remove(data.index.name)
+            coords[data.index.name] = data.index
+
+        # We want to iterate over index columns of a multi index as well
+        data = data.reset_index()
+        for col in data.columns:
+            if col not in dims:
+                continue
+            dims.remove(col)
+            category = self._col_as_category(col)
+            cat = pd.Categorical(category.categories, dtype=category)
+            coords[col] = pd.CategoricalIndex(cat, name=col)
+
+        if dims:
+            raise KeyError('Unknown columns: %s' % dims)
+
+        return coords
+
+    @property
+    def columns(self):
+        return self.data.columns

pymc3/distributions/distribution.py

@@ -56,17 +56,32 @@ def __new__(cls, name, *args, **kwargs):
                             "a 'with model:' block, or use the '.dist' syntax "
                             "for a standalone distribution.")
 
-        if isinstance(name, string_types):
-            data = kwargs.pop('observed', None)
-            cls.data = data
-            if isinstance(data, ObservedRV) or isinstance(data, FreeRV):
-                raise TypeError("observed needs to be data but got: {}".format(type(data)))
-            total_size = kwargs.pop('total_size', None)
-            dist = cls.dist(*args, **kwargs)
-            return model.Var(name, dist, data, total_size)
-        else:
+        if not isinstance(name, string_types):
             raise TypeError("Name needs to be a string but got: {}".format(name))
 
+        data = kwargs.pop('observed', None)
+        cls.data = data
+        if isinstance(data, ObservedRV) or isinstance(data, FreeRV):
+            raise TypeError("observed needs to be data but got: {}".format(type(data)))
+        total_size = kwargs.pop('total_size', None)
+
+        dims = kwargs.pop('dims', None)
+        has_shape = 'shape' in kwargs
+        shape = kwargs.pop('shape', None)
+        if dims is not None:
+            if shape is not None:
+                raise ValueError("Specify only one of 'dims' and 'shape'")
+            if isinstance(dims, string_types):
+                dims = (dims,)
+            shape = model.shape_from_dims(dims)
+
+        # Some distribution do not accept shape=None
+        if has_shape or shape is not None:
+            dist = cls.dist(*args, **kwargs, shape=shape)
+        else:
+            dist = cls.dist(*args, **kwargs)
+        return model.Var(name, dist, data, total_size, dims=dims)
+
     def __getnewargs__(self):
         return _Unpickling,
 
@@ -77,7 +92,7 @@ def dist(cls, *args, **kwargs):
         return dist
 
     def __init__(self, shape, dtype, testval=None, defaults=(),
-                 transform=None, broadcastable=None):
+                 transform=None, broadcastable=None, dims=None):
         self.shape = np.atleast_1d(shape)
         if False in (np.floor(self.shape) == self.shape):
             raise TypeError("Expected int elements in shape")

pymc3/model.py

@@ -819,8 +819,12 @@ def __new__(cls, *args, **kwargs):
         instance._theano_config = theano_config
         return instance
 
-    def __init__(self, name='', model=None, theano_config=None):
+    def __init__(self, name='', model=None, theano_config=None, coords=None):
         self.name = name
+        self.coords = {}
+        self.RV_dims = {}
+        self.add_coords(coords)
+
         if self.parent is not None:
             self.named_vars = treedict(parent=self.parent.named_vars)
             self.free_RVs = treelist(parent=self.parent.free_RVs)
@@ -972,7 +976,37 @@ def cont_vars(self):
         """All the continuous variables in the model"""
         return list(typefilter(self.vars, continuous_types))
 
-    def Var(self, name, dist, data=None, total_size=None):
+    def shape_from_dims(self, dims):
+        shape = []
+        if len(set(dims)) != len(dims):
+            raise ValueError('Can not contain the same dimension name twice.')
+        for dim in dims:
+            if dim not in self.coords:
+                raise ValueError('Unknown dimension name %s. All dimension '
+                                 'names must be specified in the `coords` '
+                                 'argument of the model or through a pm.Data '
+                                 'variable.' % dim)
+            shape.extend(self.coords[dim].shape)
+        return tuple(shape)
+
+    def add_coords(self, coords):
+        if coords is None:
+            return
+
+        for name in coords:
+            if not name.isidentifier():
+                raise ValueError('Invalid dimension name. Must be a python '
+                                 'identifier: %s' % name)
+            if name == 'draw' or name == 'chain':
+                raise ValueError('Dimensions can not be named `draw` or `chain`.')
+            if name in self.coords:
+                if not coords[name].equals(self.coords[name]):
+                    raise ValueError(
+                        'Duplicate and incompatiple coordinate: %s.' % name)
+            else:
+                self.coords[name] = coords[name]
+
+    def Var(self, name, dist, data=None, total_size=None, dims=None):
         """Create and add (un)observed random variable to the model with an
         appropriate prior distribution.
 
@@ -985,12 +1019,15 @@ def Var(self, name, dist, data=None, total_size=None):
            the variable is unobserved.
         total_size: scalar
             upscales logp of variable with ``coef = total_size/var.shape[0]``
+        dims : tuple
+            Dimension names for the variable.
 
         Returns
         -------
         FreeRV or ObservedRV
         """
         name = self.name_for(name)
+
         if data is None:
             if getattr(dist, "transform", None) is None:
                 with self:
@@ -1009,7 +1046,7 @@ def Var(self, name, dist, data=None, total_size=None):
                                 name=name,
                                 orig_name=get_transformed_name(name, dist.transform)))
                 self.deterministics.append(var)
-                self.add_random_variable(var)
+                self.add_random_variable(var, dims)
                 return var
         elif isinstance(data, dict):
             with self:
@@ -1032,14 +1069,21 @@ def Var(self, name, dist, data=None, total_size=None):
                 self.missing_values.append(var.missing_values)
                 self.named_vars[var.missing_values.name] = var.missing_values
 
-        self.add_random_variable(var)
+        self.add_random_variable(var, dims)
         return var
 
-    def add_random_variable(self, var):
+    def add_random_variable(self, var, dims=None):
         """Add a random variable to the named variables of the model."""
         if self.named_vars.tree_contains(var.name):
             raise ValueError(
                 "Variable name {} already exists.".format(var.name))
+
+        if dims is not None:
+            if isinstance(dims, str):
+                dims = (dims,)
+            assert all(dim in self.coords for dim in dims)
+            self.RV_dims[var.name] = dims
+
         self.named_vars[var.name] = var
         if not hasattr(self, self.name_of(var.name)):
             setattr(self, self.name_of(var.name), var)
@@ -1680,7 +1724,7 @@ def _latex_repr_rv(rv):
     return r'$\text{%s} \sim \text{Deterministic}(%s)$' % (rv.name, r',~'.join(_walk_up_rv(rv)))
 
 
-def Deterministic(name, var, model=None):
+def Deterministic(name, var, model=None, dims=None):
     """Create a named deterministic variable
 
     Parameters
@@ -1695,7 +1739,7 @@ def Deterministic(name, var, model=None):
     model = modelcontext(model)
     var = var.copy(model.name_for(name))
     model.deterministics.append(var)
-    model.add_random_variable(var)
+    model.add_random_variable(var, dims)
     var._repr_latex_ = functools.partial(_latex_repr_rv, var)
     var.__latex__ = var._repr_latex_
     return var