I have an xarray Dataset which contains some data defined on three coordinates (lat, lon, time) and some defined on only one (time). I'd like to separate this into two different datasets based on the number of coordinates per variable. It's not as straightforward as I'd like.
By way of an example using pandas, if I have a DataFrame which contains a different number of non-NaN values per variable (i.e. column)
df = pd.DataFrame({'a': [1, np.nan, np.nan], 'b': [1, 1, 1]})
and I wanted to split the DataFrame based on this, I would do something like
df[df.columns[df.isna().sum()==0]]
as in, I can use some logic to reduce the list of columns and access all of the columns which satisfy the logic. I can't figure out how to do this in xarray. I'm able to make a list of all the variables
list(ds.keys())
and a corresponding list of how many coordinate variables they each have
[len(ds[var].coords) for var in list(ds.keys())]
but I can't use this to variables in the Dataset because variables are not indexed. I'm sure there must be an easier way of doing this in xarray but I haven't found an answer. Can anyone help?
Update: I found a solution along the lines of my original intent
ds_a = ds[np.array(list(ds.keys()))[np.where(np.array([len(ds[var].coords) for var in list(ds.keys())]) == 3)[0]]]
ds_b = ds[np.array(list(ds.keys()))[np.where(np.array([len(ds[var].coords) for var in list(ds.keys())]) == 1)[0]]]
but although this works it's ugly as hell. I've actually always been frustrated with the inelegance of the pandas solution I posted above, so I'll leave the question open in case anyone has some beautiful insights
