Merge pull request #372 from shoyer/swap_dims

API: new methods {Dataset/DataArray}.swap_dims

Author: shoyer

doc/api.rst

@@ -65,6 +65,7 @@ Dataset contents
    Dataset.copy
    Dataset.merge
    Dataset.rename
+   Dataset.swap_dims
    Dataset.drop
    Dataset.set_coords
    Dataset.reset_coords
@@ -166,6 +167,7 @@ DataArray contents
    :toctree: generated/
 
    DataArray.rename
+   DataArray.swap_dims
    DataArray.drop
    DataArray.reset_coords
    DataArray.copy

doc/whats-new.rst

@@ -20,7 +20,7 @@ Enhancements
 
 - New documentation sections on :ref:`time-series` and
   :ref:`combining multiple files`.
-- :py:meth`~xray.Dataset.resample` lets you resample a dataset or data array to
+- :py:meth:`~xray.Dataset.resample` lets you resample a dataset or data array to
   a new temporal resolution. The syntax is the `same as pandas`_, except you
   need to supply the time dimension explicitly:
 
@@ -57,13 +57,23 @@ Enhancements
 
       array.resample('1D', dim='time', how='first')
 
+.. _same as pandas: http://pandas.pydata.org/pandas-docs/stable/timeseries.html#up-and-downsampling
+
+- :py:meth:`~xray.Dataset.swap_dims` allows for easily swapping one dimension
+  out for another:
+
+  .. ipython:: python
+
+       ds = xray.Dataset({'x': range(3), 'y': ('x', list('abc'))})
+       ds
+       ds.swap_dims({'x': 'y'})
+
+  This was possible in earlier versions of xray, but required some contortions.
 - :py:func:`~xray.open_dataset` and :py:meth:`~xray.Dataset.to_netcdf` now
   accept an ``engine`` argument to explicitly select which underlying library
   (netcdf4 or scipy) is used for reading/writing a netCDF file.
 - New documentation section on :ref:`combining multiple files`.
 
-TODO: write full docs on time-series!
-
 .. _same as pandas: http://pandas.pydata.org/pandas-docs/stable/timeseries.html#up-and-downsampling
 
 Bug fixes

xray/core/dataarray.py

@@ -309,14 +309,8 @@ def dims(self):
 
     @dims.setter
     def dims(self, value):
-        with self._set_new_dataset() as ds:
-            if not len(value) == self.ndim:
-                raise ValueError('%s dimensions supplied but data has ndim=%s'
-                                 % (len(value), self.ndim))
-            name_map = dict(zip(self.dims, value))
-            ds.rename(name_map, inplace=True)
-        if self.name in name_map:
-            self._name = name_map[self.name]
+        raise AttributeError('you cannot assign dims on a DataArray. Use '
+                             '.rename() or .swap_dims() instead.')
 
     def _item_key_to_dict(self, key):
         if utils.is_dict_like(key):
@@ -553,13 +547,23 @@ def reindex(self, method=None, copy=True, **indexers):
     def rename(self, new_name_or_name_dict):
         """Returns a new DataArray with renamed coordinates and/or a new name.
 
-        If the argument is dict-like, it it used as a mapping from old names to
-        new names for dataset variables. Otherwise, use the argument as the new
-        name for this array.
+
+        Parameters
+        ----------
+        new_name_or_name_dict : str or dict-like
+            If the argument is dict-like, it it used as a mapping from old
+            names to new names for coordinates (and/or this array itself).
+            Otherwise, use the argument as the new name for this array.
+
+        Returns
+        -------
+        renamed : DataArray
+            Renamed array or array with renamed coordinates.
 
         See Also
         --------
         Dataset.rename
+        DataArray.swap_dims
         """
         if utils.is_dict_like(new_name_or_name_dict):
             name_dict = new_name_or_name_dict
@@ -570,6 +574,32 @@ def rename(self, new_name_or_name_dict):
         renamed_dataset = self._dataset.rename(name_dict)
         return renamed_dataset[new_name]
 
+    def swap_dims(self, dims_dict):
+        """Returns a new DataArray with swapped dimensions.
+
+        Parameters
+        ----------
+        dims_dict : dict-like
+            Dictionary whose keys are current dimension names and whose values
+            are new names. Each value must already be a coordinate on this
+            array.
+        inplace : bool, optional
+            If True, swap dimensions in-place. Otherwise, return a new object.
+
+        Returns
+        -------
+        renamed : Dataset
+            DataArray with swapped dimensions.
+
+        See Also
+        --------
+
+        DataArray.rename
+        Dataset.swap_dims
+        """
+        ds = self._dataset.swap_dims(dims_dict)
+        return self._with_replaced_dataset(ds)
+
     def transpose(self, *dims):
         """Return a new DataArray object with transposed dimensions.
 

xray/core/dataset.py

@@ -607,7 +607,7 @@ def _construct_direct(cls, variables, coord_names, dims, attrs,
     __default_attrs = object()
 
     def _replace_vars_and_dims(self, variables, coord_names=None,
-                               attrs=__default_attrs):
+                               attrs=__default_attrs, inplace=False):
         """Fastpath constructor for internal use.
 
         Preserves coord names and attributes; dimensions are recalculated from
@@ -628,11 +628,21 @@ def _replace_vars_and_dims(self, variables, coord_names=None,
         new : Dataset
         """
         dims = _calculate_dims(variables)
-        if coord_names is None:
-            coord_names = self._coord_names.copy()
-        if attrs is self.__default_attrs:
-            attrs = self._attrs_copy()
-        return self._construct_direct(variables, coord_names, dims, attrs)
+        if inplace:
+            self._dims = dims
+            self._variables = variables
+            if coord_names is not None:
+                self._coord_names = coord_names
+            if attrs is not self.__default_attrs:
+                self._attrs = attrs
+            obj = self
+        else:
+            if coord_names is None:
+                coord_names = self._coord_names.copy()
+            if attrs is self.__default_attrs:
+                attrs = self._attrs_copy()
+            obj = self._construct_direct(variables, coord_names, dims, attrs)
+        return obj
 
     def copy(self, deep=False):
         """Returns a copy of this dataset.
@@ -1221,6 +1231,12 @@ def rename(self, name_dict, inplace=False):
         -------
         renamed : Dataset
             Dataset with renamed variables and dimensions.
+
+        See Also
+        --------
+
+        Dataset.swap_dims
+        DataArray.rename
         """
         for k in name_dict:
             if k not in self:
@@ -1237,14 +1253,57 @@ def rename(self, name_dict, inplace=False):
             if k in self._coord_names:
                 coord_names.add(name)
 
-        if inplace:
-            self._dims = _calculate_dims(variables)
-            self._variables = variables
-            self._coord_names = coord_names
-            obj = self
-        else:
-            obj = self._replace_vars_and_dims(variables, coord_names)
-        return obj
+        return self._replace_vars_and_dims(variables, coord_names,
+                                           inplace=inplace)
+
+    def swap_dims(self, dims_dict, inplace=False):
+        """Returns a new object with swapped dimensions.
+
+        Parameters
+        ----------
+        dims_dict : dict-like
+            Dictionary whose keys are current dimension names and whose values
+            are new names. Each value must already be a variable in the
+            dataset.
+        inplace : bool, optional
+            If True, swap dimensions in-place. Otherwise, return a new dataset
+            object.
+
+        Returns
+        -------
+        renamed : Dataset
+            Dataset with swapped dimensions.
+
+        See Also
+        --------
+
+        Dataset.rename
+        DataArray.swap_dims
+        """
+        for k, v in dims_dict.items():
+            if k not in self.dims:
+                raise ValueError('cannot swap from dimension %r because it is '
+                                 'not an existing dimension' % k)
+            if self.variables[v].dims != (k,):
+                raise ValueError('replacement dimension %r is not a 1D '
+                                 'variable along the old dimension %r'
+                                 % (v, k))
+
+        result_dims = set(dims_dict.get(dim, dim) for dim in self.dims)
+
+        variables = OrderedDict()
+
+        coord_names = self._coord_names.copy()
+        coord_names.update(dims_dict.values())
+
+        for k, v in iteritems(self.variables):
+            dims = tuple(dims_dict.get(dim, dim) for dim in v.dims)
+            var = v.to_coord() if k in result_dims else v.to_variable()
+            var.dims = dims
+            variables[k] = var
+
+        return self._replace_vars_and_dims(variables, coord_names,
+                                           inplace=inplace)
 
     def update(self, other, inplace=True):
         """Update this dataset's variables with those from another dataset.

xray/core/variable.py

@@ -332,6 +332,11 @@ def values(self, values):
                 "replacement values must match the Variable's shape")
         self._data = values
 
+    def to_variable(self):
+        """Return this variable as a base xray.Variable"""
+        return Variable(self.dims, self._data, self._attrs,
+                        encoding=self._encoding, fastpath=True)
+
     def to_coord(self):
         """Return this variable as an xray.Coordinate"""
         return Coordinate(self.dims, self._data, self._attrs,

xray/test/test_dataarray.py

@@ -69,13 +69,8 @@ def test_dims(self):
         arr = self.dv
         self.assertEqual(arr.dims, ('x', 'y'))
 
-        arr.dims = ('w', 'z')
-        self.assertEqual(arr.dims, ('w', 'z'))
-
-        x = Dataset({'x': ('x', np.arange(5))})['x']
-        x.dims = ('y',)
-        self.assertEqual(x.dims, ('y',))
-        self.assertEqual(x.name, 'y')
+        with self.assertRaisesRegexp(AttributeError, 'you cannot assign'):
+            arr.dims = ('w', 'z')
 
     def test_encoding(self):
         expected = {'foo': 'bar'}
@@ -513,6 +508,14 @@ def test_rename(self):
             renamed.to_dataset(), self.ds.rename({'foo': 'bar'}))
         self.assertEqual(renamed.name, 'bar')
 
+    def test_swap_dims(self):
+        array = DataArray(np.random.randn(3), {'y': ('x', list('abc'))}, 'x')
+        expected = DataArray(array.values,
+                             {'y': list('abc'), 'x': ('y', range(3))},
+                             dims='y')
+        actual = array.swap_dims({'x': 'y'})
+        self.assertDataArrayIdentical(expected, actual)
+
     def test_dataset_getitem(self):
         dv = self.ds['foo']
         self.assertDataArrayIdentical(dv, self.dv)

xray/test/test_dataset.py

@@ -824,6 +824,26 @@ def test_rename_inplace(self):
         # check virtual variables
         self.assertArrayEqual(data['t.dayofyear'], [1, 2, 3])
 
+    def test_swap_dims(self):
+        original = Dataset({'x': [1, 2, 3], 'y': ('x', list('abc')), 'z': 42})
+        expected = Dataset({'z': 42}, {'x': ('y', [1, 2, 3]), 'y': list('abc')})
+        actual = original.swap_dims({'x': 'y'})
+        self.assertDatasetIdentical(expected, actual)
+        self.assertIsInstance(actual.variables['y'], Coordinate)
+        self.assertIsInstance(actual.variables['x'], Variable)
+
+        roundtripped = actual.swap_dims({'y': 'x'})
+        self.assertDatasetIdentical(original.set_coords('y'), roundtripped)
+
+        actual = original.copy()
+        actual.swap_dims({'x': 'y'}, inplace=True)
+        self.assertDatasetIdentical(expected, actual)
+
+        with self.assertRaisesRegexp(ValueError, 'cannot swap'):
+            original.swap_dims({'y': 'x'})
+        with self.assertRaisesRegexp(ValueError, 'replacement dimension'):
+            original.swap_dims({'x': 'z'})
+
     def test_update(self):
         data = create_test_data(seed=0)
         expected = data.copy()