Skip to content

Commit

Permalink
rename dataset.apply to dataset.map, deprecating apply
Browse files Browse the repository at this point in the history
  • Loading branch information
max-sixty committed Oct 29, 2019
1 parent 164ce9b commit 3f0c4bd
Show file tree
Hide file tree
Showing 2 changed files with 37 additions and 16 deletions.
20 changes: 17 additions & 3 deletions xarray/core/dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -4088,14 +4088,14 @@ def reduce(
variables, coord_names=coord_names, attrs=attrs, indexes=indexes
)

def apply(
def map(
self,
func: Callable,
keep_attrs: bool = None,
args: Iterable[Any] = (),
**kwargs: Any,
) -> "Dataset":
"""Apply a function over the data variables in this dataset.
"""Apply a function over each data variable in this dataset.
Parameters
----------
Expand Down Expand Up @@ -4128,7 +4128,7 @@ def apply(
Data variables:
foo (dim_0, dim_1) float64 -0.3751 -1.951 -1.945 0.2948 0.711 -0.3948
bar (x) int64 -1 2
>>> ds.apply(np.fabs)
>>> ds.map(np.fabs)
<xarray.Dataset>
Dimensions: (dim_0: 2, dim_1: 3, x: 2)
Dimensions without coordinates: dim_0, dim_1, x
Expand All @@ -4145,6 +4145,20 @@ def apply(
attrs = self.attrs if keep_attrs else None
return type(self)(variables, attrs=attrs)

def apply(
self,
func: Callable,
keep_attrs: bool = None,
args: Iterable[Any] = (),
**kwargs: Any,
) -> "Dataset":
warnings.warn(
"Dataset.apply is deprecated in favor of Dataset.map and will be removed in a future version of xarray",
DeprecationWarning,
stacklevel=2,
)
return self.map(func, keep_attrs, args, **kwargs)

def assign(
self, variables: Mapping[Hashable, Any] = None, **variables_kwargs: Hashable
) -> "Dataset":
Expand Down
33 changes: 20 additions & 13 deletions xarray/tests/test_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -4494,29 +4494,36 @@ def test_count(self):
actual = ds.count()
assert_identical(expected, actual)

def test_apply(self):
def test_map(self):
data = create_test_data()
data.attrs["foo"] = "bar"

assert_identical(data.apply(np.mean), data.mean())
assert_identical(data.map(np.mean), data.mean())

expected = data.mean(keep_attrs=True)
actual = data.apply(lambda x: x.mean(keep_attrs=True), keep_attrs=True)
actual = data.map(lambda x: x.mean(keep_attrs=True), keep_attrs=True)
assert_identical(expected, actual)

assert_identical(data.apply(lambda x: x, keep_attrs=True), data.drop("time"))
assert_identical(data.map(lambda x: x, keep_attrs=True), data.drop("time"))

def scale(x, multiple=1):
return multiple * x

actual = data.apply(scale, multiple=2)
actual = data.map(scale, multiple=2)
assert_equal(actual["var1"], 2 * data["var1"])
assert_identical(actual["numbers"], data["numbers"])

actual = data.apply(np.asarray)
actual = data.map(np.asarray)
expected = data.drop("time") # time is not used on a data var
assert_equal(expected, actual)

def test_apply_deprecated_map(self):
data = create_test_data()
data.attrs["foo"] = "bar"

with pytest.warns(DeprecationWarning):
assert_identical(data.map(np.mean), data.mean())

def make_example_math_dataset(self):
variables = {
"bar": ("x", np.arange(100, 400, 100)),
Expand All @@ -4543,15 +4550,15 @@ def test_dataset_number_math(self):
def test_unary_ops(self):
ds = self.make_example_math_dataset()

assert_identical(ds.apply(abs), abs(ds))
assert_identical(ds.apply(lambda x: x + 4), ds + 4)
assert_identical(ds.map(abs), abs(ds))
assert_identical(ds.map(lambda x: x + 4), ds + 4)

for func in [
lambda x: x.isnull(),
lambda x: x.round(),
lambda x: x.astype(int),
]:
assert_identical(ds.apply(func), func(ds))
assert_identical(ds.map(func), func(ds))

assert_identical(ds.isnull(), ~ds.notnull())

Expand All @@ -4564,7 +4571,7 @@ def test_unary_ops(self):
def test_dataset_array_math(self):
ds = self.make_example_math_dataset()

expected = ds.apply(lambda x: x - ds["foo"])
expected = ds.map(lambda x: x - ds["foo"])
assert_identical(expected, ds - ds["foo"])
assert_identical(expected, -ds["foo"] + ds)
assert_identical(expected, ds - ds["foo"].variable)
Expand All @@ -4573,7 +4580,7 @@ def test_dataset_array_math(self):
actual -= ds["foo"]
assert_identical(expected, actual)

expected = ds.apply(lambda x: x + ds["bar"])
expected = ds.map(lambda x: x + ds["bar"])
assert_identical(expected, ds + ds["bar"])
actual = ds.copy(deep=True)
actual += ds["bar"]
Expand All @@ -4589,7 +4596,7 @@ def test_dataset_dataset_math(self):
assert_identical(ds, ds + 0 * ds)
assert_identical(ds, ds + {"foo": 0, "bar": 0})

expected = ds.apply(lambda x: 2 * x)
expected = ds.map(lambda x: 2 * x)
assert_identical(expected, 2 * ds)
assert_identical(expected, ds + ds)
assert_identical(expected, ds + ds.data_vars)
Expand Down Expand Up @@ -4686,7 +4693,7 @@ def test_dataset_transpose(self):
assert_identical(expected, actual)

actual = ds.transpose("x", "y")
expected = ds.apply(lambda x: x.transpose("x", "y", transpose_coords=True))
expected = ds.map(lambda x: x.transpose("x", "y", transpose_coords=True))
assert_identical(expected, actual)

ds = create_test_data()
Expand Down

0 comments on commit 3f0c4bd

Please sign in to comment.