Zarr Output

Zarr is a chunked, compressed, cloud-friendly format for N-dimensional arrays.

ZarrOutput is implemented as an extension that is only loaded once Zarr.jl is imported:

using SpeedyWeather
using Zarr     # this loads SpeedyWeatherZarrExt and enables ZarrOutput

spectral_grid = SpectralGrid(truncation=32, nlayers=8)
output = ZarrOutput(spectral_grid, PrimitiveWet, interval=Hour(6))
model = PrimitiveWetModel(spectral_grid; output)
simulation = initialize!(model)
run!(simulation, period=Day(10), output=true)

The constructor and option fields mirror NetCDFOutput: path, id, overwrite, interval, variables, write_restart, write_parameters_txt, write_progress_txt all behave the same way, the run folder layout (run_<id>_NNNN/) is identical, and the same AbstractOutputVariable types are used to declare which variables are written. The on-disk layout differs:

• the Zarr store is a directory (output.zarr/), not a single file

• per-variable arrays live as subdirectories with .zarray and .zattrs metadata

• coordinates lon, lat, layer, soil_layer, time are stored as 1D arrays in the same group, tagged with the conventional _ARRAY_DIMENSIONS attribute so that Xarray-compatible readers can rebuild the dataset

Two extra options are specific to ZarrOutput:

Option	Meaning
time_chunk::Int	Number of time steps per chunk along the time axis (default 1). Larger values give bigger chunks and usually better compression at the cost of higher write latency.
compressor	Any Zarr.Compressor (e.g. Zarr.BloscCompressor(clevel=3), Zarr.ZlibCompressor()); nothing (default) uses Blosc with the SpeedyWeather default compression level.

using SpeedyWeather, Zarr

spectral_grid = SpectralGrid(truncation=32, nlayers=8)
output = ZarrOutput(spectral_grid, PrimitiveWet;
    interval = Hour(1),
    time_chunk = 24,                        # bundle one day per chunk on the time axis
    compressor = Zarr.BloscCompressor(clevel=5),
)

Reading back the data only needs Zarr.jl:

using Zarr
g = Zarr.zopen(joinpath(output.run_path, output.filename))
g["time"][:]             # all stored hours since startdate
g["vor"][:, :, 1, :]     # vorticity, top layer, all time steps

Custom output variables work exactly as with NetCDFOutput (see Output variables and Customizing netCDF output): subtype AbstractOutputVariable, implement path(::MyOutputVariable, simulation) to return the AbstractField to write, and add!(output, MyOutputVariable()) it to the ZarrOutput.

Reading the store from Python with xarray

ZarrOutput writes the stores according to the conventions of xarray to ensure compability with it. Simulations can be easily opened in Python as in the following:

import xarray as xr
ds = xr.open_zarr("run_0001/output.zarr", consolidated=False)
print(ds)
# <xarray.Dataset>
# Dimensions:  (time: 41, layer: 8, lat: 32, lon: 64)
# Coordinates:
#   * time     (time) datetime64[ns]  2000-01-01 ... 2000-01-11
#   * layer    (layer) float32        0.06 0.19 ... 0.94
#   * lat      (lat) float64          85.76 ... -85.76
#   * lon      (lon) float64          0.0 ... 354.4
# Data variables:
#     vor      (time, layer, lat, lon) float32
#     u        (time, layer, lat, lon) float32
#     v        (time, layer, lat, lon) float32
#     temp     (time, layer, lat, lon) float32
#     humid    (time, layer, lat, lon) float32
#     mslp     (time, lat, lon) float32

ds["temp"].isel(time=-1, layer=0).plot()       # last step, top layer
ds["mslp"].mean(("lat", "lon")).plot.line(x="time")

xarray decodes the time axis to datetime64 automatically (from the CF-style hours since <startdate> units), so resampling and slicing by date work without extra work:

ds.sel(time=slice("2000-01-05", "2000-01-08"))["temp"].mean("time")

Ensemble output

Several runs of an ensemble can be written into a single Zarr store along an additional ensemble dimension. This is disabled by default and controlled by two options:

Option	Meaning
ensemble_index::Int	This writer's ensemble member index. 0 (default) disables ensemble output and reproduces the layout described above. A value > 0 adds an ensemble dimension and makes this writer store its data into slot ensemble_index. Members are indexed 1..ensemble_size.
ensemble_size::Int	Total number of ensemble members. It sizes the ensemble dimension and must be passed up front; it has to satisfy ensemble_size ≥ ensemble_index.
ensemble_timeout::Int	Seconds a member waits for member 1 to create the shared store before erroring (default 600).

The design targets parallel ensemble members running as separate processes, all writing into one store. This is safe because the ensemble axis is chunked with size 1, so member e only ever writes chunk files carrying its own index and no two members touch the same chunk file. Each process constructs a ZarrOutput with the same path, id, run_number, filename and ensemble_size, and a distinct ensemble_index, so all members resolve to the same run folder and store:

using SpeedyWeather, Zarr

# in each process, `member` is this process' ensemble index (1..ensemble_size)
spectral_grid = SpectralGrid(truncation=32, nlayers=8)
output = ZarrOutput(spectral_grid, PrimitiveWet;
    ensemble_index = member,        # e.g. read from an environment variable / job array id
    ensemble_size = 10,
    interval = Hour(6),
)
model = PrimitiveWetModel(spectral_grid; output)
simulation = initialize!(model)
run!(simulation, period=Day(10), output=true)

The single shared artifacts (the group and coordinate metadata, and the time axis) are written once by the creator — the member with ensemble_index == 1, which builds the store schema and then signals readiness. The remaining writer members (`ensemble_index

1`) wait for that signal, open the existing store and write only their own ensemble

slice. This assumes that all members share the same time stepping so that one common time axis is correct.

The resulting store gains an ensemble coordinate; the ensemble axis is the outermost (slowest-varying) dimension, so an xarray-compatible reader reports e.g.

import xarray as xr
ds = xr.open_zarr("run_0001/output.zarr", consolidated=False)
print(ds)
# Dimensions:  (ensemble: 10, time: 41, layer: 8, lat: 32, lon: 64)
# Data variables:
#     temp     (ensemble, time, layer, lat, lon) float32
ds["temp"].mean("ensemble")     # ensemble mean

Note that this ensemble layout is currently specific to ZarrOutput; the NetCDFOutput does not support concurrent multi-process writes.

HEALPix Output

HEALPix is an equal-area discretization of the sphere. However with its non-rectangular coordinates, it's not straight forward to save HEALPix grids in standard NetCDF files. Instead HEALPixOutput writes a Zarr store, keeping the horizontal dimension flat — one unravelled vector of npix = 12nside² cells, exactly as a Field stores its data — rather than interpolating onto a rectangular lon×lat grid the way NetCDFOutput and ZarrOutput do. Use it with:

using SpeedyWeather
using Zarr     # this loads SpeedyWeatherZarrExt and enables HEALPixOutput

spectral_grid = SpectralGrid(truncation=31, nlayers=8)
output = HEALPixOutput(spectral_grid, PrimitiveWet, nside=16, interval=Hour(6))
model = PrimitiveWetModel(spectral_grid; output)
simulation = initialize!(model)
run!(simulation, period=Day(1), output=true)

The resolution is set with nside (as in most HEALPix implementation like cuHPX), equivalently nlat_half = 2nside, and defaults to the model grid's own nlat_half (rounded up to the nearest even number, which HEALPixGrid requires).

All the other file options of ZarrOutput (path, id, overwrite, interval, variables, time_chunk, compressor, …) behave identically; lon_chunk/lat_chunk are replaced by a single cell_chunk for the flat dimension.

Store layout

array	shape	_ARRAY_DIMENSIONS
3D variable, e.g. temp	(npix, nlayers, ntime)	["time", "layer", "cell"]
2D variable, e.g. mslp	(npix, ntime)	["time", "cell"]
lat, lon, ring	(npix,)	["cell"]
cell	(npix,)	["cell"]

Alongside the data the store holds three flat vectors, one entry per cell, so that a consumer can place every data point on the sphere without knowing anything about SpeedyWeather's grids:

• lat, lon: the coordinates of each cell in degrees north/east

• ring: the latitude ring index of each cell, 1 (north) to 4nside-1 (south)

g = Zarr.zopen(joinpath(output.run_path, output.filename))
g["lat"][1:4], g["lon"][1:4], g["ring"][1:4]    # the 4 cells of the northernmost ring

([87.07581964294992, 87.07581964294992, 87.07581964294992, 87.07581964294992], [45.0, 135.0, 225.0, 315.0], [1, 1, 1, 1])

RING ordering and interoperability

Cells are written in standard HEALPix RING order of cuHPX: cell ij (1-based, as in a Field) is RING pixel ij-1 in the 0-based convention of healpy and cuHPX, so no reordering is needed on either side. The store also carries healpix_nside, healpix_npix and healpix_order as global attributes.

These conventions should be compatabile with healpy and cuHPX. Note that while HEALPixOutput writes any even nlat_half, the NESTED and earth-2 flat layouts those tools convert to require nside to be a power of two.

Our lat is latitude in degrees from the equator and lon is in [0, 360), which is exactly healpy's lonlat=True convention. Mind that this is not healpy's default: with lonlat=False (the default) healpy uses colatitude in radians from the north pole, so convert with theta = deg2rad(90 - lat), phi = deg2rad(lon) or use lonlat=True.

Mind the axis order when reading from Python: Zarr stores shape row-major while Julia is column-major, so a variable written as (cell, layer, time) from Julia is seen as (time, layer, cell) from Python — matching its _ARRAY_DIMENSIONS tag. A flat (npix,) map to hand to healpy is therefore arr.reshape(-1, npix)[-1], and the hp.mollview call above works because .isel(time=-1, layer=-1) already reduces to the trailing cell axis.

Here's an example script that loads a store and plots it with healpy:

import matplotlib 
import matplotlib.pyplot as plt
import xarray as xr # also needs zarr installed
import healpy as hp 

ds = xr.open_zarr("../run_0004/output_healpix.zarr", consolidated=True)
nside = ds.attrs["healpix_nside"]

hp.mollview(ds["temp"].isel(time=-1, layer=7).values, nest=False,
            flip="geo", rot=(180, 0, 0),
            unit=ds["temp"].attrs["units"])
# surface layer = 7 (0-indexed), rotation and flip so that it aligns with Julia CairoMakie 
hp.graticule()                        # optional lat/lon grid
plt.savefig("hpy-mollview.png", dpi=150, bbox_inches="tight")
plt.close()

Be aware that by default SpeedyWeather.jl's Makie extension together with Makie itself interpolate the data, so the output of this script might still look relatively different on first sight to directly plotting the data in Julia.

OctaHEALPix output

HEALPixOutput also writes the OctaHEALPixGrid, a SpeedyWeather-specific member of the HEALPix family (4 faces, 4nlat_half² points, no equatorial belt). It is written with exactly the same store convention: the same flat cell dimension, the same flat lat/lon/ring coordinate vectors, the same array shapes and _ARRAY_DIMENSIONS.

output = HEALPixOutput(spectral_grid, PrimitiveWet, output_grid=OctaHEALPixGrid(24))
output.field2D.grid

47-ring OctaHEALPixGrid{...}
├ nlat_half = 24 (2304 points, ~4.2˚, reduced)
└ architecture = CPU

healpy and cuHPX implement only the standard HEALPix and cannot read an OctaHEALPix map.

By default the output grid follows the model's grid when the simulation already runs on either supported grid, so both are written without interpolation.

Skipping the interpolation

If the simulation already runs on the very HEALPix grid requested for output, no interpolation is needed and none is set up — output.interpolator stays nothing and the data is copied straight out of the model state:

spectral_grid = SpectralGrid(truncation=31, nlayers=8, Grid=HEALPixGrid)
output = HEALPixOutput(spectral_grid, PrimitiveWet)
isnothing(output.interpolator)      # true, model and output share the grid

true

This saves both the per-output-step interpolation and the precomputation of the interpolator's stencil indices and weights, and makes the written data bit-identical to the model state (up to the number format and the keepbits mantissa rounding). Requesting a different nside than the model's resolution brings the interpolator back.

SpeedyWeather.HEALPixOutput Type

Output writer that writes a SpeedyWeather simulation to a Zarr store on an equal-area grid of the HEALPix family — a HEALPixGrid or an OctaHEALPixGrid — keeping the horizontal dimension flat (unravelled into a single cell dimension) exactly as a Field's data is. Data on any other grid is interpolated onto the output grid; a simulation that already runs on it is written without interpolation.

Alongside the variables the store holds one flat vector per cell with its latitude (lat, ˚N from the equator), its longitude (lon, ˚E in [0, 360)) and its ring index (ring), so consumers need no knowledge of SpeedyWeather's grid machinery to place the data on the sphere. Cells run ring by ring from north to south.

Both grids use this same store convention, but they are not equally portable:

• on a HEALPixGrid the cell order is the standard HEALPix RING order, i.e. the flat index ij is RING pixel ij - 1 (0-based), and the store carries healpix_nside, healpix_npix and healpix_order so healpy and cuHPX can read it directly. The coordinates match healpy's pix2ang(..., lonlat=true); healpy's default is colatitude in radians, i.e. θ = deg2rad(90 - lat).

• an OctaHEALPixGrid is a SpeedyWeather-specific member of the HEALPix family (4 faces, 4nlat_half² points, no equatorial belt). It is equal-area and ring-ordered, but it is not standard HEALPix and healpy/cuHPX cannot read it, so those stores deliberately carry no healpix_* attributes.

The actual implementation lives in the SpeedyWeatherZarrExt extension and is only available once Zarr.jl is loaded:

using Zarr
using SpeedyWeather
spectral_grid = SpectralGrid(truncation = 31, nlayers = 8)
output = HEALPixOutput(spectral_grid, PrimitiveWet, nside = 16)

# or on an OctaHEALPixGrid, which also takes odd nlat_half
output = HEALPixOutput(spectral_grid, PrimitiveWet, output_grid = OctaHEALPixGrid(24))

Type parameters: Field2D, Field3D are the scratch field types on the output grid, Interpolator is the interpolator type (or Nothing when the model grid already is the output grid and interpolation is skipped), DT and S are the start-date and output-step types, C is the Zarr compressor type (or Nothing for the Zarr default), Z is the Zarr group type once initialize! has been called (Nothing before) and Layers the type of the vertical output layers. Fields are

• active::Bool

• path::String: [OPTION] path to output parent folder, run folders will be created within

• run_prefix::String: [OPTION] Prefix for run folder where data is stored, e.g. 'run_'

• id::String: [OPTION] run identification, added between run_prefix and run_number

• run_number::Int64: [OPTION] run identification number, automatically determined if overwrite=false

• run_digits::Int64: [OPTION] run numbers digits

• core::OutputWriterCore: [DERIVED] shared output writer state (run folder, counters, output frequency)

• overwrite::Bool: [OPTION] Overwrite an existing run folder?

• filename::String: [OPTION] name of the output zarr store (a directory)

• write_restart::Bool: [OPTION] also write restart file if output=true?

• write_parameters_txt::Bool: [OPTION] also write parameters txt file if output=true?

• write_progress_txt::Bool: [OPTION] also write progress txt file if output=true?

• startdate::Any: [DERIVED] start date of the simulation, used for time dimension in zarr store

• interval::Any: [OPTION] output frequency, time step

• variables::Dict{Symbol, SpeedyWeather.AbstractOutputVariable}: [OPTION] dictionary of variables to output, e.g. u, v, vor, div, pres, temp, humid

• layers::Any: [OPTION] vertical layers to write 3D atmospheric variables on, ModelLayers() (default) or PressureLayers(spectral_grid), see AbstractOutputLayers

• time_chunk::Int64: [OPTION] number of time steps per chunk along the time dimension

• cell_chunk::Int64: [OPTION] chunk size along the flat cell dimension. 0 (default) means one chunk = all cells.

• vertical_chunk::Int64: [OPTION] chunk size along the vertical (layer / soil_layer). 0 (default) means full extent.

• compressor::Union{Nothing, C} where C: [OPTION] Zarr compressor (extension-typed). nothing keeps the Zarr default.

• zarr_group::Union{Nothing, Z} where Z: [DERIVED] the Zarr group to be written into, created on initialize!

• time_buffers::Dict{String, Any}: [DERIVED] per-variable time-chunk write buffers keyed by variable name, see ZarrOutput

• interpolator::Any

• land_fraction::Any

• field2D::Any

• field3D::Any

• field3Dland::Any

source