Multidimensional Data is Not Complicated — Xarray

In this article, I will demonstrate how the Xarray library makes working with multidimensional NumPy arrays intuitive and easy by abstracting away the complicated transpose operations, allowing you to perform advanced indexing simply using the dimension names.

Quick Review — Multidimensional Arrays

Multidimensional arrays are just arrays that contain subarrays.

One-dimensional array: An array of scalar values

[1, 2, 3]

Two-dimensional array: An array of arrays (a table/matrix)

[
    [1, 2, 3],
    [4, 5, 6],
    [7, 8, 9]    
]

Three-dimensional array: An array of two-dimensional arrays (a cube)

[
    [
        [1, 2, 3],
        [4, 5, 6],
        [7, 8, 9]
    ],
    [
        [1, 2, 3],
        [4, 5, 6],
        [7, 8, 9]
    ]
]

n-dimensional array: An array of (n-1)-dimensional arrays

[
    [(n-1)-dimensional array],
    [(n-1)-dimensional array],
    [(n-1)-dimensional array],
    ...
    [(n-1)-dimensional array]
]

Now that we understand what multi-dimensional arrays are, let’s see how Xarray makes them easy to work with.

Xarray Variable — NumPy Array with Named Dimensions

Imagine you have a 3D NumPy array: data = np.random.rand(2, 3, 4). What do these dimensions (of size 2, 3, and 4) represent? Is it (time, latitude, longitude), or (depth, height, width), or something else entirely? Without this basic information, even simple operations become a guessing game. The Xarray Variable object allows us to assign names to these dimensions. Consider the following example where we wrap a NumPy array into an Xarray Variable object along with its dimension names.

temperature_data = np.array([[[14.12445211, 27.70085037],
                              [11.28397733, 31.34240561],
                              [18.08515337, 29.2945892 ]],

                             [[30.53006346, 31.5863542 ],
                              [11.41459454, 20.91854694],
                              [15.43264092, 19.47200334]]])

temperature_variable = xr.Variable(
    dims=('time', 'lat', 'lon'),
    data=temperature_data
)

Now we can select data by the dimension names without knowing the order they’re in. For example, if we want to select data by the first entry in the time dimension, we could do it like so:

temperature_variable.isel(time=0)

The numpy equivalent would be:

temperature_data[0, :, :] # If time is first dimension
temperature_data[:, 0, :] # If time is second dimension
temperature_data[:, :, 0] # If time is third dimension

As you can see, the Xarray Variable object makes it much easier to index multidimensional data as we don’t have to figure out what the dimensions are or what order they’re in. However, we don’t actually know what time=0 actually means in this example. We’ll address this with the next Xarray object, the DataArray.

Xarray DataArray — Coordinates and Measurements

The Xarray DataArray lets us store NumPy arrays with their dimension names as well as the actual values corresponding to these dimensions. We call these values the coordinates. You can just think of these coordinate values as “where/when/how we took measurements”.

For example, let’s say we went out on a given day, 2023–01–01, and took some measurements of the temperature at midnight and then again at noon. Both times, we took 6 measurements — one for each combination of 3 latitudes and 2 longitudes. Since we have 3 dimensions (time, lat, lon), we can arbitrarily decide to organize our measurements by time, then latitude and longitude. So the data is stored as an array with two subarrays — the midnight measurements and noon measurements. Each of these subarrays contains 3 subarrays (one for each of the three latitudes), each of which contain 2 temperature measurements (one for each of the two longitudes).

# Coordinates
times = pd.to_datetime(['2023-01-01T00:00:00', '2023-01-01T12:00:00'])
lats = np.array([30.0, 40.0, 50.0])
lons = np.array([-100.0, -90.0])

# Measurements
temperature_data = np.array([[[14.12445211, 27.70085037],
                              [11.28397733, 31.34240561],
                              [18.08515337, 29.2945892 ]],

                             [[30.53006346, 31.5863542 ],
                              [11.41459454, 20.91854694],
                              [15.43264092, 19.47200334]]])

But what if you just downloaded this temperature_data array and imported into python, would you be able to understand what order of dimensions it’s nested in? For example, it could just as well have been ordered by longitude, latitude, and then time:

temperature_data = np.array([[[14.12445211, 30.53006346],
                              [11.28397733, 11.41459454],
                              [18.08515337, 15.43264092]],

                             [[27.70085037, 31.5863542 ],
                              [31.34240561, 20.91854694],
                              [29.2945892 , 19.47200334]]])

Or lat, lon, time:

temperature_data = np.array([[[14.12445211, 30.53006346],
                              [27.70085037, 31.5863542 ]],

                             [[11.28397733, 11.41459454],
                              [31.34240561, 20.91854694]],

                             [[18.08515337, 15.43264092],
                              [29.2945892 , 19.47200334]]])

How would you perform basic indexing operations like “all measurements from midnight”? In NumPy this would be quite difficult. You would first have to figure out what the dimensions are, which order they’re in, and then which of the subarrays within the time dimension corresponds to this timestamp. Once you had all that figured out, the query itself would be quite simple, although not very intuitive:

temperature_data[0, :, :] # First dimension is time, first entry is midnight

The beauty of Xarray is that it abstracts this multi-dimensional nesting from you and lets you index by the dimension names instead. So you would record your measurements along with the dimensions coordinates and their names, and then insert them into one Xarray DataArray object together:

# Coordinates
times = pd.to_datetime(['2023-01-01T00:00:00', '2023-01-01T12:00:00'])
lats = np.array([30.0, 40.0, 50.0])
lons = np.array([-100.0, -90.0])

# Measurements
temperature_data = np.array([[[14.12445211, 27.70085037],
                              [11.28397733, 31.34240561],
                              [18.08515337, 29.2945892 ]],

                             [[30.53006346, 31.5863542 ],
                              [11.41459454, 20.91854694],
                              [15.43264092, 19.47200334]]])
# Create the DataArray
temperature_da = xr.DataArray(
    data=temperature_data,
    coords={"time": times, "lat": lats, "lon": lons}
)

Now we can immediately start working with the data without having to figure out what the dimensions are or what order they’re in:

temperature_da.sel(time='2023-01-01T00:00:00').values

The obvious question is why do we even need to store the data as a multi-dimensional array in the first place? Couldn’t it just be a table? And the answer is of course it could be a table, here’s what it would look like:

time,lat,lon,temperature
2023-01-01 00:00:00,30.0,-100.0, 14.12445211
2023-01-01 00:00:00,30.0,-90.0,  27.70085037
2023-01-01 00:00:00,40.0,-100.0, 11.28397733
2023-01-01 00:00:00,40.0,-90.0,  31.34240561
2023-01-01 00:00:00,50.0,-100.0, 18.08515337
2023-01-01 00:00:00,50.0,-90.0,  29.2945892
2023-01-01 12:00:00,30.0,-100.0, 30.53006346
2023-01-01 12:00:00,30.0,-90.0,  31.5863542
2023-01-01 12:00:00,40.0,-100.0, 11.41459454
2023-01-01 12:00:00,40.0,-90.0,  20.91854694
2023-01-01 12:00:00,50.0,-100.0, 15.43264092
2023-01-01 12:00:00,50.0,-90.0,  19.47200334

Notice the amount of redundant data required to map each temperature measurement to its (time, lat, lon) coordinate. That’s 3 coordinate values per one measurement value! So you would imagine that n measurements stored in such a table would required 3n coordinate values! Immediately, the benefits of Xarray are obvious — we only store the coordinates once (and of course we only store the measurements once).

For example, let’s imagine a dataset with 1000 time coordinates, 500 latitude coordinates, and 500 longitude coordinates. Such a dataset would have 1000 500 500 = 250,000,000 possible values. If we were to store each measurement along with its coordinates in a table, the total number of coordinate values stored would be 250,000,000 * 3 = 750,000,000. Whereas Xarray would only store the 2000 coordinate values once.

Furthermore, storing data in contiguous blocks of memory as dense multidimensional arrays results in much faster read and write operations due to computer memory architecture.

To summarize, with Xarray, we can store datapoints as measurements along with their corresponding coordinates (and their names) in an efficient multi-dimensional array format and then perform indexing operations by simply using the dimension names without even knowing what the dimension order is.

Dataset — DataArrays Sharing Coordinates

The Dataset object is an extension of DataArray that allows you to store multiple sets of measurements with their corresponding coordinates. For example, let’s say we also took humidity measurements on our trip. We can wrap them in a DataArray as we did with the temperature measurements and group them into

times = pd.to_datetime(['2023-01-01T00:00:00', '2023-01-01T12:00:00'])
lats = np.array([30.0, 40.0, 50.0])
lons = np.array([-100.0, -90.0])

temperature_measurements = np.array([[[14.12445211, 27.70085037],
                                      [11.28397733, 31.34240561],
                                      [18.08515337, 29.2945892 ]],

                                     [[30.53006346, 31.5863542 ],
                                      [11.41459454, 20.91854694],
                                      [15.43264092, 19.47200334]]])

humidity_measurements = np.array([[[ 7.06222606, 13.85042519],
                                   [ 5.64198867, 15.67120281],
                                   [ 9.04257669, 14.6472946 ]],

                                  [[15.26503173, 15.7931771 ],
                                   [ 5.70729727, 10.45927347],
                                   [ 7.71632046,  9.73600167]]])

temperature_da = xr.DataArray(
    data=temperature_measurements,
    coords={"time": times, "lat": lats, "lon": lons}
)

humidity_da = xr.DataArray(
    data=humidity_measurements,
    coords={"time": times,"lat": lats, "lon": lons},
)

# Create a Dataset to hold both temperature and humidity
measurements_ds = xr.Dataset({
    "temperature": temperature_da,
    "humidity": humidity_da
})

And now you can query the dataset like this:

# Midnight measurements
print(measurements_ds.temperature.sel(time='2023-01-01T00:00:00').values)
print(measurements_ds.humidity.sel(time='2023-01-01T00:00:00').values)

# Geo-specific measurements
print(measurements_ds.temperature.sel(lat=40.0, lon=-90).values)
print(measurements_ds.humidity.sel(lat=40.0, lon=-90).values)

Summary

Variable: NumPy array with dimension names. DataArray: NumPy array with dimension names and coordinates Dataset: A collection of DataArrays that can share coordinates