remove bm_test changes

Author: janfb

tests/bm_test.py

@@ -11,7 +11,6 @@
 from sbi.inference.posteriors.base_posterior import NeuralPosterior
 from sbi.inference.trainers.npe import NPE_C
 from sbi.inference.trainers.nre import BNRE, NRE_A, NRE_B, NRE_C
-from sbi.neural_nets.factory import posterior_score_nn
 from sbi.utils.metrics import c2st
 
 from .mini_sbibm import get_task
@@ -34,78 +33,6 @@
     "transformer",
 ]
 
-VP_SCHEDULES = [
-    {"name": "vp_default", "beta_min": 0.01, "beta_max": 10.0},
-    {"name": "vp_wide", "beta_min": 0.1, "beta_max": 20.0},
-]
-SUBVP_SCHEDULES = [
-    {"name": "subvp_default", "beta_min": 0.01, "beta_max": 10.0},
-    {"name": "subvp_wide", "beta_min": 0.1, "beta_max": 20.0},
-]
-VE_SCHEDULES = [
-    {"name": "ve_default", "sigma_min": 1e-4, "sigma_max": 10.0},
-    {"name": "ve_wide", "sigma_min": 1e-3, "sigma_max": 50.0},
-]
-
-
-def _make_npse_builder(
-    vf_estimator: Literal["mlp", "ada_mlp", "transformer"],
-    sde_type: Literal["vp", "subvp", "ve"],
-    schedule_kwargs: dict[str, float | str],
-):
-    builder = posterior_score_nn(model=vf_estimator, sde_type=sde_type)
-
-    def build_fn(batch_theta, batch_x):
-        estimator = builder(batch_theta, batch_x)
-        if sde_type in {"vp", "subvp"}:
-            if "beta_min" in schedule_kwargs:
-                estimator.beta_min = schedule_kwargs["beta_min"]
-            if "beta_max" in schedule_kwargs:
-                estimator.beta_max = schedule_kwargs["beta_max"]
-        elif sde_type == "ve":
-            if "sigma_min" in schedule_kwargs:
-                estimator.sigma_min = schedule_kwargs["sigma_min"]
-            if "sigma_max" in schedule_kwargs:
-                estimator.sigma_max = schedule_kwargs["sigma_max"]
-
-        t_tensor = torch.as_tensor([estimator.t_max], device=estimator.mean_0.device)
-        mean_t = estimator.approx_marginal_mean(t_tensor)
-        std_t = estimator.approx_marginal_std(t_tensor)
-        mean_t = torch.broadcast_to(mean_t, (1, *estimator.input_shape))
-        std_t = torch.broadcast_to(std_t, (1, *estimator.input_shape))
-        estimator._mean_base = mean_t
-        estimator._std_base = std_t
-        return estimator
-
-    return build_fn
-
-
-def _npse_variants():
-    variants = []
-    for vf_estimator in VF_ESTIMATORS:
-        for schedule in VP_SCHEDULES:
-            label = f"{vf_estimator}/vp/{schedule['name']}"
-            variants.append({
-                "vf_estimator": _make_npse_builder(vf_estimator, "vp", schedule),
-                "sde_type": "vp",
-                "_bm_label": label,
-            })
-        for schedule in SUBVP_SCHEDULES:
-            label = f"{vf_estimator}/subvp/{schedule['name']}"
-            variants.append({
-                "vf_estimator": _make_npse_builder(vf_estimator, "subvp", schedule),
-                "sde_type": "subvp",
-                "_bm_label": label,
-            })
-        for schedule in VE_SCHEDULES:
-            label = f"{vf_estimator}/ve/{schedule['name']}"
-            variants.append({
-                "vf_estimator": _make_npse_builder(vf_estimator, "ve", schedule),
-                "sde_type": "ve",
-                "_bm_label": label,
-            })
-    return variants
-
 
 # Benchmarking method groups i.e. what to run for different --bm-mode
 METHOD_GROUPS = {
@@ -126,7 +53,11 @@ def _npse_variants():
     "nle": [{"density_estimator": de} for de in ["maf", "nsf"]],
     "nre": [{"classifier": cl} for cl in CLASSIFIERS],
     "fmpe": [{"vf_estimator": nn} for nn in VF_ESTIMATORS],
-    "npse": _npse_variants(),
+    "npse": [
+        {"vf_estimator": nn, "sde_type": sde}
+        for nn in VF_ESTIMATORS
+        for sde in ["ve", "vp"]
+    ],
     "vfpe": [{"vf_estimator": nn} for nn in VF_ESTIMATORS],
     "snpe": [{}],
     "snle": [{}],
@@ -260,8 +191,6 @@ def train_and_eval_amortized_inference(
     thetas, xs = task.get_data(benchmark_num_simulations)
     prior = task.get_prior()
 
-    extra_kwargs = dict(extra_kwargs)
-    bm_label = extra_kwargs.pop("_bm_label", None)
     inference = inference_class(prior, **extra_kwargs)
     _ = inference.append_simulations(thetas, xs).train(**TRAIN_KWARGS)
 
@@ -273,10 +202,6 @@ def train_and_eval_amortized_inference(
     results_bag.metric = mean_c2st
     results_bag.num_simulations = benchmark_num_simulations
     results_bag.task_name = task_name
-    if bm_label is None:
-        results_bag.method = inference_class.__name__ + str(extra_kwargs)
-    else:
-        results_bag.method = f"{inference_class.__name__}({bm_label})"
 
 
 def train_and_eval_sequential_inference(
@@ -305,8 +230,6 @@ def train_and_eval_sequential_inference(
     simulator = task.get_simulator()
 
     # Round 1
-    extra_kwargs = dict(extra_kwargs)
-    bm_label = extra_kwargs.pop("_bm_label", None)
     inference = inference_class(prior, **extra_kwargs)
     _ = inference.append_simulations(thetas, xs).train(**TRAIN_KWARGS)
 
@@ -330,10 +253,6 @@ def train_and_eval_sequential_inference(
     results_bag.metric = c2st_val
     results_bag.num_simulations = benchmark_num_simulations
     results_bag.task_name = task_name
-    if bm_label is None:
-        results_bag.method = inference_class.__name__ + str(extra_kwargs)
-    else:
-        results_bag.method = f"{inference_class.__name__}({bm_label})"
 
 
 @pytest.mark.benchmark