Xarray engine: writing back to GRIB

Warning: converting back Xarray to GRIB is an experimental feature and is not yet fully supported.

First, we get some example GRIB data and convert it into Xarray.

import earthkit.data as ekd

ds_fl = ekd.from_source("sample", "pl.grib").to_fieldlist()
ds_xr = ds_fl.to_xarray()
ds_xr

<xarray.Dataset> Size: 176kB
Dimensions:                  (forecast_reference_time: 4, step: 2, level: 2,
                              latitude: 19, longitude: 36)
Coordinates:
  * forecast_reference_time  (forecast_reference_time) datetime64[ns] 32B 202...
  * step                     (step) timedelta64[ns] 16B 00:00:00 06:00:00
  * level                    (level) int64 16B 500 700
  * latitude                 (latitude) float64 152B 90.0 80.0 ... -80.0 -90.0
  * longitude                (longitude) float64 288B 0.0 10.0 ... 340.0 350.0
Data variables:
    r                        (forecast_reference_time, step, level, latitude, longitude) float64 88kB ...
    t                        (forecast_reference_time, step, level, latitude, longitude) float64 88kB ...
Attributes:
    Conventions:  CF-1.8
    institution:  ECMWF

xarray.Dataset

• Dimensions:
  • forecast_reference_time: 4
  • step: 2
  • level: 2
  • latitude: 19
  • longitude: 36
• Coordinates: (5)
  • forecast_reference_time
    (forecast_reference_time)
    datetime64[ns]
    2024-06-03 ... 2024-06-04T12:00:00

    array(['2024-06-03T00:00:00.000000000', '2024-06-03T12:00:00.000000000',
           '2024-06-04T00:00:00.000000000', '2024-06-04T12:00:00.000000000'],
          dtype='datetime64[ns]')

  • step
    (step)
    timedelta64[ns]
    00:00:00 06:00:00

    standard_name :

    forecast_period

    long_name :

    time since forecast_reference_time

    array([             0, 21600000000000], dtype='timedelta64[ns]')

  • level
    (level)
    int64
    500 700

    standard_name :

    air_pressure

    long_name :

    pressure

    units :

    hectopascal

    positive :

    down

    array([500, 700])

  • latitude
    (latitude)
    float64
    90.0 80.0 70.0 ... -80.0 -90.0

    units :

    degrees_north

    standard_name :

    latitude

    long_name :

    latitude

    array([ 90.,  80.,  70.,  60.,  50.,  40.,  30.,  20.,  10.,   0., -10., -20.,
           -30., -40., -50., -60., -70., -80., -90.])

  • longitude
    (longitude)
    float64
    0.0 10.0 20.0 ... 330.0 340.0 350.0

    units :

    degrees_east

    standard_name :

    longitude

    long_name :

    longitude

    array([  0.,  10.,  20.,  30.,  40.,  50.,  60.,  70.,  80.,  90., 100., 110.,
           120., 130., 140., 150., 160., 170., 180., 190., 200., 210., 220., 230.,
           240., 250., 260., 270., 280., 290., 300., 310., 320., 330., 340., 350.])

• Data variables: (2)
  • r
    (forecast_reference_time, step, level, latitude, longitude)
    float64
    ...

    standard_name :

    relative_humidity

    long_name :

    Relative humidity

    units :

    percent

    level_type :

    pressure

    _earthkit :

    {'message': b"GRIB\x00\x00l\x01\x00\x004\x80b\x9a\xff\x80\x9dd\x01\xf4\x18\x06\x03\x00\x00\x01\x00\x00\x01\x00\x00\x00\x15\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x01\t\x04\x010001\x00\x00\x00\x00\x00 \x00\xff\x00\x00$\x00\x13\x01_\x90\x00\x00\x00\x80\x81_\x90\x05W0'\x10'\x10\x00\x00\x00\x00\x00\x00\x00\x0c\x00\x80\x02D\xb9}n\x00\x007777", 'bitsPerValue': 16, 'grid_spec': '{"grid": [10, 10]}'}

    [10944 values with dtype=float64]

  • t
    (forecast_reference_time, step, level, latitude, longitude)
    float64
    ...

    standard_name :

    air_temperature

    long_name :

    Temperature

    units :

    kelvin

    level_type :

    pressure

    _earthkit :

    {'message': b"GRIB\x00\x00l\x01\x00\x004\x80b\x9a\xff\x80\x82d\x01\xf4\x18\x06\x03\x00\x00\x01\x00\x00\x01\x00\x00\x00\x15\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x01\t\x04\x010001\x00\x00\x00\x00\x00 \x00\xff\x00\x00$\x00\x13\x01_\x90\x00\x00\x00\x80\x81_\x90\x05W0'\x10'\x10\x00\x00\x00\x00\x00\x00\x00\x0c\x00\x80\x02D\xb9}n\x00\x007777", 'bitsPerValue': 16, 'grid_spec': '{"grid": [10, 10]}'}

    [10944 values with dtype=float64]

• Attributes: (2)

  Conventions :

  CF-1.8

  institution :

  ECMWF

By default, add_earthkit_attrs=True in to_xarray() and some special earthkit attributes are added to the dataset. This is needed for the Xarray to GRIB conversion. For this reason, if the Xarray is modified we must ensure the variable attributes are copied to the new Xarray dataset. By default, variable attributes are not kept in Xarray computations so we need to set the global Xarray keep_attrs option to enable it.

# ensure earthkit attributes are set
import xarray as xr

xr.set_options(keep_attrs=True)

# modify values
ds_xr += 1

Using to_target()

It is possible to directly write the Xarray dataset created with the earthkit engine into a GRIB file with to_target(). This is a memory efficient way to write GRIB to disk since only one field is loaded into memory at a time. We can call to_target() either on the earthkit accessor or as a top level function.

First, we write a datarray into a GRIB file.

# option1: writing to GRIB file using the accessor
ds_xr["t"].earthkit.to_target("file", "_from_xr_1.grib")

# option2: writing to GRIB file using the top level function
ekd.to_target("file", "_from_xr_1a.grib", data=ds_xr["t"])

# check the results
ds_tmp1 = ekd.from_source("file", "_from_xr_1.grib").to_fieldlist()
ds_tmp1[0]

Field

data

number_of_values	684
array_type	ndarray
array_dtype	float64

parameter

variable	t
standard_name	air_temperature
long_name	Temperature
units	kelvin
chem_variable	None

time

valid_datetime	2024-06-03 00:00:00
base_datetime	2024-06-03 00:00:00
step	0:00:00

vertical

level	500
layer	None
level_type	pressure

ensemble

member

0

geography

grid_spec	{'grid': [10, 10]}
grid_type	regular_ll
shape	(19, 36)
area	(90.0, 0.0, -90.0, 350.0)

Next, we write the whole dataset into a GRIB file.

# option1: writing to GRIB file using the accessor
ds_xr.earthkit.to_target("file", "_from_xr_2.grib")

# option2: writing to GRIB file using the top level function
ekd.to_target("file", "_from_xr_2a.grib", data=ds_xr)

# check the results
ds_tmp2 = ekd.from_source("file", "_from_xr_2.grib").to_fieldlist()
ds_tmp2.head()

	parameter.variable	time.valid_datetime	time.base_datetime	time.step	vertical.level	vertical.level_type	ensemble.member	geography.grid_type
0	r	2024-06-03 00:00:00	2024-06-03 00:00:00	0 days 00:00:00	500	pressure	0	regular_ll
1	r	2024-06-03 00:00:00	2024-06-03 00:00:00	0 days 00:00:00	700	pressure	0	regular_ll
2	r	2024-06-03 06:00:00	2024-06-03 00:00:00	0 days 06:00:00	500	pressure	0	regular_ll
3	r	2024-06-03 06:00:00	2024-06-03 00:00:00	0 days 06:00:00	700	pressure	0	regular_ll
4	r	2024-06-03 12:00:00	2024-06-03 12:00:00	0 days 00:00:00	500	pressure	0	regular_ll

We check if the computation results were correctly written to the generated GRIB data.

# original GRIB data
print(ds_fl.sel({"parameter.variable": "t", "time.step": 0, "vertical.level": 500})[0].values[0])
# GRIB data converted from the modified xarray object
print(ds_tmp1.sel({"parameter.variable": "t", "time.step": 0, "vertical.level": 500})[0].values[0])

250.22500610351562
251.22500610351562

Using to_fieldlist()

We can also convert the Xarray dataset into a GRIB fieldlist by using to_fieldlist() on the earthkit accessor of the Xarray object. Please note that this will generate a fieldlist entirely stored in memory.

First, we convert a dataarray to a GRIB fieldlist.

ds_fl1 = ds_xr["t"].earthkit.to_fieldlist()
ds_fl1.head()

	parameter.variable	time.valid_datetime	time.base_datetime	time.step	vertical.level	vertical.level_type	ensemble.member	geography.grid_type
0	t	2024-06-03 00:00:00	2024-06-03 00:00:00	0 days 00:00:00	500	pressure	0	regular_ll
1	t	2024-06-03 00:00:00	2024-06-03 00:00:00	0 days 00:00:00	700	pressure	0	regular_ll
2	t	2024-06-03 06:00:00	2024-06-03 00:00:00	0 days 06:00:00	500	pressure	0	regular_ll
3	t	2024-06-03 06:00:00	2024-06-03 00:00:00	0 days 06:00:00	700	pressure	0	regular_ll
4	t	2024-06-03 12:00:00	2024-06-03 12:00:00	0 days 00:00:00	500	pressure	0	regular_ll

Next, we convert back the whole dataset into a GRIB fieldlist.

ds_fl2 = ds_xr.earthkit.to_fieldlist()
ds_fl2.head()

	parameter.variable	time.valid_datetime	time.base_datetime	time.step	vertical.level	vertical.level_type	ensemble.member	geography.grid_type
0	r	2024-06-03 00:00:00	2024-06-03 00:00:00	0 days 00:00:00	500	pressure	0	regular_ll
1	r	2024-06-03 00:00:00	2024-06-03 00:00:00	0 days 00:00:00	700	pressure	0	regular_ll
2	r	2024-06-03 06:00:00	2024-06-03 00:00:00	0 days 06:00:00	500	pressure	0	regular_ll
3	r	2024-06-03 06:00:00	2024-06-03 00:00:00	0 days 06:00:00	700	pressure	0	regular_ll
4	r	2024-06-03 12:00:00	2024-06-03 12:00:00	0 days 00:00:00	500	pressure	0	regular_ll

The generated GRIB fieldlist can be saved to disk using the to_target() method.

out_name = "_from_xr_3.grib"
ds_fl1.to_target("file", out_name)
# read back and check the saved GRIB
ds_tmp = ekd.from_source("file", out_name).to_fieldlist()
ds_tmp[0]

Field

data

number_of_values	684
array_type	ndarray
array_dtype	float64

parameter

variable	t
standard_name	air_temperature
long_name	Temperature
units	kelvin
chem_variable	None

time

valid_datetime	2024-06-03 00:00:00
base_datetime	2024-06-03 00:00:00
step	0:00:00

vertical

level	500
layer	None
level_type	pressure

ensemble

member

0

geography

grid_spec	{'grid': [10, 10]}
grid_type	regular_ll
shape	(19, 36)
area	(90.0, 0.0, -90.0, 350.0)