More climatology reductions using Cubed

This is the Cubed equivalent of More climatology reductions.

The task is to compute an hourly climatology from an hourly dataset with 744 hours in each chunk, using the “map-reduce” strategy.

import cubed
import cubed.array_api as xp
import numpy as np
import pandas as pd
import xarray as xr

import flox.xarray

Create data

Note that we use fewer lat/long points so the computation can be run locally.

spec = cubed.Spec(allowed_mem="2GB")
ds = xr.Dataset(
    {
        "tp": (
            ("time", "latitude", "longitude"),
            xp.ones((8760, 72, 144), chunks=(744, 5, 144), dtype=np.float32, spec=spec),
        )
    },
    coords={"time": pd.date_range("2021-01-01", "2021-12-31 23:59", freq="h")},
)
ds

---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
File ~/checkouts/readthedocs.org/user_builds/flox/envs/stable/lib/python3.14/site-packages/IPython/core/formatters.py:406, in BaseFormatter.__call__(self, obj)
    404     method = get_real_method(obj, self.print_method)
    405     if method is not None:
--> 406         return method()
    407     return None
    408 else:

File ~/checkouts/readthedocs.org/user_builds/flox/envs/stable/lib/python3.14/site-packages/xarray/core/dataset.py:2417, in Dataset._repr_html_(self)
   2415 if OPTIONS["display_style"] == "text":
   2416     return f"<pre>{escape(repr(self))}</pre>"
-> 2417 return formatting_html.dataset_repr(self)

File ~/checkouts/readthedocs.org/user_builds/flox/envs/stable/lib/python3.14/site-packages/xarray/core/formatting_html.py:379, in dataset_repr(ds)
    376 if ds.coords:
    377     sections.append(coord_section(ds.coords))
--> 379 sections.append(datavar_section(ds.data_vars))
    381 display_default_indexes = _get_boolean_with_default(
    382     "display_default_indexes", False
    383 )
    384 xindexes = filter_nondefault_indexes(
    385     _get_indexes_dict(ds.xindexes), not display_default_indexes
    386 )

File ~/checkouts/readthedocs.org/user_builds/flox/envs/stable/lib/python3.14/site-packages/xarray/core/formatting_html.py:225, in _mapping_section(mapping, name, details_func, max_items_collapse, expand_option_name, enabled, **kwargs)
    218 expanded = max_items_collapse is None or _get_boolean_with_default(
    219     expand_option_name, n_items < max_items_collapse
    220 )
    221 collapsed = not expanded
    223 return collapsible_section(
    224     f"{name}:",
--> 225     details=details_func(mapping, **kwargs),
    226     n_items=n_items,
    227     enabled=enabled,
    228     collapsed=collapsed,
    229 )

File ~/checkouts/readthedocs.org/user_builds/flox/envs/stable/lib/python3.14/site-packages/xarray/core/formatting_html.py:136, in summarize_vars(variables)
    135 def summarize_vars(variables) -> str:
--> 136     vars_li = "".join(
    137         f"<li class='xr-var-item'>{summarize_variable(k, v)}</li>"
    138         for k, v in variables.items()
    139     )
    141     return f"<ul class='xr-var-list'>{vars_li}</ul>"

File ~/checkouts/readthedocs.org/user_builds/flox/envs/stable/lib/python3.14/site-packages/xarray/core/formatting_html.py:137, in <genexpr>(.0)
    135 def summarize_vars(variables) -> str:
    136     vars_li = "".join(
--> 137         f"<li class='xr-var-item'>{summarize_variable(k, v)}</li>"
    138         for k, v in variables.items()
    139     )
    141     return f"<ul class='xr-var-list'>{vars_li}</ul>"

File ~/checkouts/readthedocs.org/user_builds/flox/envs/stable/lib/python3.14/site-packages/xarray/core/formatting_html.py:98, in summarize_variable(name, var, is_index, dtype)
     96 preview = escape(inline_variable_array_repr(variable, 35))
     97 attrs_ul = summarize_attrs(var.attrs)
---> 98 data_repr = short_data_repr_html(variable)
    100 attrs_icon = _icon("icon-file-text2")
    101 data_icon = _icon("icon-database")

File ~/checkouts/readthedocs.org/user_builds/flox/envs/stable/lib/python3.14/site-packages/xarray/core/formatting_html.py:45, in short_data_repr_html(array)
     43 internal_data = getattr(array, "variable", array)._data
     44 if hasattr(internal_data, "_repr_html_"):
---> 45     return internal_data._repr_html_()
     46 text = escape(short_data_repr(array))
     47 return f"<pre>{text}</pre>"

File ~/checkouts/readthedocs.org/user_builds/flox/envs/stable/lib/python3.14/site-packages/cubed/array_api/array_object.py:50, in Array._repr_html_(self)
     49 def _repr_html_(self):
---> 50     from cubed.diagnostics.widgets import get_template
     52     try:
     53         grid = self.to_svg(size=ARRAY_SVG_SIZE)

File ~/checkouts/readthedocs.org/user_builds/flox/envs/stable/lib/python3.14/site-packages/cubed/diagnostics/__init__.py:1
----> 1 from .rich import RichProgressBar as ProgressBar
      3 __all__ = ["ProgressBar"]

File ~/checkouts/readthedocs.org/user_builds/flox/envs/stable/lib/python3.14/site-packages/cubed/diagnostics/rich.py:6
      3 import time
      4 from contextlib import contextmanager
----> 6 from rich.console import RenderableType
      7 from rich.progress import (
      8     BarColumn,
      9     MofNCompleteColumn,
   (...)     15     TimeElapsedColumn,
     16 )
     17 from rich.text import Text

ModuleNotFoundError: No module named 'rich'

<xarray.Dataset> Size: 363MB
Dimensions:  (time: 8760, latitude: 72, longitude: 144)
Coordinates:
  * time     (time) datetime64[us] 70kB 2021-01-01 ... 2021-12-31T23:00:00
Dimensions without coordinates: latitude, longitude
Data variables:
    tp       (time, latitude, longitude) float32 363MB cubed.Array<chunksize=(744, 5, 144)>

Computation

hourly = flox.xarray.xarray_reduce(ds.tp, ds.time.dt.hour, func="mean", reindex=True)
hourly

---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
File ~/checkouts/readthedocs.org/user_builds/flox/envs/stable/lib/python3.14/site-packages/IPython/core/formatters.py:406, in BaseFormatter.__call__(self, obj)
    404     method = get_real_method(obj, self.print_method)
    405     if method is not None:
--> 406         return method()
    407     return None
    408 else:

File ~/checkouts/readthedocs.org/user_builds/flox/envs/stable/lib/python3.14/site-packages/xarray/core/common.py:189, in AbstractArray._repr_html_(self)
    187 if OPTIONS["display_style"] == "text":
    188     return f"<pre>{escape(repr(self))}</pre>"
--> 189 return formatting_html.array_repr(self)

File ~/checkouts/readthedocs.org/user_builds/flox/envs/stable/lib/python3.14/site-packages/xarray/core/formatting_html.py:344, in array_repr(arr)
    336 arr_name = escape(repr(arr.name)) if getattr(arr, "name", None) else ""
    338 header_components = [
    339     f"<div class='xr-obj-type'>{obj_type}</div>",
    340     f"<div class='xr-obj-name'>{arr_name}</div>",
    341     format_dims(dims, indexed_dims),
    342 ]
--> 344 sections = [array_section(arr)]
    346 if hasattr(arr, "coords"):
    347     if arr.coords:

File ~/checkouts/readthedocs.org/user_builds/flox/envs/stable/lib/python3.14/site-packages/xarray/core/formatting_html.py:250, in array_section(obj)
    248 variable = getattr(obj, "variable", obj)
    249 preview = escape(inline_variable_array_repr(variable, max_width=70))
--> 250 data_repr = short_data_repr_html(obj)
    251 data_icon = _icon("icon-database")
    253 return (
    254     "<div class='xr-array-wrap'>"
    255     f"<input id='{data_id}' class='xr-array-in' type='checkbox' {collapsed}>"
   (...)    259     "</div>"
    260 )

File ~/checkouts/readthedocs.org/user_builds/flox/envs/stable/lib/python3.14/site-packages/xarray/core/formatting_html.py:45, in short_data_repr_html(array)
     43 internal_data = getattr(array, "variable", array)._data
     44 if hasattr(internal_data, "_repr_html_"):
---> 45     return internal_data._repr_html_()
     46 text = escape(short_data_repr(array))
     47 return f"<pre>{text}</pre>"

File ~/checkouts/readthedocs.org/user_builds/flox/envs/stable/lib/python3.14/site-packages/cubed/array_api/array_object.py:50, in Array._repr_html_(self)
     49 def _repr_html_(self):
---> 50     from cubed.diagnostics.widgets import get_template
     52     try:
     53         grid = self.to_svg(size=ARRAY_SVG_SIZE)

File ~/checkouts/readthedocs.org/user_builds/flox/envs/stable/lib/python3.14/site-packages/cubed/diagnostics/__init__.py:1
----> 1 from .rich import RichProgressBar as ProgressBar
      3 __all__ = ["ProgressBar"]

File ~/checkouts/readthedocs.org/user_builds/flox/envs/stable/lib/python3.14/site-packages/cubed/diagnostics/rich.py:6
      3 import time
      4 from contextlib import contextmanager
----> 6 from rich.console import RenderableType
      7 from rich.progress import (
      8     BarColumn,
      9     MofNCompleteColumn,
   (...)     15     TimeElapsedColumn,
     16 )
     17 from rich.text import Text

ModuleNotFoundError: No module named 'rich'

<xarray.DataArray 'tp' (hour: 24, latitude: 72, longitude: 144)> Size: 995kB
cubed.Array<array-018, shape=(24, 72, 144), dtype=float32, chunks=((24,), (5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 2), (144,))>
Coordinates:
  * hour     (hour) int64 192B 0 1 2 3 4 5 6 7 8 ... 15 16 17 18 19 20 21 22 23
Dimensions without coordinates: latitude, longitude

hourly.compute()

<xarray.DataArray 'tp' (hour: 24, latitude: 72, longitude: 144)> Size: 995kB
array([[[1., 1., 1., ..., 1., 1., 1.],
        [1., 1., 1., ..., 1., 1., 1.],
        [1., 1., 1., ..., 1., 1., 1.],
        ...,
        [1., 1., 1., ..., 1., 1., 1.],
        [1., 1., 1., ..., 1., 1., 1.],
        [1., 1., 1., ..., 1., 1., 1.]],

       [[1., 1., 1., ..., 1., 1., 1.],
        [1., 1., 1., ..., 1., 1., 1.],
        [1., 1., 1., ..., 1., 1., 1.],
        ...,
        [1., 1., 1., ..., 1., 1., 1.],
        [1., 1., 1., ..., 1., 1., 1.],
        [1., 1., 1., ..., 1., 1., 1.]],

       [[1., 1., 1., ..., 1., 1., 1.],
        [1., 1., 1., ..., 1., 1., 1.],
        [1., 1., 1., ..., 1., 1., 1.],
        ...,
...
        [1., 1., 1., ..., 1., 1., 1.],
        [1., 1., 1., ..., 1., 1., 1.],
        [1., 1., 1., ..., 1., 1., 1.]],

       [[1., 1., 1., ..., 1., 1., 1.],
        [1., 1., 1., ..., 1., 1., 1.],
        [1., 1., 1., ..., 1., 1., 1.],
        ...,
        [1., 1., 1., ..., 1., 1., 1.],
        [1., 1., 1., ..., 1., 1., 1.],
        [1., 1., 1., ..., 1., 1., 1.]],

       [[1., 1., 1., ..., 1., 1., 1.],
        [1., 1., 1., ..., 1., 1., 1.],
        [1., 1., 1., ..., 1., 1., 1.],
        ...,
        [1., 1., 1., ..., 1., 1., 1.],
        [1., 1., 1., ..., 1., 1., 1.],
        [1., 1., 1., ..., 1., 1., 1.]]],
      shape=(24, 72, 144), dtype=float32)
Coordinates:
  * hour     (hour) int64 192B 0 1 2 3 4 5 6 7 8 ... 15 16 17 18 19 20 21 22 23
Dimensions without coordinates: latitude, longitude

Other climatologies: resampling by month

This uses the “blockwise” strategy.

monthly = ds.tp.resample(time="ME").sum(method="blockwise")
monthly

---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
File ~/checkouts/readthedocs.org/user_builds/flox/envs/stable/lib/python3.14/site-packages/IPython/core/formatters.py:406, in BaseFormatter.__call__(self, obj)
    404     method = get_real_method(obj, self.print_method)
    405     if method is not None:
--> 406         return method()
    407     return None
    408 else:

File ~/checkouts/readthedocs.org/user_builds/flox/envs/stable/lib/python3.14/site-packages/xarray/core/common.py:189, in AbstractArray._repr_html_(self)
    187 if OPTIONS["display_style"] == "text":
    188     return f"<pre>{escape(repr(self))}</pre>"
--> 189 return formatting_html.array_repr(self)

File ~/checkouts/readthedocs.org/user_builds/flox/envs/stable/lib/python3.14/site-packages/xarray/core/formatting_html.py:344, in array_repr(arr)
    336 arr_name = escape(repr(arr.name)) if getattr(arr, "name", None) else ""
    338 header_components = [
    339     f"<div class='xr-obj-type'>{obj_type}</div>",
    340     f"<div class='xr-obj-name'>{arr_name}</div>",
    341     format_dims(dims, indexed_dims),
    342 ]
--> 344 sections = [array_section(arr)]
    346 if hasattr(arr, "coords"):
    347     if arr.coords:

File ~/checkouts/readthedocs.org/user_builds/flox/envs/stable/lib/python3.14/site-packages/xarray/core/formatting_html.py:250, in array_section(obj)
    248 variable = getattr(obj, "variable", obj)
    249 preview = escape(inline_variable_array_repr(variable, max_width=70))
--> 250 data_repr = short_data_repr_html(obj)
    251 data_icon = _icon("icon-database")
    253 return (
    254     "<div class='xr-array-wrap'>"
    255     f"<input id='{data_id}' class='xr-array-in' type='checkbox' {collapsed}>"
   (...)    259     "</div>"
    260 )

File ~/checkouts/readthedocs.org/user_builds/flox/envs/stable/lib/python3.14/site-packages/xarray/core/formatting_html.py:45, in short_data_repr_html(array)
     43 internal_data = getattr(array, "variable", array)._data
     44 if hasattr(internal_data, "_repr_html_"):
---> 45     return internal_data._repr_html_()
     46 text = escape(short_data_repr(array))
     47 return f"<pre>{text}</pre>"

File ~/checkouts/readthedocs.org/user_builds/flox/envs/stable/lib/python3.14/site-packages/cubed/array_api/array_object.py:50, in Array._repr_html_(self)
     49 def _repr_html_(self):
---> 50     from cubed.diagnostics.widgets import get_template
     52     try:
     53         grid = self.to_svg(size=ARRAY_SVG_SIZE)

File ~/checkouts/readthedocs.org/user_builds/flox/envs/stable/lib/python3.14/site-packages/cubed/diagnostics/__init__.py:1
----> 1 from .rich import RichProgressBar as ProgressBar
      3 __all__ = ["ProgressBar"]

File ~/checkouts/readthedocs.org/user_builds/flox/envs/stable/lib/python3.14/site-packages/cubed/diagnostics/rich.py:6
      3 import time
      4 from contextlib import contextmanager
----> 6 from rich.console import RenderableType
      7 from rich.progress import (
      8     BarColumn,
      9     MofNCompleteColumn,
   (...)     15     TimeElapsedColumn,
     16 )
     17 from rich.text import Text

ModuleNotFoundError: No module named 'rich'

<xarray.DataArray 'tp' (time: 12, latitude: 72, longitude: 144)> Size: 995kB
cubed.Array<array-022, shape=(12, 72, 144), dtype=float64, chunks=((1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1), (5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 2), (144,))>
Coordinates:
  * time     (time) datetime64[us] 96B 2021-01-31 2021-02-28 ... 2021-12-31
Dimensions without coordinates: latitude, longitude

monthly.compute()

<xarray.DataArray 'tp' (time: 12, latitude: 72, longitude: 144)> Size: 995kB
array([[[744., 744., 744., ..., 744., 744., 744.],
        [744., 744., 744., ..., 744., 744., 744.],
        [744., 744., 744., ..., 744., 744., 744.],
        ...,
        [744., 744., 744., ..., 744., 744., 744.],
        [744., 744., 744., ..., 744., 744., 744.],
        [744., 744., 744., ..., 744., 744., 744.]],

       [[672., 672., 672., ..., 672., 672., 672.],
        [672., 672., 672., ..., 672., 672., 672.],
        [672., 672., 672., ..., 672., 672., 672.],
        ...,
        [672., 672., 672., ..., 672., 672., 672.],
        [672., 672., 672., ..., 672., 672., 672.],
        [672., 672., 672., ..., 672., 672., 672.]],

       [[744., 744., 744., ..., 744., 744., 744.],
        [744., 744., 744., ..., 744., 744., 744.],
        [744., 744., 744., ..., 744., 744., 744.],
        ...,
...
        ...,
        [744., 744., 744., ..., 744., 744., 744.],
        [744., 744., 744., ..., 744., 744., 744.],
        [744., 744., 744., ..., 744., 744., 744.]],

       [[720., 720., 720., ..., 720., 720., 720.],
        [720., 720., 720., ..., 720., 720., 720.],
        [720., 720., 720., ..., 720., 720., 720.],
        ...,
        [720., 720., 720., ..., 720., 720., 720.],
        [720., 720., 720., ..., 720., 720., 720.],
        [720., 720., 720., ..., 720., 720., 720.]],

       [[744., 744., 744., ..., 744., 744., 744.],
        [744., 744., 744., ..., 744., 744., 744.],
        [744., 744., 744., ..., 744., 744., 744.],
        ...,
        [744., 744., 744., ..., 744., 744., 744.],
        [744., 744., 744., ..., 744., 744., 744.],
        [744., 744., 744., ..., 744., 744., 744.]]], shape=(12, 72, 144))
Coordinates:
  * time     (time) datetime64[us] 96B 2021-01-31 2021-02-28 ... 2021-12-31
Dimensions without coordinates: latitude, longitude