refactor: improve abc implementation and tests. (#1615)

* refactor: improve implementation and consistency of inference/abc module

* refactor: add type hints to abc module

* refactor: add type hints to smcabc

* refactor method visibility and types

* docs(abc): Add docstrings to distances.py

* feat(abc): Refactor ABC module for consistent privatization and type hints

* fix abc tests, fix mmd metric

* consolidate distances and metrics

* fix string output

* refactor(abc): lra and sass speed and side effects

-  use super() in MCABC

* ignore all venvs

* add lra and sass tests

* typing and reduced comments

* refactor smcabc test for speed

* small fixes

* refactor abc tests for speed and coverage

* fix gitignore

* fix abc tests typing

* remove empty file

* small fixes

Author: janfb

.gitignore

@@ -9,7 +9,7 @@
 /*.egg
 /*.egg-info
 .venv*
-/.serena
+
 # Notebook checkpoints
 .ipynb_checkpoints
 
@@ -102,3 +102,6 @@ target/
 # uv
 uv.lock
 .python-version
+
+# Serena cache
+.serena/

sbi/inference/abc/abc_base.py

@@ -6,13 +6,13 @@
 import logging
 from typing import Callable, Dict, Optional, Union
 
-import numpy as np
 import torch
 from sklearn.linear_model import LinearRegression
 from sklearn.preprocessing import PolynomialFeatures
+from torch import Tensor
 
-from sbi.inference.abc.distances import Distance
 from sbi.simulators.simutils import simulate_in_batches
+from sbi.utils.metrics import Distance
 
 
 class ABCBASE:
@@ -84,12 +84,12 @@ def __init__(
         self.logger = logging.getLogger(__name__)
 
     @staticmethod
-    def get_sass_transform(
-        theta: torch.Tensor,
-        x: torch.Tensor,
+    def _get_sass_transform(
+        theta: Tensor,
+        x: Tensor,
         expansion_degree: int = 1,
-        sample_weight=None,
-    ) -> Callable:
+        sample_weight: Optional[Tensor] = None,
+    ) -> Callable[[Tensor], Tensor]:
         """Return semi-automatic summary statitics function.
 
         Running weighted linear regressin as in
@@ -102,17 +102,18 @@ def get_sass_transform(
         """
         expansion = PolynomialFeatures(degree=expansion_degree, include_bias=False)
         # Transform x, remove intercept.
-        x_expanded = expansion.fit_transform(x)
-        sumstats_map = np.zeros((x_expanded.shape[1], theta.shape[1]))
-
-        for parameter_idx in range(theta.shape[1]):
-            regression_model = LinearRegression(fit_intercept=True)
-            regression_model.fit(
-                X=x_expanded, y=theta[:, parameter_idx], sample_weight=sample_weight
-            )
-            sumstats_map[:, parameter_idx] = regression_model.coef_
+        x_expanded = expansion.fit_transform(x.numpy())
+
+        # Fit a single multi-output regression model for all parameters at once
+        regression_model = LinearRegression(fit_intercept=True)
+        regression_model.fit(
+            X=x_expanded,
+            y=theta.numpy(),  # All parameters at once
+            sample_weight=sample_weight.numpy() if sample_weight is not None else None,
+        )
 
-        sumstats_map = torch.tensor(sumstats_map, dtype=torch.float32)
+        # Get coefficients for all parameters (shape: [n_features, n_parameters])
+        sumstats_map = torch.tensor(regression_model.coef_.T, dtype=torch.float32)
 
         def sumstats_transform(x):
             x_expanded = torch.tensor(expansion.fit_transform(x), dtype=torch.float32)
@@ -121,28 +122,30 @@ def sumstats_transform(x):
         return sumstats_transform
 
     @staticmethod
-    def run_lra(
-        theta: torch.Tensor,
-        x: torch.Tensor,
-        observation: torch.Tensor,
-        sample_weight=None,
-    ) -> torch.Tensor:
+    def _run_lra(
+        theta: Tensor,
+        x: Tensor,
+        observation: Tensor,
+        sample_weight: Optional[Tensor] = None,
+    ) -> Tensor:
         """Return parameters adjusted with linear regression adjustment.
 
         Implementation as in Beaumont et al. 2002: https://arxiv.org/abs/1707.01254
         """
 
-        theta_adjusted = theta
-        for parameter_idx in range(theta.shape[1]):
-            regression_model = LinearRegression(fit_intercept=True)
-            regression_model.fit(
-                X=x,
-                y=theta[:, parameter_idx],
-                sample_weight=sample_weight,
-            )
-            theta_adjusted[:, parameter_idx] += regression_model.predict(
-                observation.reshape(1, -1)
-            )
-            theta_adjusted[:, parameter_idx] -= regression_model.predict(x)
+        # Fit a single multi-output regression model
+        regression_model = LinearRegression(fit_intercept=True)
+        regression_model.fit(
+            X=x.numpy(),
+            y=theta.numpy(),  # All parameters at once
+            sample_weight=sample_weight.numpy() if sample_weight is not None else None,
+        )
+
+        # Predict for observation and simulated data
+        pred_obs = regression_model.predict(observation.reshape(1, -1).numpy())
+        pred_sim = regression_model.predict(x.numpy())
+
+        # Apply adjustment: theta + m(x_o) - m(x)
+        theta_adjusted = theta + torch.from_numpy(pred_obs - pred_sim)
 
         return theta_adjusted

sbi/inference/abc/distances.py

@@ -8,6 +8,7 @@
 import torch
 from numpy import ndarray
 from torch import Tensor
+from torch.distributions import Distribution
 
 from sbi.inference.abc.abc_base import ABCBASE
 from sbi.utils.kde import KDEWrapper, get_kde
@@ -20,9 +21,9 @@ class MCABC(ABCBASE):
     def __init__(
         self,
         simulator: Callable,
-        prior,
+        prior: Distribution,
         distance: Union[str, Callable] = "l2",
-        requires_iid_data: Optional[None] = None,
+        requires_iid_data: Optional[bool] = None,
         distance_kwargs: Optional[Dict] = None,
         num_workers: int = 1,
         simulation_batch_size: int = 1,
@@ -147,15 +148,15 @@ def __call__(
             pilot_theta = self.prior.sample((num_pilot_simulations,))
             pilot_x = self._batched_simulator(pilot_theta)
 
-            sass_transform = self.get_sass_transform(
+            sass_transform = super()._get_sass_transform(
                 pilot_theta, pilot_x, sass_expansion_degree
             )
 
             # Add sass transform to simulator and x_o.
             def simulator(theta):
                 return sass_transform(self._batched_simulator(theta))
 
-            x_o = sass_transform(x_o)
+            x_o = sass_transform(process_x(x_o))
         else:
             simulator = self._batched_simulator
 
@@ -173,11 +174,10 @@ def simulator(theta):
         if not self.distance.requires_iid_data:
             x = x.squeeze(1)
             self.x_shape = x[0].shape
-            self.x_o = process_x(x_o, self.x_shape)
         else:
             self.x_shape = x[0, 0].shape
-            self.x_o = process_x(x_o, self.x_shape)
 
+        self.x_o = process_x(x_o, self.x_shape)
         distances = self.distance(self.x_o, x)
 
         # Select based on acceptance threshold epsilon.
@@ -204,16 +204,18 @@ def simulator(theta):
         # Maybe adjust theta with LRA.
         if lra:
             self.logger.info("Running Linear regression adjustment.")
-            final_theta = self.run_lra(theta_accepted, x_accepted, observation=self.x_o)
+            final_theta = super()._run_lra(
+                theta_accepted, x_accepted, observation=self.x_o
+            )
         else:
             final_theta = theta_accepted
 
         if kde:
             self.logger.info(
-                """KDE on %s samples with bandwidth option
-                {kde_kwargs["bandwidth"] if "bandwidth" in kde_kwargs else "cv"}.
-                Beware that KDE can give unreliable results when used with too few
-                samples and in high dimensions.""",
+                "KDE on %s samples with bandwidth option "
+                f"{kde_kwargs.get('bandwidth', 'cv')}. "
+                "Beware that KDE can give unreliable results when used with too few"
+                " samples and in high dimensions.",
                 final_theta.shape[0],
             )