Pythonic Way to Perform Statistics Across Multiple Variables with Xarray | by Luke Gloege, Ph.D. | Towards Data Science

Consider this, you have an ensemble of models stored in a Dataset (ds), with each ensemble member stored as a different variable. How would you take the ensemble average? If you have three members (ens1, ens2, ens3), one approach is to do this:
ens_mean = (ds['ens1'] + ds['ens2'] + ds['ens3']) / 3
Although this is fine for a small ensemble, there is a more Pythonic way to do this. Let’s pretend each member represents the monthly temperature at each location on the globe, what we really want is code that looks like this:
ens_mean = ds['temperature'].mean('ensemble')
This is clean, has fewer characters to type, and it is clear what the code is accomplishing.
In this post, I will demonstrate how to manipulate a Xarray Dataset so you can average across a new categorical variable. I expect anyone who uses Xarray to analyze geospatial data or ensemble models to find relevance in this post.
If you want to follow along, I will be using a Dataset populated with synthetic data to explain this concept. The code to this Dataset is below. Let’s pretend this Dataset represents monthly temperature output from different models with a 1-degree spatial resolution.
The Dataset’s info() is shown below. Notice that each ensemble member (ens1, ens2, ens3) is a different variable.
Now that we have an example Dataset, let’s merge all the members into a single variable named temperature with dimension (ensemble, time, lat, lon), where the new ensemble coordinate will provide access to each member. The code below accomplishes this task.
There is alot happening in this, let’s walk through it line-by-line.
variables = list(ds) → this creates a list of all the variable names in your Dataset. If you don’t need all the variables, then you can simply hardcode a list of variable names.
The second line does three things:
expand_dims(‘ensemble’, axis=0) → Takes each variable and adds a new dimension at the first axis position.assign_coords(ensemble=[var]) → Creates a new coordinate with the same name as the new dimension from the previous step. This also assigns the variable name to the coordinaterename("temperature") → Renames the variable to temperature
expand_dims(‘ensemble’, axis=0) → Takes each variable and adds a new dimension at the first axis position.
assign_coords(ensemble=[var]) → Creates a new coordinate with the same name as the new dimension from the previous step. This also assigns the variable name to the coordinate
rename("temperature") → Renames the variable to temperature
Notices this is done for each variable via list comprehension. If this is confusing, I encourage you to run this on on a single variable and look at how the object changes after each step. The code below can help with this.
var = 'ens1'expand = ds[var].expand_dims('ensemble',axis=0)assign = expand.assign_coords(ensemble=[var])rename = assign.rename('temperature')
You can then output expand, assign, and rename to see how the object changes after each step.
Once we have a list of DataArrays, we need to stitch them together along the ensemble dimension using xr.merge().
xr.merge(da_list) → This merges all the DataArrays together into a single Dataset. The new Dataset only has one variable (temperature) with all the DataArrays merged along the ensemble coordinate
The last two lines may not be necessary. Those lines preserve the global attributes in ds.
global_attrs = ds.attrsds_merged.attrs = global_attrs
The final Dataset will look like this:
There are multiple ways to merge variables onto a new dimension. Here is another approach.
Now that we have our merged Dataset we can easily perform statistics across ensemble members.
ds_merged[‘temperature’].mean(‘ensemble’)
ds_merged[‘temperature’].var(‘ensemble’)
If we want to plot just ens1, we can select it like this
ds_merged.sel(ensemble='ens1')
Although accessing individual ensemble members is slightly more cumbersome, I think the benefit of being able to perform statistics across the entire ensemble is more important.
In this section, I will describe two methods to generalize this technique. The code I present is merely meant to give you ideas on how to extend this post.
The first approach is to pipe() the function to a Dataset. This is more aptly applied to chaining multiple functions together. However, depending on the situation, you may want to pipe() a single function.
Another option is to extend Xarray by registering an accessor. In practice, this is done by simply adding a class decorator. This approach is appropriate if you are working on a set of special-purpose functions. For example, maybe you’re working on a class of methods for processing ensemble models. The code below shows how to decorate a class that has the var_merger() function.
Keep in mind that once you register an accessor you can not change it. For example, if you run this code in a notebook and then later decide to add another function to the class and re-run it, that new function will not be registered with the accessor. You would first have to restart the kernel for it to take effect. For that reason, registering an accessor is more appropriately placed in the code library.
You can use your new accessor on a Dataset (ds) like this:
ds_merged = ds.ensemble.var_merger()
Notice that the name of the class is irrelevant since it is the accessor name we supply to the Dataset.
Hopefully, this tutorial helps you work more efficiently with ensembles. At the time of writing, I was unable to find a one-liner approach to merge variables like this in the documentation as well as their “How do I...” section.
Xarray is extremely powerful and many “how do I...” questions can be solved with existing functions. Having said that, it can be tricky to find the right functions and chain them together in the right way to achieve the desired object.
Please let me know if you have comments/suggestions on ways to improve this post. Also, please let me know if you are aware of a more elegant solution to merging variables onto a common coordinate.
As always, I am happy to help troubleshoot any issues you may come across.
Thank you for reading and supporting Medium writers