Skip to content

Commit

Permalink
Allow cubed arrays to be passed to flox groupby (#7941)
Browse files Browse the repository at this point in the history 
* fix error raised by ruff

* whatsnew
  • Loading branch information
@TomNicholas
TomNicholas authored Jun 26, 2023
1 parent 86f9933 commit e5071aa
Show file tree
Hide file tree
Showing 4 changed files with 52 additions and 48 deletions.
3 changes: 3 additions & 0 deletions doc/whats-new.rst
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,9 @@ Documentation
Internal Changes
~~~~~~~~~~~~~~~~

- Allow chunked non-dask arrays (i.e. Cubed arrays) in groupby operations. (:pull:`7941`)
By `Tom Nicholas <https://github.com/TomNicholas>`_.


.. _whats-new.2023.06.0:

Expand Down
82 changes: 41 additions & 41 deletions xarray/core/_aggregations.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
from xarray.core import duck_array_ops
from xarray.core.options import OPTIONS
from xarray.core.types import Dims
from xarray.core.utils import contains_only_dask_or_numpy, module_available
from xarray.core.utils import contains_only_chunked_or_numpy, module_available

if TYPE_CHECKING:
from xarray.core.dataarray import DataArray
Expand Down Expand Up @@ -2413,7 +2413,7 @@ def count(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="count",
Expand Down Expand Up @@ -2511,7 +2511,7 @@ def all(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="all",
Expand Down Expand Up @@ -2609,7 +2609,7 @@ def any(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="any",
Expand Down Expand Up @@ -2723,7 +2723,7 @@ def max(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="max",
Expand Down Expand Up @@ -2839,7 +2839,7 @@ def min(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="min",
Expand Down Expand Up @@ -2957,7 +2957,7 @@ def mean(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="mean",
Expand Down Expand Up @@ -3092,7 +3092,7 @@ def prod(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="prod",
Expand Down Expand Up @@ -3229,7 +3229,7 @@ def sum(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="sum",
Expand Down Expand Up @@ -3363,7 +3363,7 @@ def std(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="std",
Expand Down Expand Up @@ -3497,7 +3497,7 @@ def var(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="var",
Expand Down Expand Up @@ -3927,7 +3927,7 @@ def count(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="count",
Expand Down Expand Up @@ -4025,7 +4025,7 @@ def all(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="all",
Expand Down Expand Up @@ -4123,7 +4123,7 @@ def any(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="any",
Expand Down Expand Up @@ -4237,7 +4237,7 @@ def max(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="max",
Expand Down Expand Up @@ -4353,7 +4353,7 @@ def min(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="min",
Expand Down Expand Up @@ -4471,7 +4471,7 @@ def mean(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="mean",
Expand Down Expand Up @@ -4606,7 +4606,7 @@ def prod(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="prod",
Expand Down Expand Up @@ -4743,7 +4743,7 @@ def sum(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="sum",
Expand Down Expand Up @@ -4877,7 +4877,7 @@ def std(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="std",
Expand Down Expand Up @@ -5011,7 +5011,7 @@ def var(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="var",
Expand Down Expand Up @@ -5436,7 +5436,7 @@ def count(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="count",
Expand Down Expand Up @@ -5527,7 +5527,7 @@ def all(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="all",
Expand Down Expand Up @@ -5618,7 +5618,7 @@ def any(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="any",
Expand Down Expand Up @@ -5723,7 +5723,7 @@ def max(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="max",
Expand Down Expand Up @@ -5830,7 +5830,7 @@ def min(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="min",
Expand Down Expand Up @@ -5939,7 +5939,7 @@ def mean(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="mean",
Expand Down Expand Up @@ -6063,7 +6063,7 @@ def prod(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="prod",
Expand Down Expand Up @@ -6189,7 +6189,7 @@ def sum(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="sum",
Expand Down Expand Up @@ -6312,7 +6312,7 @@ def std(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="std",
Expand Down Expand Up @@ -6435,7 +6435,7 @@ def var(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="var",
Expand Down Expand Up @@ -6842,7 +6842,7 @@ def count(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="count",
Expand Down Expand Up @@ -6933,7 +6933,7 @@ def all(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="all",
Expand Down Expand Up @@ -7024,7 +7024,7 @@ def any(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="any",
Expand Down Expand Up @@ -7129,7 +7129,7 @@ def max(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="max",
Expand Down Expand Up @@ -7236,7 +7236,7 @@ def min(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="min",
Expand Down Expand Up @@ -7345,7 +7345,7 @@ def mean(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="mean",
Expand Down Expand Up @@ -7469,7 +7469,7 @@ def prod(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="prod",
Expand Down Expand Up @@ -7595,7 +7595,7 @@ def sum(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="sum",
Expand Down Expand Up @@ -7718,7 +7718,7 @@ def std(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="std",
Expand Down Expand Up @@ -7841,7 +7841,7 @@ def var(
if (
flox_available
and OPTIONS["use_flox"]
and contains_only_dask_or_numpy(self._obj)
and contains_only_chunked_or_numpy(self._obj)
):
return self._flox_reduce(
func="var",
Expand Down
Loading

0 comments on commit e5071aa

Please sign in to comment.