loop-dimensions

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.

def square(item):
    return item ** 2
    
xr.apply_ufunc(square,
               xda,
               input_core_dims=["y"]("y"), 
               output_core_dims=["y"]("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.

References

1. https://tutorial.xarray.dev/advanced/apply_ufunc/automatic-vectorizing-numpy.html#core-dimensions-and-looping