Add methods for combining variables of differing dimensionality (#1597)

* Add stack_cat and unstack_cat methods

This partially resolves #1317.

Change names of methods

stack_cat -> to_stacked_array
unstack_cat -> to_unstacked_dataset

Test that the dtype of the stacked dimensions is preserved

This is not passing at the moment because concatenating None with
a dimension that has values upcasts the combined dtype to object

Fix dtypes of stacked dimensions

This commit ensures that the dtypes of the stacked coordinate match the input dimensions.

Use new index variable rather than patching the old one

I didn't like the inplace modification of a private member.

Handle variable_dim correctly

I also fixed

1. f-string formatting issue
2. Use an OrderedDict as @jhamman recommends

Add documentation to api.rst and reshaping.rst

I also added appropriate See Also sections to the docstrings for
to_stacked_array and to_unstacked_dataset.

Add changes to whats-new

Fixing style errors.

Split up lengthy test

Remove "physical variable" from docs

This is in response to Joe's "nit"

* Fix to_stacked_array with new master

An error arose when checking for the precence of a dimension in array.
The code 'dim in data' no longer works.
Replaced this with 'dim in data.dims'

* Move entry in whats-new to most recent release

* Fix code styling errors

It needs to pass `pycodestyle xarray`

* Improve docstring of to_stacked_array

Added example and additional description.

* Move "See Also" section to end of docstring

* Doc and comment improvements.

* Improve documented example

@benbovy pointed out that the old example was confusing.

* Add name argument to to_stacked_array and test

* Allow level argument to be an int or str

* Remove variable_dim argument of to_unstacked_array

* Actually removed variable_dim

* Change function signature of to_stacked_array

Previously, this function was passed a list of dimensions which should be
stacked together. However, @benbovy found that the function failed when the
_non-stacked_ dimensions were not shared across all variables. Thus, it is
easier to specify the dimensions which should remain unchanged, rather than the
dimensions to be stacked.

The function to_stacked_array now takes an argument ''sample_dim'' which defines
these non-stacked dimensions. If these dims are not shared accross all
variables than an error is raised.

* Fix lint error

The line was too long

* Fix validation and failing tests

1. the test which stacks a scalar and an array doesn't make sense anymore given
the new API.

2. Fixed a bug in the validation code which raised an error almost always.

* Fix typo

* Improve docs and error messages

* Remove extra spaces

* Test warning in to_unstacked_dataset

* Improve formatting and naming

* Fix flake8 error

* Respond to @max-sixty's suggestions

Author: nbren12

doc/api.rst

@@ -208,6 +208,7 @@ Reshaping and reorganizing
    Dataset.transpose
    Dataset.stack
    Dataset.unstack
+   Dataset.to_stacked_array
    Dataset.shift
    Dataset.roll
    Dataset.sortby
@@ -381,6 +382,7 @@ Reshaping and reorganizing
    DataArray.transpose
    DataArray.stack
    DataArray.unstack
+   DataArray.to_unstacked_dataset
    DataArray.shift
    DataArray.roll
    DataArray.sortby

doc/reshaping.rst

@@ -133,6 +133,48 @@ pandas, it does not automatically drop missing values. Compare:
 We departed from pandas's behavior here because predictable shapes for new
 array dimensions is necessary for :ref:`dask`.
 
+Stacking different variables together
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+These stacking and unstacking operations are particularly useful for reshaping
+xarray objects for use in machine learning packages, such as `scikit-learn
+<http://scikit-learn.org/stable/>`_, that usually require two-dimensional numpy
+arrays as inputs. For datasets with only one variable, we only need ``stack``
+and ``unstack``, but combining multiple variables in a
+:py:class:`xarray.Dataset` is more complicated. If the variables in the dataset
+have matching numbers of dimensions, we can call
+:py:meth:`~xarray.Dataset.to_array` and then stack along the the new coordinate.
+But :py:meth:`~xarray.Dataset.to_array` will broadcast the dataarrays together,
+which will effectively tile the lower dimensional variable along the missing
+dimensions. The method :py:meth:`xarray.Dataset.to_stacked_array` allows
+combining variables of differing dimensions without this wasteful copying while
+:py:meth:`xarray.DataArray.to_unstacked_dataset` reverses this operation.
+Just as with :py:meth:`xarray.Dataset.stack` the stacked coordinate is
+represented by a :py:class:`pandas.MultiIndex` object. These methods are used
+like this:
+
+.. ipython:: python
+        data = xr.Dataset(
+            data_vars={'a': (('x', 'y'), [[0, 1, 2], [3, 4, 5]]),
+                      'b': ('x', [6, 7])},
+            coords={'y': ['u', 'v', 'w']}
+        )
+        stacked = data.to_stacked_array("z", sample_dims=['x'])
+        stacked
+        unstacked = stacked.to_unstacked_dataset("z")
+        unstacked
+
+In this example, ``stacked`` is a two dimensional array that we can easily pass to a scikit-learn or another generic
+numerical method.
+
+.. note::
+
+    Unlike with ``stack``,  in ``to_stacked_array``, the user specifies the dimensions they **do not** want stacked.
+    For a machine learning task, these unstacked dimensions can be interpreted as the dimensions over which samples are
+    drawn, whereas the stacked coordinates are the features. Naturally, all variables should possess these sampling
+    dimensions.
+
+
 .. _reshape.set_index:
 
 Set and reset index

doc/whats-new.rst

@@ -18,6 +18,12 @@ What's New
 v0.12.3 (unreleased)
 --------------------
 
+New functions/methods
+~~~~~~~~~~~~~~~~~~~~~
+
+- New methods for reshaping Datasets of variables with different dimensions
+  (:issue:`1317`). By `Noah Brenowitz <https://github.com/nbren12>`_.
+
 Enhancements
 ~~~~~~~~~~~~
 

xarray/core/dataarray.py

@@ -1540,6 +1540,72 @@ def unstack(self, dim: Union[Hashable, Sequence[Hashable], None] = None
         ds = self._to_temp_dataset().unstack(dim)
         return self._from_temp_dataset(ds)
 
+    def to_unstacked_dataset(self, dim, level=0):
+        """Unstack DataArray expanding to Dataset along a given level of a
+        stacked coordinate.
+
+        This is the inverse operation of Dataset.to_stacked_array.
+
+        Parameters
+        ----------
+        dim : str
+            Name of existing dimension to unstack
+        level : int or str
+            The MultiIndex level to expand to a dataset along. Can either be
+            the integer index of the level or its name.
+        label : int, default 0
+            Label of the level to expand dataset along. Overrides the label
+            argument if given.
+
+        Returns
+        -------
+        unstacked: Dataset
+
+        Examples
+        --------
+        >>> import xarray as xr
+        >>> arr = DataArray(np.arange(6).reshape(2, 3),
+        ...                 coords=[('x', ['a', 'b']), ('y', [0, 1, 2])])
+        >>> data = xr.Dataset({'a': arr, 'b': arr.isel(y=0)})
+        >>> data
+        <xarray.Dataset>
+        Dimensions:  (x: 2, y: 3)
+        Coordinates:
+          * x        (x) <U1 'a' 'b'
+          * y        (y) int64 0 1 2
+        Data variables:
+            a        (x, y) int64 0 1 2 3 4 5
+            b        (x) int64 0 3
+        >>> stacked = data.to_stacked_array("z", ['y'])
+        >>> stacked.indexes['z']
+        MultiIndex(levels=[['a', 'b'], [0, 1, 2]],
+                labels=[[0, 0, 0, 1], [0, 1, 2, -1]],
+                names=['variable', 'y'])
+        >>> roundtripped = stacked.to_unstacked_dataset(dim='z')
+        >>> data.identical(roundtripped)
+        True
+
+        See Also
+        --------
+        Dataset.to_stacked_array
+        """
+
+        idx = self.indexes[dim]
+        if not isinstance(idx, pd.MultiIndex):
+            raise ValueError("'{}' is not a stacked coordinate".format(dim))
+
+        level_number = idx._get_level_number(level)
+        variables = idx.levels[level_number]
+        variable_dim = idx.names[level_number]
+
+        # pull variables out of datarray
+        data_dict = OrderedDict()
+        for k in variables:
+            data_dict[k] = self.sel({variable_dim: k}).squeeze(drop=True)
+
+        # unstacked dataset
+        return Dataset(data_dict)
+
     def transpose(self,
                   *dims: Hashable,
                   transpose_coords: Optional[bool] = None) -> 'DataArray':

xarray/core/dataset.py

@@ -2767,6 +2767,119 @@ def stack(self, dimensions=None, **dimensions_kwargs):
             result = result._stack_once(dims, new_dim)
         return result
 
+    def to_stacked_array(self, new_dim, sample_dims, variable_dim='variable',
+                         name=None):
+        """Combine variables of differing dimensionality into a DataArray
+        without broadcasting.
+
+        This method is similar to Dataset.to_array but does not broadcast the
+        variables.
+
+        Parameters
+        ----------
+        new_dim : str
+            Name of the new stacked coordinate
+        sample_dims : Sequence[str]
+            Dimensions that **will not** be stacked. Each array in the dataset
+            must share these dimensions. For machine learning applications,
+            these define the dimensions over which samples are drawn.
+        variable_dim : str, optional
+            Name of the level in the stacked coordinate which corresponds to
+            the variables.
+        name : str, optional
+            Name of the new data array.
+
+        Returns
+        -------
+        stacked : DataArray
+            DataArray with the specified dimensions and data variables
+            stacked together. The stacked coordinate is named ``new_dim``
+            and represented by a MultiIndex object with a level containing the
+            data variable names. The name of this level is controlled using
+            the ``variable_dim`` argument.
+
+        See Also
+        --------
+        Dataset.to_array
+        Dataset.stack
+        DataArray.to_unstacked_dataset
+
+        Examples
+        --------
+        >>> data = Dataset(
+        ...     data_vars={'a': (('x', 'y'), [[0, 1, 2], [3, 4, 5]]),
+        ...                'b': ('x', [6, 7])},
+        ...     coords={'y': ['u', 'v', 'w']}
+        ... )
+
+        >>> data
+        <xarray.Dataset>
+        Dimensions:  (x: 2, y: 3)
+        Coordinates:
+        * y        (y) <U1 'u' 'v' 'w'
+        Dimensions without coordinates: x
+        Data variables:
+            a        (x, y) int64 0 1 2 3 4 5
+            b        (x) int64 6 7
+
+        >>> data.to_stacked_array("z", sample_dims=['x'])
+        <xarray.DataArray (x: 2, z: 4)>
+        array([[0, 1, 2, 6],
+            [3, 4, 5, 7]])
+        Coordinates:
+        * z         (z) MultiIndex
+        - variable  (z) object 'a' 'a' 'a' 'b'
+        - y         (z) object 'u' 'v' 'w' nan
+        Dimensions without coordinates: x
+
+        """
+        stacking_dims = tuple(dim for dim in self.dims
+                              if dim not in sample_dims)
+
+        for variable in self:
+            dims = self[variable].dims
+            dims_include_sample_dims = set(sample_dims) <= set(dims)
+            if not dims_include_sample_dims:
+                raise ValueError(
+                    "All variables in the dataset must contain the "
+                    "dimensions {}.".format(dims)
+                )
+
+        def ensure_stackable(val):
+            assign_coords = {variable_dim: val.name}
+            for dim in stacking_dims:
+                if dim not in val.dims:
+                    assign_coords[dim] = None
+
+            expand_dims = set(stacking_dims).difference(set(val.dims))
+            expand_dims.add(variable_dim)
+            # must be list for .expand_dims
+            expand_dims = list(expand_dims)
+
+            return (val.assign_coords(**assign_coords)
+                    .expand_dims(expand_dims)
+                    .stack({new_dim: (variable_dim,) + stacking_dims}))
+
+        # concatenate the arrays
+        stackable_vars = [ensure_stackable(self[key])
+                          for key in self.data_vars]
+        data_array = xr.concat(stackable_vars, dim=new_dim)
+
+        # coerce the levels of the MultiIndex to have the same type as the
+        # input dimensions. This code is messy, so it might be better to just
+        # input a dummy value for the singleton dimension.
+        idx = data_array.indexes[new_dim]
+        levels = ([idx.levels[0]]
+                  + [level.astype(self[level.name].dtype)
+                     for level in idx.levels[1:]])
+        new_idx = idx.set_levels(levels)
+        data_array[new_dim] = IndexVariable(new_dim, new_idx)
+
+        if name is not None:
+            data_array.name = name
+
+        return data_array
+
     def _unstack_once(self, dim):
         index = self.get_index(dim)
         # GH2619. For MultiIndex, we need to call remove_unused.

xarray/tests/test_dataarray.py

@@ -1798,6 +1798,12 @@ def test_stack_nonunique_consistency(self):
         expected = DataArray(orig.to_pandas().stack(), dims='z')
         assert_identical(expected, actual)
 
+    def test_to_unstacked_dataset_raises_value_error(self):
+        data = DataArray([0, 1], dims='x', coords={'x': [0, 1]})
+        with pytest.raises(
+                ValueError, match="'x' is not a stacked coordinate"):
+            data.to_unstacked_dataset('x', 0)
+
     def test_transpose(self):
         da = DataArray(np.random.randn(3, 4, 5), dims=('x', 'y', 'z'),
                        coords={'x': range(3), 'y': range(4), 'z': range(5),

xarray/tests/test_dataset.py

@@ -110,6 +110,14 @@ def create_test_multiindex():
     return Dataset({}, {'x': mindex})
 
 
+def create_test_stacked_array():
+    x = DataArray(pd.Index(np.r_[:10], name='x'))
+    y = DataArray(pd.Index(np.r_[:20], name='y'))
+    a = x * y
+    b = x * y * y
+    return a, b
+
+
 class InaccessibleVariableDataStore(backends.InMemoryDataStore):
     def __init__(self):
         super().__init__()
@@ -2482,6 +2490,61 @@ def test_stack_unstack_slow(self):
         actual = stacked.isel(z=slice(None, None, -1)).unstack('z')
         assert actual.identical(ds[['b']])
 
+    def test_to_stacked_array_invalid_sample_dims(self):
+        data = xr.Dataset(
+            data_vars={'a': (('x', 'y'), [[0, 1, 2], [3, 4, 5]]),
+                       'b': ('x', [6, 7])},
+            coords={'y': ['u', 'v', 'w']}
+        )
+        with pytest.raises(ValueError):
+            data.to_stacked_array("features", sample_dims=['y'])
+
+    def test_to_stacked_array_name(self):
+        name = 'adf9d'
+
+        # make a two dimensional dataset
+        a, b = create_test_stacked_array()
+        D = xr.Dataset({'a': a, 'b': b})
+        sample_dims = ['x']
+
+        y = D.to_stacked_array('features', sample_dims, name=name)
+        assert y.name == name
+
+    def test_to_stacked_array_dtype_dims(self):
+        # make a two dimensional dataset
+        a, b = create_test_stacked_array()
+        D = xr.Dataset({'a': a, 'b': b})
+        sample_dims = ['x']
+        y = D.to_stacked_array('features', sample_dims)
+        assert y.indexes['features'].levels[1].dtype == D.y.dtype
+        assert y.dims == ('x', 'features')
+
+    def test_to_stacked_array_to_unstacked_dataset(self):
+        # make a two dimensional dataset
+        a, b = create_test_stacked_array()
+        D = xr.Dataset({'a': a, 'b': b})
+        sample_dims = ['x']
+        y = D.to_stacked_array('features', sample_dims)\
+            .transpose("x", "features")
+
+        x = y.to_unstacked_dataset("features")
+        assert_identical(D, x)
+
+        # test on just one sample
+        x0 = y[0].to_unstacked_dataset("features")
+        d0 = D.isel(x=0)
+        assert_identical(d0, x0)
+
+    def test_to_stacked_array_to_unstacked_dataset_different_dimension(self):
+        # test when variables have different dimensionality
+        a, b = create_test_stacked_array()
+        sample_dims = ['x']
+        D = xr.Dataset({'a': a, 'b': b.isel(y=0)})
+
+        y = D.to_stacked_array('features', sample_dims)
+        x = y.to_unstacked_dataset('features')
+        assert_identical(D, x)
+
     def test_update(self):
         data = create_test_data(seed=0)
         expected = data.copy()