xarray-contrib / xwrf

A lightweight interface for working with the Weather Research and Forecasting (WRF) model output in Xarray.
https://xwrf.readthedocs.io/
Apache License 2.0
56 stars 16 forks source link

Fix destagger attrs #97

Closed lpilz closed 1 year ago

lpilz commented 1 year ago

Change Summary

When using the destagger method of Dataset and DataArray accessors, variable attributes were getting removed. Attributes like units or grid_mapping are however important for processing of data, e.g. with metpy.

Related issue number

Closing #96

Checklist

jthielen commented 1 year ago

Very good catch, and an important thing to fix before releasing this! I think I missed this because of keep_attrs settings (which, I'm going to go look if there is a context-based way of having different keep_attrs settings for use in our tests without changing global state).

I would recommend, however, a different implementation approach than used here. Attrs (and encoding...that's one I forgot about too) are able to exist on the xarray.Variable level, so the issue really lies in the low-level _destag_variable function rather than the accessor methods. Particularly,

https://github.com/xarray-contrib/xwrf/blob/290c2e7cd42c6d0a42872857fd2ebc421f504c34/xwrf/destagger.py#L73-L79 should get replaced with

return xr.Variable(
    dims=tuple(str(unstag_dim_name) if dim == stagger_dim else dim for dim in center_mean.dims),
    data=center_mean.data,
    attrs=_drop_attrs(datavar.attrs, ('stagger',)),
    encoding=datavar.encoding,
    fastpath=True,
)

(i.e., use original datavar for attrs and encoding instead of center_mean) and the accessor methods left unchanged.

@lpilz Would you want to implement these changes in this PR? Otherwise, I could put together a replacement PR if not?

Also, if you want to go ahead on implementing this, I'd recommend against modifying the data attrs in the tests. For example, instead of

if 'stagger' in data.attrs:
    data.attrs.pop('stagger')
assert set(destaggered.attrs.keys()) == set(data.attrs.keys())

have

assert set(destaggered.attrs) == set(data.attrs) - set(('stagger',))
lpilz commented 1 year ago

Thanks for the hints, I don't know why I didn't see the obvious fix... I applied the changes you requested, hope this works.

I implemented the set subtraction with - {'stagger'} instead of the tuple to set cast. Is there any reason why one would like to have that?

jthielen commented 1 year ago

I implemented the set subtraction with - {'stagger'} instead of the tuple to set cast. Is there any reason why one would like to have that?

No, that's definitely better! The reason I did it the way I did instead of thinking of the more obvious construction is that I've gotten into the bad habit of over-using set() instead of {} to not get confused with dicts.