Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Switch all methods to dim #8982

Merged
merged 7 commits into from
Apr 29, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions doc/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,10 @@ Internal Changes
``xarray/testing/assertions`` for ``DataTree``. (:pull:`8967`)
By `Owen Littlejohns <https://github.com/owenlittlejohns>`_ and
`Tom Nicholas <https://github.com/TomNicholas>`_.
- ``transpose``, ``set_dims``, ``stack`` & ``unstack`` now use a ``dim`` kwarg
rather than ``dims`` or ``dimensions``. This is the final change to make xarray methods
consistent with their use of ``dim``. Using the existing kwarg will raise a
warning. By `Maximilian Roos <https://github.com/max-sixty>`_


.. _whats-new.2024.03.0:
Expand Down
31 changes: 17 additions & 14 deletions xarray/core/dataarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import datetime
import warnings
from collections.abc import Hashable, Iterable, Mapping, MutableMapping, Sequence
from functools import partial
from os import PathLike
from typing import (
TYPE_CHECKING,
Expand Down Expand Up @@ -2808,12 +2809,13 @@ def reorder_levels(
ds = self._to_temp_dataset().reorder_levels(dim_order, **dim_order_kwargs)
return self._from_temp_dataset(ds)

@partial(deprecate_dims, old_name="dimensions")
def stack(
self,
dimensions: Mapping[Any, Sequence[Hashable]] | None = None,
dim: Mapping[Any, Sequence[Hashable]] | None = None,
create_index: bool | None = True,
index_cls: type[Index] = PandasMultiIndex,
**dimensions_kwargs: Sequence[Hashable],
**dim_kwargs: Sequence[Hashable],
) -> Self:
"""
Stack any number of existing dimensions into a single new dimension.
Expand All @@ -2823,7 +2825,7 @@ def stack(

Parameters
----------
dimensions : mapping of Hashable to sequence of Hashable
dim : mapping of Hashable to sequence of Hashable
Mapping of the form `new_name=(dim1, dim2, ...)`.
Names of new dimensions, and the existing dimensions that they
replace. An ellipsis (`...`) will be replaced by all unlisted dimensions.
Expand All @@ -2837,9 +2839,9 @@ def stack(
index_cls: class, optional
Can be used to pass a custom multi-index type. Must be an Xarray index that
implements `.stack()`. By default, a pandas multi-index wrapper is used.
**dimensions_kwargs
The keyword arguments form of ``dimensions``.
One of dimensions or dimensions_kwargs must be provided.
**dim_kwargs
The keyword arguments form of ``dim``.
One of dim or dim_kwargs must be provided.

Returns
-------
Expand Down Expand Up @@ -2874,10 +2876,10 @@ def stack(
DataArray.unstack
"""
ds = self._to_temp_dataset().stack(
dimensions,
dim,
create_index=create_index,
index_cls=index_cls,
**dimensions_kwargs,
**dim_kwargs,
)
return self._from_temp_dataset(ds)

Expand Down Expand Up @@ -3011,17 +3013,18 @@ def to_unstacked_dataset(self, dim: Hashable, level: int | Hashable = 0) -> Data
# unstacked dataset
return Dataset(data_dict)

@deprecate_dims
def transpose(
self,
*dims: Hashable,
*dim: Hashable,
transpose_coords: bool = True,
missing_dims: ErrorOptionsWithWarn = "raise",
) -> Self:
"""Return a new DataArray object with transposed dimensions.

Parameters
----------
*dims : Hashable, optional
*dim : Hashable, optional
By default, reverse the dimensions. Otherwise, reorder the
dimensions to this order.
transpose_coords : bool, default: True
Expand Down Expand Up @@ -3049,13 +3052,13 @@ def transpose(
numpy.transpose
Dataset.transpose
"""
if dims:
dims = tuple(infix_dims(dims, self.dims, missing_dims))
variable = self.variable.transpose(*dims)
if dim:
dim = tuple(infix_dims(dim, self.dims, missing_dims))
variable = self.variable.transpose(*dim)
if transpose_coords:
coords: dict[Hashable, Variable] = {}
for name, coord in self.coords.items():
coord_dims = tuple(dim for dim in dims if dim in coord.dims)
coord_dims = tuple(d for d in dim if d in coord.dims)
coords[name] = coord.variable.transpose(*coord_dims)
return self._replace(variable, coords)
else:
Expand Down
39 changes: 21 additions & 18 deletions xarray/core/dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
MutableMapping,
Sequence,
)
from functools import partial
from html import escape
from numbers import Number
from operator import methodcaller
Expand Down Expand Up @@ -124,7 +125,7 @@
from xarray.namedarray.parallelcompat import get_chunked_array_type, guess_chunkmanager
from xarray.namedarray.pycompat import array_type, is_chunked_array
from xarray.plot.accessor import DatasetPlotAccessor
from xarray.util.deprecation_helpers import _deprecate_positional_args
from xarray.util.deprecation_helpers import _deprecate_positional_args, deprecate_dims

if TYPE_CHECKING:
from dask.dataframe import DataFrame as DaskDataFrame
Expand Down Expand Up @@ -5264,12 +5265,13 @@ def _stack_once(
new_variables, coord_names=new_coord_names, indexes=indexes
)

@partial(deprecate_dims, old_name="dimensions")
def stack(
self,
dimensions: Mapping[Any, Sequence[Hashable | ellipsis]] | None = None,
dim: Mapping[Any, Sequence[Hashable | ellipsis]] | None = None,
create_index: bool | None = True,
index_cls: type[Index] = PandasMultiIndex,
**dimensions_kwargs: Sequence[Hashable | ellipsis],
**dim_kwargs: Sequence[Hashable | ellipsis],
) -> Self:
"""
Stack any number of existing dimensions into a single new dimension.
Expand All @@ -5279,7 +5281,7 @@ def stack(

Parameters
----------
dimensions : mapping of hashable to sequence of hashable
dim : mapping of hashable to sequence of hashable
Mapping of the form `new_name=(dim1, dim2, ...)`. Names of new
dimensions, and the existing dimensions that they replace. An
ellipsis (`...`) will be replaced by all unlisted dimensions.
Expand All @@ -5295,9 +5297,9 @@ def stack(
index_cls: Index-class, default: PandasMultiIndex
Can be used to pass a custom multi-index type (must be an Xarray index that
implements `.stack()`). By default, a pandas multi-index wrapper is used.
**dimensions_kwargs
The keyword arguments form of ``dimensions``.
One of dimensions or dimensions_kwargs must be provided.
**dim_kwargs
The keyword arguments form of ``dim``.
One of dim or dim_kwargs must be provided.

Returns
-------
Expand All @@ -5308,9 +5310,9 @@ def stack(
--------
Dataset.unstack
"""
dimensions = either_dict_or_kwargs(dimensions, dimensions_kwargs, "stack")
dim = either_dict_or_kwargs(dim, dim_kwargs, "stack")
result = self
for new_dim, dims in dimensions.items():
for new_dim, dims in dim.items():
result = result._stack_once(dims, new_dim, index_cls, create_index)
return result

Expand Down Expand Up @@ -6218,9 +6220,10 @@ def drop_dims(
drop_vars = {k for k, v in self._variables.items() if set(v.dims) & drop_dims}
return self.drop_vars(drop_vars)

@deprecate_dims
def transpose(
self,
*dims: Hashable,
*dim: Hashable,
missing_dims: ErrorOptionsWithWarn = "raise",
) -> Self:
"""Return a new Dataset object with all array dimensions transposed.
Expand All @@ -6230,7 +6233,7 @@ def transpose(

Parameters
----------
*dims : hashable, optional
*dim : hashable, optional
By default, reverse the dimensions on each array. Otherwise,
reorder the dimensions to this order.
missing_dims : {"raise", "warn", "ignore"}, default: "raise"
Expand All @@ -6257,20 +6260,20 @@ def transpose(
numpy.transpose
DataArray.transpose
"""
# Raise error if list is passed as dims
if (len(dims) > 0) and (isinstance(dims[0], list)):
list_fix = [f"{repr(x)}" if isinstance(x, str) else f"{x}" for x in dims[0]]
# Raise error if list is passed as dim
if (len(dim) > 0) and (isinstance(dim[0], list)):
list_fix = [f"{repr(x)}" if isinstance(x, str) else f"{x}" for x in dim[0]]
raise TypeError(
f'transpose requires dims to be passed as multiple arguments. Expected `{", ".join(list_fix)}`. Received `{dims[0]}` instead'
f'transpose requires dim to be passed as multiple arguments. Expected `{", ".join(list_fix)}`. Received `{dim[0]}` instead'
)

# Use infix_dims to check once for missing dimensions
if len(dims) != 0:
_ = list(infix_dims(dims, self.dims, missing_dims))
if len(dim) != 0:
_ = list(infix_dims(dim, self.dims, missing_dims))

ds = self.copy()
for name, var in self._variables.items():
var_dims = tuple(dim for dim in dims if dim in (var.dims + (...,)))
var_dims = tuple(d for d in dim if d in (var.dims + (...,)))
ds._variables[name] = var.transpose(*var_dims)
return ds

Expand Down
Loading
Loading