Computations and Masks with Xarray
Overview
In this tutorial, we will cover the following topics:
Performing basic arithmetic on
DataArrays
andDatasets
Performing aggregation (i.e., reduction) along single or multiple dimensions of a
DataArray
orDataset
Computing climatologies and anomalies of data using Xarray’s “split-apply-combine” approach, via the
.groupby()
methodPerforming weighted-reduction operations along single or multiple dimensions of a
DataArray
orDataset
Providing a broad overview of Xarray’s data-masking capability
Using the
.where()
method to mask Xarray data
Prerequisites
Concepts |
Importance |
Notes |
---|---|---|
Necessary |
Time to learn: 60 minutes
Imports
In order to work with data and plotting, we must import NumPy, Matplotlib, and Xarray. These packages are covered in greater detail in earlier tutorials. We also import a package that allows quick download of Pythia example datasets.
import matplotlib.pyplot as plt
import numpy as np
import xarray as xr
from pythia_datasets import DATASETS
Data Setup
The bulk of the examples in this tutorial make use of a single dataset. This dataset contains monthly sea surface temperature (SST, call ‘tos’ here) data, and is obtained from the Community Earth System Model v2 (CESM2). (For this tutorial, however, the dataset will be retrieved from the Pythia example data repository.) The following example illustrates the process of retrieving this Global Climate Model dataset:
filepath = DATASETS.fetch('CESM2_sst_data.nc')
ds = xr.open_dataset(filepath)
ds
/free1/ktyle/miniforge3/envs/foundations_book24/lib/python3.12/site-packages/xarray/conventions.py:440: SerializationWarning: variable 'tos' has multiple fill values {1e+20, 1e+20}, decoding all values to NaN.
new_vars[k] = decode_cf_variable(
<xarray.Dataset> Size: 47MB Dimensions: (time: 180, d2: 2, lat: 180, lon: 360) Coordinates: * time (time) object 1kB 2000-01-15 12:00:00 ... 2014-12-15 12:00:00 * lat (lat) float64 1kB -89.5 -88.5 -87.5 -86.5 ... 86.5 87.5 88.5 89.5 * lon (lon) float64 3kB 0.5 1.5 2.5 3.5 4.5 ... 356.5 357.5 358.5 359.5 Dimensions without coordinates: d2 Data variables: time_bnds (time, d2) object 3kB ... lat_bnds (lat, d2) float64 3kB ... lon_bnds (lon, d2) float64 6kB ... tos (time, lat, lon) float32 47MB ... Attributes: (12/45) Conventions: CF-1.7 CMIP-6.2 activity_id: CMIP branch_method: standard branch_time_in_child: 674885.0 branch_time_in_parent: 219000.0 case_id: 972 ... ... sub_experiment_id: none table_id: Omon tracking_id: hdl:21.14100/2975ffd3-1d7b-47e3-961a-33f212ea4eb2 variable_id: tos variant_info: CMIP6 20th century experiments (1850-2014) with C... variant_label: r11i1p1f1
Arithmetic Operations
In a similar fashion to NumPy arrays, performing an arithmetic operation on a DataArray
will automatically perform the operation on all array values; this is known as vectorization. To illustrate the process of vectorization, the following example converts the air temperature data from units of degrees Celsius to units of Kelvin:
ds.tos + 273.15
<xarray.DataArray 'tos' (time: 180, lat: 180, lon: 360)> Size: 47MB array([[[ nan, nan, nan, ..., nan, nan, nan], [ nan, nan, nan, ..., nan, nan, nan], [ nan, nan, nan, ..., nan, nan, nan], ..., [271.3552 , 271.3553 , 271.3554 , ..., 271.35495, 271.355 , 271.3551 ], [271.36005, 271.36014, 271.36023, ..., 271.35986, 271.35992, 271.36 ], [271.36447, 271.36453, 271.3646 , ..., 271.3643 , 271.36435, 271.3644 ]], [[ nan, nan, nan, ..., nan, nan, nan], [ nan, nan, nan, ..., nan, nan, nan], [ nan, nan, nan, ..., nan, nan, nan], ... [271.40677, 271.40674, 271.4067 , ..., 271.40695, 271.4069 , 271.40683], [271.41296, 271.41293, 271.41293, ..., 271.41306, 271.413 , 271.41296], [271.41772, 271.41772, 271.41772, ..., 271.41766, 271.4177 , 271.4177 ]], [[ nan, nan, nan, ..., nan, nan, nan], [ nan, nan, nan, ..., nan, nan, nan], [ nan, nan, nan, ..., nan, nan, nan], ..., [271.39386, 271.39383, 271.3938 , ..., 271.39407, 271.394 , 271.39392], [271.39935, 271.39932, 271.39932, ..., 271.39948, 271.39944, 271.39938], [271.40372, 271.40372, 271.40375, ..., 271.4037 , 271.4037 , 271.40372]]], dtype=float32) Coordinates: * time (time) object 1kB 2000-01-15 12:00:00 ... 2014-12-15 12:00:00 * lat (lat) float64 1kB -89.5 -88.5 -87.5 -86.5 ... 86.5 87.5 88.5 89.5 * lon (lon) float64 3kB 0.5 1.5 2.5 3.5 4.5 ... 356.5 357.5 358.5 359.5
In addition, there are many other arithmetic operations that can be performed on DataArrays
. In this example, we demonstrate squaring the original Celsius values of our air temperature data:
ds.tos**2
<xarray.DataArray 'tos' (time: 180, lat: 180, lon: 360)> Size: 47MB array([[[ nan, nan, nan, ..., nan, nan, nan], [ nan, nan, nan, ..., nan, nan, nan], [ nan, nan, nan, ..., nan, nan, nan], ..., [3.2213385, 3.2209656, 3.220537 , ..., 3.2221622, 3.221913 , 3.2216525], [3.203904 , 3.203617 , 3.2032912, ..., 3.2045207, 3.2043478, 3.2041442], [3.1881146, 3.1879027, 3.1876712, ..., 3.188714 , 3.1885312, 3.1883302]], [[ nan, nan, nan, ..., nan, nan, nan], [ nan, nan, nan, ..., nan, nan, nan], [ nan, nan, nan, ..., nan, nan, nan], ... [3.0388296, 3.0389647, 3.0390673, ..., 3.038165 , 3.0383828, 3.0386322], [3.0173173, 3.0173445, 3.0173297, ..., 3.0169601, 3.0171173, 3.0172386], [3.000791 , 3.0007784, 3.0007539, ..., 3.000933 , 3.000896 , 3.0008452]], [[ nan, nan, nan, ..., nan, nan, nan], [ nan, nan, nan, ..., nan, nan, nan], [ nan, nan, nan, ..., nan, nan, nan], ..., [3.0839543, 3.0841148, 3.0842566, ..., 3.0832636, 3.0834875, 3.0837412], [3.064733 , 3.0648024, 3.0648358, ..., 3.0642793, 3.0644639, 3.0646174], [3.0494578, 3.0494475, 3.0494263, ..., 3.049596 , 3.0495603, 3.0495107]]], dtype=float32) Coordinates: * time (time) object 1kB 2000-01-15 12:00:00 ... 2014-12-15 12:00:00 * lat (lat) float64 1kB -89.5 -88.5 -87.5 -86.5 ... 86.5 87.5 88.5 89.5 * lon (lon) float64 3kB 0.5 1.5 2.5 3.5 4.5 ... 356.5 357.5 358.5 359.5
Aggregation Methods
A common practice in the field of data analysis is aggregation. Aggregation is the process of reducing data through methods such as sum()
, mean()
, median()
, min()
, and max()
, in order to gain greater insight into the nature of large datasets. In this set of examples, we demonstrate correct usage of a select group of aggregation methods:
Compute the mean:
ds.tos.mean()
<xarray.DataArray 'tos' ()> Size: 4B array(14.250171, dtype=float32)
Notice that we did not specify the dim
keyword argument; this means that the function was applied over all of the dataset’s dimensions. In other words, the aggregation method computed the mean of every element of the temperature dataset across every temporal and spatial data point. However, if a dimension name is used with the dim
keyword argument, the aggregation method computes an aggregation along the given dimension. In this next example, we use aggregation to calculate the temporal mean across all spatial data; this is performed by providing the dimension name 'time'
to the dim
keyword argument:
ds.tos.mean(dim='time').plot(size=7);

There are many other combinations of aggregation methods and dimensions on which to perform these methods. In this example, we compute the temporal minimum:
ds.tos.min(dim=['time'])
<xarray.DataArray 'tos' (lat: 180, lon: 360)> Size: 259kB array([[ nan, nan, nan, ..., nan, nan, nan], [ nan, nan, nan, ..., nan, nan, nan], [ nan, nan, nan, ..., nan, nan, nan], ..., [-1.8083605, -1.8083031, -1.8082187, ..., -1.8083988, -1.8083944, -1.8083915], [-1.8025414, -1.8024837, -1.8024155, ..., -1.8026428, -1.8026177, -1.8025846], [-1.7984415, -1.7983989, -1.7983514, ..., -1.7985678, -1.7985296, -1.7984871]], dtype=float32) Coordinates: * lat (lat) float64 1kB -89.5 -88.5 -87.5 -86.5 ... 86.5 87.5 88.5 89.5 * lon (lon) float64 3kB 0.5 1.5 2.5 3.5 4.5 ... 356.5 357.5 358.5 359.5
This example computes the spatial sum. Note that this dataset contains no altitude data; as such, the required spatial dimensions passed to the method consist only of latitude and longitude.
ds.tos.sum(dim=['lat', 'lon'])
<xarray.DataArray 'tos' (time: 180)> Size: 720B array([603767. , 607702.5 , 603976.5 , 599373.56, 595119.94, 595716.75, 598177.3 , 600670.6 , 597825.56, 591869. , 590507.7 , 597189.2 , 605954.06, 609151. , 606868.9 , 602329.9 , 599465.75, 601205.5 , 605144.4 , 608588.5 , 604046.9 , 598927.75, 597519.75, 603876.9 , 612424.44, 615765.2 , 612615.44, 606310.6 , 602034.4 , 600784.9 , 602013.5 , 603142.2 , 598850.9 , 591917.44, 589234.56, 596162.5 , 602942.06, 607196.9 , 604928.2 , 601735.6 , 599011.8 , 599490.9 , 600801.44, 602786.94, 598867.2 , 594081.8 , 593736.25, 598995.6 , 607285.25, 611901.06, 609562.75, 603527.3 , 600215.4 , 601372.6 , 604144.5 , 605376.75, 601256.2 , 595245.2 , 594002.06, 600490.4 , 611878.6 , 616563. , 613050.8 , 605734. , 600808.75, 600898.06, 603930.56, 605644.7 , 599917.5 , 592048.06, 590082.8 , 596950.7 , 607701.94, 610844.7 , 609509.6 , 603380.94, 599838.1 , 600334.25, 604386.6 , 607848.1 , 602155.2 , 594949.06, 593815.06, 598365.3 , 608730.8 , 612056.5 , 609922.5 , 603077.1 , 600134.1 , 602821.2 , 606152.75, 610257.8 , 604685.8 , 596858. , 592894.8 , 599944.9 , 609764.44, 614610.75, 611434.75, 605606.4 , 603790.94, 605750.2 , 609250.06, 612935.7 , 609645.06, 601706.4 , 598896.5 , 605349.75, 614671.8 , 618686.7 , 615895.2 , 609438.2 , 605399.56, 606126.75, 607942.3 , 609680.4 , 604814.25, 595841.94, 591908.44, 595638.7 , 604798.94, 611327.1 , 609765.7 , 603727.56, 600970. , 602514. , 606303.7 , 609225.25, 603724.3 , 595944.8 , 594477.4 , 597807.4 , 607379.06, 611808.56, 610112.94, 607196.3 , 604733.06, 605488.25, 610048.3 , 612655.75, 608906.25, 602349.7 , 601754.2 , 609220.4 , 619367.1 , 623783.2 , 619949.7 , 613369.06, 610190.8 , 611091.2 , 614213.44, 615665.06, 611722.2 , 606259.56, 605970.2 , 611463.3 , 619794.6 , 626036.5 , 623085.44, 616295.9 , 611886.3 , 611881.9 , 614420.75, 616853.56, 610375.44, 603471.5 , 602108.25, 608094.3 , 617450.7 , 623508.7 , 619830.2 , 612033.3 , 608737.2 , 610105.25, 613692.7 , 616360.44, 611735.4 , 606512.7 , 604249.44, 608777.44], dtype=float32) Coordinates: * time (time) object 1kB 2000-01-15 12:00:00 ... 2014-12-15 12:00:00
For the last example in this set of aggregation examples, we compute the temporal median:
ds.tos.median(dim='time')
<xarray.DataArray 'tos' (lat: 180, lon: 360)> Size: 259kB array([[ nan, nan, nan, ..., nan, nan, nan], [ nan, nan, nan, ..., nan, nan, nan], [ nan, nan, nan, ..., nan, nan, nan], ..., [-1.7648907, -1.7648032, -1.7647004, ..., -1.7650614, -1.7650102, -1.7649589], [-1.7590305, -1.7589546, -1.7588665, ..., -1.7591925, -1.7591486, -1.759095 ], [-1.7536805, -1.753602 , -1.7535168, ..., -1.753901 , -1.753833 , -1.7537591]], dtype=float32) Coordinates: * lat (lat) float64 1kB -89.5 -88.5 -87.5 -86.5 ... 86.5 87.5 88.5 89.5 * lon (lon) float64 3kB 0.5 1.5 2.5 3.5 4.5 ... 356.5 357.5 358.5 359.5
In addition, there are many other commonly used aggregation methods in Xarray. Some of the more popular aggregation methods are summarized in the following table:
Aggregation |
Description |
---|---|
|
Total number of items |
|
Mean and median |
|
Minimum and maximum |
|
Standard deviation and variance |
|
Compute product of elements |
|
Compute sum of elements |
|
Find index of minimum and maximum value |
GroupBy: Split, Apply, Combine
While we can obtain useful summaries of datasets using simple aggregation methods, it is more often the case that aggregation must be performed over coordinate labels or groups. In order to perform this type of aggregation, it is helpful to use the split-apply-combine workflow. Fortunately, Xarray provides this functionality for DataArrays
and Datasets
by means of the groupby
operation. The following figure illustrates the split-apply-combine workflow in detail:

Based on the above figure, you can understand the split-apply-combine process performed by groupby
. In detail, the steps of this process are:
The split step involves breaking up and grouping an xarray
Dataset
orDataArray
depending on the value of the specified group key.The apply step involves computing some function, usually an aggregate, transformation, or filtering, within the individual groups.
The combine step merges the results of these operations into an output xarray
Dataset
orDataArray
.
In this set of examples, we will remove the seasonal cycle (also known as a climatology) from our dataset using groupby
. There are many types of input that can be provided to groupby
; a full list can be found in Xarray’s groupby
user guide.
In this first example, we plot data to illustrate the annual cycle described above. We first select the grid point closest to a specific latitude-longitude point. Once we have this grid point, we can plot a temporal series of sea-surface temperature (SST) data at that location. Reviewing the generated plot, the annual cycle of the data becomes clear.
ds.tos.sel(lon=310, lat=50, method='nearest').plot();

Split
The first step of the split-apply-combine process is splitting. As described above, this step involves splitting a dataset into groups, with each group matching a group key. In this example, we split the SST data using months as a group key. Therefore, there is one resulting group for January data, one for February data, etc. This code illustrates how to perform such a split:
ds.tos.groupby(ds.time.dt.month)
DataArrayGroupBy, grouped over 'month'
12 groups with labels 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12.
In the above code example, we are extracting components of date/time data by way of the time coordinate’s .dt
attribute. This attribute is a DatetimeAccessor
object that contains additional attributes for units of time, such as hour, day, and year. Since we are splitting the data into monthly data, we use the month
attribute of .dt in this example. (In addition, there exists similar functionality in Pandas; see the official documentation for details.)