Add choose_any and choose_all methods

[andersy005]

In [1]: import xarray as xr, xcollection

In [2]: ds = xr.tutorial.open_dataset('rasm')

In [3]: ds.attrs = {}

In [4]: dsa = xr.tutorial.open_dataset('air_temperature')

In [5]: dsa.attrs = {}

In [6]: c = xcollection.Collection({'foo': ds, 'bar': dsa})

In [7]: c
Out[7]: 
<Collection (2 keys)>
🔑 foo
<xarray.Dataset>
Dimensions:  (time: 36, y: 205, x: 275)
Coordinates:
  * time     (time) object 1980-09-16 12:00:00 ... 1983-08-17 00:00:00
    xc       (y, x) float64 ...
    yc       (y, x) float64 ...
Dimensions without coordinates: y, x
Data variables:
    Tair     (time, y, x) float64 ...

🔑 bar
<xarray.Dataset>
Dimensions:  (lat: 25, time: 2920, lon: 53)
Coordinates:
  * lat      (lat) float32 75.0 72.5 70.0 67.5 65.0 ... 25.0 22.5 20.0 17.5 15.0
  * lon      (lon) float32 200.0 202.5 205.0 207.5 ... 322.5 325.0 327.5 330.0
  * time     (time) datetime64[ns] 2013-01-01 ... 2014-12-31T18:00:00
Data variables:
    air      (time, lat, lon) float32 ...

• Choose ANY

In [8]: c.choose_any('air')
Out[8]: 
<Collection (1 keys)>
🔑 bar
<xarray.Dataset>
Dimensions:  (lat: 25, time: 2920, lon: 53)
Coordinates:
  * lat      (lat) float32 75.0 72.5 70.0 67.5 65.0 ... 25.0 22.5 20.0 17.5 15.0
  * lon      (lon) float32 200.0 202.5 205.0 207.5 ... 322.5 325.0 327.5 330.0
  * time     (time) datetime64[ns] 2013-01-01 ... 2014-12-31T18:00:00
Data variables:
    air      (time, lat, lon) float32 ...


In [9]: c.choose_any(['air', 'Tair'])
Out[9]: <Collection (0 keys)>

• Choose ALL

In [10]: c.choose_all('air')
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
~/.mambaforge/envs/xcollection-dev/lib/python3.9/site-packages/xarray/core/dataset.py in _copy_listed(self, names)
   1362             try:
-> 1363                 variables[name] = self._variables[name]
   1364             except KeyError:

KeyError: 'air'

During handling of the above exception, another exception occurred:

KeyError                                  Traceback (most recent call last)
~/devel/ncar/xcollection/xcollection/main.py in _select_vars(dset)
     92             try:
---> 93                 return dset[data_vars]
     94             except KeyError:

~/.mambaforge/envs/xcollection-dev/lib/python3.9/site-packages/xarray/core/dataset.py in __getitem__(self, key)
   1503         else:
-> 1504             return self._copy_listed(key)
   1505 

~/.mambaforge/envs/xcollection-dev/lib/python3.9/site-packages/xarray/core/dataset.py in _copy_listed(self, names)
   1364             except KeyError:
-> 1365                 ref_name, var_name, var = _get_virtual_variable(
   1366                     self._variables, name, self._level_coords, self.dims

~/.mambaforge/envs/xcollection-dev/lib/python3.9/site-packages/xarray/core/dataset.py in _get_virtual_variable(variables, key, level_vars, dim_sizes)
    172     else:
--> 173         ref_var = variables[ref_name]
    174 

KeyError: 'air'

During handling of the above exception, another exception occurred:

KeyError                                  Traceback (most recent call last)
<ipython-input-10-056950f9daff> in <module>
----> 1 c.choose_all('air')

~/devel/ncar/xcollection/xcollection/main.py in choose_all(self, data_vars)
     95                 raise KeyError(f'No data variables: `{data_vars}` found in dataset: {dset!r}')
     96 
---> 97         result = toolz.valmap(_select_vars, self.datasets)
     98         return type(self)(datasets=result)
     99 

~/.mambaforge/envs/xcollection-dev/lib/python3.9/site-packages/toolz/dicttoolz.py in valmap(func, d, factory)
     81     """
     82     rv = factory()
---> 83     rv.update(zip(d.keys(), map(func, d.values())))
     84     return rv
     85 

~/devel/ncar/xcollection/xcollection/main.py in _select_vars(dset)
     93                 return dset[data_vars]
     94             except KeyError:
---> 95                 raise KeyError(f'No data variables: `{data_vars}` found in dataset: {dset!r}')
     96 
     97         result = toolz.valmap(_select_vars, self.datasets)

KeyError: "No data variables: `['air']` found in dataset: <xarray.Dataset>\nDimensions:  (time: 36, y: 205, x: 275)\nCoordinates:\n  * time     (time) object 1980-09-16 12:00:00 ... 1983-08-17 00:00:00\n    xc       (y, x) float64 ...\n    yc       (y, x) float64 ...\nDimensions without coordinates: y, x\nData variables:\n    Tair     (time, y, x) float64 ..."

[andersy005]

Note that with choose_any and choose_all, the returned value is always a Collection of datasets due to this promotion:

https://github.com/NCAR/xcollection/blob/973126a20e011f3030ceffab0d6c2c2ba8431365/xcollection/main.py#L11-L12

[dcherian]

Note that with choose_any and choose_all, the returned value is always a Collection of datasets

I think this is perfect. We can add to_dataset and to_array methods if users want something else.

[dcherian]

Shall we do

def choose(self, data_vars, *, mode="all")

and mode could be "any".

We should add an equivalent to Dataset.filter_by_attrs but with a different name potentially to allow users to filter by standard_name etc. iris may also have some related functionality that we could copy.

[mnlevy1981]

Should the default be mode="any" and then the user needs to explicitly specify when they expect the variable to be in every dataset?

[andersy005]

Should the default be mode="any" and then the user needs to explicitly specify when they expect the variable to be in every dataset?

Done!

[kmpaul]

@andersy005 I think this looks good. Thanks for putting this together.