From fef7a68e859c903c73de77d29e059a8348ca83c6 Mon Sep 17 00:00:00 2001 From: Giacomo Caria Date: Thu, 24 Apr 2025 17:53:21 -0400 Subject: [PATCH 1/9] extract data from dataarray --- xarray/computation/apply_ufunc.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/xarray/computation/apply_ufunc.py b/xarray/computation/apply_ufunc.py index 50c3ed4bbd8..5789e1d220d 100644 --- a/xarray/computation/apply_ufunc.py +++ b/xarray/computation/apply_ufunc.py @@ -30,10 +30,7 @@ from xarray.core.formatting import limit_lines from xarray.core.indexes import Index, filter_indexes_from_coords from xarray.core.options import _get_keep_attrs -from xarray.core.utils import ( - is_dict_like, - result_name, -) +from xarray.core.utils import is_dict_like, result_name from xarray.core.variable import Variable from xarray.namedarray.parallelcompat import get_chunked_array_type from xarray.namedarray.pycompat import is_chunked_array @@ -1212,6 +1209,8 @@ def apply_ufunc( dask_gufunc_kwargs.setdefault("output_sizes", output_sizes) if kwargs: + if "where" in kwargs and isinstance(kwargs["where"], DataArray): + kwargs["where"] = kwargs["where"].data func = functools.partial(func, **kwargs) if keep_attrs is None: From 84670c9bc569391308bdff6f2525047612489a72 Mon Sep 17 00:00:00 2001 From: Giacomo Caria Date: Sat, 26 Apr 2025 10:48:36 -0400 Subject: [PATCH 2/9] add test --- xarray/tests/test_computation.py | 44 ++++++++++++++++++++------------ 1 file changed, 28 insertions(+), 16 deletions(-) diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py index dc7016238df..f19ef8865c5 100644 --- a/xarray/tests/test_computation.py +++ b/xarray/tests/test_computation.py @@ -10,24 +10,17 @@ from numpy.testing import assert_allclose, assert_array_equal import xarray as xr -from xarray.computation.apply_ufunc import ( - _UFuncSignature, - apply_ufunc, - broadcast_compat_data, - collect_dict_values, - join_dict_keys, - ordered_set_intersection, - ordered_set_union, - unified_dim_sizes, -) +from xarray.computation.apply_ufunc import (_UFuncSignature, apply_ufunc, + broadcast_compat_data, + collect_dict_values, + join_dict_keys, + ordered_set_intersection, + ordered_set_union, + unified_dim_sizes) from xarray.core.utils import result_name from xarray.structure.alignment import broadcast -from xarray.tests import ( - has_dask, - raise_if_dask_computes, - requires_cftime, - requires_dask, -) +from xarray.tests import (has_dask, raise_if_dask_computes, requires_cftime, + requires_dask) def assert_identical(a, b): @@ -2627,3 +2620,22 @@ def test_complex_number_reduce(compute_backend): # Check that xarray doesn't call into numbagg, which doesn't compile for complex # numbers at the moment (but will when numba supports dynamic compilation) da.min() + + +@pytest.mark.parametrize("use_dask", [False, True]) +def test_fix(use_dask: bool) -> None: + + val = 3.0 + val_fixed = np.fix(val) + + da = xr.DataArray([val]) + expected = xr.DataArray([val_fixed]) + + if use_dask: + if not has_dask: + pytest.skip("test for dask.") + da = da.chunk() + expected = expected.chunk() + + actual = np.fix(da) + xr.testing.assert_duckarray_allclose(expected, actual) From b4063b6f8a14dd0ae6c790a711c9b17cb4684589 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 28 Apr 2025 14:51:03 +0000 Subject: [PATCH 3/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- xarray/tests/test_computation.py | 26 ++++++++++++++++---------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py index f19ef8865c5..65c95f66c4a 100644 --- a/xarray/tests/test_computation.py +++ b/xarray/tests/test_computation.py @@ -10,17 +10,24 @@ from numpy.testing import assert_allclose, assert_array_equal import xarray as xr -from xarray.computation.apply_ufunc import (_UFuncSignature, apply_ufunc, - broadcast_compat_data, - collect_dict_values, - join_dict_keys, - ordered_set_intersection, - ordered_set_union, - unified_dim_sizes) +from xarray.computation.apply_ufunc import ( + _UFuncSignature, + apply_ufunc, + broadcast_compat_data, + collect_dict_values, + join_dict_keys, + ordered_set_intersection, + ordered_set_union, + unified_dim_sizes, +) from xarray.core.utils import result_name from xarray.structure.alignment import broadcast -from xarray.tests import (has_dask, raise_if_dask_computes, requires_cftime, - requires_dask) +from xarray.tests import ( + has_dask, + raise_if_dask_computes, + requires_cftime, + requires_dask, +) def assert_identical(a, b): @@ -2624,7 +2631,6 @@ def test_complex_number_reduce(compute_backend): @pytest.mark.parametrize("use_dask", [False, True]) def test_fix(use_dask: bool) -> None: - val = 3.0 val_fixed = np.fix(val) From 6c85acfb000e977fd030f436da26132cc128712f Mon Sep 17 00:00:00 2001 From: Giacomo Caria Date: Mon, 28 Apr 2025 11:31:17 -0400 Subject: [PATCH 4/9] fix test --- xarray/tests/test_computation.py | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py index f19ef8865c5..402603e0ed0 100644 --- a/xarray/tests/test_computation.py +++ b/xarray/tests/test_computation.py @@ -2622,8 +2622,7 @@ def test_complex_number_reduce(compute_backend): da.min() -@pytest.mark.parametrize("use_dask", [False, True]) -def test_fix(use_dask: bool) -> None: +def test_fix() -> None: val = 3.0 val_fixed = np.fix(val) @@ -2631,11 +2630,5 @@ def test_fix(use_dask: bool) -> None: da = xr.DataArray([val]) expected = xr.DataArray([val_fixed]) - if use_dask: - if not has_dask: - pytest.skip("test for dask.") - da = da.chunk() - expected = expected.chunk() - actual = np.fix(da) xr.testing.assert_duckarray_allclose(expected, actual) From 5b82c96a4c9046e9bf5e0594dcf9bb31831d5531 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 28 Apr 2025 15:34:47 +0000 Subject: [PATCH 5/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- xarray/tests/test_computation.py | 25 ++++++++++++++++--------- 1 file changed, 16 insertions(+), 9 deletions(-) diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py index 4242dfbb1b0..546b6c06a6d 100644 --- a/xarray/tests/test_computation.py +++ b/xarray/tests/test_computation.py @@ -10,17 +10,24 @@ from numpy.testing import assert_allclose, assert_array_equal import xarray as xr -from xarray.computation.apply_ufunc import (_UFuncSignature, apply_ufunc, - broadcast_compat_data, - collect_dict_values, - join_dict_keys, - ordered_set_intersection, - ordered_set_union, - unified_dim_sizes) +from xarray.computation.apply_ufunc import ( + _UFuncSignature, + apply_ufunc, + broadcast_compat_data, + collect_dict_values, + join_dict_keys, + ordered_set_intersection, + ordered_set_union, + unified_dim_sizes, +) from xarray.core.utils import result_name from xarray.structure.alignment import broadcast -from xarray.tests import (has_dask, raise_if_dask_computes, requires_cftime, - requires_dask) +from xarray.tests import ( + has_dask, + raise_if_dask_computes, + requires_cftime, + requires_dask, +) def assert_identical(a, b): From 5277b4bc055fb5e4364abb7a956d196d58dab4a3 Mon Sep 17 00:00:00 2001 From: Deepak Cherian Date: Mon, 28 Apr 2025 10:43:40 -0600 Subject: [PATCH 6/9] Update xarray/tests/test_computation.py --- xarray/tests/test_computation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py index 546b6c06a6d..cb12f3df534 100644 --- a/xarray/tests/test_computation.py +++ b/xarray/tests/test_computation.py @@ -2637,4 +2637,4 @@ def test_fix() -> None: expected = xr.DataArray([val_fixed]) actual = np.fix(da) - xr.testing.assert_duckarray_allclose(expected, actual) + assert_identical(expected, actual) From c6ac784b58edd571633fd23847530688f7632239 Mon Sep 17 00:00:00 2001 From: Deepak Cherian Date: Mon, 28 Apr 2025 10:51:18 -0600 Subject: [PATCH 7/9] Update xarray/computation/apply_ufunc.py --- xarray/computation/apply_ufunc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/computation/apply_ufunc.py b/xarray/computation/apply_ufunc.py index 5789e1d220d..980127cba62 100644 --- a/xarray/computation/apply_ufunc.py +++ b/xarray/computation/apply_ufunc.py @@ -1210,7 +1210,7 @@ def apply_ufunc( if kwargs: if "where" in kwargs and isinstance(kwargs["where"], DataArray): - kwargs["where"] = kwargs["where"].data + kwargs["where"] = kwargs["where"].data. #+type: ignore[index] func = functools.partial(func, **kwargs) if keep_attrs is None: From a73965f2fca9775e1861dbb3c0a6b9c3e9a796b9 Mon Sep 17 00:00:00 2001 From: Deepak Cherian Date: Mon, 28 Apr 2025 10:51:44 -0600 Subject: [PATCH 8/9] Update xarray/computation/apply_ufunc.py --- xarray/computation/apply_ufunc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/computation/apply_ufunc.py b/xarray/computation/apply_ufunc.py index 980127cba62..1694f9c1fec 100644 --- a/xarray/computation/apply_ufunc.py +++ b/xarray/computation/apply_ufunc.py @@ -1210,7 +1210,7 @@ def apply_ufunc( if kwargs: if "where" in kwargs and isinstance(kwargs["where"], DataArray): - kwargs["where"] = kwargs["where"].data. #+type: ignore[index] + kwargs["where"] = kwargs["where"].data # type:ignore[index] func = functools.partial(func, **kwargs) if keep_attrs is None: From ab9a69773166b41b9db0f84bf0d9c009e248b61d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 28 Apr 2025 16:52:03 +0000 Subject: [PATCH 9/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- xarray/computation/apply_ufunc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/computation/apply_ufunc.py b/xarray/computation/apply_ufunc.py index 1694f9c1fec..41f4adb5fd5 100644 --- a/xarray/computation/apply_ufunc.py +++ b/xarray/computation/apply_ufunc.py @@ -1210,7 +1210,7 @@ def apply_ufunc( if kwargs: if "where" in kwargs and isinstance(kwargs["where"], DataArray): - kwargs["where"] = kwargs["where"].data # type:ignore[index] + kwargs["where"] = kwargs["where"].data # type:ignore[index] func = functools.partial(func, **kwargs) if keep_attrs is None: