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

---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/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 /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/xarray/core/dataset.py:2404, in Dataset._repr_html_(self)
   2402 if OPTIONS["display_style"] == "text":
   2403     return f"<pre>{escape(repr(self))}</pre>"
-> 2404 return formatting_html.dataset_repr(self)

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/xarray/core/formatting_html.py:374, in dataset_repr(ds)
    371 if ds.coords:
    372     sections.append(coord_section(ds.coords))
--> 374 sections.append(datavar_section(ds.data_vars))
    376 display_default_indexes = _get_boolean_with_default(
    377     "display_default_indexes", False
    378 )
    379 xindexes = filter_nondefault_indexes(
    380     _get_indexes_dict(ds.xindexes), not display_default_indexes
    381 )

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/xarray/core/formatting_html.py:224, in _mapping_section(mapping, name, details_func, max_items_collapse, expand_option_name, enabled, max_option_name)
    218     if n_items > max_items:
    219         inline_details = f"({max_items}/{n_items})"
    221 return collapsible_section(
    222     name,
    223     inline_details=inline_details,
--> 224     details=details_func(mapping),
    225     n_items=n_items,
    226     enabled=enabled,
    227     collapsed=collapsed,
    228 )

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/xarray/core/formatting_html.py:134, in summarize_vars(variables)
    133 def summarize_vars(variables) -> str:
--> 134     vars_li = "".join(
    135         f"<li class='xr-var-item'>{summarize_variable(k, v)}</li>"
    136         for k, v in variables.items()
    137     )
    139     return f"<ul class='xr-var-list'>{vars_li}</ul>"

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/xarray/core/formatting_html.py:135, in <genexpr>(.0)
    133 def summarize_vars(variables) -> str:
    134     vars_li = "".join(
--> 135         f"<li class='xr-var-item'>{summarize_variable(k, v)}</li>"
    136         for k, v in variables.items()
    137     )
    139     return f"<ul class='xr-var-list'>{vars_li}</ul>"

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/xarray/core/formatting_html.py:96, in summarize_variable(name, var, is_index, dtype)
     94 preview = escape(inline_variable_array_repr(variable, 35))
     95 attrs_ul = summarize_attrs(var.attrs)
---> 96 data_repr = short_data_repr_html(variable)
     98 attrs_icon = _icon("icon-file-text2")
     99 data_icon = _icon("icon-database")

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/xarray/core/formatting_html.py:43, in short_data_repr_html(array)
     41 internal_data = getattr(array, "variable", array)._data
     42 if hasattr(internal_data, "_repr_html_"):
---> 43     return internal_data._repr_html_()
     44 text = escape(short_data_repr(array))
     45 return f"<pre>{text}</pre>"

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/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 /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/cubed/diagnostics/__init__.py:1
----> 1 from .rich import RichProgressBar as ProgressBar
      3 __all__ = ["ProgressBar"]

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/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: 31MB
Dimensions:  (time: 2920, lat: 25, lon: 53)
Coordinates:
  * time     (time) datetime64[ns] 23kB 2013-01-01 ... 2014-12-31T18:00:00
  * 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      (time, lat, lon) float64 31MB cubed.Array<chunksize=(2920, 25, 53)>
Attributes: (5)

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

/tmp/ipykernel_2388/172254778.py:1: SerializationWarning: saving variable None with floating point data as an integer dtype without any _FillValue to use for NaNs
  ds.to_zarr("air_temperature_cubed.zarr", mode="w", consolidated=True);
/opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/zarr/api/asynchronous.py:247: ZarrUserWarning: Consolidated metadata is currently not part in the Zarr format 3 specification. It may not be supported by other zarr implementations and may change in the future.
  warnings.warn(

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

---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/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 /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/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 /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/xarray/core/formatting_html.py:339, in array_repr(arr)
    331 arr_name = escape(repr(arr.name)) if getattr(arr, "name", None) else ""
    333 header_components = [
    334     f"<div class='xr-obj-type'>{obj_type}</div>",
    335     f"<div class='xr-obj-name'>{arr_name}</div>",
    336     format_dims(dims, indexed_dims),
    337 ]
--> 339 sections = [array_section(arr)]
    341 if hasattr(arr, "coords"):
    342     if arr.coords:

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/xarray/core/formatting_html.py:249, in array_section(obj)
    247 variable = getattr(obj, "variable", obj)
    248 preview = escape(inline_variable_array_repr(variable, max_width=70))
--> 249 data_repr = short_data_repr_html(obj)
    250 data_icon = _icon("icon-database")
    252 return (
    253     "<div class='xr-array-wrap'>"
    254     f"<input id='{data_id}' class='xr-array-in' type='checkbox' {collapsed}>"
   (...)    258     "</div>"
    259 )

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/xarray/core/formatting_html.py:43, in short_data_repr_html(array)
     41 internal_data = getattr(array, "variable", array)._data
     42 if hasattr(internal_data, "_repr_html_"):
---> 43     return internal_data._repr_html_()
     44 text = escape(short_data_repr(array))
     45 return f"<pre>{text}</pre>"

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/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 /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/cubed/diagnostics/__init__.py:1
----> 1 from .rich import RichProgressBar as ProgressBar
      3 __all__ = ["ProgressBar"]

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/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 '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
Attributes: (11)

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
Attributes: (11)

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)

• Attributes: (11)

  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 ]

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

/opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/zarr/api/asynchronous.py:247: ZarrUserWarning: Consolidated metadata is currently not part in the Zarr format 3 specification. It may not be supported by other zarr implementations and may change in the future.
  warnings.warn(

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={}
)

---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/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 /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/xarray/core/dataset.py:2404, in Dataset._repr_html_(self)
   2402 if OPTIONS["display_style"] == "text":
   2403     return f"<pre>{escape(repr(self))}</pre>"
-> 2404 return formatting_html.dataset_repr(self)

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/xarray/core/formatting_html.py:374, in dataset_repr(ds)
    371 if ds.coords:
    372     sections.append(coord_section(ds.coords))
--> 374 sections.append(datavar_section(ds.data_vars))
    376 display_default_indexes = _get_boolean_with_default(
    377     "display_default_indexes", False
    378 )
    379 xindexes = filter_nondefault_indexes(
    380     _get_indexes_dict(ds.xindexes), not display_default_indexes
    381 )

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/xarray/core/formatting_html.py:224, in _mapping_section(mapping, name, details_func, max_items_collapse, expand_option_name, enabled, max_option_name)
    218     if n_items > max_items:
    219         inline_details = f"({max_items}/{n_items})"
    221 return collapsible_section(
    222     name,
    223     inline_details=inline_details,
--> 224     details=details_func(mapping),
    225     n_items=n_items,
    226     enabled=enabled,
    227     collapsed=collapsed,
    228 )

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/xarray/core/formatting_html.py:134, in summarize_vars(variables)
    133 def summarize_vars(variables) -> str:
--> 134     vars_li = "".join(
    135         f"<li class='xr-var-item'>{summarize_variable(k, v)}</li>"
    136         for k, v in variables.items()
    137     )
    139     return f"<ul class='xr-var-list'>{vars_li}</ul>"

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/xarray/core/formatting_html.py:135, in <genexpr>(.0)
    133 def summarize_vars(variables) -> str:
    134     vars_li = "".join(
--> 135         f"<li class='xr-var-item'>{summarize_variable(k, v)}</li>"
    136         for k, v in variables.items()
    137     )
    139     return f"<ul class='xr-var-list'>{vars_li}</ul>"

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/xarray/core/formatting_html.py:96, in summarize_variable(name, var, is_index, dtype)
     94 preview = escape(inline_variable_array_repr(variable, 35))
     95 attrs_ul = summarize_attrs(var.attrs)
---> 96 data_repr = short_data_repr_html(variable)
     98 attrs_icon = _icon("icon-file-text2")
     99 data_icon = _icon("icon-database")

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/xarray/core/formatting_html.py:43, in short_data_repr_html(array)
     41 internal_data = getattr(array, "variable", array)._data
     42 if hasattr(internal_data, "_repr_html_"):
---> 43     return internal_data._repr_html_()
     44 text = escape(short_data_repr(array))
     45 return f"<pre>{text}</pre>"

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/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 /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/cubed/diagnostics/__init__.py:1
----> 1 from .rich import RichProgressBar as ProgressBar
      3 __all__ = ["ProgressBar"]

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/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: 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)>