Gridded Data with Xarray¶

Xarray is a very powerful tool for exploring multi-dimensional (esp. geospatial) data in a way that is efficient and robust to making coding mistakes. Pandas provided us with a way to look at tabular data, Xarray takes this further and provides a framework for N-dimensional data. One of the coolest Xarray features, is the integration with Dask. This will allow us to easily parallelize our Xarray analysis!

When using Xarray, our data will be stored in DataArrays which are collected together into a DataSet. A nice example of this is in the context of a climate model:

DataSet - contains all possible coordinates on the model grid and provides a list of all model variables (DataArrays)
DataArray - an individual model variable (e.g., sea surface temperature), the variable’s coordinates on the model grid, and any additional meta data about that specific variable

In the graphic below, you can see that temperature and precipitation are both variables with coordinates lat and lon. In this simple example temperature and precip each have 3 dimensions. Not only does the DataSet store all of this information, it also relates these two variables to each other by understanding that they share coordinates lat and lon. We will use some test data to inspect this further. Any NetCDF can be read into an Xarray dataset. Many popular models facilitate reading raw output directly into Xarray.

38447245c3224dccaa665182c279df8f

[1]:

import matplotlib.pyplot as plt
import xarray as xr

# load a sample dataset from the xarray library
ds = xr.tutorial.load_dataset("air_temperature")

TIP If you realize you need to use Pandas for a particular problem, never fear! You can convert between Pandas and Xarray with a single line of code.

[2]:

df = ds.isel(time=slice(10)).to_dataframe() # to dataframe from xarray
df.to_xarray() # to xarray from dataframe

Interpolating (or Resampling) Data¶

[11]:

# interpolate the air temperature to a new longitude (or new set of longitudes)

Plotting in Xarray (so easy!)¶

[12]:

ds.air.isel(time=0).plot()

[12]:

<matplotlib.collections.QuadMesh at 0x177c91dc0>

../../_images/in_workshop_python_earth_science_xarray_20_1.png

[13]:

import cmocean.cm as cmo

# Let's use a non-default colormap
ds.air.isel(time=0).plot(cmap=cmo.thermal)

[13]:

<matplotlib.collections.QuadMesh at 0x177c52870>

../../_images/in_workshop_python_earth_science_xarray_21_1.png

[14]:

ds.air.sel(lon=220.0,method='nearest').plot(y='lat',x='time',cmap=cmo.thermal)

[14]:

<matplotlib.collections.QuadMesh at 0x318691dc0>

../../_images/in_workshop_python_earth_science_xarray_22_1.png

Statistics in Xarray¶

[15]:

ds.air.sel(lon=220.0,method='nearest').mean(dim='time').plot()

[15]:

[<matplotlib.lines.Line2D at 0x319317230>]

../../_images/in_workshop_python_earth_science_xarray_24_1.png

[16]:

del ds

A More Exciting Xarray Example¶

Let’s look at a more interesting and realistic example:

[17]:

import xarray as xr
ds = xr.open_dataset('TPOSE6_Daily_2012_surface.nc')

What about something harder? Gradients? Integrals? … Enter xgcm.

EOF decomposition (a.k.a. PCA) with xeofs¶

Xarray automatically sped up with Dask¶

Let’s load a larger dataset (~2.5GB). Here is where Xarray really shines! The data can be easily chunked into Dask arrays. If we spin up a Dask client with the 3 lines of code we learned earlier, our Xarray data analysis will be parallelized over the chunks. We can see the size of the chunks through inspection of the DataSet.

[1]:

import xarray as xr
import xgcm

ds = xr.open_dataset('TPOSE6_Daily_2012.nc')
grid = xgcm.Grid(ds, periodic=['X','Y'])
grid

[1]:

<xgcm.Grid>
Y Axis (periodic, boundary=None):
  * center   YC --> outer
  * outer    YG --> center
Z Axis (not periodic, boundary=None):
  * center   Z
T Axis (not periodic, boundary=None):
  * center   time
X Axis (periodic, boundary=None):
  * center   XC --> outer
  * outer    XG --> center

[2]:

%%time

# calculate the divergence of the flow at all depths
u_transport = ds.UVEL * ds.dyG * ds.hFacW * ds.drF
v_transport = ds.VVEL * ds.dxG * ds.hFacS * ds.drF

div_uv = (grid.diff(u_transport, 'X') + grid.diff(v_transport, 'Y')) / ds.rA  # calculate the divergence of the flow
div_uv.compute()

CPU times: user 321 ms, sys: 585 ms, total: 907 ms
Wall time: 926 ms

[3]:

del ds

[ ]:

# spin up our Dask client
from dask.distributed import LocalCluster, Client
cluster = LocalCluster()
client = Client(cluster)
client

[4]:

ds = xr.open_dataset('TPOSE6_Daily_2012.nc').chunk({'time':1})
grid = xgcm.Grid(ds, periodic=['X','Y'])
grid

[4]:

<xgcm.Grid>
Y Axis (periodic, boundary=None):
  * center   YC --> outer
  * outer    YG --> center
Z Axis (not periodic, boundary=None):
  * center   Z
T Axis (not periodic, boundary=None):
  * center   time
X Axis (periodic, boundary=None):
  * center   XC --> outer
  * outer    XG --> center

[5]:

ds.THETA

[5]:

<xarray.DataArray 'THETA' (time: 366, Z: 22, YC: 84, XC: 240)>
dask.array<xarray-THETA, shape=(366, 22, 84, 240), dtype=float32, chunksize=(1, 22, 84, 240), chunktype=numpy.ndarray>
Coordinates: (12/14)
    iter     (time) int64 dask.array<chunksize=(1,), meta=np.ndarray>
  * time     (time) timedelta64[ns] 00:01:12 00:02:24 ... 01:10:48 01:12:00
  * YC       (YC) float64 -3.917 -3.75 -3.583 -3.417 ... 9.417 9.583 9.75 9.917
  * Z        (Z) float64 -1.0 -3.0 -5.0 -7.0 -9.0 ... -68.5 -76.0 -85.0 -95.0
    drF      (Z) float32 dask.array<chunksize=(22,), meta=np.ndarray>
    PHrefC   (Z) float32 dask.array<chunksize=(22,), meta=np.ndarray>
    ...       ...
    rA       (YC, XC) float32 dask.array<chunksize=(84, 240), meta=np.ndarray>
    Depth    (YC, XC) float32 dask.array<chunksize=(84, 240), meta=np.ndarray>
    hFacC    (Z, YC, XC) float32 dask.array<chunksize=(22, 84, 240), meta=np.ndarray>
    maskC    (Z, YC, XC) bool dask.array<chunksize=(22, 84, 240), meta=np.ndarray>
    dxF      (YC, XC) float32 dask.array<chunksize=(84, 240), meta=np.ndarray>
    dyF      (YC, XC) float32 dask.array<chunksize=(84, 240), meta=np.ndarray>
Attributes:
    standard_name:  THETA
    long_name:      Potential Temperature
    units:          degC

[ ]:

%%time

# calculate the divergence of the flow at all depths
u_transport = ds.UVEL * ds.dyG * ds.hFacW * ds.drF
v_transport = ds.VVEL * ds.dxG * ds.hFacS * ds.drF

div_uv = (grid.diff(u_transport, 'X') + grid.diff(v_transport, 'Y')) / ds.rA  # calculate the divergence of the flow
div_uv.compute()

[ ]:

client.close()

See this page for a more detailed discussion of best practices with Xarray and Dask.

NOTE Take a look at all of the cool Xarray related projects! Most of them are for geoscience applications. There is also an extensive set of Xarray tutorials.