ENH Add the fused CSR dense case for Euclidean Specializations

[jjerphan]

Reference Issues/PRs

Towards #22587.
Follow up of #24556.

What does this implement/fix? Explain your changes.

The CSR-dense and the dense-CSR cases were chosen not to be supported for the Euclidean specialization of all PairwiseDistancesReductions (for more context see #23585 (comment)).

This PR implements SparseDenseMiddleTermComputer allows computing the middle term of the distance matrix decomposition for the Euclidean specializations, covering those two missing cases.

Hence this completes all the combinations for the Euclidean specialisations (:tada:).

Any other comments?

Different designs have been explored in other Pull Requests to factor some logic altogether or rethink DatasetsPairs w.r.t. MiddleTermComputer for the Euclidean specialisations.

In overall, this PR seems to have the best tradeoff regarding performance and duplication of code.

Benchmarks

This makes using PairwiseDistancesReductions on the CSR-dense and the dense-CSR for euclidean competitive w.r.t to the previous implementation relying on joblib.

One can get up to ×2 on a laptop.

Details

In [1]: from sklearn.neighbors import NearestNeighbors
   ...: from sklearn.datasets import make_classification
   ...: from scipy.sparse import csr_matrix
   ...: 
   ...: import sklearn

In [2]: X_train, _ = make_classification(n_samples=100_000, n_features=100)
   ...: X_test, _ = make_classification(n_samples=256, n_features=100)
   ...: 
   ...: X_test = csr_matrix(X_test)

In [3]: nn = NearestNeighbors().fit(X_train)

In [4]: %%timeit
   ...: nn.kneighbors(X_test)
   ...: 
   ...: 
899 ms ± 13.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [5]: %%timeit
   ...: with sklearn.config_context(enable_cython_pairwise_dist=False):
   ...:     nn.kneighbors(X_test)
   ...: 
2.24 s ± 13.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [6]: X_train, _ = make_classification(n_samples=256, n_features=100)
   ...: X_test, _ = make_classification(n_samples=100_000, n_features=100)
   ...: 
   ...: X_test = csr_matrix(X_test)

In [7]: nn = NearestNeighbors().fit(X_train)

In [8]: %%timeit
   ...: nn.kneighbors(X_test)
   ...: 
   ...: 
760 ms ± 21.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [9]: %%timeit
   ...: with sklearn.config_context(enable_cython_pairwise_dist=False):
   ...:     nn.kneighbors(X_test)
   ...: 
1.22 s ± 14.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

TODO:

• Refactor the logic to:
  • Remove the cast on the entire dense array and data array and only upcast chunks using dedicated buffers (as done for DenseDenseMiddleTermComputer64
  • Redesign DatasetsPair w.r.t the new MiddleTermComputer to remove the duplicated logic and to clarify responsabilities (esp. squared euclidean norm computations)
    • Comment: the redesign was tried in ENH Euclidean specialization of DatasetsPair instead of ArgKmin and RadiusNeighbors #25170, works, but suffers from a lot of method dispatches overhead
• Add more tests?

[jjerphan]

Hi @OmarManzoor, @Vincent-Maladiere, @Micky774 and @adam2392: might some of you might be interested in reviewing this PR? 🙂

[OmarManzoor]

Hi @OmarManzoor, @Vincent-Maladiere and @adam2392: might some of you might be interested in reviewing this PR? 🙂

Sure! I might need some time to understand the context of the PR though.

[jjerphan]

I am wondering about the kind of tests we could add. Should we add tests on combinations of sparse and dense datasets for all the interfaces? Should we add more of them?

[jjerphan]

I might need some time to understand the context of the PR though.

Yes, feel free to take your time and do not feel pressured (the class hierarchy is rather dense and its implementations Tempita-heavy).

[jjerphan]

Should I write more about the design?

[OmarManzoor]

LGTM. Maybe we can enhance the tests by including the combinations within test_sqeuclidean_row_norms and test_pairwise_distances_argkmin.

[ogrisel]

Great work!

and an extra +1 for the opportunity to remove the lines with the XOR operator which were not easy to reason about :)

[ogrisel]

Add more tests?

I think there are enough tests in this PR and the already existing tests in main.

[lorentzenchr]

LGTM
PS: I had to find something:smirk:

[lorentzenchr]

Is this "trick" still documented somewhere in the pairwise distances codebase?

[jjerphan]

Yes, it is given here:

scikit-learn/sklearn/metrics/_pairwise_distances_reduction/_middle_term_computer.pyx.tp

Lines 77 to 93 in dfe9e2e

	cdef class MiddleTermComputer{{name_suffix}}:
	"""Helper class to compute a Euclidean distance matrix in chunks.
	This is an abstract base class that is further specialized depending
	on the type of data (dense or sparse).
	`EuclideanDistance` subclasses relies on the squared Euclidean
	distances between chunks of vectors X_c and Y_c using the
	following decomposition for the (i,j) pair :
	||X_c_i - Y_c_j||² = ||X_c_i||² - 2 X_c_i.Y_c_j^T + ||Y_c_j||²
	This helper class is in charge of wrapping the common logic to compute
	the middle term, i.e. `- 2 X_c_i.Y_c_j^T`.
	"""

I am thinking of writing a couple of notes for the design of sklearn.metrics._pairwise_distances_reduction. What do you think of it?

[lorentzenchr]

That would be great. Not too many details, but giving a good overview.

[lorentzenchr]

Which co-authors should I keep when merging?

[jjerphan]

Which co-authors should I keep when merging?

All, I guess? I let this choice to the maintainer merging.

[jjerphan]

Thank you for your reviews, @OmarManzoor, @ogrisel, and @lorentzenchr.