Xarray: using CuPy

This notebook demonstrates how to use Xarray on a GPU with CuPy. Since CuPy is not a dependency for earthkit-data it has to be installed separately. Also a CUDA-based GPU environment has to be up and running for the notebook to work.

# Get GRIB data on pressure levels
import earthkit.data as ekd

ds = ekd.from_source("sample", "pl.grib").to_fieldlist()

# Create a lazy loaded Xarray with Numpy arrays
r = ds.to_xarray()
r

<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:
    class:        od
    stream:       oper
    levtype:      pl
    type:         fc
    expver:       0001
    date:         20240603
    time:         0
    domain:       g
    number:       0
    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

    standard_name :

    forecast_reference_time

    long_name :

    initial time of forecast

    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

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

  • level
    (level)
    int64
    500 700

    units :

    hPa

    positive :

    down

    stored_direction :

    decreasing

    standard_name :

    air_pressure

    long_name :

    pressure

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

    param :

    r

    standard_name :

    relative_humidity

    long_name :

    Relative humidity

    paramId :

    157

    units :

    %

    _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}

    [10944 values with dtype=float64]

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

    param :

    t

    standard_name :

    air_temperature

    long_name :

    Temperature

    paramId :

    130

    units :

    K

    _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}

    [10944 values with dtype=float64]

• Indexes: (5)
  • forecast_reference_time
    PandasIndex

    PandasIndex(DatetimeIndex(['2024-06-03 00:00:00', '2024-06-03 12:00:00',
                   '2024-06-04 00:00:00', '2024-06-04 12:00:00'],
                  dtype='datetime64[ns]', name='forecast_reference_time', freq=None))

  • step
    PandasIndex

    PandasIndex(TimedeltaIndex(['0 days 00:00:00', '0 days 06:00:00'], dtype='timedelta64[ns]', name='step', freq=None))

  • level
    PandasIndex

    PandasIndex(Index([500, 700], dtype='int64', name='level'))

  • latitude
    PandasIndex

    PandasIndex(Index([ 90.0,  80.0,  70.0,  60.0,  50.0,  40.0,  30.0,  20.0,  10.0,   0.0,
           -10.0, -20.0, -30.0, -40.0, -50.0, -60.0, -70.0, -80.0, -90.0],
          dtype='float64', name='latitude'))

  • longitude
    PandasIndex

    PandasIndex(Index([  0.0,  10.0,  20.0,  30.0,  40.0,  50.0,  60.0,  70.0,  80.0,  90.0,
           100.0, 110.0, 120.0, 130.0, 140.0, 150.0, 160.0, 170.0, 180.0, 190.0,
           200.0, 210.0, 220.0, 230.0, 240.0, 250.0, 260.0, 270.0, 280.0, 290.0,
           300.0, 310.0, 320.0, 330.0, 340.0, 350.0],
          dtype='float64', name='longitude'))

• Attributes: (11)

  class :

  od

  stream :

  oper

  levtype :

  pl

  type :

  fc

  expver :

  0001

  date :

  20240603

  time :

  0

  domain :

  g

  number :

  0

  Conventions :

  CF-1.8

  institution :

  ECMWF

type(r.t.data)

numpy.ndarray

Move to the GPU as CuPy

We use the to_device() method, which is available on the earthkit Xarray accessor. The first argument specifies the device. When the device is not “cpu” and the array_backend keyword argument is not specified it is automatically set to “cupy”.

r_cp = r.earthkit.to_device("cuda:0")
# equivalent code:
# r_cp = r.earthkit.to_device("cuda:0", array_backend="cupy")

type(r_cp.t.data)

cupy.ndarray

# Xarray computations work
r_cp.t.mean()

<xarray.DataArray 't' ()> Size: 8B
array(261.56490497)

xarray.DataArray

't'

• array(261.56490497)

  array(261.56490497)

• Coordinates: (0)
• Indexes: (0)
• Attributes: (0)

# Alter the values
r_cp += 1
type(r_cp.t.data)

cupy.ndarray

Move back to the CPU as Numpy

We use to_device() again to move back the dataset to the cpu. When the device is “cpu” and the array_backend keyword argument is not specified it is automatically set to “numpy”.

r_np = r_cp.earthkit.to_device("cpu")
# equivalent code:
# r_np = r.earthkit.to_device("cpu", array_backend="numpy")

type(r_np.t.data)

numpy.ndarray

# The dataset contains the values altered on the GPU
r_np.t.mean()

<xarray.DataArray 't' ()> Size: 8B
array(262.56490497)

xarray.DataArray

't'

• 262.6

  array(262.56490497)

• Coordinates: (0)
• Indexes: (0)
• Attributes: (0)