Skip to content

Commit e70d9a3

Browse files
shoyerWeatherbench2 authors
authored and
Weatherbench2 authors
committed
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
1 parent bb1713d commit e70d9a3

File tree

1 file changed

+4
-2
lines changed

1 file changed

+4
-2
lines changed

weatherbench2/derived_variables.py

+4-2
Original file line numberDiff line numberDiff line change
@@ -664,7 +664,9 @@ def interpolate_spectral_frequencies(
664664

665665
def interp_at_one_lat(da: xr.DataArray) -> xr.DataArray:
666666
da = (
667-
da.swap_dims({wavenumber_dim: 'frequency'}) # pytype: disable=wrong-arg-types
667+
da
668+
.squeeze('latitude')
669+
.swap_dims({wavenumber_dim: 'frequency'}) # pytype: disable=wrong-arg-types
668670
.drop_vars(wavenumber_dim)
669671
.interp(frequency=frequencies, method=method, **interp_kwargs)
670672
)
@@ -673,7 +675,7 @@ def interp_at_one_lat(da: xr.DataArray) -> xr.DataArray:
673675
da['wavelength'] = da['wavelength'].assign_attrs(units='m')
674676
return da
675677

676-
return spectrum.groupby('latitude').apply(interp_at_one_lat)
678+
return spectrum.groupby('latitude', squeeze=False).apply(interp_at_one_lat)
677679

678680

679681
@dataclasses.dataclass

0 commit comments

Comments
 (0)