Neural batch size error more informative (#152)

src/bmi/estimators/neural/_estimators.py

@@ -138,6 +138,15 @@ def estimate_with_info(self, x: ArrayLike, y: ArrayLike) -> EstimateResult:
             xs, ys, train_size=self._params.train_test_split, key=key_split
         )
 
+        if self._params.batch_size > len(xs_train):
+            if self._verbose:
+                print("ERROR: Batch size larger than train dataset.")
+
+            return EstimateResult(
+                mi_estimate=float("nan"),
+                additional_information={"batch_size_larger_than_train": True},
+            )
+
         # initialize critic
         critic = self._critic_factory(key_init, xs_train.shape[-1], ys_train.shape[-1])
 

src/bmi/estimators/neural/_mine_estimator.py

@@ -329,6 +329,15 @@ def estimate_with_info(self, x: ArrayLike, y: ArrayLike) -> EstimateResult:
             xs, ys, train_size=self._params.train_test_split, key=key_split
         )
 
+        if self._params.batch_size > len(xs_train):
+            if self._verbose:
+                print("ERROR: Batch size larger than train dataset.")
+
+            return EstimateResult(
+                mi_estimate=float("nan"),
+                additional_information={"batch_size_larger_than_train": True},
+            )
+
         # initialize critic
         critic = self._create_critic(dim_x=space.dim_x, dim_y=space.dim_y, key=key_init)
 

tests/estimators/neural/test_neural.py

@@ -53,3 +53,25 @@ def test_critic_saved(estimator_class, n_points: int = 100) -> None:
 
     estimator.estimate(xs, ys)
     assert isinstance(estimator.trained_critic, eqx.Module)
+
+
+@pytest.mark.parametrize(
+    "estimator_class",
+    [
+        neural.DonskerVaradhanEstimator,
+        neural.NWJEstimator,
+        neural.InfoNCEEstimator,
+        neural.MINEEstimator,
+    ],
+)
+def test_batch_size_error(estimator_class, n_points: int = 100) -> None:
+    """Tests whether the critic is saved after estimation."""
+    batch_size = n_points + 1
+    estimator = estimator_class(batch_size=batch_size, learning_rate=0.15, max_n_steps=100)
+    sampler = samplers.SplitMultinormal(dim_x=2, dim_y=1, covariance=np.eye(3))
+    xs, ys = sampler.sample(n_points=n_points, rng=0)
+
+    result = estimator.estimate_with_info(xs, ys)
+    assert np.isnan(result.mi_estimate)
+    assert "batch_size_larger_than_train" in result.additional_information
+    assert result.additional_information["batch_size_larger_than_train"]

workflows/benchmark/_utils.py

@@ -105,7 +105,7 @@ def make_pretty(styler):
         )
         if converged_only:
             styler.set_caption(
-                "Estimates smaller than 10% of the maximal estimate are excluded."
+                "Estimates smaller than 10% of the maximal estimate are excluded. "
                 "This can help neural estimators which sometimes fail to converge."
             )
         return styler
@@ -154,7 +154,7 @@ def make_pretty(styler):
             cmap="gray",
         )
         styler.set_caption(
-            "Estimates higher than 10% of the maximal estimate are considered"
+            "Estimates higher than 10% of the maximal estimate are considered "
             "converged. This table shows the percentage of converged estimates."
         )
         return styler