Fixing numerical error from self-looping

rapidsai · Oct 17, 2023 · 18743cc · 18743cc
1 parent 3ec6dcc
commit 18743cc
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 12 deletions.
diff --git a/cpp/include/raft/distance/detail/distance_ops/l2_exp.cuh b/cpp/include/raft/distance/detail/distance_ops/l2_exp.cuh
@@ -31,10 +31,20 @@ struct l2_exp_cutlass_op {
   __device__ AccT operator()(DataT& aNorm, const DataT& bNorm, DataT& accVal) const noexcept
   {
     AccT outVal = aNorm + bNorm - DataT(2.0) * accVal;
-    // outVal could be negative due to numerical instability, especially when
-    // calculating self distance.
-    // clamp to 0 to avoid potential NaN in sqrt
-    outVal = outVal * (outVal >= DataT(1e-6));
+
+    if (aNorm == bNorm) {
+      printf("aNorm: %f, bNorm:%f, acc: %f, outVal: %f\n",
+             aNorm,
+             bNorm,
+             accVal,
+             outVal * (outVal >= 1e-6));
+    }
+
+    /**
+     * Self-neighboring points should have (aNorm == bNorm) == accVal and the dot product (accVal)
+     * can sometimes have round-off errors, which will cause (aNorm == bNorm) ~ accVal instead.
+     */
+    outVal = outVal * (outVal > 1e-4 && !(aNorm == bNorm && accVal > 0.0));
     return sqrt ? raft::sqrt(outVal) : outVal;
   }
 
@@ -86,10 +96,15 @@ struct l2_exp_distance_op {
     for (int i = 0; i < Policy::AccRowsPerTh; ++i) {
 #pragma unroll
       for (int j = 0; j < Policy::AccColsPerTh; ++j) {
-        DataT val = regxn[i] + regyn[j] - (DataT)2.0 * acc[i][j];
-        // val could be negative due to numerical instability, especially when
-        // calculating self distance. Clamp to 0 to avoid potential NaN in sqrt
-        acc[i][j] = val * (val >= DataT(1e-6));
+        DataT accVal = acc[i][j];
+        DataT val    = regxn[i] + regyn[j] - (DataT)2.0 * accVal;
+
+        /**
+         * Self-neighboring points should have (aNorm == bNorm) == accVal and the dot product
+         * (accVal) can sometimes have round-off errors, which will cause (aNorm == bNorm) ~ accVal
+         * instead.
+         */
+        acc[i][j] = val * (val >= 1e-4 && !(regxn[i] == regyn[j] && accVal > 0.0));
       }
     }
     if (sqrt) {

diff --git a/python/pylibraft/pylibraft/test/test_distance.py b/python/pylibraft/pylibraft/test/test_distance.py
@@ -21,8 +21,8 @@
 from pylibraft.distance import pairwise_distance
 
 
-@pytest.mark.parametrize("n_rows", [32, 100])
-@pytest.mark.parametrize("n_cols", [40, 100])
+@pytest.mark.parametrize("n_rows", [50, 100])
+@pytest.mark.parametrize("n_cols", [10, 50])
 @pytest.mark.parametrize(
     "metric",
     [
@@ -63,7 +63,7 @@ def test_distance(n_rows, n_cols, inplace, metric, order, dtype):
     else:
         expected = cdist(input1, input1, metric)
 
-    expected[expected <= 1e-5] = 0.0
+    # expected[expected <= 1e-5] = 0.0
 
     input1_device = device_ndarray(input1)
     output_device = device_ndarray(output) if inplace else None
@@ -79,6 +79,8 @@ def test_distance(n_rows, n_cols, inplace, metric, order, dtype):
 
     actual = output_device.copy_to_host()
 
-    actual[actual <= 1e-5] = 0.0
+    # actual[actual <= 1e-5] = 0.0
+    # if metric == "euclidean":
+    #     np.fill_diagonal(actual, 0.0)
 
     assert np.allclose(expected, actual, atol=1e-3, rtol=1e-3)