Xarray

Cubed can work with Xarray datasets via the cubed-xarray package.

Install by running the following:

pip install cubed cubed-xarray xarray pooch netCDF4

Note that pooch and netCDF4 are needed to access the Xarray tutorial datasets that we use in the example below.

Open dataset

Start by importing Xarray - note that we don’t need to import Cubed or cubed-xarray, since they will be picked up automatically.

import xarray as xr

xr.set_options(display_expand_attrs=False, display_expand_data=True);

We open an Xarray dataset (in netCDF format) using the usual open_dataset function. By specifying chunks={} we ensure that the dataset is chunked using the on-disk chunking (here it is the netCDF file chunking). The chunked_array_type argument specifies which chunked array type to use - Cubed in this case.

ds = xr.tutorial.open_dataset(
    "air_temperature", chunked_array_type="cubed", chunks={}
)
ds

<xarray.Dataset> Size: 31MB
Dimensions:  (lat: 25, time: 2920, lon: 53)
Coordinates:
  * lat      (lat) float32 100B 75.0 72.5 70.0 67.5 65.0 ... 22.5 20.0 17.5 15.0
  * lon      (lon) float32 212B 200.0 202.5 205.0 207.5 ... 325.0 327.5 330.0
  * time     (time) datetime64[ns] 23kB 2013-01-01 ... 2014-12-31T18:00:00
Data variables:
    air      (time, lat, lon) float64 31MB cubed.Array<chunksize=(2920, 25, 53)>
Attributes: (5)

xarray.Dataset

• Dimensions:
  • lat: 25
  • time: 2920
  • lon: 53
• Coordinates: (3)
  • lat
    (lat)
    float32
    75.0 72.5 70.0 ... 20.0 17.5 15.0

    standard_name :

    latitude

    long_name :

    Latitude

    units :

    degrees_north

    axis :

    Y

    array([75. , 72.5, 70. , 67.5, 65. , 62.5, 60. , 57.5, 55. , 52.5, 50. , 47.5,
           45. , 42.5, 40. , 37.5, 35. , 32.5, 30. , 27.5, 25. , 22.5, 20. , 17.5,
           15. ], dtype=float32)

  • lon
    (lon)
    float32
    200.0 202.5 205.0 ... 327.5 330.0

    standard_name :

    longitude

    long_name :

    Longitude

    units :

    degrees_east

    axis :

    X

    array([200. , 202.5, 205. , 207.5, 210. , 212.5, 215. , 217.5, 220. , 222.5,
           225. , 227.5, 230. , 232.5, 235. , 237.5, 240. , 242.5, 245. , 247.5,
           250. , 252.5, 255. , 257.5, 260. , 262.5, 265. , 267.5, 270. , 272.5,
           275. , 277.5, 280. , 282.5, 285. , 287.5, 290. , 292.5, 295. , 297.5,
           300. , 302.5, 305. , 307.5, 310. , 312.5, 315. , 317.5, 320. , 322.5,
           325. , 327.5, 330. ], dtype=float32)

  • time
    (time)
    datetime64[ns]
    2013-01-01 ... 2014-12-31T18:00:00

    standard_name :

    time

    long_name :

    Time

    array(['2013-01-01T00:00:00.000000000', '2013-01-01T06:00:00.000000000',
           '2013-01-01T12:00:00.000000000', ..., '2014-12-31T06:00:00.000000000',
           '2014-12-31T12:00:00.000000000', '2014-12-31T18:00:00.000000000'],
          shape=(2920,), dtype='datetime64[ns]')

• Data variables: (1)
  • air
    (time, lat, lon)
    float64
    cubed.Array<chunksize=(2920, 25, 53)>

    long_name :

    4xDaily Air temperature at sigma level 995

    units :

    degK

    precision :

    2

    GRIB_id :

    11

    GRIB_name :

    TMP

    var_desc :

    Air temperature

    dataset :

    NMC Reanalysis

    level_desc :

    Surface

    statistic :

    Individual Obs

    parent_stat :

    Other

    actual_range :

    [185.16 322.1 ]

    Array Chunk Bytes 31.0 MB 31.0 MB Shape (2920, 25, 53) (2920, 25, 53) Count 3 arrays in Plan 1 Chunks Type float64 np.ndarray

• Indexes: (3)
  • lat
    PandasIndex

    PandasIndex(Index([75.0, 72.5, 70.0, 67.5, 65.0, 62.5, 60.0, 57.5, 55.0, 52.5, 50.0, 47.5,
           45.0, 42.5, 40.0, 37.5, 35.0, 32.5, 30.0, 27.5, 25.0, 22.5, 20.0, 17.5,
           15.0],
          dtype='float32', name='lat'))

  • lon
    PandasIndex

    PandasIndex(Index([200.0, 202.5, 205.0, 207.5, 210.0, 212.5, 215.0, 217.5, 220.0, 222.5,
           225.0, 227.5, 230.0, 232.5, 235.0, 237.5, 240.0, 242.5, 245.0, 247.5,
           250.0, 252.5, 255.0, 257.5, 260.0, 262.5, 265.0, 267.5, 270.0, 272.5,
           275.0, 277.5, 280.0, 282.5, 285.0, 287.5, 290.0, 292.5, 295.0, 297.5,
           300.0, 302.5, 305.0, 307.5, 310.0, 312.5, 315.0, 317.5, 320.0, 322.5,
           325.0, 327.5, 330.0],
          dtype='float32', name='lon'))

  • time
    PandasIndex

    PandasIndex(DatetimeIndex(['2013-01-01 00:00:00', '2013-01-01 06:00:00',
                   '2013-01-01 12:00:00', '2013-01-01 18:00:00',
                   '2013-01-02 00:00:00', '2013-01-02 06:00:00',
                   '2013-01-02 12:00:00', '2013-01-02 18:00:00',
                   '2013-01-03 00:00:00', '2013-01-03 06:00:00',
                   ...
                   '2014-12-29 12:00:00', '2014-12-29 18:00:00',
                   '2014-12-30 00:00:00', '2014-12-30 06:00:00',
                   '2014-12-30 12:00:00', '2014-12-30 18:00:00',
                   '2014-12-31 00:00:00', '2014-12-31 06:00:00',
                   '2014-12-31 12:00:00', '2014-12-31 18:00:00'],
                  dtype='datetime64[ns]', name='time', length=2920, freq=None))

• Attributes: (5)

  Conventions :

  COARDS

  title :

  4x daily NMC reanalysis (1948)

  description :

  Data is from NMC initialized reanalysis (4x/day). These are the 0.9950 sigma level values.

  platform :

  Model

  references :

  http://www.esrl.noaa.gov/psd/data/gridded/data.ncep.reanalysis.html

Notice that the air data variable is a cubed.Array. Since Cubed has a lazy computation model, this array is not loaded from disk until a computation is run.

Convert to Zarr

We can use Cubed to convert the dataset to Zarr format by calling to_zarr on the dataset:

ds.to_zarr("air_temperature_cubed.zarr", mode="w", consolidated=True);

/opt/hostedtoolcache/Python/3.10.16/x64/lib/python3.10/site-packages/xarray/core/dataset.py:2622: SerializationWarning: saving variable None with floating point data as an integer dtype without any _FillValue to use for NaNs
  return to_zarr(  # type: ignore[call-overload,misc]

This will run a computation that loads the input data and writes it out to a Zarr store on the local filesystem.

Compute the mean

We can also use Xarray’s API to run computations on the dataset using Cubed. Here we find the mean air temperature over time, for each location:

mean = ds.air.mean("time", skipna=False)
mean

<xarray.DataArray 'air' (lat: 25, lon: 53)> Size: 11kB
cubed.Array<array-012, shape=(25, 53), dtype=float64, chunks=((25,), (53,))>
Coordinates:
  * lat      (lat) float32 100B 75.0 72.5 70.0 67.5 65.0 ... 22.5 20.0 17.5 15.0
  * lon      (lon) float32 212B 200.0 202.5 205.0 207.5 ... 325.0 327.5 330.0

xarray.DataArray

'air'

• lat: 25
• lon: 53

• cubed.Array<chunksize=(25, 53)>

  Array Chunk Bytes 10.6 KB 10.6 KB Shape (25, 53) (25, 53) Count 3 arrays in Plan 1 Chunks Type float64 np.ndarray

• Coordinates: (2)
  • lat
    (lat)
    float32
    75.0 72.5 70.0 ... 20.0 17.5 15.0

    standard_name :

    latitude

    long_name :

    Latitude

    units :

    degrees_north

    axis :

    Y

    array([75. , 72.5, 70. , 67.5, 65. , 62.5, 60. , 57.5, 55. , 52.5, 50. , 47.5,
           45. , 42.5, 40. , 37.5, 35. , 32.5, 30. , 27.5, 25. , 22.5, 20. , 17.5,
           15. ], dtype=float32)

  • lon
    (lon)
    float32
    200.0 202.5 205.0 ... 327.5 330.0

    standard_name :

    longitude

    long_name :

    Longitude

    units :

    degrees_east

    axis :

    X

    array([200. , 202.5, 205. , 207.5, 210. , 212.5, 215. , 217.5, 220. , 222.5,
           225. , 227.5, 230. , 232.5, 235. , 237.5, 240. , 242.5, 245. , 247.5,
           250. , 252.5, 255. , 257.5, 260. , 262.5, 265. , 267.5, 270. , 272.5,
           275. , 277.5, 280. , 282.5, 285. , 287.5, 290. , 292.5, 295. , 297.5,
           300. , 302.5, 305. , 307.5, 310. , 312.5, 315. , 317.5, 320. , 322.5,
           325. , 327.5, 330. ], dtype=float32)

• Indexes: (2)
  • lat
    PandasIndex

    PandasIndex(Index([75.0, 72.5, 70.0, 67.5, 65.0, 62.5, 60.0, 57.5, 55.0, 52.5, 50.0, 47.5,
           45.0, 42.5, 40.0, 37.5, 35.0, 32.5, 30.0, 27.5, 25.0, 22.5, 20.0, 17.5,
           15.0],
          dtype='float32', name='lat'))

  • lon
    PandasIndex

    PandasIndex(Index([200.0, 202.5, 205.0, 207.5, 210.0, 212.5, 215.0, 217.5, 220.0, 222.5,
           225.0, 227.5, 230.0, 232.5, 235.0, 237.5, 240.0, 242.5, 245.0, 247.5,
           250.0, 252.5, 255.0, 257.5, 260.0, 262.5, 265.0, 267.5, 270.0, 272.5,
           275.0, 277.5, 280.0, 282.5, 285.0, 287.5, 290.0, 292.5, 295.0, 297.5,
           300.0, 302.5, 305.0, 307.5, 310.0, 312.5, 315.0, 317.5, 320.0, 322.5,
           325.0, 327.5, 330.0],
          dtype='float32', name='lon'))

• Attributes: (0)

To run the computation we need to call compute:

mean.compute()

<xarray.DataArray 'air' (lat: 25, lon: 53)> Size: 11kB
array([[260.37644178, 260.18305137, 259.88662671, ..., 250.81590068,
        251.93811644, 253.43804795],
       [262.73439384, 262.79397603, 262.74933904, ..., 249.75590411,
        251.58575685, 254.35926027],
       [264.7687637 , 264.32730822, 264.06169521, ..., 250.60789041,
        253.58351027, 257.71559932],
       ...,
       [297.64986301, 296.95333219, 296.62931507, ..., 296.81092466,
        296.28796233, 295.81645548],
       [298.12920205, 297.93700685, 297.47039384, ..., 296.85954795,
        296.7770274 , 296.44383562],
       [298.36615068, 298.38573973, 298.11414384, ..., 297.33820548,
        297.28144521, 297.30510274]], shape=(25, 53))
Coordinates:
  * lat      (lat) float32 100B 75.0 72.5 70.0 67.5 65.0 ... 22.5 20.0 17.5 15.0
  * lon      (lon) float32 212B 200.0 202.5 205.0 207.5 ... 325.0 327.5 330.0

xarray.DataArray

'air'

• lat: 25
• lon: 53

• 260.4 260.2 259.9 259.5 259.0 258.6 ... 297.9 297.8 297.3 297.3 297.3

  array([[260.37644178, 260.18305137, 259.88662671, ..., 250.81590068,
          251.93811644, 253.43804795],
         [262.73439384, 262.79397603, 262.74933904, ..., 249.75590411,
          251.58575685, 254.35926027],
         [264.7687637 , 264.32730822, 264.06169521, ..., 250.60789041,
          253.58351027, 257.71559932],
         ...,
         [297.64986301, 296.95333219, 296.62931507, ..., 296.81092466,
          296.28796233, 295.81645548],
         [298.12920205, 297.93700685, 297.47039384, ..., 296.85954795,
          296.7770274 , 296.44383562],
         [298.36615068, 298.38573973, 298.11414384, ..., 297.33820548,
          297.28144521, 297.30510274]], shape=(25, 53))

• Coordinates: (2)
  • lat
    (lat)
    float32
    75.0 72.5 70.0 ... 20.0 17.5 15.0

    standard_name :

    latitude

    long_name :

    Latitude

    units :

    degrees_north

    axis :

    Y

    array([75. , 72.5, 70. , 67.5, 65. , 62.5, 60. , 57.5, 55. , 52.5, 50. , 47.5,
           45. , 42.5, 40. , 37.5, 35. , 32.5, 30. , 27.5, 25. , 22.5, 20. , 17.5,
           15. ], dtype=float32)

  • lon
    (lon)
    float32
    200.0 202.5 205.0 ... 327.5 330.0

    standard_name :

    longitude

    long_name :

    Longitude

    units :

    degrees_east

    axis :

    X

    array([200. , 202.5, 205. , 207.5, 210. , 212.5, 215. , 217.5, 220. , 222.5,
           225. , 227.5, 230. , 232.5, 235. , 237.5, 240. , 242.5, 245. , 247.5,
           250. , 252.5, 255. , 257.5, 260. , 262.5, 265. , 267.5, 270. , 272.5,
           275. , 277.5, 280. , 282.5, 285. , 287.5, 290. , 292.5, 295. , 297.5,
           300. , 302.5, 305. , 307.5, 310. , 312.5, 315. , 317.5, 320. , 322.5,
           325. , 327.5, 330. ], dtype=float32)

• Indexes: (2)
  • lat
    PandasIndex

    PandasIndex(Index([75.0, 72.5, 70.0, 67.5, 65.0, 62.5, 60.0, 57.5, 55.0, 52.5, 50.0, 47.5,
           45.0, 42.5, 40.0, 37.5, 35.0, 32.5, 30.0, 27.5, 25.0, 22.5, 20.0, 17.5,
           15.0],
          dtype='float32', name='lat'))

  • lon
    PandasIndex

    PandasIndex(Index([200.0, 202.5, 205.0, 207.5, 210.0, 212.5, 215.0, 217.5, 220.0, 222.5,
           225.0, 227.5, 230.0, 232.5, 235.0, 237.5, 240.0, 242.5, 245.0, 247.5,
           250.0, 252.5, 255.0, 257.5, 260.0, 262.5, 265.0, 267.5, 270.0, 272.5,
           275.0, 277.5, 280.0, 282.5, 285.0, 287.5, 290.0, 292.5, 295.0, 297.5,
           300.0, 302.5, 305.0, 307.5, 310.0, 312.5, 315.0, 317.5, 320.0, 322.5,
           325.0, 327.5, 330.0],
          dtype='float32', name='lon'))

• Attributes: (0)

This is fine for outputs that fit in memory like the example here, but sometimes we want to write the output of the computation to Zarr, which we do by calling to_zarr on the dataset instead of compute:

mean.to_zarr("mean_air_temperature.zarr", mode="w", consolidated=True);

We can check that the Zarr file was created by loading it from disk using xarray.open_dataset:

xr.open_dataset(
    "mean_air_temperature.zarr", chunked_array_type="cubed", chunks={}
)

<xarray.Dataset> Size: 11kB
Dimensions:  (lat: 25, lon: 53)
Coordinates:
  * lat      (lat) float32 100B 75.0 72.5 70.0 67.5 65.0 ... 22.5 20.0 17.5 15.0
  * lon      (lon) float32 212B 200.0 202.5 205.0 207.5 ... 325.0 327.5 330.0
Data variables:
    air      (lat, lon) float64 11kB cubed.Array<chunksize=(25, 53)>

xarray.Dataset

• Dimensions:
  • lat: 25
  • lon: 53
• Coordinates: (2)
  • lat
    (lat)
    float32
    75.0 72.5 70.0 ... 20.0 17.5 15.0

    axis :

    Y

    long_name :

    Latitude

    standard_name :

    latitude

    units :

    degrees_north

    array([75. , 72.5, 70. , 67.5, 65. , 62.5, 60. , 57.5, 55. , 52.5, 50. , 47.5,
           45. , 42.5, 40. , 37.5, 35. , 32.5, 30. , 27.5, 25. , 22.5, 20. , 17.5,
           15. ], dtype=float32)

  • lon
    (lon)
    float32
    200.0 202.5 205.0 ... 327.5 330.0

    axis :

    X

    long_name :

    Longitude

    standard_name :

    longitude

    units :

    degrees_east

    array([200. , 202.5, 205. , 207.5, 210. , 212.5, 215. , 217.5, 220. , 222.5,
           225. , 227.5, 230. , 232.5, 235. , 237.5, 240. , 242.5, 245. , 247.5,
           250. , 252.5, 255. , 257.5, 260. , 262.5, 265. , 267.5, 270. , 272.5,
           275. , 277.5, 280. , 282.5, 285. , 287.5, 290. , 292.5, 295. , 297.5,
           300. , 302.5, 305. , 307.5, 310. , 312.5, 315. , 317.5, 320. , 322.5,
           325. , 327.5, 330. ], dtype=float32)

• Data variables: (1)
  • air
    (lat, lon)
    float64
    cubed.Array<chunksize=(25, 53)>

    Array Chunk Bytes 10.6 KB 10.6 KB Shape (25, 53) (25, 53) Count 3 arrays in Plan 1 Chunks Type float64 np.ndarray

• Indexes: (2)
  • lat
    PandasIndex

    PandasIndex(Index([75.0, 72.5, 70.0, 67.5, 65.0, 62.5, 60.0, 57.5, 55.0, 52.5, 50.0, 47.5,
           45.0, 42.5, 40.0, 37.5, 35.0, 32.5, 30.0, 27.5, 25.0, 22.5, 20.0, 17.5,
           15.0],
          dtype='float32', name='lat'))

  • lon
    PandasIndex

    PandasIndex(Index([200.0, 202.5, 205.0, 207.5, 210.0, 212.5, 215.0, 217.5, 220.0, 222.5,
           225.0, 227.5, 230.0, 232.5, 235.0, 237.5, 240.0, 242.5, 245.0, 247.5,
           250.0, 252.5, 255.0, 257.5, 260.0, 262.5, 265.0, 267.5, 270.0, 272.5,
           275.0, 277.5, 280.0, 282.5, 285.0, 287.5, 290.0, 292.5, 295.0, 297.5,
           300.0, 302.5, 305.0, 307.5, 310.0, 312.5, 315.0, 317.5, 320.0, 322.5,
           325.0, 327.5, 330.0],
          dtype='float32', name='lon'))

• Attributes: (0)