Skip to content

Commit

Permalink
Misc. fixes for newer versions of Xarray
Browse files Browse the repository at this point in the history
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
  • Loading branch information
shoyer authored and Weatherbench2 authors committed Sep 6, 2024
1 parent bb1713d commit 5461f2c
Showing 1 changed file with 3 additions and 2 deletions.
5 changes: 3 additions & 2 deletions weatherbench2/derived_variables.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
)
Expand All @@ -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
Expand Down

0 comments on commit 5461f2c

Please sign in to comment.