Introduction to CuPy-Xarray#

Negin Sobhani, Deepak Cherian, and Max Jones
negins@ucar.edu, dcherian@ucar.edu, and max@carbonplan.org


Overview#

In this tutorial, you learn:#

  • Basics of Cupy-Xarray

  • Creating and handling Xarray DataArrays on GPUs

  • Data Transfer Between Host and Device

Prerequisites#

Concepts

Importance

Notes

Familiarity with NumPy

Necessary

Basics of Cupy

Necessary

Familiarity with Xarray

Necessary

  • Time to learn: 10 minutes

Introduction#

Xarray is a powerful library for working with labeled multi-dimensional arrays in Python. It provides a convenient and intuitive way to manipulate large and complex datasets, and is built on top of NumPy. CuPy, on the other hand, is a library that allows for GPU-accelerated computing with Python and is compatible with NumPy.

When used together, Xarray and CuPy can provide an easy way to take advantage of GPU acceleration for scientific computing tasks.

Xarray can wrap custom duck array objects (i.e. NumPy-like arrays) that follow specific protocols.

CuPy-Xarray provides an interface for using CuPy in Xarray, providing accessors on the Xarray objects.

This tutorial showcases the use of cupy-xarray, which offers a cupy accessor that allows access to cupy-specific features.

First, let’s import our packages

## Import NumPy and CuPy
import cupy as cp
import numpy as np
import xarray as xr
import cupy_xarray  # Adds .cupy to Xarray objects

Creating Xarray DataArray with CuPy#

In the previous tutorial, we learned how to create a NumPy and CuPy array:

arr_cpu = np.random.rand(10, 10, 10)
arr_gpu = cp.random.rand(10, 10, 10)

We can create the Xarray DataArray using the CuPy array or NumPy array as the data source in a similar fashion:

# create a DataArray using CuPy array with three dimensions and 10 elements along each dimension
da_np = xr.DataArray(arr_cpu, dims=["x", "y", "time"])

da_np
<xarray.DataArray (x: 10, y: 10, time: 10)>
array([[[9.86345808e-01, 4.17660665e-01, 1.88327552e-01, 8.90280942e-01,
         2.21689274e-01, 3.17943609e-01, 5.00880587e-01, 7.53337533e-01,
         3.59031996e-01, 1.91030893e-01],
        [1.33144043e-02, 5.02596284e-01, 5.42913172e-01, 5.23846968e-01,
         4.05313585e-01, 9.46118837e-01, 1.06548298e-01, 3.47524404e-01,
         1.52843324e-01, 8.48253778e-01],
        [4.33649929e-01, 6.23728979e-01, 6.47445402e-01, 9.03623126e-01,
         9.05012105e-01, 4.06989322e-03, 4.19896664e-01, 8.60406673e-02,
         7.41788571e-01, 6.22665340e-01],
        [1.74267952e-01, 6.14897148e-01, 5.01242328e-01, 6.66759345e-01,
         8.44182632e-01, 3.19620076e-01, 7.90701915e-01, 2.43897985e-01,
         8.25956047e-01, 6.06534832e-01],
        [5.28111326e-01, 7.42343565e-01, 8.05094324e-02, 8.84691476e-01,
         1.69856723e-02, 3.32512453e-01, 6.67738160e-01, 7.06905069e-01,
         5.16369945e-01, 2.71965903e-01],
        [2.81638568e-01, 2.89389278e-01, 8.19006807e-01, 3.53878654e-01,
         9.21084664e-02, 5.69411698e-01, 8.56797393e-01, 3.24107223e-01,
         8.15087813e-01, 4.70350855e-01],
        [5.70124339e-01, 7.92088214e-01, 9.33540441e-01, 9.88027072e-01,
         9.05585677e-01, 5.28417548e-01, 4.40169554e-01, 1.40924601e-01,
...
         5.29797743e-01, 9.87589722e-01, 9.18635655e-01, 8.68580278e-01,
         4.71548324e-01, 3.64257635e-01],
        [6.42229124e-01, 2.33643023e-02, 5.85033551e-01, 8.80436137e-02,
         7.07996956e-01, 4.40586920e-01, 3.10391741e-01, 1.22763638e-02,
         8.02412664e-01, 4.33761051e-01],
        [1.24780820e-01, 3.53875474e-01, 8.36031716e-01, 2.84138174e-02,
         3.57476794e-01, 2.44890794e-02, 1.47504786e-02, 3.19465404e-01,
         2.91984584e-01, 3.39490525e-01],
        [2.04021642e-01, 4.71267303e-01, 9.03187717e-02, 3.83928128e-01,
         5.96265409e-01, 3.17287239e-01, 3.22413673e-01, 8.38235070e-01,
         9.58316554e-01, 9.73589612e-01],
        [6.13802208e-01, 8.70356525e-01, 5.17350919e-01, 7.72374828e-03,
         5.35340510e-01, 8.89268388e-01, 6.93943330e-01, 6.29953006e-01,
         1.70230716e-01, 7.16573680e-01],
        [8.44214598e-01, 3.35186917e-01, 8.78891352e-01, 1.98027834e-01,
         6.36005433e-01, 1.21753118e-01, 6.48103717e-01, 8.68341345e-01,
         7.81023406e-01, 4.45064620e-01],
        [3.85731750e-01, 8.02230895e-01, 6.41415045e-01, 7.60886886e-01,
         2.00746550e-01, 3.76787007e-01, 6.68073723e-01, 7.87222270e-01,
         6.75273015e-01, 8.63705777e-01]]])
Dimensions without coordinates: x, y, time
# create a DataArray using NumPy array with three dimensions and 10 elements along each dimension
da_cp = xr.DataArray(arr_gpu, dims=["x", "y", "time"])

da_cp
<xarray.DataArray (x: 10, y: 10, time: 10)>
array([[[2.88196731e-01, 3.71102840e-01, 8.22413516e-01, 7.61603373e-01,
         2.14247694e-01, 6.08972260e-01, 6.35605124e-01, 4.51735394e-02,
         3.56580833e-01, 3.33245593e-01],
        [9.56686233e-02, 3.09634487e-01, 5.72034429e-01, 8.64203361e-01,
         5.44551902e-01, 4.54445926e-01, 1.21606888e-01, 2.30160410e-01,
         6.14639953e-01, 7.73246535e-01],
        [8.03011705e-01, 2.69969912e-01, 2.03781951e-01, 6.64806547e-01,
         4.93709552e-01, 2.60248353e-01, 6.82195033e-01, 6.75837492e-01,
         5.07293067e-01, 6.45924343e-01],
        [1.03968071e-01, 1.31787260e-01, 2.31666523e-02, 2.90727455e-01,
         6.22514068e-02, 9.54996781e-01, 1.38868633e-01, 3.18043546e-01,
         9.94141764e-01, 6.52825114e-01],
        [6.72144360e-01, 9.25109790e-01, 9.24907616e-01, 9.97835547e-01,
         1.30089788e-01, 3.28381980e-01, 9.47761645e-01, 2.15451004e-01,
         1.55072912e-01, 2.84564825e-01],
        [5.32157180e-01, 4.05812774e-01, 6.65152077e-01, 1.62793186e-01,
         8.38375837e-01, 4.38498164e-01, 3.93970103e-01, 3.25181026e-01,
         8.43314943e-01, 6.37218468e-01],
        [9.47935236e-01, 1.39071514e-01, 3.34994498e-01, 7.42907508e-01,
         1.13865457e-01, 3.69531071e-01, 6.58907523e-01, 4.10997683e-01,
...
         5.01101857e-01, 6.76530919e-01, 6.01550513e-01, 1.91761020e-01,
         2.01591335e-01, 3.73443454e-01],
        [8.72935075e-01, 9.28175014e-01, 7.03819938e-01, 4.25757273e-01,
         6.80355431e-01, 1.22351044e-01, 8.22086635e-03, 9.23118431e-01,
         8.00040998e-02, 3.51963004e-01],
        [5.30917733e-01, 1.73025731e-03, 5.46551386e-01, 3.41904305e-01,
         6.11276326e-01, 7.83903426e-01, 7.67650251e-01, 9.27383669e-02,
         5.99146336e-01, 1.44674661e-02],
        [9.32478257e-02, 6.51279678e-01, 3.40032365e-01, 6.66761485e-02,
         3.88243075e-01, 3.06181721e-02, 5.58666002e-01, 3.10356676e-01,
         6.46523629e-01, 1.19013418e-01],
        [1.81940990e-01, 3.89650142e-01, 9.98204973e-01, 4.39178186e-02,
         6.88137446e-02, 7.61541679e-02, 6.26075251e-01, 9.14708720e-01,
         4.45414011e-01, 5.16678456e-01],
        [8.51618677e-01, 6.81900815e-01, 6.66821786e-01, 8.75685884e-01,
         2.90499242e-01, 3.25977864e-01, 3.67627054e-01, 3.93770674e-01,
         7.40898577e-01, 3.50451112e-02],
        [7.06374026e-01, 7.19519511e-01, 1.79160522e-01, 8.81425785e-01,
         3.51431945e-01, 4.11507382e-01, 6.86088790e-01, 3.04671156e-01,
         5.70729870e-01, 7.76584760e-01]]])
Dimensions without coordinates: x, y, time

But how are these two DataArrays different from each other? How do we know which array is on CPU vs. GPU?

Checking for CuPy Arrays#

The cupy accessor provides the is_cupy method to check if these arrays are on the host or device. For example:

da_np.cupy.is_cupy
False
da_cp.cupy.is_cupy
True

Accessing Device Information of the DataArray#

To access the underlying CuPy array, use the data property of the DataArray. It returns the CuPy array:

cupy_array = da_cp.data
type(cupy_array)
cupy.ndarray

In the previous tutorial, we learned about CuPy’s introduction of the notion of a current device. We also learned that to identify the device assigned to a CuPy array, the cupy.ndarray.device attribute can be used. Similar concept can be applied to a DataArray:

da_cp.data.device
<CUDA Device 0>

Data Transfer#

Transferring DataArrays to another Device#

In the previous lesson we learned that by default, code execution is carried out on Device 0. However, with CuPy, we have the ability to transfer arrays to other devices using cp.cuda.Device(). This feature becomes particularly valuable when your code is designed to leverage the capabilities of multiple GPUs. Similar concept applies to DataArrays that include Cupy Arrays:

with cp.cuda.Device(1):
    x_on_gpu1 = cp.array([5, 7, 8, 5, 5])
    da_cp1 = xr.DataArray(x_on_gpu1, dims=['time'])

da_cp1.data.device

Transferring Data between Host and Device#

Xarray provides DataArray.as_numpy to convert all kinds of arrays to NumPy arrays.

# Move data to host
da_np = da_cp.as_numpy()
da_np
<xarray.DataArray (x: 10, y: 10, time: 10)>
array([[[2.88196731e-01, 3.71102840e-01, 8.22413516e-01, 7.61603373e-01,
         2.14247694e-01, 6.08972260e-01, 6.35605124e-01, 4.51735394e-02,
         3.56580833e-01, 3.33245593e-01],
        [9.56686233e-02, 3.09634487e-01, 5.72034429e-01, 8.64203361e-01,
         5.44551902e-01, 4.54445926e-01, 1.21606888e-01, 2.30160410e-01,
         6.14639953e-01, 7.73246535e-01],
        [8.03011705e-01, 2.69969912e-01, 2.03781951e-01, 6.64806547e-01,
         4.93709552e-01, 2.60248353e-01, 6.82195033e-01, 6.75837492e-01,
         5.07293067e-01, 6.45924343e-01],
        [1.03968071e-01, 1.31787260e-01, 2.31666523e-02, 2.90727455e-01,
         6.22514068e-02, 9.54996781e-01, 1.38868633e-01, 3.18043546e-01,
         9.94141764e-01, 6.52825114e-01],
        [6.72144360e-01, 9.25109790e-01, 9.24907616e-01, 9.97835547e-01,
         1.30089788e-01, 3.28381980e-01, 9.47761645e-01, 2.15451004e-01,
         1.55072912e-01, 2.84564825e-01],
        [5.32157180e-01, 4.05812774e-01, 6.65152077e-01, 1.62793186e-01,
         8.38375837e-01, 4.38498164e-01, 3.93970103e-01, 3.25181026e-01,
         8.43314943e-01, 6.37218468e-01],
        [9.47935236e-01, 1.39071514e-01, 3.34994498e-01, 7.42907508e-01,
         1.13865457e-01, 3.69531071e-01, 6.58907523e-01, 4.10997683e-01,
...
         5.01101857e-01, 6.76530919e-01, 6.01550513e-01, 1.91761020e-01,
         2.01591335e-01, 3.73443454e-01],
        [8.72935075e-01, 9.28175014e-01, 7.03819938e-01, 4.25757273e-01,
         6.80355431e-01, 1.22351044e-01, 8.22086635e-03, 9.23118431e-01,
         8.00040998e-02, 3.51963004e-01],
        [5.30917733e-01, 1.73025731e-03, 5.46551386e-01, 3.41904305e-01,
         6.11276326e-01, 7.83903426e-01, 7.67650251e-01, 9.27383669e-02,
         5.99146336e-01, 1.44674661e-02],
        [9.32478257e-02, 6.51279678e-01, 3.40032365e-01, 6.66761485e-02,
         3.88243075e-01, 3.06181721e-02, 5.58666002e-01, 3.10356676e-01,
         6.46523629e-01, 1.19013418e-01],
        [1.81940990e-01, 3.89650142e-01, 9.98204973e-01, 4.39178186e-02,
         6.88137446e-02, 7.61541679e-02, 6.26075251e-01, 9.14708720e-01,
         4.45414011e-01, 5.16678456e-01],
        [8.51618677e-01, 6.81900815e-01, 6.66821786e-01, 8.75685884e-01,
         2.90499242e-01, 3.25977864e-01, 3.67627054e-01, 3.93770674e-01,
         7.40898577e-01, 3.50451112e-02],
        [7.06374026e-01, 7.19519511e-01, 1.79160522e-01, 8.81425785e-01,
         3.51431945e-01, 4.11507382e-01, 6.86088790e-01, 3.04671156e-01,
         5.70729870e-01, 7.76584760e-01]]])
Dimensions without coordinates: x, y, time

Let’s confirm this isn’t a CuPy array anymore:

da_np.cupy.is_cupy
False

We also can convert an Xarray DataArray that include NumPy array to a CuPy array (move data to Device) use cupy.as_cupy():

# Move data to GPU
da_cp = da_np.cupy.as_cupy()
da_cp.as_cupy().cupy.is_cupy
True

Plotting#

Plotting DataArrays with underlying data as CuPy arrays work in the same way as DataArrays with Numpy Arrays; however, data is first transferred to CPU before being plotted.

da_cp.plot()
(array([111.,  93.,  96., 112.,  88.,  93., 115.,  88.,  98., 106.]),
 array([5.01237631e-05, 9.98656087e-02, 1.99681094e-01, 2.99496578e-01,
        3.99312063e-01, 4.99127548e-01, 5.98943033e-01, 6.98758518e-01,
        7.98574003e-01, 8.98389488e-01, 9.98204973e-01]),
 <BarContainer object of 10 artists>)
../_images/6faacf616085285a967889590b3056ef23765105ef12a9d38fcc8ab8b9e9af77.png

Summary#

In this notebook, we have learned about:

  • CuPy-Xarray Basics

  • Data Transfer between Device to Host

Additional Resources#

CuPy User Guide
Xarray User Guide
Cupy-Xarray Github