Speed up quantiles with sorting

[SarahG-579462]

Pull Request Checklist:

• This PR addresses an already opened issue (for bug fixes / features)
  • This PR will help Efficient Application of Bias Correction #1255
• Tests for the changes have been added (for bug fixes / features)
  • (If applicable) Documentation has been added / updated (for bug fixes / features)
• CHANGES.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?

• nbutils.quantile has a speed-up of more than 2.5x by a combination of changes in nbutils.quantile and nbutils._quantile
• This does not cover nbutils.vec_quantiles (used for adapt_freq) but similar principles could be used
• It adds the possibility of using fastnanquantile module which is very fast

Does this PR introduce a breaking change?

No

Other information:

• The new low-level function to compute quantiles nbutils._quantile is a 1d jitted version of xclim.core.utils._nan_quantile
• Manual benchmarking can be performed in the notebook benchmarks/sdba_quantile.ipynb, attached to this PR.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[github-actions]

Welcome, new contributor!

It appears that this is your first Pull Request. To give credit where it's due, we ask that you add your information to the AUTHORS.rst and .zenodo.json:

• The relevant author information has been added to AUTHORS.rst and .zenodo.json

Please make sure you've read our contributing guide. We look forward to reviewing your Pull Request shortly ✨

[juliettelavoie]

For a 30 years period, over QC, the bias adjustment (first initial DQM + Npdf + restore) with these parameters:

n_iter: 30
nquantiles:50
group: time
n_escore: 1

took 8 hours for this branch and 8h40 for the master. (time for branch faster_mbcn is coming.)

EDIT: Not sure if it makes a difference but, the test for this branch ran mostly during work hours and the test for the master ran during the night.

UPDATE: It took 6h30 with branch faster_mbcn. @coxipi your branch does help in my case!

[Zeitsperre]

@SarahG-579462

checking consistency... /home/docs/checkouts/readthedocs.org/user_builds/xclim/checkouts/1513/docs/notebooks/sdba-speed.ipynb: WARNING: document isn't included in any toctree

Small FYI, you need to add the notebook to the exclude_patterns in docs/conf.py if you don't want it rendered.

[coxipi]

So I hijacked the PR a bit to test general benchmark tools. I needed to install pytest-benchmarkand pytest-benchmark[histogram]. The idea would be to have a benchmarking suite, similar to the structure of the testing suite. For a test similar to your current notebook, I run:

pytest  -n0 --dist no benchmarks/sdba/test_nbutils.py

where test_nbutils.py is

from xclim.sdba import nbutils
import pytest 
import importlib
import os
import numpy as np

nq = 50
i = 0
@pytest.mark.parametrize(
    "use_nanquantile,size",
    [
    (True, 100),
    (True, 200),
    (True, 500),
    (True, 1000),
    (False, 100),
    (False, 200),
    (False, 500),
    (False, 1000),
    ],
)
def test_nanquantile_simple(benchmark,use_nanquantile, size):
    np.random.seed(0)
    arr = np.random.randn(size)
    os.environ["USE_NANQUANTILE"] = "True" if use_nanquantile else ""
    importlib.reload(nbutils)   
    def func():
        return nbutils._choosequantile(arr, np.linspace(0, 1, nq))
    # trigger jit compilation on first run
    global i
    if i==0: 
        func()
        i = i + 1
    benchmark(func)

(I should have named it test__choosequantile but anyways)

I can then have an histogram with the results:

It's a bit annoying that the order of parameters in pytest.mark.parametrize is not respected. I could sort the names alphabetically, but then "1000" points comes before the other smaller samples. Anyways.

Anyways, it seems pretty plug and play. I find the visualization with the notebook is better, so we might want to keep a notebook for benchmarks too? Simpler tests that are not comparing multiple methods would benefit of simple numerical tests just to assert the performance doesn't drop below a certain level perhaps?

[coxipi]

I did a test with EmpiricalQuantileMapping, 30 years, 3 locations, dayofyear-31.
Sorting is still advantageous, but the effect is marginal. There is maybe an overhead which is just much greater than the actual computation of quantiles? False/True refers to use_nanquantile

In comparison, group="time" has a better performance for sortquantile, it's about 85-90% of the nanquantile running time. In that case, there must be less overhead (complicated map_blocks structure has a less important role in this case?)

[SarahG-579462]

I did a test with EmpiricalQuantileMapping, 30 years, 3 locations, dayofyear-31. Sorting is still advantageous, but the effect is marginal. There is maybe an overhead which is just much greater than the actual computation of quantiles? False/True refers to use_nanquantile

This is really interesting, thanks. I totally support you taking over this PR, I don't really have time to work on it these days.

[SarahG-579462]

The present implementation is not right, the call to nan_quantile is not properly "jitterized". Among other things, some operations in nan_quantile use the argument axis in nan_quantile which is unsupported for many functions used with njit. The library fastnanquantile shows how those parts working on axis must be done outside the njit calls. Also, np.asanyarray doesn't seem to be supported. I'm not sure why this is triggered or not in certain contexts.

The function would either need to be properly jitterized, or keep it unjittered and study again the performance. Since sdba and xclim are to be separated, I feel some redundancy would not be bad on this front, I was considering to simply revert some commits and continue to use sortquantile. Thoughts?

During testing, I had made a version of Abel's nan_quantile that supports only 1 dimension, that is jitt'able. Would it be preferable to use this? Abel's unjitted version is really not much slower than the jitted version, because he's using quite pure numpy functions.

[coxipi]

Should be good to go, I added a few tests

[coxipi]

Would it be possible for you to add fastnanquantile as a dependency that can be installed via tox (e.g. under deps: fastnanquantile: fastnanquantile) and add that to one of the PyPI/tox CI builds (in main.yml) so that we can test against it?

@Zeitsperre under test-pypi?
EDIT: I will wait your return, exploring this part of the codebase is not good for my cardiac rhythm

[github-actions]

Note

It appears that this Pull Request modifies the main.yml workflow.

On inspection, the XCLIM_TESTDATA_BRANCH environment variable is set to the most recent tag (v2023.12.14).

No further action is required.

[coveralls]

coverage: 90.417% (-0.3%) from 90.679%
when pulling df6963b on speed-up-quantile
into 1b83536 on main.

[Zeitsperre]

@aulemahal I leave you to do the final review and merge. I won't pretend to be the expert here.

[coxipi]

@SarahG-579462 can you inspect a last time? Should be good now

[SarahG-579462]

Looks good to me! a note about fastmath below.