Selecting the value at a given date for each lat/lon point in xarray

user_16662504 0
Viewed 164

I have a xr.DataArray object that has a day of 2015 (as a cftime.DateTimeNoLeap object) for each lat-lon point on the grid.

date_matrix2015

    <xarray.DataArray (lat: 160, lon: 320)>
array([[cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0), ...,
        cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0)],
       [cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0), ...,
        cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0)],
       [cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0), ...,
        cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 12, 11, 12, 0, 0, 0)],
       ...,
       [cftime.DatetimeNoLeap(2015, 3, 14, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 3, 14, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 3, 14, 12, 0, 0, 0), ...,
        cftime.DatetimeNoLeap(2015, 9, 16, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 9, 16, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 9, 16, 12, 0, 0, 0)],
       [cftime.DatetimeNoLeap(2015, 9, 15, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 9, 15, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 9, 15, 12, 0, 0, 0), ...,
        cftime.DatetimeNoLeap(2015, 9, 16, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 9, 15, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 9, 15, 12, 0, 0, 0)],
       [cftime.DatetimeNoLeap(2015, 9, 16, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 9, 16, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 9, 16, 12, 0, 0, 0), ...,
        cftime.DatetimeNoLeap(2015, 9, 16, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 9, 16, 12, 0, 0, 0),
        cftime.DatetimeNoLeap(2015, 9, 16, 12, 0, 0, 0)]], dtype=object)
Coordinates:
    year     int64 2015
  * lat      (lat) float64 -89.14 -88.03 -86.91 -85.79 ... 86.91 88.03 89.14
  * lon      (lon) float64 0.0 1.125 2.25 3.375 4.5 ... 355.5 356.6 357.8 358.9

I have another xr.DataArray on the same lat-lon grid for vertical velocity (omega) that has data for every day in 2015. At each lat-lon point I would like to select the velocity value on the corresponding day given in date_matrix2015. Ideally I would like to do something like this: omega.sel(time=date_matrix2015)

I have tried constructing the new dataarray manually with iteration, but I haven't had much luck.

Does anyone have any ideas? Thank you in advance!

------------EDIT---------------

Here is a minimal reproducible example for the problem. To clarify what I am looking for: I have two DataArrays, one for daily precipitation values, and one for daily omega values. I want to determine for each lat/lon point the day that saw the maximum precipitation (I think I have done this part correctly). From there I want to select at each lat/lon point the omega value that occurred on the day of maximum precipitation. So ultimately I would like to end up with a DataArray of omega values that has two dimensions, lat and lon, where the value at each lat/lon point is the omega value on the day of maximum rainfall at that location.

import numpy as np
import xarray as xr
import pandas as pd

precip = np.abs(8*np.random.randn(10,10,10))
omega = 15*np.random.randn(10,10,10)
lat = np.arange(0,10)
lon = np.arange(0, 10)
##Note: actual data resolution is 160x360
dates = pd.date_range('01-01-2015', '01-10-2015')

precip_da = xr.DataArray(precip).rename({'dim_0':'time', 'dim_1':'lat', 'dim_2':'lon'}).assign_coords({'time':dates, 'lat':lat, 'lon':lon})
omega_da = xr.DataArray(omega).rename({'dim_0':'time', 'dim_1':'lat', 'dim_2':'lon'}).assign_coords({'time':dates, 'lat':lat, 'lon':lon})

#Find Date of maximum precip for each lat lon point and store in an array
maxDateMatrix = precip_da.idxmax(dim='time')


#For each lat lon point, select the value from omega_da on the day of maximum precip (i.e. the date given at that location in the maxDateMatrix)
python-xarray indexing
1 Answers

You can pair da.sel with da.idxmax to select the index of the maxima along any number of dimensions:

In [10]: omega_da.sel(time=precip_da.idxmax(dim='time'))
Out[10]:
<xarray.DataArray (lat: 10, lon: 10)>
array([[ 17.72211193, -16.20781517,   9.65493368, -28.16691093,
         18.8756182 ,  16.81924325, -20.55251804, -18.36625778,
        -19.57938236, -10.77385357],
       [  3.95402784,  -5.28478105,  -8.6632994 ,   2.46787932,
         20.53981254,  -4.74908659,   9.5274101 ,  -1.08191372,
          9.4637305 , -10.91884369],
       [-31.30033085,   6.6284144 ,   8.15945444,   5.74849304,
         12.49505739,   2.11797825, -18.12861347,   7.27497695,
          5.16197504, -32.99882591],
       ...
       [-34.73945635,  24.40515233,  14.56982584,  12.16550083,
         -8.3558104 , -20.16328749, -33.89051472,  -0.09599935,
          2.65689584,  29.54056082],
       [-18.8660847 ,  -7.58120994,  15.57632568,   4.19142695,
          8.71046261,   9.05684805,   8.48128361,   0.34166869,
          8.41090015,  -2.31386572],
       [ -4.38999926,  17.00411671,  16.66619606,  24.99390669,
        -14.01424591,  19.85606151, -16.87897   ,  12.84205521,
        -16.78824975,  -6.33920671]])
Coordinates:
    time     (lat, lon) datetime64[ns] 2015-01-01 2015-01-01 ... 2015-01-10
  * lat      (lat) int64 0 1 2 3 4 5 6 7 8 9
  * lon      (lon) int64 0 1 2 3 4 5 6 7 8 9

See the great section of the xarray docs on Indexing and Selecting Data for more info, especially the section on Advanced Indexing, which goes into using DataArrays as indexers for powerful reshaping operations.

user_38887190
Related