xarray.Dataset conditionally indexing variables-CodePudding

Starting with a hrrr file downloaded from ncep.

read into a xarray.Dataset like...

ds: xr.Dataset = xr.open_dataset(file, engine="pynio")

`Dataset`

<xarray.Dataset>
Dimensions:                        (ygrid_0: 1059, xgrid_0: 1799, lv_HYBL0: 50,
                                    lv_HTGL1: 2, lv_HTGL2: 2, lv_TMPL3: 2,
                                    lv_SPDL4: 3, lv_HTGL5: 2, lv_HTGL6: 2,
                                    lv_DBLL7: 2, lv_HTGL8: 2, lv_HTGL9: 3)
Coordinates:
  * lv_HTGL6                       (lv_HTGL6) float32 1e 03 4e 03
  * lv_TMPL3                       (lv_TMPL3) float32 253.0 263.0
  * lv_HTGL1                       (lv_HTGL1) float32 10.0 80.0
  * lv_HYBL0                       (lv_HYBL0) float32 1.0 2.0 3.0 ... 49.0 50.0
    gridlat_0                      (ygrid_0, xgrid_0) float32 ...
    gridlon_0                      (ygrid_0, xgrid_0) float32 ...
Dimensions without coordinates: ygrid_0, xgrid_0, lv_HTGL2, lv_SPDL4, lv_HTGL5,
                                lv_DBLL7, lv_HTGL8, lv_HTGL9
Data variables: (12/149)
    TMP_P0_L1_GLC0                 (ygrid_0, xgrid_0) float32 ...
    TMP_P0_L103_GLC0               (ygrid_0, xgrid_0) float32 ...
    TMP_P0_L105_GLC0               (lv_HYBL0, ygrid_0, xgrid_0) float32 ...
    POT_P0_L103_GLC0               (ygrid_0, xgrid_0) float32 ...
    DPT_P0_L103_GLC0               (ygrid_0, xgrid_0) float32 ...
    LHTFL_P0_L1_GLC0               (ygrid_0, xgrid_0) float32 ...
    ...                             ...
    lv_HTGL5_l0                    (lv_HTGL5) float32 ...
    lv_SPDL4_l1                    (lv_SPDL4) float32 ...
    lv_SPDL4_l0                    (lv_SPDL4) float32 ...
    lv_HTGL2_l1                    (lv_HTGL2) float32 ...
    lv_HTGL2_l0                    (lv_HTGL2) float32 ...
    gridrot_0                      (ygrid_0, xgrid_0) float32 ...

for the time being I am only concerned with Variables that contain these 3 common Coordinates [lv_HYBL0, gridlat_0, gridlon_0]

I can manually select/index those Variables that have the Coordinates that I want, like....

ds[["TMP_P0_L105_GLC0",...]]

but I would prefer a more abstract method. In pandas I would do some sort of bool indexing along the lines of ... ds[ds.variables[ds.coords.isin(["gridlat_0","gridlon_0","lv_HYBL0"])]]

this unfortunately does not work.

How can I select Variables based on a condition where the Variable is tied to a Coordinate?

CodePudding user response：

You can still do something similar. You can filter a dataset’s variables using a list of keys, and determine the dimensions by testing the elements of each array’s dims attribute, which is a tuple.

In this case:

required_dims = ['lv_HYBL0', 'gridlat_0', 'gridlon_0']

#sorted tuple
required_dims = tuple(sorted(required_dims))

subset = ds[[
    k for k, v in ds.data_vars.keys()
    if tuple(sorted(v.dims)) == required_dims
]]

CodePudding user response：

I found that the drop_dims method worked sufficiently

def dont_drop(dims: Mapping, *args: str):
    a = np.array(tuple(dims.keys()))
    mask = np.all(a == np.array(args)[:, np.newaxis], axis=0)
    return a[~mask]


ds.drop_dims(dont_drop(ds.dims, "lv_HYBL0", "ygrid_0", "xgrid_0"))