zxdawn / S5P-WRFChem

Core code for the TROPOMI NO2 retrieval based on WRF-Chem outputs
GNU General Public License v3.0
13 stars 2 forks source link

Better method of sorting pressure levels #3

Closed zxdawn closed 3 years ago

zxdawn commented 4 years ago

Feature Request

Currently, the method of sorting pressure levels is slow and Series indexes are reassigned again. This is related to https://github.com/pandas-dev/pandas/issues/33420 and https://github.com/pydata/xarray/issues/3957.

zxdawn commented 3 years ago

A simple method of using pure xarray:

    sort_index = (-1*s5p_pcld.load()).argsort(axis=0) # (layer, y, x)
    s5p_pcld = xr.DataArray(np.take_along_axis(s5p_pcld.values, sort_index, axis=0),
                            dims=['plevel', 'y', 'x'])

However, the result is different from the pandas one ... I need to check it in detail.

image

zxdawn commented 3 years ago

The difference mentioned above is caused by the missing coordinates of plevel in the xarray method. Instead, we should add it:

    # get the sorting index
    sort_index = (-1*s5p_pcld.load()).argsort(axis=0) # (layer, y, x)
    # sort pressures
    s5p_pcld = xr.DataArray(np.take_along_axis(s5p_pcld.values, sort_index, axis=0),
                            dims=['plevel', 'y', 'x'])
    # assign plevel coordinates
    s5p_pcld = s5p_pcld.assign_coords(plevel=range(s5p_pcld.sizes['plevel']))

image