Fix purify stopping before the tensor is pure with a random order

python/interpret-core/tests/utils/test_purify.py

@@ -268,3 +268,46 @@ def test_purify_multiclass_5():
     scores -= np.average(scores, axis=-1, keepdims=True)
 
     assert np.allclose(refilled, scores)
+
+
+def _random_empty_cell_inputs(rng, shape, p_empty, n_classes=None):
+    weights = rng.uniform(0.1, 3.0, shape) * (rng.random(shape) > p_empty)
+    score_shape = shape if n_classes is None else (*shape, n_classes)
+    scores = rng.normal(size=score_shape)
+    return scores, weights
+
+
+def test_purify_randomized_with_empty_cells_regression():
+    # A fresh random visiting order on each sweep used to end the purification at the first
+    # sweep that did not shrink the total impurity, which left such tensors impure.
+    rng = np.random.default_rng(0)
+    for shape, p_empty in (((5, 6), 0.6), ((3, 3, 3), 0.3), ((4, 5, 6), 0.3)):
+        for _ in range(10):
+            scores, weights = _random_empty_cell_inputs(rng, shape, p_empty)
+            purified, impurities, intercept = purify(
+                scores, weights, tolerance=0.0, is_randomized=True
+            )
+
+            for axis in range(len(shape)):
+                assert np.allclose((purified * weights).sum(axis=axis), 0.0, atol=1e-8)
+
+            refilled = purified + intercept
+            for idxes, impurity in impurities:
+                new_shape = [1] * len(shape)
+                for i in idxes:
+                    new_shape[i] = shape[i]
+                refilled += impurity.reshape(new_shape)
+            assert np.allclose(refilled, scores)
+
+
+def test_purify_randomized_multiclass_with_empty_cells_regression():
+    rng = np.random.default_rng(1)
+    shape = (5, 6)
+    n_classes = 3
+    for _ in range(10):
+        scores, weights = _random_empty_cell_inputs(rng, shape, 0.6, n_classes)
+        purified, _, _ = purify(scores, weights, tolerance=0.0, is_randomized=True)
+
+        weighted = purified * weights[..., np.newaxis]
+        for axis in range(len(shape)):
+            assert np.allclose(weighted.sum(axis=axis), 0.0, atol=1e-8)

shared/libebm/Purify.cpp

@@ -27,6 +27,53 @@ namespace DEFINED_ZONE_NAME {
 #error DEFINED_ZONE_NAME must be defined
 #endif // DEFINED_ZONE_NAME
 
+// When the surfaces are visited in a new random order on each sweep, the total impurity of a sweep does not shrink
+// monotonically. Empty cells make the convergence slow, so a sweep that fails to beat the previous one usually comes
+// long before the tensor is pure. Such a sweep only ends the purification once the total impurity has dropped to
+// floating point noise relative to the first sweep. Otherwise the purification continues until k_cSweepsPatience
+// sweeps in a row have failed to set a new minimum, which also guarantees that it terminates.
+static constexpr double k_impurityNoiseRatio = 1e-12;
+static constexpr size_t k_cSweepsPatience = 20;
+
+struct PurifyProgress {
+   double m_impurityNoise;
+   double m_impurityBest;
+   size_t m_cSweepsSinceBest;
+};
+
+static void ResetPurifyProgress(PurifyProgress* const pProgress) {
+   pProgress->m_impurityNoise = std::numeric_limits<double>::quiet_NaN();
+   pProgress->m_impurityBest = std::numeric_limits<double>::infinity();
+   pProgress->m_cSweepsSinceBest = 0;
+}
+
+// Returns true when no further sweep is needed. impurityPrev is +inf for the first sweep.
+static bool IsPurifySweepingDone(
+      const bool bRandomized, const double impurityPrev, const double impurityCur, PurifyProgress* const pProgress) {
+   if(!bRandomized) {
+      // the order is fixed, so the total impurity shrinks until it reaches floating point noise
+      return impurityPrev <= impurityCur;
+   }
+   if(std::isnan(impurityCur)) {
+      // a score became infinite, so the total is not comparable with earlier sweeps. Track progress anew.
+      ResetPurifyProgress(pProgress);
+      return false;
+   }
+   if(std::isnan(pProgress->m_impurityNoise)) {
+      pProgress->m_impurityNoise = impurityCur * k_impurityNoiseRatio;
+   }
+   if(impurityCur < pProgress->m_impurityBest) {
+      pProgress->m_impurityBest = impurityCur;
+      pProgress->m_cSweepsSinceBest = 0;
+   } else {
+      ++pProgress->m_cSweepsSinceBest;
+   }
+   if(impurityPrev <= impurityCur && impurityCur <= pProgress->m_impurityNoise) {
+      return true;
+   }
+   return k_cSweepsPatience <= pProgress->m_cSweepsSinceBest;
+}
+
 EBM_API_BODY double EBM_CALLING_CONVENTION MeasureImpurity(IntEbm countMultiScores,
       IntEbm indexMultiScore,
       IntEbm countDimensions,
@@ -507,6 +554,8 @@ extern ErrorEbm PurifyInternal(const double tolerance,
             EBM_ASSERT(std::numeric_limits<double>::min() <= impuritySumOverflowPreventer);
             double impurityPrev;
             double impurityCur = std::numeric_limits<double>::infinity();
+            PurifyProgress progress;
+            ResetPurifyProgress(&progress);
             bool bRetry;
             do {
                // if any non-infinite value was flipped to an infinite value, it could increase the impurity
@@ -842,7 +891,7 @@ extern ErrorEbm PurifyInternal(const double tolerance,
                   ++iRandom;
                } while(cSurfaceBins != iRandom);
                // this loops on std::isnan(impurityCur)
-            } while(bRetry && !(impurityPrev <= impurityCur));
+            } while(bRetry && !IsPurifySweepingDone(nullptr != aRandomize, impurityPrev, impurityCur, &progress));
             EBM_ASSERT(!std::isnan(impurityCur));
 
             if(nullptr != pIntercept) {
@@ -1514,6 +1563,8 @@ static ErrorEbm PurifyNormalizedMulticlass(const size_t cScores,
       const double impuritySumOverflowPreventer = 0.5 / static_cast<double>(cScores * cSurfaceBins);
       double impurityPrev;
       double impurityCur = std::numeric_limits<double>::infinity();
+      PurifyProgress progress;
+      ResetPurifyProgress(&progress);
       do {
          // if any non-infinite value was flipped to an infinite value, it could increase the impurity
          // so we set impurityCur to NaN. We need to reset it to +inf to avoid stopping early
@@ -1765,7 +1816,7 @@ static ErrorEbm PurifyNormalizedMulticlass(const size_t cScores,
             ++iRandom;
          } while(cSurfaceBins != iRandom);
          // this loops on std::isnan(impurityCur)
-      } while(!(impurityPrev <= impurityCur));
+      } while(!IsPurifySweepingDone(nullptr != aRandomize, impurityPrev, impurityCur, &progress));
       EBM_ASSERT(!std::isnan(impurityCur));
 
       if(nullptr != aInterceptOut) {

shared/libebm/tests/PurifyTest.cpp

@@ -603,6 +603,115 @@ TEST_CASE("Purify simple 3x4x5") {
    CHECK(-1e-20 < impurity && impurity < 1e-20);
 }
 
+// a repeatable pseudo random number in [0, 1)
+static double NextUniform(uint64_t& state) {
+   state = state * 6364136223846793005ULL + 1442695040888963407ULL;
+   return static_cast<double>(state >> 11) / 9007199254740992.0;
+}
+
+// The largest absolute weighted sum over any row or column of a 2 dimensional tensor. It is zero when the tensor is
+// pure. MeasureImpurity cannot be used here since it divides by the weight of each row and column, which is NaN when
+// a row or column has no weight. The first dimension varies fastest in the tensor, so cell (iRow, iCol) is stored at
+// iRow + cRows * iCol.
+static double MaxWeightedFiberSum(const size_t cRows,
+      const size_t cCols,
+      const double* const aWeights,
+      const double* const aScores,
+      const size_t cClasses,
+      const size_t iClass) {
+   double maxSum = 0.0;
+   for(size_t iRow = 0; iRow < cRows; ++iRow) {
+      double sum = 0.0;
+      for(size_t iCol = 0; iCol < cCols; ++iCol) {
+         const size_t iBin = iRow + cRows * iCol;
+         sum += aWeights[iBin] * aScores[iBin * cClasses + iClass];
+      }
+      maxSum = std::max(maxSum, std::abs(sum));
+   }
+   for(size_t iCol = 0; iCol < cCols; ++iCol) {
+      double sum = 0.0;
+      for(size_t iRow = 0; iRow < cRows; ++iRow) {
+         const size_t iBin = iRow + cRows * iCol;
+         sum += aWeights[iBin] * aScores[iBin * cClasses + iClass];
+      }
+      maxSum = std::max(maxSum, std::abs(sum));
+   }
+   return maxSum;
+}
+
+TEST_CASE("Purify randomized 5x6 with empty cells, zero tolerance") {
+   // Visiting the surfaces in a new random order on each sweep used to end the purification at the first sweep that
+   // did not shrink the total impurity, which left tensors with empty cells impure.
+   constexpr size_t cBins = 5 * 6;
+   const IntEbm dimensionLengths[]{5, 6};
+   const size_t cDimensions = sizeof(dimensionLengths) / sizeof(dimensionLengths[0]);
+
+   uint64_t state = 1;
+   for(size_t iTrial = 0; iTrial < 50; ++iTrial) {
+      double weights[cBins];
+      double scores[cBins];
+      for(size_t iBin = 0; iBin < cBins; ++iBin) {
+         const double isEmpty = NextUniform(state);
+         const double weight = 0.1 + 2.9 * NextUniform(state);
+         weights[iBin] = isEmpty < 0.6 ? 0.0 : weight;
+         scores[iBin] = 2.0 * NextUniform(state) - 1.0;
+      }
+      double impurities[5 + 6];
+      double residualIntercept;
+
+      ErrorEbm error = Purify(0.0,
+            k_isRandomizedDefault,
+            k_isMulticlassNormalization,
+            IntEbm{1},
+            cDimensions,
+            dimensionLengths,
+            weights,
+            scores,
+            impurities,
+            &residualIntercept);
+      CHECK(Error_None == error);
+      CHECK(MaxWeightedFiberSum(5, 6, weights, scores, 1, 0) < 1e-9);
+   }
+}
+
+TEST_CASE("Purify randomized multiclass 5x6 with empty cells, zero tolerance") {
+   constexpr IntEbm cClasses = 3;
+   constexpr size_t cBins = 5 * 6;
+   const IntEbm dimensionLengths[]{5, 6};
+   const size_t cDimensions = sizeof(dimensionLengths) / sizeof(dimensionLengths[0]);
+
+   uint64_t state = 1;
+   for(size_t iTrial = 0; iTrial < 50; ++iTrial) {
+      double weights[cBins];
+      double scores[cBins * cClasses];
+      for(size_t iBin = 0; iBin < cBins; ++iBin) {
+         const double isEmpty = NextUniform(state);
+         const double weight = 0.1 + 2.9 * NextUniform(state);
+         weights[iBin] = isEmpty < 0.6 ? 0.0 : weight;
+      }
+      for(size_t iScore = 0; iScore < cBins * cClasses; ++iScore) {
+         scores[iScore] = 2.0 * NextUniform(state) - 1.0;
+      }
+      double impurities[(5 + 6) * cClasses];
+      double residualIntercept[cClasses];
+
+      ErrorEbm error = Purify(0.0,
+            k_isRandomizedDefault,
+            k_isMulticlassNormalization,
+            cClasses,
+            cDimensions,
+            dimensionLengths,
+            weights,
+            scores,
+            impurities,
+            &residualIntercept[0]);
+      CHECK(Error_None == error);
+      for(size_t iClass = 0; iClass < static_cast<size_t>(cClasses); ++iClass) {
+         CHECK(MaxWeightedFiberSum(5, 6, weights, scores, static_cast<size_t>(cClasses), iClass) < 1e-9);
+      }
+   }
+}
+
 TEST_CASE("Purify simple multiclass 3x4") {
    constexpr IntEbm cClasses = 2;
    const IntEbm dimensionLengths[]{3, 4};