From 5461f2cc33337e84736f26234b12b86f55c4beda Mon Sep 17 00:00:00 2001 From: Stephan Hoyer Date: Wed, 4 Sep 2024 22:59:23 -0700 Subject: [PATCH] Misc. fixes for newer versions of Xarray These fixes fall into three categories: 1. Xarray's groupby no longer supports squeezing out size 1 dimensions. Instead, I've set `squeeze=False` and added explicit calls to `.squeeze()`. 2. Slight differences in dtypes/rounding (I believe these are indirectly due to changes required for NumPy 2.0 compatibility). 3. Various cases where type checking is stricter (I've added `type: ignore` to these lines) PiperOrigin-RevId: 671238424 --- weatherbench2/derived_variables.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/weatherbench2/derived_variables.py b/weatherbench2/derived_variables.py index 774299c..09a22d2 100644 --- a/weatherbench2/derived_variables.py +++ b/weatherbench2/derived_variables.py @@ -664,7 +664,8 @@ def interpolate_spectral_frequencies( def interp_at_one_lat(da: xr.DataArray) -> xr.DataArray: da = ( - da.swap_dims({wavenumber_dim: 'frequency'}) # pytype: disable=wrong-arg-types + da.squeeze('latitude') + .swap_dims({wavenumber_dim: 'frequency'}) # pytype: disable=wrong-arg-types .drop_vars(wavenumber_dim) .interp(frequency=frequencies, method=method, **interp_kwargs) ) @@ -673,7 +674,7 @@ def interp_at_one_lat(da: xr.DataArray) -> xr.DataArray: da['wavelength'] = da['wavelength'].assign_attrs(units='m') return da - return spectrum.groupby('latitude').apply(interp_at_one_lat) + return spectrum.groupby('latitude', squeeze=False).apply(interp_at_one_lat) @dataclasses.dataclass