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)>

xarray.Dataset

• Dimensions:
  • time: 8760
  • latitude: 72
  • longitude: 144
• Coordinates: (1)
  • time
    (time)
    datetime64[ns]
    2021-01-01 ... 2021-12-31T23:00:00

    array(['2021-01-01T00:00:00.000000000', '2021-01-01T01:00:00.000000000',
           '2021-01-01T02:00:00.000000000', ..., '2021-12-31T21:00:00.000000000',
           '2021-12-31T22:00:00.000000000', '2021-12-31T23:00:00.000000000'],
          shape=(8760,), dtype='datetime64[ns]')

• Data variables: (1)
  • tp
    (time, latitude, longitude)
    float32
    cubed.Array<chunksize=(744, 5, 144)>

    Array Chunk Bytes 363.3 MB 2.1 MB Shape (8760, 72, 144) (744, 5, 144) Count 1 arrays in Plan 180 Chunks Type float32 np.ndarray

• Indexes: (1)
  • time
    PandasIndex

    PandasIndex(DatetimeIndex(['2021-01-01 00:00:00', '2021-01-01 01:00:00',
                   '2021-01-01 02:00:00', '2021-01-01 03:00:00',
                   '2021-01-01 04:00:00', '2021-01-01 05:00:00',
                   '2021-01-01 06:00:00', '2021-01-01 07:00:00',
                   '2021-01-01 08:00:00', '2021-01-01 09:00:00',
                   ...
                   '2021-12-31 14:00:00', '2021-12-31 15:00:00',
                   '2021-12-31 16:00:00', '2021-12-31 17:00:00',
                   '2021-12-31 18:00:00', '2021-12-31 19:00:00',
                   '2021-12-31 20:00:00', '2021-12-31 21:00:00',
                   '2021-12-31 22:00:00', '2021-12-31 23:00:00'],
                  dtype='datetime64[ns]', name='time', length=8760, freq='h'))

• Attributes: (0)

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

xarray.DataArray

'tp'

• hour: 24
• latitude: 72
• longitude: 144

• cubed.Array<chunksize=(24, 5, 144)>

  Array Chunk Bytes 995.3 KB 69.1 KB Shape (24, 72, 144) (24, 5, 144) Count 9 arrays in Plan 15 Chunks Type float32 np.ndarray

• Coordinates: (1)
  • hour
    (hour)
    int64
    0 1 2 3 4 5 6 ... 18 19 20 21 22 23

    array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17,
           18, 19, 20, 21, 22, 23])

• Indexes: (1)
  • hour
    PandasIndex

    PandasIndex(Index([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17,
           18, 19, 20, 21, 22, 23],
          dtype='int64', name='hour'))

• Attributes: (0)

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

xarray.DataArray

'tp'

• hour: 24
• latitude: 72
• longitude: 144

• 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 ... 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0

  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: (1)
  • hour
    (hour)
    int64
    0 1 2 3 4 5 6 ... 18 19 20 21 22 23

    array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17,
           18, 19, 20, 21, 22, 23])

• Indexes: (1)
  • hour
    PandasIndex

    PandasIndex(Index([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17,
           18, 19, 20, 21, 22, 23],
          dtype='int64', name='hour'))

• Attributes: (0)

Other climatologies: resampling by month

This uses the “blockwise” strategy.

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

<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[ns] 96B 2021-01-31 2021-02-28 ... 2021-12-31
Dimensions without coordinates: latitude, longitude

xarray.DataArray

'tp'

• time: 12
• latitude: 72
• longitude: 144

• cubed.Array<chunksize=(1, 5, 144)>

  Array Chunk Bytes 995.3 KB 5.8 KB Shape (12, 72, 144) (1, 5, 144) Count 3 arrays in Plan 180 Chunks Type float64 np.ndarray

• Coordinates: (1)
  • time
    (time)
    datetime64[ns]
    2021-01-31 ... 2021-12-31

    array(['2021-01-31T00:00:00.000000000', '2021-02-28T00:00:00.000000000',
           '2021-03-31T00:00:00.000000000', '2021-04-30T00:00:00.000000000',
           '2021-05-31T00:00:00.000000000', '2021-06-30T00:00:00.000000000',
           '2021-07-31T00:00:00.000000000', '2021-08-31T00:00:00.000000000',
           '2021-09-30T00:00:00.000000000', '2021-10-31T00:00:00.000000000',
           '2021-11-30T00:00:00.000000000', '2021-12-31T00:00:00.000000000'],
          dtype='datetime64[ns]')

• Indexes: (1)
  • time
    PandasIndex

    PandasIndex(DatetimeIndex(['2021-01-31', '2021-02-28', '2021-03-31', '2021-04-30',
                   '2021-05-31', '2021-06-30', '2021-07-31', '2021-08-31',
                   '2021-09-30', '2021-10-31', '2021-11-30', '2021-12-31'],
                  dtype='datetime64[ns]', name='time', freq='ME'))

• Attributes: (0)

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[ns] 96B 2021-01-31 2021-02-28 ... 2021-12-31
Dimensions without coordinates: latitude, longitude

xarray.DataArray

'tp'

• time: 12
• latitude: 72
• longitude: 144

• 744.0 744.0 744.0 744.0 744.0 744.0 ... 744.0 744.0 744.0 744.0 744.0

  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: (1)
  • time
    (time)
    datetime64[ns]
    2021-01-31 ... 2021-12-31

    array(['2021-01-31T00:00:00.000000000', '2021-02-28T00:00:00.000000000',
           '2021-03-31T00:00:00.000000000', '2021-04-30T00:00:00.000000000',
           '2021-05-31T00:00:00.000000000', '2021-06-30T00:00:00.000000000',
           '2021-07-31T00:00:00.000000000', '2021-08-31T00:00:00.000000000',
           '2021-09-30T00:00:00.000000000', '2021-10-31T00:00:00.000000000',
           '2021-11-30T00:00:00.000000000', '2021-12-31T00:00:00.000000000'],
          dtype='datetime64[ns]')

• Indexes: (1)
  • time
    PandasIndex

    PandasIndex(DatetimeIndex(['2021-01-31', '2021-02-28', '2021-03-31', '2021-04-30',
                   '2021-05-31', '2021-06-30', '2021-07-31', '2021-08-31',
                   '2021-09-30', '2021-10-31', '2021-11-30', '2021-12-31'],
                  dtype='datetime64[ns]', name='time', freq='ME'))

• Attributes: (0)