Calculating BIOCLIM variables using Xarray and UKCP18 - Multivariable indexing

Question

I am currently generating several bioclimatic variables (climatic derivatives) to apply to some biodiversity work using UKCP18 data. I am generating bioclimatic variable "Bio 19": Precipitation of the Coldest Quarter (https://pubs.usgs.gov/ds/691/ds691.pdf) which uses tas and pr.

The task involves finding the 3-month rolling sum of mean temperatures, identifying the minimum of these (constituting the coldest quarter) and then extracting the total precipitation over that 3-month rolling period to obtain "Bio 19".

My issue: I can find the coldest quarter (using tas) without problem, but pr data is dropped by Xarray in the action alongside the time indexing. It means I cannot know which period to extract rainfall from, because the data is not linked in that way across variables (using my method).

Example code:

# previous code here
...
# calculate rolling climatological mean over 3 month period(s)
ds_rolling_3mos_sum = ds_monthly.rolling(time=3, center=True).sum().dropna('time')
ds_rolling_3mos_sum


Dimensions:                     (time: 12, grid_latitude: 606, grid_longitude: 484)
Coordinates:
  * time                        (time) object 1981-01-01 00:00:00 ... 1981-12...
  * grid_latitude               (grid_latitude) float64 -4.683 -4.647 ... 8.063
  * grid_longitude              (grid_longitude) float64 353.9 354.0 ... 364.3
    latitude                    (grid_latitude, grid_longitude) float64 dask.array
    longitude                   (grid_latitude, grid_longitude) float64 dask.array
Data variables:
    pr                          (time, grid_latitude, grid_longitude) float32 dask.array
    rotated_latitude_longitude  (time) float64 -6.442e+09 ... -6.442e+09
    grid_latitude_bnds          (time, grid_latitude) float64 dask.array
    grid_longitude_bnds         (time, grid_longitude) float64 dask.array
    tas                         (time, grid_latitude, grid_longitude) float32 dask.array

Find the coldest quarter:

ds_rolling_3mos_sum_tas_min = ds_rolling_3mos_sum.tas.min('time')

I now have neither time index information - which I could use to obtain the correct month for rainfall in a .sel(), nor do I have any connected pr data to access for this variable.


dask.array
Coordinates:
  * grid_latitude   (grid_latitude) float64 -4.683 -4.647 -4.611 ... 8.027 8.063
  * grid_longitude  (grid_longitude) float64 353.9 354.0 354.0 ... 364.3 364.3
    latitude        (grid_latitude, grid_longitude) float64 dask.array
    longitude       (grid_latitude, grid_longitude) float64 dask.array

I've tried using a .where statement and a few other hands on techniques but nothing working to date. I feel like there's a recognised Xarray action I'm not aware of! Please help!

user_3888719 · Answer

This sounds like a case for xarray's advanced indexing! Get excited - this is one of the most fun & powerful features of xarray in my opinion :)

Here's a quick MRE I'll use for this example:

import xarray as xr, numpy as np, pandas as pd
time = pd.date_range('1980-12-01', '1982-01-01', freq='MS')
x = np.arange(-4, 8, 0.25)
y = np.arange(353, 364, 0.25)
tas = np.random.random(size=(len(time), len(y), len(x))) * 100 + 260
pr = np.exp(np.random.random(size=(len(time), len(y), len(x)))) * 100
ds = xr.Dataset(
    {'pr': (('time', 'lat', 'lon'), pr), 'tas': (('time', 'lat', 'lon'), tas)},
    coords={'lat': y, 'lon': x, 'time': time},
)

Your dataset of rolling 3mo sums contains the temperature information about which locations within the precip array you want to select. So what we can do is use the idxmin method to build a 2D DataArray of times, indexed by lat/lon, which give the center of the coldest quarter:

In [24]: # find the time of the min rolled tas for each lat/lon
    ...: time_of_min_tas_by_cell = ds_rolling_3mos_sum.tas.idxmin(dim='time')

In [25]: time_of_min_tas_by_cell
Out[25]:

array([['1981-06-01T00:00:00.000000000', '1981-12-01T00:00:00.000000000',
        '1981-06-01T00:00:00.000000000', ...,
        '1981-03-01T00:00:00.000000000', '1981-04-01T00:00:00.000000000',
        '1981-10-01T00:00:00.000000000'],
       ...,
       ['1981-07-01T00:00:00.000000000', '1981-09-01T00:00:00.000000000',
        '1981-11-01T00:00:00.000000000', ...,
        '1981-07-01T00:00:00.000000000', '1981-12-01T00:00:00.000000000',
        '1981-09-01T00:00:00.000000000']], dtype='datetime64[ns]')
Coordinates:
  * lat      (lat) float64 353.0 353.2 353.5 353.8 ... 363.0 363.2 363.5 363.8
  * lon      (lon) float64 -4.0 -3.75 -3.5 -3.25 -3.0 ... 6.75 7.0 7.25 7.5 7.75

This can be used to index into the precip array directly to find the rolling 3mo precip total at the center of the min temp quarter:

In [30]: ds_rolling_3mos_sum.pr.sel(time=time_of_min_tas_by_cell)
Out[30]:

array([[449.50779525, 531.90472182, 747.26749901, ..., 405.24357679,
        610.08658199, 488.83666056],
       [487.65599173, 567.01802137, 380.9979117 , ..., 613.84289448,
        442.8228211 , 629.50269312],
       [432.48761645, 444.76568124, 480.11564481, ..., 464.74424834,
        543.97169369, 491.91926534],
       ...,
       [488.68368642, 455.70782431, 363.25961252, ..., 457.72558376,
        529.17600183, 438.6763931 ],
       [370.4485618 , 491.65565156, 391.47992765, ..., 689.95878533,
        585.65987576, 407.78032041],
       [576.67281438, 551.36298132, 389.643589  , ..., 366.8810199 ,
        526.52862773, 593.30879779]])
Coordinates:
  * lat      (lat) float64 353.0 353.2 353.5 353.8 ... 363.0 363.2 363.5 363.8
  * lon      (lon) float64 -4.0 -3.75 -3.5 -3.25 -3.0 ... 6.75 7.0 7.25 7.5 7.75
    time     (lat, lon) datetime64[ns] 1981-06-01 1981-12-01 ... 1981-09-01

xarray knows to reshape the result into an array indexed by (lat, lon) because the time indices are also indexed by (lat, lon). So it collapses across time, matching the indexer's lat/lon values to the source array's dims. Cool right?

Calculating BIOCLIM variables using Xarray and UKCP18 - Multivariable indexing

1 Answers