Allow continuation of MCMC chains when using latest_sample

CHANGELOG.md

@@ -1,6 +1,7 @@
 # v0.13.1
 
 - Make logging of vectorized numpy slice sampler slightly less verbose and address NumPy future warning (#347)
+- Allow continuation of MCMC chains (#348)
 
 
 # v0.13.0

sbi/inference/posteriors/base_posterior.py

@@ -22,7 +22,14 @@
 from torch import nn
 
 from sbi import utils as utils
-from sbi.mcmc import Slice, SliceSampler, SliceSamplerVectorized, prior_init, sir
+from sbi.mcmc import (
+    Slice,
+    SliceSampler,
+    SliceSamplerVectorized,
+    IterateParameters,
+    prior_init,
+    sir,
+)
 from sbi.types import Array, Shape
 from sbi.user_input.user_input_checks import process_x
 from sbi.utils.torchutils import (
@@ -594,7 +601,8 @@ def _build_mcmc_init_fn(
         elif init_strategy == "sir":
             return lambda: sir(prior, potential_fn, **kwargs)
         elif init_strategy == "latest_sample":
-            return lambda: self._mcmc_init_params
+            latest_sample = IterateParameters(self._mcmc_init_params, **kwargs)
+            return latest_sample
         else:
             raise NotImplementedError
 

sbi/inference/snle/snle_base.py

@@ -158,6 +158,10 @@ def __call__(
                 mcmc_parameters=self._mcmc_parameters,
             )
 
+        # Copy MCMC init parameters for latest sample init
+        if hasattr(proposal, "_mcmc_init_params"):
+            self._posterior._mcmc_init_params = proposal._mcmc_init_params
+
         # Fit neural likelihood to newly aggregated dataset.
         self._train(
             training_batch_size=training_batch_size,

sbi/inference/snpe/snpe_base.py

@@ -198,6 +198,10 @@ def __call__(
                 rejection_sampling_parameters=self._rejection_sampling_parameters,
             )
 
+        # Copy MCMC init parameters for latest sample init
+        if hasattr(proposal, "_mcmc_init_params"):
+            self._posterior._mcmc_init_params = proposal._mcmc_init_params
+
         # Fit posterior using newly aggregated data set.
         self._train(
             proposal=proposal,

sbi/inference/snre/snre_base.py

@@ -164,6 +164,10 @@ def __call__(
                 mcmc_parameters=self._mcmc_parameters,
             )
 
+        # Copy MCMC init parameters for latest sample init
+        if hasattr(proposal, "_mcmc_init_params"):
+            self._posterior._mcmc_init_params = proposal._mcmc_init_params
+
         # Fit posterior using newly aggregated data set.
         self._train(
             num_atoms=num_atoms,

sbi/mcmc/__init__.py

@@ -1,4 +1,4 @@
 from sbi.mcmc.slice_numpy import SliceSampler
 from sbi.mcmc.slice_numpy_vectorized import SliceSamplerVectorized
 from sbi.mcmc.slice import Slice
-from sbi.mcmc.init_strategy import prior_init, sir
+from sbi.mcmc.init_strategy import IterateParameters, prior_init, sir

sbi/mcmc/init_strategy.py

@@ -9,6 +9,22 @@
 from torch import Tensor
 
 
+class IterateParameters:
+    """Iterates through parameters by rows
+    """
+
+    def __init__(self, parameters: torch.Tensor, **kwargs):
+        self.iter = self._make_iterator(parameters)
+
+    @staticmethod
+    def _make_iterator(t):
+        for i in range(t.shape[0]):
+            yield t[i, :].reshape(1, -1)
+
+    def __call__(self):
+        return next(self.iter)
+
+
 def prior_init(prior: Any, **kwargs: Any) -> Tensor:
     """Return a sample from the prior."""
     return prior.sample((1,)).detach()