CODEWITHSUNDEEP
python-3.xpython-xarray
Shawn
Shawn

Reputation: 382

How to use xr.apply_ufunc with changing dimensions

I have a climate dataset with 3 dimensions loaded with xarray

climate = xr.open_dataset(data_file)
climate
<xarray.Dataset>
Dimensions:  (lat: 621, lon: 1405, time: 424)
Coordinates:
  * time     (time) datetime64[ns] 2017-11-01 2017-11-02 2017-11-03 ...
  * lon      (lon) float64 -125.0 -125.0 -124.9 -124.9 -124.8 -124.8 -124.7 ...
  * lat      (lat) float64 49.92 49.87 49.83 49.79 49.75 49.71 49.67 49.62 ...
Data variables:
  tmean    (time, lat, lon) float64 nan nan nan nan nan nan nan nan nan ...
  status   (time) object 'provisional' 'provisional' 'provisional' ...

I have a model which gets applied to the time dimension, and returns a 2d array of just the lat,lon.

apply_model(climate.tmean.values).shape
(621, 1405)

How can I use this within xr.apply_ufunc()? I've tried several different things but it always complains of some dimension error.

For example:

def apply_model_ufunc(climate):
    return xr.apply_ufunc(
            apply_model, climate,
            input_core_dims=[['lat','lon']])

apply_model_ufunc(climate)
ValueError: dimensions ('time',) must have the same length as the number of data dimensions, ndim=2

Or

def apply_model_ufunc(climate):
    return xr.apply_ufunc(
            apply_model, climate,
            input_core_dims=[['time','lat','lon']],
            output_core_dims=[['lat','lon']])

apply_model_ufunc(climate)
ValueError: operand to apply_ufunc has required core dimensions ['time', 'lat', 'lon'], but some of these are missing on the input variable:  ['lat', 'lon']

Upvotes: 2

Views: 1818

Answers (1)

Shawn
Shawn

Reputation: 382

After some more fiddling I think I figured out out. The issue was apply_ufunc will apply the function over all data variables. The "status" variable in my dataset was causing issues as it only has the time dimension. The working code was

def apply_model_ufunc(climate):
    return xr.apply_ufunc(
            apply_model, climate,
            input_core_dims=[['time']],
            output_dtypes=[float])

apply_model_ufunc(climate['tmean'])

Upvotes: 1

Related Questions