<jupyter_start><jupyter_text>Xarray 6: Computations and Masks ---  Overview1. Do basic arithmetic with DataArrays and Datasets2. Perform aggregation (reduction) along one or multiple dimensions of a DataArray or Dataset3. Compute climatology and anomaly using xarray's "split-apply-combine" approach via `.groupby()`4. Perform weighted reductions along one or multiple dimensions of a DataArray or Dataset5. Provide an overview of masking data in xarray6. Mask data using `.where()` method  Prerequisites| Concepts | Importance | Notes || --- | --- | --- || [Xarray Intro](../week7/01_Xarray_Intro.ipynb) | Necessary | |- **Time to learn**: 45 minutes ---  Imports<jupyter_code>import matplotlib.pyplot as plt
import numpy as np
import xarray as xr
from pythia_datasets import DATASETS<jupyter_output><empty_output><jupyter_text>Let's open the monthly sea surface temperature dataset from the Community Earth System Model v2 (CESM2), which is a Global Climate Model:<jupyter_code>filepath = DATASETS.fetch('CESM2_sst_data.nc')
ds = xr.open_dataset(filepath)
ds<jupyter_output><empty_output><jupyter_text>Arithmetic OperationsArithmetic operations with a single DataArray automatically apply over all array values (like NumPy). This process is called vectorization.  Let's convert the air temperature from degrees Celsius to kelvins:<jupyter_code>ds.tos + 273.15<jupyter_output><empty_output><jupyter_text>Let's square all values in `tos`:<jupyter_code>ds.tos ** 2<jupyter_output><empty_output><jupyter_text>Aggregation Methods A very common step during data analysis is to summarize the data in question by computing aggregations like `sum()`, `mean()`, `median()`, `min()`, `max()`. This *reduced data* provide insight into the nature of large datasets. Let's explore some of these aggregation methods. Compute the mean:<jupyter_code>ds.tos.mean()<jupyter_output><empty_output><jupyter_text>Because we specified no `dim` argument the function was applied over all dimensions, computing the mean of every element of `tos` across time and space. It is possible to specify a dimension along which to compute an aggregation. For example, to calculate the mean in time for all locations, specify the time dimension as the dimension along which the mean should be calculated. Using Xarray's `plot` function, visualize the mean SST.<jupyter_code>ds.tos.mean(dim='time').plot(size=7);<jupyter_output><empty_output><jupyter_text>Compute the temporal min:<jupyter_code>ds.tos.min(dim=['time'])<jupyter_output><empty_output><jupyter_text>Compute the spatial sum:<jupyter_code>ds.tos.sum(dim=['lat', 'lon'])<jupyter_output><empty_output><jupyter_text>Compute the temporal median:<jupyter_code>ds.tos.median(dim='time')<jupyter_output><empty_output><jupyter_text>The following table summarizes some other built-in xarray aggregations:| Aggregation              | Description                     ||--------------------------|---------------------------------|| ``count()``              | Total number of items           || ``mean()``, ``median()`` | Mean and median                 || ``min()``, ``max()``     | Minimum and maximum             || ``std()``, ``var()``     | Standard deviation and variance || ``prod()``               | Compute product of elements            || ``sum()``                | Compute sum of elements                || ``argmin()``, ``argmax()``| Find index of minimum and maximum value |  GroupBy: Split, Apply, CombineSimple aggregations can give useful summaries of our dataset, but often we would prefer to aggregate conditionally on some coordinate labels or groups. Xarray provides the so-called `groupby` operation which enables the **split-apply-combine** workflow on Xarray DataArrays and Datasets. The split-apply-combine operation is illustrated in this figure:  First, let's select a gridpoint closest to a specified lat-lon, and plot a time series of SST at that point. The annual cycle will be quite evident.<jupyter_code>ds.tos.sel(lon=310, lat=50, method='nearest').plot();<jupyter_output><empty_output><jupyter_text>SplitLet's group data by month, i.e. all Januaries in one group, all Februaries in one group, etc.<jupyter_code>ds.tos.groupby(ds.time.dt.month)<jupyter_output><empty_output><jupyter_text>In the above example, we are using datetime's DatetimeAccessor function to extract specific components of dates/times in our time coordinate dimension. For example, we can extract the year with: ds.time.dt.year. Xarray also offers a more concise syntax when the variable you’re grouping on is already present in the dataset. This is identical to `ds.tos.groupby(ds.time.dt.month)`:<jupyter_code>ds.tos.groupby('time.month')<jupyter_output><empty_output><jupyter_text>Apply & Combine Now that we have groups defined, it’s time to “apply” a calculation to the group. These calculations can either be:- aggregation: reduces the size of the group- transformation: preserves the group’s full sizeAt then end of the apply step, xarray will automatically combine the aggregated/transformed groups back into a single object.  Compute climatology Let's calculate the climatology at every point in the dataset:<jupyter_code>tos_clim = ds.tos.groupby('time.month').mean()
tos_clim<jupyter_output><empty_output><jupyter_text>Plot climatology at a specific point:<jupyter_code>tos_clim.sel(lon=310, lat=50, method='nearest').plot();<jupyter_output><empty_output><jupyter_text>Plot zonal mean climatology:<jupyter_code>tos_clim.mean(dim='lon').transpose().plot.contourf(levels=12, robust=True, cmap='turbo');<jupyter_output><empty_output><jupyter_text>Calculate and plot the difference between January and December climatologies:<jupyter_code>(tos_clim.sel(month=1) - tos_clim.sel(month=12)).plot(size=6, robust=True);<jupyter_output><empty_output><jupyter_text>Compute anomalyNow let's combine the previous steps to compute climatology and use xarray's `groupby` arithmetic to remove this climatology from our original data:<jupyter_code>gb = ds.tos.groupby('time.month')
tos_anom = gb - gb.mean(dim='time')
tos_anom
tos_anom.sel(lon=310, lat=50, method='nearest').plot();<jupyter_output><empty_output><jupyter_text>Let's compute and visualize the mean global anomaly over time. We need to specify both `lat` and `lon` dimensions in the `dim` argument to `mean()`:<jupyter_code>unweighted_mean_global_anom = tos_anom.mean(dim=['lat', 'lon'])
unweighted_mean_global_anom.plot();<jupyter_output><empty_output><jupyter_text>Note: An operation which combines grid cells of different size is not scientifically valid unless each cell is weighted by the size of the grid cell. xarray has a convenient weighted method to accomplish this. Let's first load the cell area data, from another CESM2 dataset that contains the weights for the grid cells:<jupyter_code>filepath2 = DATASETS.fetch('CESM2_grid_variables.nc')
areacello = xr.open_dataset(filepath2).areacello
areacello<jupyter_output><empty_output><jupyter_text>As before, let's calculate area-weighted mean global anomaly:<jupyter_code>weighted_mean_global_anom = tos_anom.weighted(areacello).mean(dim=['lat', 'lon'])<jupyter_output><empty_output><jupyter_text>Let's plot both unweighted and weighted means:<jupyter_code>unweighted_mean_global_anom.plot(size=7)
weighted_mean_global_anom.plot()
plt.legend(['unweighted', 'weighted']);<jupyter_output><empty_output><jupyter_text>Other high level computation functionality- `resample`: [Groupby-like functionality specifialized for time dimensions. Can be used for temporal upsampling and downsampling](https://xarray.pydata.org/en/stable/user-guide/time-series.htmlresampling-and-grouped-operations)- `rolling`: [Useful for computing aggregations on moving windows of your dataset e.g. computing moving averages](https://xarray.pydata.org/en/stable/user-guide/computation.htmlrolling-window-operations)- `coarsen`: [Generic functionality for downsampling data](https://xarray.pydata.org/en/stable/user-guide/computation.htmlcoarsen-large-arrays) For example, resample to annual frequency:<jupyter_code>r = ds.tos.resample(time='AS')
r
r.mean()<jupyter_output><empty_output><jupyter_text>Compute a 5-month moving average:<jupyter_code>m_avg = ds.tos.rolling(time=5, center=True).mean()
m_avg
lat = 50
lon = 310

m_avg.isel(lat=lat, lon=lon).plot(size=6)
ds.tos.isel(lat=lat, lon=lon).plot()
plt.legend(['5-month moving average', 'monthly data']);<jupyter_output><empty_output><jupyter_text>Masking Data Using the `xr.where()` or `.where()` method, elements of an xarray Dataset or xarray DataArray that satisfy a given condition or multiple conditions can be replaced/masked. To demonstrate this, we are going to use the `.where()` method on the `tos` DataArray.  We will use the same sea surface temperature dataset:<jupyter_code>ds<jupyter_output><empty_output><jupyter_text>Using `where` with one condition Imagine we wish to analyze just the last time in the dataset. We could of course use `isel` for this:<jupyter_code>sample = ds.tos.isel(time=-1)
sample<jupyter_output><empty_output><jupyter_text>Unlike `.isel()` and `.sel()` that change the shape of the returned results, `.where()` preserves the shape of the original data. It accomplishes this by returning values from the original DataArray or Dataset if the `condition` is `True`, and fills in values (by default `nan`) wherever the `condition` is `False`. Before applying it, let's look at the [`.where()` documentation](http://xarray.pydata.org/en/stable/generated/xarray.DataArray.where.html). As the documention points out, the conditional expression in `.where()` can be: - a DataArray- a Dataset- a functionFor demonstration purposes, let's use `.where()` to mask locations with temperature values greater than `0`:<jupyter_code>masked_sample = sample.where(sample < 0.0)
masked_sample<jupyter_output><empty_output><jupyter_text>Let's plot both our original sample, and the masked sample:<jupyter_code>fig, axes = plt.subplots(ncols=2, figsize=(19, 6))
sample.plot(ax=axes[0], robust=True)
masked_sample.plot(ax=axes[1], robust=True);
fig, axes = plt.subplots(ncols=2, figsize=(19, 6))
sample.plot(ax=axes[0])
masked_sample.plot(ax=axes[1]);<jupyter_output><empty_output><jupyter_text>Using `where` with multiple conditions `.where()` allows providing multiple conditions. To do this, we need to make sure each conditional expression is enclosed in `()`. To combine conditions, we use the `bit-wise and` (`&`) operator and/or the `bit-wise or` (`|`). Let's use `.where()` to mask locations with temperature values less than 25 and greater than 30:<jupyter_code>sample.where((sample > 25) & (sample < 30)).plot(size=6, robust=True);
sample.where((sample > 25) & (sample < 30)).plot(size=6);<jupyter_output><empty_output><jupyter_text>We can use coordinates to apply a mask as well. Below, we use the `latitude` and `longitude` coordinates to mask everywhere outside of the [Niño 3.4 region](https://www.ncdc.noaa.gov/teleconnections/enso/indicators/sst/):![](https://www.ncdc.noaa.gov/monitoring-content/teleconnections/nino-regions.gif)<jupyter_code>sample.where(
    (sample.lat < 5) & (sample.lat > -5) & (sample.lon > 190) & (sample.lon < 240)
).plot(size=6, robust=True);<jupyter_output><empty_output><jupyter_text>Using `where` with a custom fill value `.where()` can take a second argument, which, if supplied, defines a *fill value* for the masked region. Below we fill masked regions with a constant `0`:<jupyter_code>sample.where((sample > 25) & (sample < 30), 0).plot(size=6, robust=False);<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11770833333333333, 'fraction_numerical': 0.009114583333333334, 'mean_word_length': 5.932009626955475, 'pattern_counts': {'":': 0, '<': 127, '<?xml version=': 0, '>': 126, 'https://': 5, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '30', 'dir': 'jupyter-structured-clean-dedup', 'id': '304501', 'license': 'CC0-1.0', 'n_tokens_mistral': 3674, 'n_tokens_neox': 3496, 'n_words': 1089, 'path': '06_Xarray_computation-masking.ipynb', 'repo_name': 'ktyle/binder_xarray', 'size': '11520'}