Repository: sbi-dev/sbi

Issue:
sampling without inferring the prior transform
Dear SBI team, 

Im absolutely loving the package! Very well written and documented.

I need to fit a distribution that has convex polytope support, meaning that all samples `theta` need to satisfy the following: `A @ theta <= b`. I made the following support class for my prior:

```
from torch.distributions.constraints import Constraint, _Dependent

class _CannonicalPolytopeSupport(_Dependent):
    def __init__(
            self,
            A: torch.Tensor,
            b: torch.Tensor,
            validation_tol = 1e-13,
    ):
        self._A = A 
        self._b = b + validation_tol 
        super().__init__(is_discrete=False, event_dim=self._A.shape[1])

    def check(self, value):
        vape    = value.shape
        viewlue = value.permute((-1, *range(value.ndim -1))).view(vape[-1], vape[:-1].numel())
        valid   = (self._A @ viewlue <= self._b).T
        return valid.view(*vape[:-1], self._A.shape[0])
```

Now I run into the following issue. Once a round of inference is done and I try to build a posterior, `sbi` assumes that if you have a bounded and dependent support, that there must be a bijective transform to an  unconstrained space:

```
---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
Input In [32], in <module>
----> 1 posterior = snpec.build_posterior(density_estimator)

File C:\Miniconda3\envs\pta\lib\site-packages\sbi\inference\snpe\snpe_base.py:420, in PosteriorEstimator.build_posterior(self, density_estimator, prior, sample_with, mcmc_method, vi_method, mcmc_parameters, vi_parameters, rejection_sampling_parameters)
    417     # Otherwise, infer it from the device of the net parameters.
    418     device = next(density_estimator.parameters()).device.type
--> 420 potential_fn, theta_transform = posterior_estimator_based_potential(
    421     posterior_estimator=posterior_estimator, prior=prior, x_o=None
    422 )
    424 if sample_with == "rejection":
    425     if "proposal" in rejection_sampling_parameters.keys():

File C:\Miniconda3\envs\pta\lib\site-packages\sbi\inference\potentials\posterior_based_potential.py:46, in posterior_estimator_based_potential(posterior_estimator, prior, x_o)
     41 device = str(next(posterior_estimator.parameters()).device)
     43 potential_fn = PosteriorBasedPotential(
     44     posterior_estimator, prior, x_o, device=device
     45 )
---> 46 theta_transform = mcmc_transform(prior, device=device)
     48 return potential_fn, theta_transform

File C:\Miniconda3\envs\pta\lib\site-packages\sbi\utils\sbiutils.py:615, in mcmc_transform(prior, num_prior_samples_for_zscoring, enable_transform, device, **kwargs)
    613 # Prior with bounded support, e.g., uniform priors.
    614 if has_support and support_is_bounded:
--> 615     transform = biject_to(prior.support)
    616 # For all other cases build affine transform with mean and std.
    617 else:
    618     if hasattr(prior, "mean") and hasattr(prior, "stddev"):

File C:\Miniconda3\envs\pta\lib\site-packages\torch\distributions\constraint_registry.py:142, in ConstraintRegistry.__call__(self, constraint)
    140     factory = self._registry[type(constraint)]
    141 except KeyError:
--> 142     raise NotImplementedError(
    143         f'Cannot transform {type(constraint).__name__} constraints') from None
    144 return factory(constraint)

NotImplementedError: Cannot transform _CannonicalPolytopeSupport constraints
```

Is there any way around this issue? I've thought of such a transform (have not implemented it yet, since its a real headache), but for now I would like to just check whether a sample is within the support without transforming it to unconstrained space. 

Thank you for you efforts!


Discussion:
Perhaps I could just use rejection sampling by `.build_posterior(density_estimator, sample_with='rejection')`, and then we do not use the `theta_transform` at all, so we just have to build the `potential_fun` in `posterior_estimator_based_potential`. Or perhaps we can pass the `enable_transform=False` kwarg to `mcmc_transform` from `build_posterior` somehow.
Hi there! Thanks for the compliments, we're happy that you're enjoying the package :)

If you want to turn off the tranfsformation, you should probably use the [sampler interface](https://www.mackelab.org/sbi/tutorial/11_sampler_interface/).

Could you try this (I did not try it yet):
```python
from sbi.inference import DirectPosterior

inference = SNPE()
posterior_estimator = inference.append_simulations(theta, x).train()
posterior = DirectPosterior(posterior_estimator, prior)
```
So at the instantiation of `DirectPosterior`, there is also a call to `posterior_estimator_based_potential` and thus to `mcmc_transform`. For now, I inserted a `enable_theta_transform=False` keyword to the constructor of `DirectPosterior` which is passed further and this works! Thanks for thinking along.
Ah, great that you found this fix! I'll leave this issue open because we should probably make this more easy...

Resolving patch:
diff --git a/sbi/inference/posteriors/base_posterior.py b/sbi/inference/posteriors/base_posterior.py
--- a/sbi/inference/posteriors/base_posterior.py
+++ b/sbi/inference/posteriors/base_posterior.py
@@ -3,7 +3,6 @@
 
 from abc import ABC, abstractmethod
 from typing import Any, Callable, Dict, Optional, Union
-from warnings import warn
 
 import torch
 import torch.distributions.transforms as torch_tf
diff --git a/sbi/inference/posteriors/direct_posterior.py b/sbi/inference/posteriors/direct_posterior.py
--- a/sbi/inference/posteriors/direct_posterior.py
+++ b/sbi/inference/posteriors/direct_posterior.py
@@ -12,7 +12,7 @@
     posterior_estimator_based_potential,
 )
 from sbi.samplers.rejection.rejection import rejection_sample_posterior_within_prior
-from sbi.types import Shape
+from sbi.types import Shape, TorchTransform
 from sbi.utils import check_prior, match_theta_and_x_batch_shapes, within_support
 from sbi.utils.torchutils import ensure_theta_batched
 
@@ -35,6 +35,7 @@ def __init__(
         self,
         posterior_estimator: flows.Flow,
         prior: Distribution,
+        theta_transform: Optional[TorchTransform] = None,
         max_sampling_batch_size: int = 10_000,
         device: Optional[str] = None,
         x_shape: Optional[torch.Size] = None,
@@ -43,6 +44,12 @@ def __init__(
         Args:
             prior: Prior distribution with `.log_prob()` and `.sample()`.
             posterior_estimator: The trained neural posterior.
+            theta_transform: Custom transform to perform MAP optimization in
+                unconstrained space. If None (default), a suitable transform is
+                built from the prior support. In order to not use a transform at all,
+                pass an identity transform, e.g., `theta_transform=torch.distrbutions.
+                transforms`.
+                identity_transform()`.
             max_sampling_batch_size: Batchsize of samples being drawn from
                 the proposal at every iteration.
             device: Training device, e.g., "cpu", "cuda" or "cuda:0". If None,
@@ -55,7 +62,7 @@ def __init__(
         # obtaining the MAP.
         check_prior(prior)
         potential_fn, theta_transform = posterior_estimator_based_potential(
-            posterior_estimator, prior, None
+            posterior_estimator, prior, x_o=None, theta_transform=theta_transform
         )
 
         super().__init__(
diff --git a/sbi/inference/potentials/posterior_based_potential.py b/sbi/inference/potentials/posterior_based_potential.py
--- a/sbi/inference/potentials/posterior_based_potential.py
+++ b/sbi/inference/potentials/posterior_based_potential.py
@@ -4,6 +4,7 @@
 from typing import Callable, Optional, Tuple
 
 import torch
+import torch.distributions.transforms as torch_tf
 from pyknos.nflows import flows
 from torch import Tensor, nn
 from torch.distributions import Distribution
@@ -19,6 +20,7 @@ def posterior_estimator_based_potential(
     posterior_estimator: nn.Module,
     prior: Distribution,
     x_o: Optional[Tensor],
+    theta_transform: Optional[TorchTransform] = None,
 ) -> Tuple[Callable, TorchTransform]:
     r"""Returns the potential for posterior-based methods.
 
@@ -32,6 +34,11 @@ def posterior_estimator_based_potential(
         posterior_estimator: The neural network modelling the posterior.
         prior: The prior distribution.
         x_o: The observed data at which to evaluate the posterior.
+        theta_transform: Transform to map the parameters to an
+            unconstrained space. If None (default), a suitable transform is
+            built from the prior support. In order to not use a transform at all,
+            pass an identity transform, e.g., `theta_transform=torch.distrbutions.
+            transforms`.
 
     Returns:
         The potential function and a transformation that maps
@@ -43,7 +50,9 @@ def posterior_estimator_based_potential(
     potential_fn = PosteriorBasedPotential(
         posterior_estimator, prior, x_o, device=device
     )
-    theta_transform = mcmc_transform(prior, device=device)
+
+    if theta_transform is None:
+        theta_transform = mcmc_transform(prior, device=device)
 
     return potential_fn, theta_transform
 


Test patch: