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
<xarray.Dataset> Size: 363MB
Dimensions:  (time: 8760, latitude: 72, longitude: 144)
Coordinates:
  * time     (time) datetime64[ns] 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
<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.]]], 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