如何使用 xr.apply_ufunc 改变维度

How to use xr.apply_ufunc with changing dimensions

我有一个加载了 xarray 的 3 维气候数据集

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' ...

我有一个应用于时间维度的模型,returns 一个只有纬度、经度的二维数组。

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

如何在 xr.apply_ufunc() 中使用它?我尝试了几种不同的方法,但它总是抱怨尺寸错误。

例如:

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

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']

经过更多的摆弄,我想我明白了。问题是 apply_ufunc 会将函数应用于所有数据变量。我的数据集中的 "status" 变量导致了问题,因为它只有时间维度。工作代码是

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

apply_model_ufunc(climate['tmean'])