Example of vectorizing ruptures with xarray.apply_ufunc

[spencerahill]

(originally Issue #136)

In case others come across a similar need (whether or not it makes sense to implement this into ruptures proper), here's a way of using xarray.apply_ufunc to perform calculations on rupture on xarray.DataArrays.

This particular example is for detecting a single breakpoint; note that if you want to return multiple, the apply_ufunc call signature will have to be modified, because then each call is returning an array rather than a scalar.

import ruptures as rpt
import xarray as xr

def detect_breakpoint(arr, dim, rpt_class=rpt.Binseg, model="l2", n_bkps=1):
    """Use xr.apply_ufunc to broadcast breakpoint detections from ruptures"""

    def _detect_bp(arr):
        """Wrapper to use in apply_ufunc."""
        return rpt_class(model=model).fit(arr).predict(n_bkps=n_bkps)[0]

    inds_bp = xr.apply_ufunc(
        _detect_bp,
        arr,
        input_core_dims=[[dim]],
        vectorize=True,
        dask="parallelized",
    )
    return arr[dim][inds_bp]

(FYI This came up when attempting to detect changepoints at each point of a lat-lon gridded rainfall dataset.)