Public API to create lazy variables after function application

[dcherian]

Is your feature request related to a problem?

With flexible indexes, a lot more people are going to want lazy coordinate variables, that compose function application with slicing.

We've provided one set of adapters with CoordinateTransformIndex, but there is certainly a need for more general coordinate variables. For example, I have used the following to create a lazy array: valid_time[time,step] = time[time] + step[step]

def create_lazy_valid_time_variable(
    *, reference_time: xr.DataArray, period: xr.DataArray
) -> xr.DataArray:
    # TODO: work to make this public API upstream
    from xarray.coding.variables import lazy_elemwise_func


    assert reference_time.ndim == 1
    assert period.ndim == 1


    shape = (reference_time.size, period.size)
    time_broadcast = np.broadcast_to(reference_time.data[:, np.newaxis], shape)
    step_broadcast = np.broadcast_to(period.data[np.newaxis, :], shape)


    valid_time = lazy_elemwise_func(
        time_broadcast, partial(operator.add, step_broadcast), dtype=reference_time.dtype
    )
    return xr.DataArray(
        valid_time, dims=(*reference_time.dims, *period.dims), attrs={"standard_name": "time"}
    )

Describe the solution you'd like

It would be nice to have some kind of public API for this.

One possibility is supporting lazy=True in apply_ufunc: #2302.

Describe alternatives you've considered

We could ask users to use dask/cubed/etc but that seems unnecessary

[dcherian]

Next step would be to look in to how apply_ufunc does broadcasting; and how broadcasting interacts with lazy arrays