Add log_post_var strategy arg (facebookresearch#822)

Summary:
Pull Request resolved: facebookresearch#822

A new strategy option that uses the same computation as min_post_range but instead of deciding when to finish just logs prediction variance to see the progress of the model.

Reviewed By: tymmsc

Differential Revision: D83012919

fbshipit-source-id: 65097f611d9563742adda3d377ab62c3e7308ec8

Author: JasonKChow

aepsych/strategy/strategy.py

@@ -53,6 +53,7 @@ def __init__(
         max_asks: int | None = None,
         keep_most_recent: int | None = None,
         min_post_range: float | None = None,
+        log_post_var: bool = False,
         name: str = "",
         run_indefinitely: bool = False,
         transforms: ChainedInputTransform = ChainedInputTransform(**{}),
@@ -82,6 +83,7 @@ def __init__(
                 as data collected from later trials. When None, the model is fitted on all data.
             min_post_range (float, optional): Experimental. The required difference between the posterior's minimum and maximum value in
                 probablity space before the strategy will finish. Ignored if None (default).
+            log_post_var (bool): Whether to log the posterior prediction variance (as it would be used for min_post_range). Defaults to False.
             name (str): The name of the strategy. Defaults to the empty string.
             run_indefinitely (bool): If true, the strategy will run indefinitely until finish() is explicitly called. Other stopping criteria will
                 be ignored. Defaults to False.
@@ -161,8 +163,11 @@ def __init__(
         self.transforms = transforms
 
         self.min_post_range = min_post_range
-        if self.min_post_range is not None:
-            assert model is not None, "min_post_range must be None if model is None!"
+        self.log_post_var = log_post_var
+        if self.min_post_range is not None or self.log_post_var:
+            assert (
+                model is not None
+            ), "posterior range cannot be evaluated if model is None!"
             self.eval_grid = make_scaled_sobol(
                 lb=self.lb, ub=self.ub, size=self._n_eval_points
             )
@@ -429,16 +434,23 @@ def finished(self) -> bool:
         else:
             sufficient_outcomes = True
 
-        if self.min_post_range is not None:
+        if self.min_post_range is not None or self.log_post_var:
             assert (
                 self.model is not None
             ), "model is None! Cannot predict without a model!"
-            fmean, _ = self.model.predict(self.eval_grid, probability_space=True)
-            meets_post_range = bool(
-                ((fmean.max() - fmean.min()) >= self.min_post_range).item()
-            )
+            fmean, fvar = self.model.predict(self.eval_grid, probability_space=True)
+            post_range = fmean.max() - fmean.min()
+        else:
+            post_range = None
+
+        if post_range is not None:
+            logger.info(f"Mean posterior variance = {fvar.mean().item()}")
+
+            if self.min_post_range:
+                meets_post_range = bool((post_range >= self.min_post_range).item())
         else:
             meets_post_range = True
+
         finished = (
             self._count >= self.min_asks
             and self.n >= self.min_total_tells

tests/test_strategy.py

@@ -227,6 +227,42 @@ def test_run_indefinitely(self):
         self.strat.finish()
         self.assertTrue(self.strat.finished)
 
+    def test_log_post_var(self):
+        seed = 1
+        torch.manual_seed(seed)
+        np.random.seed(seed)
+        lb = [0]
+        ub = [1]
+
+        self.strat = Strategy(
+            model=GPClassificationModel(
+                dim=1,
+            ),
+            generator=SobolGenerator(lb=lb, ub=ub),
+            min_asks=10,
+            lb=lb,
+            ub=ub,
+            stimuli_per_trial=1,
+            outcome_types=["binary"],
+            log_post_var=True,
+        )
+
+        # Add some initial data
+        for _ in range(5):
+            points = self.strat.gen(1)
+            response = int(np.random.rand() < 0.5)
+            self.strat.add_data(points, torch.tensor([response]))
+
+        # Check that the log prints the expected message when checking finished status
+        with self.assertLogs() as log:
+            _ = self.strat.finished
+
+        # Look for the log message about mean posterior variance
+        log_found = any("Mean posterior variance" in msg for msg in log.output)
+        self.assertTrue(
+            log_found, "Expected log message about mean posterior variance not found"
+        )
+
     def test_batchsobol_pairwise(self):
         lb = [1, 2, 3]
         ub = [2, 3, 4]