Improve SPI performance

[coxipi]

Pull Request Checklist:

• This PR addresses an already opened issue (for bug fixes / features)
  • This PR fixes Xclim Standardised Precipitation Index always throws error #1270 and fixes xclim.indices.standardized_precipitation_index() return infinite values #1416 and fixes _agro.standardized_precipitation_evapotranspiration_index creates NaNs  #1474
• Tests for the changes have been added (for bug fixes / features)
  • (If applicable) Documentation has been added / updated (for bug fixes / features)
• HISTORY.rst has been updated (with summary of main changes)
  • Link to issue (:issue:number) and pull request (:pull:number) has been added

What kind of change does this PR introduce?

• Make SPI/SPEI faster
• fit params are now modular, can be computed before computing SPI/SPEI. This allows more options to segment computations and allow to obtain the fitting params if troubleshooting is needed.
• time indexing now possible
• dist_method now avoids vectorize=True in its xr.apply_ufunc. This is the main improvement in SPI/SPEI.
• Better document the limits of usage of standardized indices. Now standardized indices are capped at extreme values ±8.21. The upper bound is a limit resulting of the use of float64.

Does this PR introduce a breaking change?

Yes.

• pr_cal or wb_cal will not be input options in the future:

Inputing pr_cal will be deprecated in xclim==0.46.0. If pr_cal is a subset of pr, then instead of:
standardized_precipitation_index(pr=pr,pr_cal=pr.sel(time=slice(t0,t1)),...), one can call:
standardized_precipitation_index(pr=pr,cal_range=(t0,t1),...).
If for some reason pr_cal is not a subset of pr, then the following approach will still be possible:
params = standardized_index_fit_params(da=pr_cal, freq=freq, window=window, dist=dist, method=method).
spi = standardized_precipitation_index(pr=pr, params=params).
This approach can be used in both scenarios to break up the computations in two, i.e. get params, then compute
standardized indice

I could revert this breaking change if we prefer. This was a first attempt to make the computation faster, but the improvements are now independent of this change. We could also keep the modular structure for params, but revert to pr_cal instead of cal_range. It's a bit less efficient when pr_cal is simply a subset of pr, because you end up doing resampling/rolling two times on the calibration range for nothing. When first computing params, then obtaining spi in two steps, then it makes no difference

Other information:

[coxipi]

I think the computation is still quite heavy. A few things should help. There is less redundant computation (we can either reuse params, or at least we don't resample/roll more or less twice (like it used to do for pr and pr_cal)).

The user can now dissect the computation more easily:

• Produce params before hand, reuse them for many computations
• Use of indexer to obtain SPI only for certain periods. This not like selecting, say month of June, from the onset. The indexing must be done after the rolling. Said another way, you still need months "Jan to June" even if you only want SPI-6 for June. So the indexer must be used mid-computation, this is what is done.

[aulemahal]

But yes, params needs an entry in declare_units if it is a Quantified.

[aulemahal]

I'm confused, how does this work with dask ? Won't da.isnull().all() as a conditional trigger the computation ?

And why indexer != {} ? Why would this fastpath only be used with indexing ?

[coxipi]

Good point about Dask. There is no way then to achieve a fastpath with lazy computing, no?

For the other question: Yes, indeed, spatial regions with only NaNs independently of time selecting would also benefit from this speedup (ignoring the Dask issue). I wanted to reserve this check for cases when I was sure there was a potential of all-NaN slices, e.g. when there was time-selecting. I'm not sure how costly it is to check .isnull.all, but probably very small in comparison to the whole algorithm.

[aulemahal]

According to the doc (https://github.com/ml31415/numpy-groupies/), numpy groupies accepts a callable as func would it make sense to use that ?

[coxipi]

I had tried but got a confusing error message leading me to believe this claim was maybe false. Maybe I just messed up. But I will retry, it would be worth it to clean the function.

[aulemahal]

Suggestion:

Suggested change

	offset: Quantified
	offset: Quantified | None = None,

# pseudo-code
if offset is None:
  if params is not None:
      if dist in bounded_distributions:
          offset = 1000 mm / d
      else: 
        offset = 0
    else:
        offset = params.offset
else:
    if params is not None and params.offset != offset:
        warning

C'est à dire : put a default of None and clearly explain in the docstring what the default behaviour is. I know I just suggested something different for the snow thing, but that was in a "temporary" perspective!

[aulemahal]

Also, could there some option where offset = -wb.min() ?

[coxipi]

I will check, but I think actually, there is not a restriction in scipy to use gamma/fisk distrubutions with 0 bounded data. A parameter can also be fitted for this offset. This is was an idea for another PR though. I think the problem could be: if there are more negative values outside of the calibration period, then the offset could be too weak.

The problem with offset = -wb.min() is that you can become sensitive to what data is included. Imagine your reference data is 1980-2010 and you fit two cases:

1. 1980-2020
2. 1980-2050

You could have a different minimum in cases 1 & 2, so different offset. Ideally, since this is a trick, we would like that our computation are not too sensitive on this. But in any case, it would be nice that the results of 1. are en exact subset of 2. if you use the same calibration data / methods for both computations.

[coxipi]

Ideally, I would eventually like to get rid of the default 1 mm/day default. It was just copied on monocongo's climate indices. It's a good rule of thumb, but users have found cases where a bigger offset was needed. Maybe we can keep it in but warn users about having NaNs? Maybe just in the docstring. Not sure. I could check the R implementation to see if I can dig ideas.

[coxipi]

Showing problems with dist_method that explain the previous slow behaviour

I think it may be worthwhile to benchmark a bit the previous implementation, where is the bottleneck. The main culprit I believe is dist_method. For instance, let us first obtain some parameters.

from xclim.testing import open_dataset
import xclim
import scipy

pr = open_dataset("sdba/CanESM2_1950-2100.nc").pr
freq, window, dist, method = "MS", 1, "gamma", "APP"
params = xclim.indices.stats.standardized_index_fit_params(
    pr.sel(time=slice("1950", "1980")), freq=freq, window=window, dist=dist, method=method
)
# broadcast params like pr
params = params.rename(month="time").reindex(time=pr.time.dt.month)
params["time"] = pr.time

Now, we either use a numpy/ scipy approach (fast):

# ~0.5 seconds
import xarray as xr
from xclim.indices.stats import get_dist


dist = get_dist(params.attrs["scipy_dist"])
def wrap_cdf(da, pars):
    return dist.cdf(da[:], *pars)

out1 = xr.apply_ufunc(
    wrap_cdf,
    pr.where(pr>0), 
    params,
    input_core_dims=[["time"], ["dparams","time"]],
    output_core_dims=[["time"]],
    vectorize=True,
)
out1.values

or we use dist_method (50 times slower)

# about 25 seconds
from xclim.indices.stats import dist_method
out2 = dist_method("cdf", params, pr.where(pr > 0))
out2.values

Both methods give the same results. It remains to be seen if the speed-up persists as we scale things up, but this is already a weird result.

full spi computation

I can use the approach outlined above for the full spi computation. It is much faster than the previous implementation before this PR. It looks like:

    probs_of_zero = da.groupby(group).map(lambda x: (x == 0).sum("time") / x.notnull().sum("time"))
    params, probs_of_zero = [resample_to_time(dax, da) for dax in [params, probs_of_zero]]

    def wrap_cdf_ppf(da, pars, probs_of_zero):
        dist_probs = get_dist(params.attrs["scipy_dist"]).cdf(da[:], *pars)
        probs = probs_of_zero + ((1 - probs_of_zero) * dist_probs)
        return norm.ppf(probs)

    std_index = xr.apply_ufunc(
        wrap_cdf_ppf,
        da, 
        params,
        probs_of_zero,
        input_core_dims=[["time"], ["dparams","time"],["time"]],
        output_core_dims=[["time"]],
        vectorize=True,
        dask="parallelized",
    )

It remains about 2x slower than the approach with group_idxs, I think mainly because of the probs_of_zero computation, I will try a different approach

2nd try:

Using flox, I almost get as fast as the weird loop on group_idx I implemented:

probs_of_zero = flox.xarray.xarray_reduce(da, idxs,func="sum")/floxx.xarray_reduce(da.notnull(), idxs,func="sum")

• weird loop with group_idx : 45 s
• using floxx : 53 s
• no floxx : 80 s

I don't think flox supports custom function. That would be ideal because then I could write:

def func(da):
    return (da==0).sum(dim="time")/da.notnull().sum(dim="time")

and have only one flox call. There is this Aggregation thing I'm trying to get the hang of...

3rd try

A little dirty trick that seems performant, but a bit more difficult to understand

floxx.xarray_reduce((pr+1)*(pr==0), idxs,func="mean")

It's equivalent to :

floxx.xarray_reduce((pr==0), idxs,func="sum")/floxx.xarray_reduce(pr.notnull(), idxs,func="sum")
floxx.xarray_reduce((pr==0).where(pr.notnull()), idxs,func="mean")

but avoids the double-flox, or the "where" call. Doesn't seem much faster than the 2nd try though.

By the way ...

I should try to improve dist_method instead of having a specific method for SPI/SPEI computations

[coxipi]

Sabotaging my previous example to mimick dist_method performance

The performance issue seems to stem from the input_core_dims used in the xr.apply_ufunc call of dist_method. With my little example, I can reproduce the slowness of dist_method:

    # like `dist_method`, about 25 s
    def wrap_cdf(da, pars):
        return dist.cdf(da, *pars)

    out = xr.apply_ufunc(
        wrap_cdf,
        pr.where(pr>0), 
        params,
        input_core_dims=[[], ["dparams"]],
        output_core_dims=[[]],
        vectorize=True,
    )

while using "time" in the core dims goes much faster. This remains true if I try using dask etc., but maybe I'm not using this well. At this point, we could chunk the time dimension, but earlier steps in stats usually require unchunked time dimension. If using "time" as a core dimension is always appropriate in this scenario, this could be way to improve dist_method easily

[Zeitsperre]

Well done!

[aulemahal]

Two small comments, but it looks good!

[Zeitsperre]

Ignore the cancelled builds, merging when docs clear.