Merge pull request #661 from aai-institute/feature/max-samples

Feature/max samples

Author: mdbenito

CHANGELOG.md

@@ -5,6 +5,8 @@
 
 ### Added
 
+- New stopping criterion `MaxSamples`
+  [PR #661](https://github.com/aai-institute/pyDVL/pull/661)
 - Introduced `UtilityModel` and two implementations `IndicatorUtilityModel`
   and `DeepSetsUtilityModel` for data utility learning
   [PR #650](https://github.com/aai-institute/pyDVL/pull/650)

src/pydvl/valuation/stopping.py

@@ -187,9 +187,10 @@
     "AbsoluteStandardError",
     "HistoryDeviation",
     "MaxChecks",
+    "MaxSamples",
+    "MaxTime",
     "MaxUpdates",
     "MinUpdates",
-    "MaxTime",
     "NoStopping",
     "RankCorrelation",
     "StoppingCriterion",
@@ -614,9 +615,52 @@ def completion(self) -> float:
             return 0.0
 
     def __str__(self) -> str:
+        if self.sampler is not None:
+            return f"NoStopping({self.sampler.__class__.__name__})"
         return "NoStopping()"
 
 
+class MaxSamples(StoppingCriterion):
+    """Run until the sampler has sampled the given number of samples.
+
+    !!! warning
+        If the sampler is batched, and the valuation method runs in parallel, the check
+        might be off by the sampler's batch size.
+
+    Args:
+        sampler: The sampler to check.
+        n_samples: The number of samples to run until.
+        modify_result: If `True` the status of the input
+            [ValuationResult][pydvl.valuation.result.ValuationResult] is modified in
+            place after the call.
+    """
+
+    def __init__(
+        self, sampler: IndexSampler, n_samples: int, modify_result: bool = True
+    ):
+        if n_samples <= 0:
+            raise ValueError("n_samples must be positive")
+        super().__init__(modify_result=modify_result)
+        self.sampler = sampler
+        self.n_samples = n_samples
+        self._completion = 0.0
+
+    def _check(self, result: ValuationResult) -> Status:
+        self._completion = np.clip(self.sampler.n_samples / self.n_samples, 0.0, 1.0)
+        if self.sampler.n_samples >= self.n_samples:
+            self._converged = np.full_like(result.indices, True, dtype=bool)
+            return Status.Converged
+        return Status.Pending
+
+    def completion(self) -> float:
+        return self._completion
+
+    def __str__(self) -> str:
+        return (
+            f"MaxSamples({self.sampler.__class__.__name__}, n_samples={self.n_samples})"
+        )
+
+
 class MinUpdates(StoppingCriterion):
     """Terminate as soon as all value updates exceed or equal the given threshold.
 
@@ -965,4 +1009,4 @@ def reset(self) -> Self:
         return super().reset()
 
     def __str__(self):
-        return f"RankCorrelation({self.rtol=}, {self.burn_in=}, {self.fraction=})"
+        return f"RankCorrelation(rtol={self.rtol}, burn_in={self.burn_in}, fraction={self.fraction})"

tests/valuation/test_stopping.py

@@ -12,6 +12,7 @@
     AbsoluteStandardError,
     HistoryDeviation,
     MaxChecks,
+    MaxSamples,
     MaxTime,
     MaxUpdates,
     MinUpdates,
@@ -349,25 +350,41 @@ def test_no_stopping_without_sampler():
     assert str(no_stop) == "NoStopping()"
 
 
-def test_no_stopping_with_finite_sampler():
-    class DummyFiniteSampler(IndexSampler):
-        def __init__(self, total_samples: int = 10, batch_size: int = 1):
-            super().__init__(batch_size=batch_size)
-            self.total_samples = total_samples
+class DummyFiniteSampler(IndexSampler):
+    def __init__(self, total_samples: int = 10, batch_size: int = 1):
+        super().__init__(batch_size=batch_size)
+        self.total_samples = total_samples
+
+    def sample_limit(self, indices):
+        return self.total_samples
+
+    def generate(self, indices):
+        for i in range(self.total_samples):
+            yield i, set()
+
+    def log_weight(self, n, subset_len):
+        return 0.0
+
+    def make_strategy(self, utility, log_coefficient=None):
+        return None
+
+
+class DummyInfiniteSampler(IndexSampler):
+    def sample_limit(self, indices):
+        return None  # Indicates an infinite sampler.
 
-        def sample_limit(self, indices):
-            return self.total_samples
+    def generate(self, indices):
+        while True:
+            yield (0, set())
 
-        def generate(self, indices):
-            for i in range(self.total_samples):
-                yield i, set()
+    def log_weight(self, n, subset_len):
+        return 0.0
 
-        def log_weight(self, n, subset_len):
-            return 0.0
+    def make_strategy(self, utility, log_coefficient=None):
+        return None
 
-        def make_strategy(self, utility, log_coefficient=None):
-            return None
 
+def test_no_stopping_with_finite_sampler():
     r = ValuationResult.from_random(5)
     total_samples = 10
     batch_size = 3
@@ -397,28 +414,59 @@ def make_strategy(self, utility, log_coefficient=None):
 
 
 def test_no_stopping_infinite_sampler():
-    class DummyInfiniteSampler(IndexSampler):
-        def sample_limit(self, indices):
-            return None  # Indicates an infinite sampler.
-
-        def generate(self, indices):
-            while True:
-                yield (0, set())
-
-        def log_weight(self, n, subset_len):
-            return 0.0
-
-        def make_strategy(self, utility, log_coefficient=None):
-            return None
-
     sampler = DummyInfiniteSampler(batch_size=1)
     no_stop = NoStopping(sampler=sampler)
 
-    batches = list(islice(sampler.generate_batches(np.array([0])), 10))
-    assert sampler.n_samples == len(batches)
+    _ = list(islice(sampler.generate_batches(np.array([0])), 10))
 
     # Verify that calling the criterion still returns Pending and marks no index as converged.
     result = ValuationResult.from_random(5)
     status = no_stop(result)
     assert status == Status.Pending
+    assert no_stop.completion() == 0.0
     np.testing.assert_equal(no_stop.converged, False)
+
+
+def test_max_samples_pending_and_convergence():
+    sampler = DummyInfiniteSampler(batch_size=1)
+    threshold = 10
+    max_samples = MaxSamples(sampler, n_samples=threshold)
+    result = ValuationResult.from_random(5)  # Create a result with 5 indices
+
+    status = max_samples(result)
+    assert status == Status.Pending
+    np.testing.assert_allclose(max_samples.completion(), 0.0)
+    assert not max_samples.converged.all()
+
+    # Set sampler.n_samples below threshold.
+    _ = list(islice(sampler.generate_batches(np.array([0])), 5))
+    status = max_samples(result)
+    assert status == Status.Pending
+    np.testing.assert_allclose(max_samples.completion(), 5 / threshold)
+    assert not max_samples.converged.all()
+
+    # Set sampler.n_samples exactly equal to threshold.
+    _ = list(islice(sampler.generate_batches(np.array([0])), 10))
+    status = max_samples(result)
+    assert status == Status.Converged
+    np.testing.assert_allclose(max_samples.completion(), 1.0)
+    assert max_samples.converged.all()
+
+    # Set sampler.n_samples above threshold.
+    _ = list(islice(sampler.generate_batches(np.array([0])), 15))
+    status = max_samples(result)
+    assert status == Status.Converged
+    np.testing.assert_allclose(max_samples.completion(), 1.0)
+    assert max_samples.converged.all()
+
+
+def test_max_samples_str_and_invalid():
+    sampler = DummyFiniteSampler(total_samples=0)
+    max_samples = MaxSamples(sampler, 10)
+    expected_str = f"MaxSamples({sampler.__class__.__name__}, n_samples=10)"
+    assert str(max_samples) == expected_str
+
+    with pytest.raises(ValueError):
+        MaxSamples(sampler, 0)
+    with pytest.raises(ValueError):
+        MaxSamples(sampler, -5)