Dask friendly check in `.weighted()` #4559

jbusecke · 2020-10-31T19:11:37Z

Closes Option to skip tests in weighted() #4541
Tests added
Passes isort . && black . && mypy . && flake8
User visible changes (including notable bug fixes) are documented in whats-new.rst

pep8speaks · 2020-10-31T19:11:40Z

Hello @jbusecke! Thanks for updating this PR. We checked the lines you've touched for PEP 8 issues, and found:

There are currently no PEP 8 issues detected in this Pull Request. Cheers! 🍻

Comment last updated at 2020-11-09 15:24:48 UTC

max-sixty

Looks good! Thanks @jbusecke . I'll let @dcherian or @mathause comment too.

xarray/core/weighted.py

jbusecke · 2020-11-01T03:16:06Z

The Ci environments without dask are failing. Should I add some pytest skip logic, or what is the best way to handle this?

max-sixty · 2020-11-01T03:46:37Z

Yes, @requires_dask on the test

keewis · 2020-11-01T13:42:04Z

since the test you added requires dask it definitely should get a @requires_dask.

That won't fix all the failing tests, though: map_blocks is only valid if we operate on dask arrays. I guess you either have to apply the function directly on other arrays, or we have to make DataArray.map_blocks on non-dask arrays call to DataArray.map.

mathause

For simplicity I would use if dask_duck_array(weights):.

mathause · 2020-11-01T16:10:32Z

xarray/tests/test_weighted.py

+
+    weights = DataArray(weights).chunk({"dim_0": -1})
+
+    weighted = data.weighted(weights)


You test that dask does not comoute:

xarray/xarray/tests/test_dask.py

Lines 189 to 190 in 83884a1

with raise_if_dask_computes():

actual = v.argmax(dim="x")

xarray/core/weighted.py

Co-authored-by: Deepak Cherian <[email protected]>

Co-authored-by: Maximilian Roos <[email protected]>

mathause · 2020-11-01T22:53:24Z

I think you need to do someting along the lines of:

if dask_duck_array(weights):
    import dask.array as dsa
    dsa.map_blocks(_weight_check, weights.data, dtype=weights.dtype)
else:
    _weight_check()

jbusecke · 2020-11-01T23:11:51Z

I did have to fiddle with this a bit. I did change .isnull() to np.isnan() and replace the data of weights, otherwise the error would not be raised for the dask array at all.
This does not look terribly elegant to me, but passes tests locally. Waiting for the CI to see if the others pass aswell. Happy to make further changes, and thanks for all the input already.

keewis · 2020-11-02T01:59:46Z

xarray/core/weighted.py

-                "`weights` cannot contain missing values. "
-                "Missing values can be replaced by `weights.fillna(0)`."
+        def _weight_check(w):
+            if np.isnan(w).any():


.isnull() does a bit more than that: np.isnan won't detect NaT. @mathause, how likely is it to get datetime-like arrays here? They don't make much sense as weights, but as far as I can tell we don't check (I might be missing something, though)

There is no check for that. A TimeDelta may make some sense as weights. DateTime not so much. I think we can get away with using np.isnan. A Date* array as weights containing NaT should be super uncommon.

We could still operate on the dataarray instead of the dask/numpy array, but as @dcherian suggesred, that would be less efficient. I would be curious as to what penalties would actually occur when we use the weights.map_blocks compared to dask.array.map_blocks?

Could use duck_array_ops.isnull to account for timedelta64? It is weird to have it as a weight though. Does that work?

Re map_blocks: the xarray version adds tasks that create xarray objects wrapping every block in a dask array. That adds overhead which is totally unneccesary here.

jbusecke · 2020-11-03T20:39:31Z

Do you think this works or are further changes needed? Many thanks for the guidance so far!

max-sixty

Looks great to me. Want to add a whatsnew?

xarray/core/weighted.py

dcherian · 2020-11-03T23:37:19Z

xarray/core/weighted.py

+        if is_duck_dask_array(weights.data):
+            import dask.array as dsa
+
+            weights.data = dsa.map_blocks(


Suggested change

weights.data = dsa.map_blocks(

weights = weights.copy(data=dsa.map_blocks(

so we don't modify the original object. Could even do weights.data.map_blocks(...) to save some typing...

Co-authored-by: Maximilian Roos <[email protected]>

jbusecke · 2020-11-04T16:21:35Z

I am getting some failures for py38-upstream (TestDataset.test_polyfit_warnings) that seem unrelated?

jbusecke · 2020-11-04T16:22:43Z

Similarly on MacOSX py38 this fails: TestDask.test_save_mfdataset_compute_false_roundtrip

keewis · 2020-11-04T16:26:05Z

the TestDask.test_save_mfdataset_compute_false_roundtrip failure should be #4539, and the failing upstream-dev CI shouldn't be related, either (my guess is that a dependency – most likely numpy – added a new warning). I'll investigate and open a new issue for that (edit: see #4565).

jbusecke · 2020-11-05T17:12:06Z

Ok I think this should be good to go. I have implemented all the requested changes. The remaining failures are related to other problems upstream (I think). Anything else I should add here?

xarray/core/weighted.py

Co-authored-by: Mathias Hauser <[email protected]>

jbusecke · 2020-11-06T16:20:04Z

I am not understanding why that MinimumVersionPolicy test is failing (or what it does haha). Is this something upstream, or should I fix?

keewis · 2020-11-06T16:25:46Z

this CI is sometimes flaky, but it's usually enough to just rerun it. I'll do that for you once the other CI finished.

jbusecke · 2020-11-06T17:15:33Z

Seems like all the other test are passing (minus the two upstream problems discussed before).

dcherian

Thanks @jbusecke great contribution!

xarray/core/weighted.py

mathause · 2020-11-09T16:22:51Z

Thanks @jbusecke

* initial changes * Using map_blocks to lazily mask input arrays, following #4559 * Adding lazy corr cov test with `raise_if_dask_computes` * adding test for one da without nans * checking ordering of arrays doesnt matter Co-authored-by: Deepak Cherian <[email protected]> * adjust inputs to test * add test for array with no missing values * added whatsnew * fixing format issues Co-authored-by: Deepak Cherian <[email protected]>

Use map_blocks for weighted init checks

c3b47db

added dask test

36f83bc

jbusecke marked this pull request as draft October 31, 2020 20:14

max-sixty reviewed Oct 31, 2020

View reviewed changes

xarray/core/weighted.py Outdated Show resolved Hide resolved

mathause reviewed Nov 1, 2020

View reviewed changes

dcherian reviewed Nov 1, 2020

View reviewed changes

xarray/core/weighted.py Outdated Show resolved Hide resolved

jbusecke and others added 4 commits November 1, 2020 16:43

Update xarray/core/weighted.py

5cabe4a

Co-authored-by: Deepak Cherian <[email protected]>

Update xarray/core/weighted.py

dc129bb

Co-authored-by: Maximilian Roos <[email protected]>

implement requires_dask

a3b51e0

Implement raise_if_dask_computes

2f0cc53

Added logic to check for dask arrays

feb0f1c

applied isort

4da7120

keewis reviewed Nov 2, 2020

View reviewed changes

jbusecke marked this pull request as ready for review November 2, 2020 19:17

max-sixty approved these changes Nov 3, 2020

View reviewed changes

xarray/core/weighted.py Outdated Show resolved Hide resolved

dcherian reviewed Nov 3, 2020

View reviewed changes

jbusecke and others added 3 commits November 4, 2020 09:37

Update xarray/core/weighted.py

8886baa

Co-authored-by: Maximilian Roos <[email protected]>

Refactor dask mapping

dc92dbe

Try duck_array_ops.isnull

681f34d

jbusecke changed the title ~~Dask friendly check for nans in .weighted()~~ Dask friendly check in .weighted() Nov 4, 2020

jbusecke added 6 commits November 4, 2020 11:30

black formatting

0852564

Add whatsnew

64bef42

Merge branch 'master' into weighted_nancheck_dask

e99b068

Remove numpy

c81a829

apply isort

6a85409

Merge branch 'master' into weighted_nancheck_dask

07f5623

dcherian reviewed Nov 5, 2020

View reviewed changes

xarray/core/weighted.py Outdated Show resolved Hide resolved

Update xarray/core/weighted.py

a5b4cf3

mathause reviewed Nov 5, 2020

View reviewed changes

xarray/core/weighted.py Outdated Show resolved Hide resolved

dcherian and others added 2 commits November 5, 2020 12:28

Update xarray/core/weighted.py

98c4e32

Co-authored-by: Mathias Hauser <[email protected]>

black formatting

8964608

dcherian approved these changes Nov 6, 2020

View reviewed changes

mathause reviewed Nov 9, 2020

View reviewed changes

xarray/core/weighted.py Show resolved Hide resolved

mathause and others added 2 commits November 9, 2020 14:42

Update xarray/core/weighted.py

c49d312

Merge branch 'master' into weighted_nancheck_dask

b5d5a30

mathause reviewed Nov 9, 2020

View reviewed changes

xarray/core/weighted.py Outdated Show resolved Hide resolved

Update xarray/core/weighted.py

1d59e0f

mathause merged commit 5a60354 into pydata:master Nov 9, 2020

jbusecke mentioned this pull request Nov 30, 2020

Non lazy behavior for weighted average when using resampled data #4625

Closed

mathause mentioned this pull request Jan 13, 2021

Improve performance of xarray.corr() on big datasets #4804

Open

AndrewILWilliams mentioned this pull request May 9, 2021

Dask-friendly nan check in xr.corr() and xr.cov() #5284

Merged

4 tasks

AndrewILWilliams added a commit to AndrewILWilliams/xarray that referenced this pull request May 16, 2021

Using map_blocks to lazily mask input arrays, following pydata#4559

227b622

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Dask friendly check in `.weighted()` #4559

Dask friendly check in `.weighted()` #4559

jbusecke commented Oct 31, 2020 •

edited

Loading

pep8speaks commented Oct 31, 2020 •

edited

Loading

max-sixty left a comment

jbusecke commented Nov 1, 2020

max-sixty commented Nov 1, 2020

keewis commented Nov 1, 2020

mathause left a comment

mathause Nov 1, 2020

mathause commented Nov 1, 2020

jbusecke commented Nov 1, 2020

keewis Nov 2, 2020

mathause Nov 2, 2020

jbusecke Nov 2, 2020

dcherian Nov 3, 2020

jbusecke commented Nov 3, 2020

max-sixty left a comment

dcherian Nov 3, 2020 •

edited

Loading

jbusecke commented Nov 4, 2020

jbusecke commented Nov 4, 2020

keewis commented Nov 4, 2020 •

edited

Loading

jbusecke commented Nov 5, 2020

jbusecke commented Nov 6, 2020

keewis commented Nov 6, 2020 •

edited

Loading

jbusecke commented Nov 6, 2020

dcherian left a comment

mathause commented Nov 9, 2020


		weights = DataArray(weights).chunk({"dim_0": -1})

		weighted = data.weighted(weights)

	weights.data = dsa.map_blocks(
	weights = weights.copy(data=dsa.map_blocks(

Dask friendly check in .weighted() #4559

Dask friendly check in .weighted() #4559

Conversation

jbusecke commented Oct 31, 2020 • edited Loading

pep8speaks commented Oct 31, 2020 • edited Loading

Comment last updated at 2020-11-09 15:24:48 UTC

max-sixty left a comment

Choose a reason for hiding this comment

jbusecke commented Nov 1, 2020

max-sixty commented Nov 1, 2020

keewis commented Nov 1, 2020

mathause left a comment

Choose a reason for hiding this comment

mathause Nov 1, 2020

Choose a reason for hiding this comment

mathause commented Nov 1, 2020

jbusecke commented Nov 1, 2020

keewis Nov 2, 2020

Choose a reason for hiding this comment

mathause Nov 2, 2020

Choose a reason for hiding this comment

jbusecke Nov 2, 2020

Choose a reason for hiding this comment

dcherian Nov 3, 2020

Choose a reason for hiding this comment

jbusecke commented Nov 3, 2020

max-sixty left a comment

Choose a reason for hiding this comment

dcherian Nov 3, 2020 • edited Loading

Choose a reason for hiding this comment

jbusecke commented Nov 4, 2020

jbusecke commented Nov 4, 2020

keewis commented Nov 4, 2020 • edited Loading

jbusecke commented Nov 5, 2020

jbusecke commented Nov 6, 2020

keewis commented Nov 6, 2020 • edited Loading

jbusecke commented Nov 6, 2020

dcherian left a comment

Choose a reason for hiding this comment

mathause commented Nov 9, 2020

Dask friendly check in `.weighted()` #4559

Dask friendly check in `.weighted()` #4559

jbusecke commented Oct 31, 2020 •

edited

Loading

pep8speaks commented Oct 31, 2020 •

edited

Loading

dcherian Nov 3, 2020 •

edited

Loading

keewis commented Nov 4, 2020 •

edited

Loading

keewis commented Nov 6, 2020 •

edited

Loading