<jupyter_start><jupyter_text>Computation with XarrayIn this lesson, we discuss how to do scientific computations with xarrayobjects. Our learning goals are as follows. By the end of the lesson, we will beable to:- Apply basic arithmetic and numpy functions to xarray DataArrays / Dataset.- Use Xarray's label-aware reduction operations (e.g. `mean`, `sum`) weighted  reductions.- Apply arbitrary functions to Xarray data via `apply_ufunc`.- Use Xarray's broadcasting to compute on arrays of different dimensionality.- Perform "split / apply / combine" workflows in Xarray using `groupby`,  including  - reductions within groups  - transformations on groups- Use the `resample`, `rolling` and `coarsen` functions to manipulate data.<jupyter_code>import expectexception
import numpy as np
import xarray as xr
from matplotlib import pyplot as plt<jupyter_output><empty_output><jupyter_text>Example DatasetFirst we load a dataset. We will use the[NOAA Extended Reconstructed Sea Surface Temperature (ERSST) v5](https://www.ncdc.noaa.gov/data-access/marineocean-data/extended-reconstructed-sea-surface-temperature-ersst-v5)product, a widely used and trusted gridded compilation of of historical datagoing back to 1854.Since the data is provided via an[OPeNDAP](https://en.wikipedia.org/wiki/OPeNDAP) server, we can load it directlywithout downloading anything:<jupyter_code>### NOTE: If hundreds of people connect to this server at once and download the same dataset,
###       things might not go so well! Recommended to use the Google Cloud copy instead.

# url = "http://www.esrl.noaa.gov/psd/thredds/dodsC/Datasets/noaa.ersst.v5/sst.mnmean.nc"
# # drop an unnecessary variable which complicates some operations
# ds = xr.open_dataset(url, drop_variables=["time_bnds"])
# # will take a minute or two to complete
# ds = ds.sel(time=slice("1960", "2018")).load()
# ds
import gcsfs

fs = gcsfs.GCSFileSystem(token="anon")
ds = xr.open_zarr(
    fs.get_mapper("gs://pangeo-noaa-ncei/noaa.ersst.v5.zarr"), consolidated=True
).load()
ds<jupyter_output><empty_output><jupyter_text>Let's do some basic visualizations of the data, just to make sure it looksreasonable.<jupyter_code>ds.sst[0].plot(vmin=-2, vmax=30)<jupyter_output><empty_output><jupyter_text>Basic ArithmeticXarray dataarrays and datasets work seamlessly with arithmetic operators andnumpy array functions.For example, imagine we want to convert the temperature (given in Celsius) toKelvin:<jupyter_code>sst_kelvin = ds.sst + 273.15
sst_kelvin<jupyter_output><empty_output><jupyter_text>The dimensions and coordinates were preseved following the operation.    Warning: Although many xarray datasets have a units attribute, which is used in plotting,    Xarray does not inherently understand units. However, work is underway to integrate xarray    with pint, which will provide full unit-aware operations.We can apply more complex functions, including numpy ufuncs, to Xarray objects.Imagine we wanted to compute the following expression as a function of SST($\Theta$) in Kelvin:$$ f(\Theta) =  0.5 \ln(\Theta^2) $$<jupyter_code>f = 0.5 * np.log(sst_kelvin ** 2)
f<jupyter_output><empty_output><jupyter_text>Applying Aribtrary FunctionsIt's awesome that we can call `np.log(ds)` and have it "just work". However, notall third party libraries work this way.In this example, we will use functions from the[Gibbs Seawater Toolkit](https://teos-10.github.io/GSW-Python/), a package forthe thermodynamics of seawater. This package provides ufuncs that operate onnumpy arrays.<jupyter_code>import gsw

# an example function
# http://www.teos-10.org/pubs/gsw/html/gsw_t90_from_t68.html
gsw.t90_from_t68?
gsw.t90_from_t68(ds.sst)  # -> returns a numpy array<jupyter_output><empty_output><jupyter_text>It would be nice to keep our dimensions and coordinates. We can accomplish thiswith `xr.apply_ufunc`.<jupyter_code>xr.apply_ufunc(gsw.t90_from_t68, ds.sst)<jupyter_output><empty_output><jupyter_text>Note: apply_ufunc is a powerful and mysterious function.    It has many options for doing more complicated things.    Unfortunately, we don't have time to go into more depth here.    Please consult the [Xarray docs](http://xarray.pydata.org/en/latest/generated/xarray.apply_ufunc.html) for more details.  ReductionsJust like in numpy, we can reduce xarray DataArrays along any number of axes:<jupyter_code>sst = ds.sst
sst.mean(axis=0)
sst.mean(axis=(1, 2))
sst.mean()<jupyter_output><empty_output><jupyter_text>However, rather than performing reductions on axes (as in numpy), we can performthem on dimensions. This turns out to be a huge convenience, particularly incomplex calculations when you can't easily remember which axis corresponds towhich dimension:<jupyter_code>sst.mean(dim="time")<jupyter_output><empty_output><jupyter_text>All of the standard numpy reductions (e.g. `min`, `max`, `sum`, `std`, etc.) areavailable  ExerciseTake the mean of `sst` in both longitude and latitude. Make a simple timeseriesplot:<jupyter_code># your code here<jupyter_output><empty_output><jupyter_text>BroadcastingBroadcasting refers to the alignmed of arrays with different numbers ofdimensions. Numpy's broadcasting rules, based on array shape, can sometimes bedifficult to understand and remember. Xarray does brodcasting by dimension name,rather than array shape. This is a huge convenience.Let's now create two arrays with some dimensions in common. For this example, wewill create a "weights" array proportional to cosine of latitude. Modulo anormalization, this is the correct area-weighting factor for data on a regularlat-lon grid.<jupyter_code>weights = np.cos(np.deg2rad(ds.lat))
weights.dims<jupyter_output><empty_output><jupyter_text>If we multiply this by SST, it "just works," and the arrays are broadcastedproperly:<jupyter_code>(ds.sst * weights).dims<jupyter_output><empty_output><jupyter_text>Warning: If the arrays being broadcasted share a dimension name, but have different coordinates,    they will first be aligned using Xarray's default align settings (including filling missing values with NaNs).    If that's not what you want, it's best to call align explicitly before broadcasting.  Weighted ReductionsWe could imagine computing the weighted spatial mean of SST manually, like this:<jupyter_code>sst_mean = (ds.sst * weights).sum(dim=("lon", "lat")) / weights.sum(dim="lat")
sst_mean.plot()
plt.title("This is wrong!")<jupyter_output><empty_output><jupyter_text>That would be wrong, however, because the denominator (`weights.sum(dim='lat')`)needs to be expanded to include the `lon` dimension and modified to account forthe missing values (land points).In general, weighted reductions on multidimensional arrays are complicated. Tomake it a bit easier, Xarray provides a mechanism for weighted reductions. Itdoes this by creating a special intermediate `DataArrayWeighted` object, towhich different reduction operations can applied.<jupyter_code>sst_weighted = ds.sst.weighted(weights)
sst_weighted
sst_weighted.mean(dim=("lon", "lat")).plot()
plt.title("Correct Global Mean SST")<jupyter_output><empty_output><jupyter_text>GroupbyXarray copies Pandas' very useful groupby functionality, enabling the "split /apply / combine" workflow on xarray DataArrays and Datasets.To provide a physically motivated example, let's examine a timeseries of SST ata single point.<jupyter_code>ds.sst.sel(lon=300, lat=50).plot()<jupyter_output><empty_output><jupyter_text>As we can see from the plot, the timeseries at any one point is totallydominated by the seasonal cycle. We would like to remove this seasonal cycle(called the "climatology") in order to better see the long-term variaitions intemperature. We can accomplish this using **groupby**.Before moving forward, w note that xarray correctly parsed the time index,resulting in a Pandas datetime index on the time dimension.<jupyter_code>ds.time<jupyter_output><empty_output><jupyter_text>The syntax of Xarray's groupby is almost identical to Pandas.<jupyter_code>ds.groupby?<jupyter_output><empty_output><jupyter_text>Split StepThe most important argument is `group`: this defines the unique values we willus to "split" the data for grouped analysis. We can pass either a DataArray or aname of a variable in the dataset. Lets first use a DataArray. Just like withPandas, we can use the time indexe to extract specific components of dates andtimes. Xarray uses a special syntax for this `.dt`, called the`DatetimeAccessor`.<jupyter_code>ds.time.dt
ds.time.dt.month<jupyter_output><empty_output><jupyter_text>ds.time.dt.year We can use these arrays in a groupby operation:<jupyter_code>gb = ds.groupby(ds.time.dt.month)
gb<jupyter_output><empty_output><jupyter_text>Xarray also offers a more concise syntax when the variable you're grouping on isalready present in the dataset. This is identical to the previous line:<jupyter_code>gb = ds.groupby("time.month")
gb<jupyter_output><empty_output><jupyter_text>Now that the data are split, we can manually iterate over the group. Theiterator returns the key (group name) and the value (the actual datasetcorresponding to that group) for each group.<jupyter_code>for group_name, group_ds in gb:
    # stop iterating after the first loop
    break
print(group_name)
group_ds<jupyter_output><empty_output><jupyter_text>Apply & CombineNow that we have groups defined, it's time to "apply" a calculation to thegroup. Like in Pandas, 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.The most fundamental way to apply is with the `.map` method.<jupyter_code>gb.map?<jupyter_output><empty_output><jupyter_text>Aggregations`.apply` accepts as its argument a function. We can pass an existing function:<jupyter_code>gb.map(np.mean)<jupyter_output><empty_output><jupyter_text>Because we specified no extra arguments (like `axis`) the function was appliedover all space and time dimensions. This is not what we wanted. Instead, wecould define a custom function. This function takes a single argument--the groupdataset--and returns a new dataset to be combined:<jupyter_code>def time_mean(a):
    return a.mean(dim="time")


gb.map(time_mean)<jupyter_output><empty_output><jupyter_text>Like Pandas, xarray's groupby object has many built-in aggregation operations(e.g. `mean`, `min`, `max`, `std`, etc):<jupyter_code># this does the same thing as the previous cell
ds_mm = gb.mean(dim="time")
ds_mm<jupyter_output><empty_output><jupyter_text>So we did what we wanted to do: calculate the climatology at every point in thedataset. Let's look at the data a bit._Climatlogy at a specific point in the North Atlantic_<jupyter_code>ds_mm.sst.sel(lon=300, lat=50).plot()<jupyter_output><empty_output><jupyter_text>_Zonal Mean Climatolgoy_<jupyter_code>ds_mm.sst.mean(dim="lon").plot.contourf(x="month", levels=12, vmin=-2, vmax=30)<jupyter_output><empty_output><jupyter_text>_Difference between January and July Climatology_<jupyter_code>(ds_mm.sst.sel(month=1) - ds_mm.sst.sel(month=7)).plot(vmax=10)<jupyter_output><empty_output><jupyter_text>TransformationsNow we want to _remove_ this climatology from the dataset, to examine theresidual, called the _anomaly_, which is the interesting part from a climateperspective. Removing the seasonal climatology is a perfect example of atransformation: it operates over a group, but doesn't change the size of thedataset. Here is one way to code it<jupyter_code>def remove_time_mean(x):
    return x - x.mean(dim="time")


ds_anom = ds.groupby("time.month").map(remove_time_mean)
ds_anom<jupyter_output><empty_output><jupyter_text>Xarray makes these sorts of transformations easy by supporting _groupbyarithmetic_. This concept is easiest explained with an example:<jupyter_code>gb = ds.groupby("time.month")
ds_anom = gb - gb.mean(dim="time")
ds_anom<jupyter_output><empty_output><jupyter_text>Now we can view the climate signal without the overwhelming influence of theseasonal cycle._Timeseries at a single point in the North Atlantic_<jupyter_code>ds_anom.sst.sel(lon=300, lat=50).plot()<jupyter_output><empty_output><jupyter_text>_Difference between Jan. 1 2018 and Jan. 1 1960_<jupyter_code>(ds_anom.sel(time="2018-01-01") - ds_anom.sel(time="1960-01-01")).sst.plot()<jupyter_output><empty_output><jupyter_text>Grouby-Related: Resample, Rolling, CoarsenResample in xarray is nearly identical to Pandas. It is effectively a group-byoperation, and uses the same basic syntax. It can be applied only to time-indexdimensions. Here we compute the five-year mean.<jupyter_code>resample_obj = ds_anom.resample(time="5Y")
resample_obj
ds_anom_resample = resample_obj.mean(dim="time")
ds_anom_resample
ds_anom.sst.sel(lon=300, lat=50).plot()
ds_anom_resample.sst.sel(lon=300, lat=50).plot(marker="o")<jupyter_output><empty_output><jupyter_text>Note: resample only works with proper datetime indexes.Rolling is also similar to pandas, but can be applied along any dimension. Itworks with logical coordinates.<jupyter_code>ds_anom_rolling = ds_anom.rolling(time=12, center=True).mean()
ds_anom_rolling
ds_anom.sst.sel(lon=300, lat=50).plot(label="monthly anom")
ds_anom_resample.sst.sel(lon=300, lat=50).plot(
    marker="o", label="5 year resample"
)
ds_anom_rolling.sst.sel(lon=300, lat=50).plot(label="12 month rolling mean")
plt.legend()<jupyter_output><empty_output><jupyter_text>`coarsen` does something similar to `resample`, but without being aware of time.It operates on logical coordinates only but can work on multiple dimensions at atime.<jupyter_code>ds_anom_coarsen_time = ds_anom.coarsen(time=12).mean()

ds_anom_rolling.sst.sel(lon=300, lat=50).plot(label="12 month rolling mean")
ds_anom_coarsen_time.sst.sel(lon=300, lat=50).plot(
    marker="^", label="12 item coarsen"
)
plt.legend()
%%expect_exception
ds_anom_coarsen_space = ds_anom.coarsen(lon=4, lat=4).mean()
ds_anom_coarsen_space = (
    ds_anom.isel(lat=slice(0, -1)).coarsen(lon=4, lat=4).mean()
)
ds_anom_coarsen_space
ds_anom_coarsen_space.sst.isel(time=0).plot()<jupyter_output><empty_output><jupyter_text>ExerciseLoad the following "basin mask" dataset, and use it to take a weighted averageof SST in each ocean basin. Figure out which ocean basins are the warmest andcoldest.**Hint:** you will first need to align this dataset with the SST dataset. Usewhat you learned in the "indexing and alignment" lesson.<jupyter_code>basin = xr.open_dataset(
    "http://iridl.ldeo.columbia.edu/SOURCES/.NOAA/.NODC/.WOA09/.Masks/.basin/dods"
)
basin<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.09210796036897848, 'fraction_numerical': 0.011001024940211821, 'mean_word_length': 6.358471593765711, 'pattern_counts': {'":': 0, '<': 145, '<?xml version=': 0, '>': 146, 'https://': 3, 'lorem ipsum': 0, 'www.': 3, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '36', 'dir': 'jupyter-structured-clean-dedup', 'id': '243080', 'license': 'Apache-2.0', 'n_tokens_mistral': 4768, 'n_tokens_neox': 4558, 'n_words': 1481, 'path': 'scipy-tutorial/03_computation_with_xarray.ipynb', 'repo_name': 'nkorinek/xarray-tutorial', 'size': '14635'}