Loop Dimensions
Functions that performs operations on data with core-dimensions and taking axis as argument, can be easily used with apply_ufunc.
However, functions which don’t take axis argument and need to work on core-dimensions need loop dimensions. Loop dimension is something where the loop happens and function gets the sub-array as input.
For example, a 2d data array with dimensions x and y and square needs to work on only y, we need to loop over x so that function can get only 1d array on dimension y.
xda = xr.DataArray(np.arange(2*5).reshape(2, 5), dims=("x", "y"), coords={"x": np.arange(2), "y": np.arange(5)})
def square(item):
return item ** 2
xr.apply_ufunc(square,
xda,
input_core_dims=[["y"]],
output_core_dims=[["y"]],
vectorize=True)square gets 1d array of dimension y. input_core_dims provides the core dimensions function will take and dimensions which are not provided will be taken as loop dimensions (which is x) here.