DOC Better document implementation

jjerphan · jjerphan · commit 100949fb17c6 · 2023-03-09T10:28:52.000+01:00
Signed-off-by: Julien Jerphanion &lt;git@jjerphan.xyz&gt;
diff --git a/doc/whats_new/v1.3.rst b/doc/whats_new/v1.3.rst
@@ -80,8 +80,8 @@ Changes impacting all modules
   Euclidean distances (sometimes followed by a fused reduction operation) for a
   pair of datasets consisting of a sparse CSR matrix and a dense NumPy.
 
-  This can improve the performance following functions and estimators.
-  
+  This can improve the performance of following functions and estimators:
+
   - :func:`sklearn.metrics.pairwise_distances_argmin`
   - :func:`sklearn.metrics.pairwise_distances_argmin_min`
   - :class:`sklearn.cluster.AffinityPropagation`
@@ -105,7 +105,7 @@ Changes impacting all modules
 
   A typical example of this performance improvement happens when passing a sparse
   CSR matrix to the `predict` or `transform` method of estimators that rely on
-  a dense numpy representation to store their fitted parameters (or the reverse).
+  a dense NumPy representation to store their fitted parameters (or the reverse).
 
   For instance, :meth:`sklearn.NearestNeighbors.kneighbors` is up to 2 times faster
   for this case now on commonly available laptops.
diff --git a/sklearn/metrics/_pairwise_distances_reduction/_middle_term_computer.pxd.tp b/sklearn/metrics/_pairwise_distances_reduction/_middle_term_computer.pxd.tp
@@ -194,6 +194,10 @@ cdef class SparseDenseMiddleTermComputer{{name_suffix}}(MiddleTermComputer{{name
 
         const {{INPUT_DTYPE_t}}[:, ::1] Y
 
+        # We treat the dense-sparse case with the sparse-dense case by simply
+        # treating the dist_middle_terms as F-ordered and by swapping arguments.
+        # This attribute is meant to encode the case and adapt the logic
+        # accordingly.
         bint c_ordered_middle_term
 
     cdef void _parallel_on_X_pre_compute_and_reduce_distances_on_chunks(
diff --git a/sklearn/metrics/_pairwise_distances_reduction/_middle_term_computer.pyx.tp b/sklearn/metrics/_pairwise_distances_reduction/_middle_term_computer.pyx.tp
@@ -49,9 +49,9 @@ cdef void _middle_term_sparse_sparse_64(
     const SPARSE_INDEX_TYPE_t[:] Y_indptr,
     ITYPE_t Y_start,
     ITYPE_t Y_end,
-    DTYPE_t * D,
+    DTYPE_t * dist_middle_terms,
 ) noexcept nogil:
-    # This routine assumes that D points to the first element of a
+    # This routine assumes that D is a pointer to the first element of a
     # zeroed buffer of length at least equal to n_X × n_Y, conceptually
     # representing a 2-d C-ordered array.
     cdef:
@@ -68,7 +68,7 @@ cdef void _middle_term_sparse_sparse_64(
                 for y_ptr in range(Y_indptr[Y_start+j], Y_indptr[Y_start+j+1]):
                     y_col = Y_indices[y_ptr]
                     if x_col == y_col:
-                        D[k] += -2 * X_data[x_ptr] * Y_data[y_ptr]
+                        dist_middle_terms[k] += -2 * X_data[x_ptr] * Y_data[y_ptr]
 
 
 {{for name_suffix, upcast_to_float64, INPUT_DTYPE_t, INPUT_DTYPE in implementation_specific_values}}
@@ -83,11 +83,11 @@ cdef void _middle_term_sparse_dense_{{name_suffix}}(
     ITYPE_t Y_start,
     ITYPE_t Y_end,
     bint c_ordered_middle_term,
-    DTYPE_t * D,
+    DTYPE_t * dist_middle_terms,
 ) nogil:
-    # This routine assumes that D points to the first element of a
-    # zeroed buffer of length at least equal to n_X × n_Y, conceptually
-    # representing a 2-d C-ordered array.
+    # This routine assumes that dist_middle_terms is a pointer to the first element
+    # of a zeroed buffer of length at least equal to n_X × n_Y, conceptually
+    # representing a 2-d C-ordered of F-ordered array.
     cdef:
         ITYPE_t i, j, k
         ITYPE_t n_X = X_end - X_start
@@ -99,7 +99,7 @@ cdef void _middle_term_sparse_dense_{{name_suffix}}(
             k = i * n_Y + j if c_ordered_middle_term else j * n_X + i
             for X_i_ptr in range(X_indptr[X_start+i], X_indptr[X_start+i+1]):
                 X_i_col_idx = X_indices[X_i_ptr]
-                D[k] += -2 * X_data[X_i_ptr] * Y[Y_start + j, X_i_col_idx]
+                dist_middle_terms[k] += -2 * X_data[X_i_ptr] * Y[Y_start + j, X_i_col_idx]
 
 
 cdef class MiddleTermComputer{{name_suffix}}:
@@ -183,9 +183,16 @@ cdef class MiddleTermComputer{{name_suffix}}:
                 c_ordered_middle_term=True
             )
         if not X_is_sparse and Y_is_sparse:
+            # NOTE: The Dense-Sparse case is implement via the Sparse-Dense case.
+            #
+            # To do so:
+            #    - X (dense) and Y (sparse) are swapped
+            #    - the distance middle term is seen as F-ordered for consistency
+            #      (c_ordered_middle_term = False)
             return SparseDenseMiddleTermComputer{{name_suffix}}(
-                Y,
-                X,
+                # Mind that X and Y are swapped here.
+                X=Y,
+                Y=X,
                 effective_n_threads,
                 chunks_n_threads,
                 dist_middle_terms_chunks_size,
@@ -572,7 +579,8 @@ cdef class SparseDenseMiddleTermComputer{{name_suffix}}(MiddleTermComputer{{name
         ITYPE_t Y_end,
         ITYPE_t thread_num,
     ) noexcept nogil:
-        # Flush the thread dist_middle_terms_chunks to 0.0
+        # Fill the thread's dist_middle_terms_chunks with 0.0 before
+        # computing its elements in _compute_dist_middle_terms.
         fill(
             self.dist_middle_terms_chunks[thread_num].begin(),
             self.dist_middle_terms_chunks[thread_num].end(),
@@ -587,7 +595,8 @@ cdef class SparseDenseMiddleTermComputer{{name_suffix}}(MiddleTermComputer{{name
         ITYPE_t Y_end,
         ITYPE_t thread_num,
     ) noexcept nogil:
-        # Flush the thread dist_middle_terms_chunks to 0.0
+        # Fill the thread's dist_middle_terms_chunks with 0.0 before
+        # computing its elements in _compute_dist_middle_terms.
         fill(
             self.dist_middle_terms_chunks[thread_num].begin(),
             self.dist_middle_terms_chunks[thread_num].end(),
@@ -607,15 +616,22 @@ cdef class SparseDenseMiddleTermComputer{{name_suffix}}(MiddleTermComputer{{name
                 self.dist_middle_terms_chunks[thread_num].data()
             )
 
+        # For the dense-sparse case, we use the sparse-dense case
+        # with dist_middle_terms seen as F-ordered.
+        # Hence we swap indices pointers here.
+        if not self.c_ordered_middle_term:
+            X_start, Y_start = Y_start, X_start
+            X_end, Y_end = Y_end, X_end
+
         _middle_term_sparse_dense_{{name_suffix}}(
             self.X_data,
             self.X_indices,
             self.X_indptr,
-            X_start if self.c_ordered_middle_term else Y_start,
-            X_end if self.c_ordered_middle_term else Y_end,
+            X_start,
+            X_end,
             self.Y,
-            Y_start if self.c_ordered_middle_term else X_start,
-            Y_end if self.c_ordered_middle_term else X_end,
+            Y_start,
+            Y_end,
             self.c_ordered_middle_term,
             dist_middle_terms,
         )
