#!/usr/bin/env python # coding: utf-8 # # Xarray for multidimensional gridded data # # In last week's lecture, we saw how Pandas provided a way to keep track of additional "metadata" surrounding tabular datasets, including "indexes" for each row and labels for each column. These features, together with Pandas' many useful routines for all kinds of data munging and analysis, have made Pandas one of the most popular python packages in the world. # # However, not all Earth science datasets easily fit into the "tabular" model (i.e. rows and columns) imposed by Pandas. In particular, we often deal with _multidimensional data_. By _multidimensional data_ (also often called _N-dimensional_), I mean data with many independent dimensions or axes. For example, we might represent Earth's surface temperature $T$ as a three dimensional variable # # $$ T(x, y, t) $$ # # where $x$ is longitude, $y$ is latitude, and $t$ is time. # # The point of xarray is to provide pandas-level convenience for working with this type of data. # # # # ![xarray data model](https://github.com/pydata/xarray/raw/master/doc/_static/dataset-diagram.png) # **NOTE:** In order to run this tutorial, you need `xarray` and `netCDF4` packages installed. The best thing to do is to create a custom conda environment, as described on the [python installation page](https://rabernat.github.io/research_computing/python.html) (scroll to Geosciences Python Environment). To test whether your environment is set up properly, try the following imports: # In[2]: import xarray import netCDF4 # ## Xarray data structures # # Like Pandas, xarray has two fundamental data structures: # * a `DataArray`, which holds a single multi-dimensional variable and its coordinates # * a `Dataset`, which holds multiple variables that potentially share the same coordinates # # ### DataArray # # A `DataArray` has four essential attributes: # * `values`: a `numpy.ndarray` holding the array’s values # * `dims`: dimension names for each axis (e.g., `('x', 'y', 'z')`) # * `coords`: a dict-like container of arrays (coordinates) that label each point (e.g., 1-dimensional arrays of numbers, datetime objects or strings) # * `attrs`: an `OrderedDict` to hold arbitrary metadata (attributes) # # Let's start by constructing some DataArrays manually # In[3]: import numpy as np import xarray as xr from matplotlib import pyplot as plt get_ipython().run_line_magic('matplotlib', 'inline') plt.rcParams['figure.figsize'] = (8,5) # A simple DataArray without dimensions or coordinates isn't much use. # In[4]: da = xr.DataArray([9, 0, 2, 1, 0]) da # We can add a dimension name... # In[5]: da = xr.DataArray([9, 0, 2, 1, 0], dims=['x']) da # But things get most interesting when we add a coordinate: # In[6]: da = xr.DataArray([9, 0, 2, 1, 0], dims=['x'], coords={'x': [10, 20, 30, 40, 50]}) da # Xarray has built-in plotting, like pandas. # In[7]: da.plot(marker='o') # ### Multidimensional DataArray # # If we are just dealing with 1D data, Pandas and Xarray have very similar capabilities. Xarray's real potential comes with multidimensional data. # # Let's go back to the multidimensional ARGO data we loaded in the numpy lession. If you haven't already downloaded it, you can do so at the command line with # # ```shell # curl -O http://www.ldeo.columbia.edu/~rpa/argo_float_4901412.npz # ``` # # We reload this data and examine its keys. # In[8]: argo_data = np.load('argo_float_4901412.npz') argo_data.keys() # The values of the `argo_data` object are numpy arrays. # In[9]: S = argo_data.f.S T = argo_data.f.T P = argo_data.f.P levels = argo_data.f.levels lon = argo_data.f.lon lat = argo_data.f.lat date = argo_data.f.date print(S.shape, lon.shape, date.shape) # Let's organize the data and coordinates of the salinity variable into a DataArray. # In[10]: da_salinity = xr.DataArray(S, dims=['level', 'date'], coords={'level': levels, 'date': date},) da_salinity # In[11]: da_salinity.plot(yincrease=False) # Attributes can be used to store metadata. What metadata should you store? The [CF Conventions](http://cfconventions.org/Data/cf-conventions/cf-conventions-1.7/cf-conventions.html#_description_of_the_data) are a great resource for thinking about climate metadata. Below we define two of the required CF-conventions attributes. # In[11]: da_salinity.attrs['units'] = 'PSU' da_salinity.attrs['standard_name'] = 'sea_water_salinity' da_salinity # ### Datasets # # A Dataset holds many DataArrays which potentially can share coordinates. In analogy to pandas: # # pandas.Series : pandas.Dataframe :: xarray.DataArray : xarray.Dataset # # Constructing Datasets manually is a bit more involved in terms of syntax. The Dataset constructor takes three arguments: # # * `data_vars` should be a dictionary with each key as the name of the variable and each value as one of: # * A `DataArray` or Variable # * A tuple of the form `(dims, data[, attrs])`, which is converted into arguments for Variable # * A pandas object, which is converted into a `DataArray` # * A 1D array or list, which is interpreted as values for a one dimensional coordinate variable along the same dimension as it’s name # * `coords` should be a dictionary of the same form as data_vars. # * `attrs` should be a dictionary. # # Let's put together a Dataset with temperature, salinity and pressure all together # In[12]: argo = xr.Dataset( data_vars={'salinity': (('level', 'date'), S), 'temperature': (('level', 'date'), T), 'pressure': (('level', 'date'), P)}, coords={'level': levels, 'date': date}) argo # What about lon and lat? We forgot them in the creation process, but we can add them after the fact. # In[13]: argo['lon'] = lon argo # That was not quite right...we want lon to have dimension `date`: # In[23]: del argo['lon'] argo['lon'] = ('date', lon) argo['lat'] = ('date', lat) # ### Coordinates vs. Data Variables # # Data variables can be modified through arithmentic operations or other functions. Coordinates are always keept the same. # In[15]: argo * 10000 # Clearly lon and lat are coordinates rather than data variables. We can change their status as follows: # In[16]: argo = argo.set_coords(['lon', 'lat']) argo # The `*` symbol in the representation above indicates that `level` and `date` are "dimension coordinates" (they describe the coordinates associated with data variable axes) while `lon` and `lat` are "non-dimension coordinates". We can make any variable a non-dimension coordiante. # ## Working with Labeled Data # # Xarray's labels make working with multidimensional data much easier. # # ### Selecting Data (Indexing) # # We can always use regular numpy indexing and slicing on DataArrays # In[17]: argo.salinity[2].plot() # In[18]: argo.salinity[:, 10].plot() # However, it is often much more powerful to use xarray's `.sel()` method to use label-based indexing. # In[19]: argo.salinity.sel(level=2).plot() # In[20]: argo.salinity.sel(date='2012-10-22').plot() # `.sel()` also supports slicing. Unfortunately we have to use a somewhat awkward syntax, but it still works. # In[21]: argo.salinity.sel(date=slice('2012-10-01', '2012-12-01')).plot() # `.sel()` also works on the whole Dataset # In[22]: argo.sel(date='2012-10-22') # ### Computation # # Xarray dataarrays and datasets work seamlessly with arithmetic operators and numpy array functions. # In[23]: α = np.cos(argo.temperature) * np.sin(argo.salinity)**2 α.plot() # ### Broadcasting # # Broadcasting arrays in numpy is a nightmare. It is much easier when the data axes are labeled! # # This is a useless calculation, but it illustrates how perfoming an operation on arrays with differenty coordinates will result in automatic broadcasting # In[24]: level_times_lat = argo.level * argo.lat level_times_lat.plot() # ### Reductions # # Rather than performing reductions on axes (as in numpy), we can perform them on dimensions. # In[25]: argo_mean = argo.mean(dim='date') argo_mean # In[26]: argo_mean.salinity.plot() # ### Aside: Swapping Dims # # Now we can fix a pesky problem with this dataset: the fact that it uses level (rather than pressure) as the vertical coordinate # In[27]: argo['pres_mean'] = argo_mean.pressure argo_pcoords = argo.swap_dims({'level': 'pres_mean'}) argo_pcoords.salinity.plot(yincrease=False) # ## Groupby: Example with SST Climatology # # Here will we work with SST data from NOAA's NERSST project. Download it by running # # ```shell # curl -O http://ldeo.columbia.edu/~rpa/NOAA_NCDC_ERSST_v3b_SST.nc # ``` # In[50]: ds = xr.open_dataset('NOAA_NCDC_ERSST_v3b_SST.nc').chunk({'time':1}) ds # In[51]: sst = ds.sst # In[30]: sst.mean(dim='time').plot(vmin=-2, vmax=30) # In[31]: sst.mean(dim=('time', 'lon')).plot() # In[32]: sst_zonal_time_mean = sst.mean(dim=('time', 'lon')) # In[33]: (sst.mean(dim='lon') - sst_zonal_time_mean).T.plot() # In[34]: #sst.sel(lon=200, lat=0).plot() sst.sel(lon=230, lat=0, method='nearest').plot() sst.sel(lon=230, lat=45, method='nearest').plot() # In[35]: # climatologies sst_clim = sst.groupby('time.month').mean(dim='time') sst_clim.sel(lon=230, lat=45, method='nearest').plot() # In[36]: sst_clim.mean(dim='lon').T.plot.contourf(levels=np.arange(-2,30)) # In[37]: sst_anom = sst.groupby('time.month') - sst_clim sst_anom.sel(lon=230, lat=45, method='nearest').plot() # In[38]: sst_anom.std(dim='time').plot() # https://www.ncdc.noaa.gov/teleconnections/enso/indicators/sst.php # # El Niño (La Niña) is a phenomenon in the equatorial Pacific Ocean characterized by a five consecutive 3-month running mean of sea surface temperature (SST) anomalies in the Niño 3.4 region that is above (below) the threshold of +0.5°C (-0.5°C). This standard of measure is known as the Oceanic Niño Index (ONI). # # ![nino-regions](https://www.ncdc.noaa.gov/monitoring-content/teleconnections/nino-regions.gif) # In[39]: sst_anom_nino34 = sst_anom.sel(lat=slice(-5, 5), lon=slice(190, 240)) sst_anom_nino34[0].plot() # In[40]: sst_anom_nino34_mean = sst_anom_nino34.mean(dim=('lon', 'lat')) oni = sst_anom_nino34_mean.rolling(time=3).mean(dim='time') # In[41]: fig, ax = plt.subplots() sst_anom_nino34_mean.plot(ax=ax, label='raw') oni.plot(ax=ax, label='smoothed') ax.grid() # ![nino3.4](https://www.ncdc.noaa.gov/monitoring-content/teleconnections/eln-f-pg.gif) # In[42]: # create a categorical dataarray nino34 = xr.full_like(oni, 'none', dtype='U4') nino34[oni >= 0.5] = 'nino' nino34[oni <= -0.5] = 'nina' nino34 # In[43]: sst_nino_composite = sst_anom.groupby(nino34.rename('nino34')).mean(dim='time') # In[44]: sst_nino_composite.sel(nino34='nino').plot() # In[45]: sst_nino_composite.sel(nino34='nina').plot() # In[46]: nino_ds = xr.Dataset({'nino34': nino34, 'oni': oni}).drop('month') nino_ds.to_netcdf('nino34_index.nc') # In[48]: # In[ ]: