feat(vf): implement transform_to_unconstrained z-scoring for vector field estimators

sbi/inference/posteriors/vector_field_posterior.py

@@ -395,6 +395,11 @@ def _sample_via_diffusion(
         # Concatenate all batches and ensure we return exactly the requested number
         samples = torch.cat(all_samples, dim=0)[:total_samples_needed]
 
+        # Inverse-transform samples from unconstrained back to original space
+        input_transform = self.vector_field_estimator.input_transform
+        if input_transform is not None:
+            samples = input_transform.inv(samples)
+
         if torch.isnan(samples).all():
             raise RuntimeError(
                 "All samples NaN after diffusion sampling. "
@@ -412,7 +417,9 @@ def sample_via_ode(
         Return samples from posterior distribution with probability flow ODE.
 
         This builds the probability flow ODE and then samples from the corresponding
-        flow.
+        flow. When ``input_transform`` is set on the estimator, the ODE integrates
+        in unconstrained space and samples are inverse-transformed back to the
+        original constrained space.
 
         Args:
             sample_shape: The shape of the samples to be returned.
@@ -429,6 +436,11 @@ def sample_via_ode(
             torch.Size((num_samples,))
         )
 
+        # Inverse-transform samples from unconstrained back to original space
+        input_transform = self.vector_field_estimator.input_transform
+        if input_transform is not None:
+            samples = input_transform.inv(samples)
+
         return samples
 
     def log_prob(

sbi/inference/potentials/vector_field_potential.py

@@ -136,6 +136,10 @@ def __call__(
         """
         Return the potential (posterior log prob) via probability flow ODE.
 
+        When ``input_transform`` is set on the estimator, theta is transformed
+        to unconstrained space before the flow evaluation, and the log-det-Jacobian
+        of the transform is added to account for the change of variables.
+
         Args:
             theta: The parameters at which to evaluate the potential.
             track_gradients: Whether to track gradients. Default is False.
@@ -155,6 +159,8 @@ def __call__(
         )
         self.vector_field_estimator.eval()
 
+        input_transform = self.vector_field_estimator.input_transform
+
         with torch.set_grad_enabled(track_gradients):
             if self.x_is_iid:
                 assert self.prior is not None, (
@@ -164,10 +170,14 @@ def __call__(
                     "Flows for each iid x are required for evaluating log_prob."
                 )
                 num_iid = self.x_o.shape[0]  # number of iid samples
+                if input_transform is not None:
+                    theta_z = input_transform(theta_density_estimator)
+                else:
+                    theta_z = theta_density_estimator
                 iid_posteriors_prob = torch.sum(
                     torch.stack(
                         [
-                            flow.log_prob(theta_density_estimator).squeeze(-1)
+                            flow.log_prob(theta_z).squeeze(-1)
                             for flow in self.flows
                         ],
                         dim=0,
@@ -180,7 +190,18 @@ def __call__(
                     theta_density_estimator
                 ).squeeze(-1)
             else:
-                log_probs = self.flow.log_prob(theta_density_estimator).squeeze(-1)
+                if input_transform is not None:
+                    theta_z = input_transform(theta_density_estimator)
+                else:
+                    theta_z = theta_density_estimator
+                log_probs = self.flow.log_prob(theta_z).squeeze(-1)
+
+            # Add log-det-Jacobian correction for change of variables
+            if input_transform is not None:
+                log_probs = log_probs + input_transform.log_abs_det_jacobian(
+                    theta_density_estimator, theta_z
+                ).sum(dim=-1)
+
             # Force probability to be zero outside prior support.
             in_prior_support = within_support(self.prior, theta)
 

sbi/neural_nets/estimators/base.py

@@ -6,6 +6,7 @@
 
 import torch
 from torch import Tensor, nn
+from torch.distributions import Transform as TorchTransform
 
 ConditionalEstimatorType = TypeVar(
     'ConditionalEstimatorType',
@@ -350,6 +351,7 @@ def __init__(
         embedding_net: Optional[nn.Module] = None,
         mean_base: Union[float, Tensor] = 0.0,
         std_base: Union[float, Tensor] = 1.0,
+        input_transform: Optional[TorchTransform] = None,
     ) -> None:
         r"""Base class for vector field estimators.
 
@@ -364,6 +366,9 @@ def __init__(
                 condition.
             mean_base: Mean of the base distribution.
             std_base: Standard deviation of the base distribution.
+            input_transform: Optional transform mapping constrained -> unconstrained
+                space. When set, ``forward()`` operates in unconstrained space and
+                ``ode_fn()`` applies this transform before calling ``forward()``.
         """
         super().__init__(input_shape, condition_shape)
         self.net = net
@@ -385,12 +390,18 @@ def __init__(
         self._embedding_net = (
             embedding_net if embedding_net is not None else nn.Identity()
         )
+        self._input_transform = input_transform
 
     @property
     def embedding_net(self) -> nn.Module:
         r"""Return the embedding network if it exists."""
         return self._embedding_net
 
+    @property
+    def input_transform(self) -> Optional[TorchTransform]:
+        r"""Return the input transform (constrained -> unconstrained) if set."""
+        return self._input_transform
+
     @abstractmethod
     def forward(self, input: Tensor, condition: Tensor, **kwargs) -> Tensor:
         r"""Forward pass of the score estimator.

sbi/neural_nets/estimators/flowmatching_estimator.py

@@ -7,6 +7,7 @@
 import torch
 import torch.nn as nn
 from torch import Tensor
+from torch.distributions import Transform as TorchTransform
 
 from sbi.neural_nets.estimators.base import ConditionalVectorFieldEstimator
 from sbi.utils.vector_field_utils import VectorFieldNet
@@ -65,6 +66,7 @@ def __init__(
         mean_0: float = 0.0,
         std_0: float = 1.0,
         gaussian_baseline: bool = False,
+        input_transform: Optional[TorchTransform] = None,
         **kwargs,
     ) -> None:
         r"""Creates a vector field estimator for Flow Matching.
@@ -81,8 +83,12 @@ def __init__(
             std_0: Std of the data distribution at t=0 (used for time-dependent
                 z-scoring).
             gaussian_baseline: If True, use analytical Gaussian baseline velocity
-                derived from Bayes' rule: v = factor * (x - μ_true) - mean.
-                The network then only learns the residual. Default: False.
+                derived from Bayes' rule. The network then only learns the residual.
+                Default: False.
+            input_transform: Optional transform mapping constrained -> unconstrained
+                space. When set, ``ode_fn()`` applies this transform before calling
+                ``forward()``, so the ODE integrates in unconstrained space, and
+                ``loss()`` transforms training data accordingly.
         """
 
         if "num_freqs" in kwargs:
@@ -99,6 +105,7 @@ def __init__(
             input_shape=input_shape,
             condition_shape=condition_shape,
             embedding_net=embedding_net,
+            input_transform=input_transform,
         )
         self.noise_scale = noise_scale
         self.gaussian_baseline = gaussian_baseline
@@ -196,7 +203,13 @@ def _compute_velocity_baseline(self, input: Tensor, time: Tensor) -> Tensor:
     def forward(self, input: Tensor, condition: Tensor, time: Tensor) -> Tensor:
         """Forward pass of the FlowMatchingEstimator.
 
-        Returns velocity in ORIGINAL SPACE for ODE integration.
+        Returns velocity in internal space (unconstrained z-space when
+        ``input_transform`` is set).
+
+        Note:
+            This method does **not** apply ``input_transform`` internally. The
+            caller (e.g. ``ode_fn()``, ``loss()``) is responsible for transforming
+            inputs if needed.
 
         Args:
             input: Inputs to evaluate the vector field on of shape
@@ -207,7 +220,7 @@ def forward(self, input: Tensor, condition: Tensor, time: Tensor) -> Tensor:
                 `(batch_dim_time, *event_shape_time)`.
 
         Returns:
-            The estimated vector field in original space.
+            The estimated vector field in internal space.
         """
         # Continue with standard processing (broadcast shapes etc.)
         batch_shape_input = input.shape[: -len(self.input_shape)]
@@ -274,6 +287,9 @@ def loss(
             \theta_t = t \cdot \theta_1 + (1 - t) \cdot \theta_0,
             \left[ \| v(\theta_t, t; x_o = x_o) - (\theta_1 - \theta_0) \|^2 \right]
 
+        When ``input_transform`` is set, the input is first transformed to
+        unconstrained space, and all computations happen in that space.
+
         Args:
             input: Parameters (:math:`\theta_0`).
             condition: Observed data (:math:`x_o`).
@@ -288,12 +304,16 @@ def loss(
             times = torch.rand(input.shape[:-1], device=input.device, dtype=input.dtype)
         times_ = times[..., None]
 
+        # Transform to unconstrained space if needed
+        if self._input_transform is not None:
+            input = self._input_transform(input)
+
         # Sample from probability path at time t
         # θ_t = (1-t) * θ_data + t * θ_noise, where θ_noise ~ N(0,1)
         theta_1 = torch.randn_like(input)
         theta_t = (1 - times_) * input + (times_ + self.noise_scale) * theta_1
 
-        # Target vector field in original space
+        # Target vector field in internal space
         vector_field = theta_1 - input
 
         # Embed condition
@@ -348,6 +368,10 @@ def ode_fn(self, input: Tensor, condition: Tensor, times: Tensor) -> Tensor:
         flow matching neural network (see Equation 1 in [1]_ with added
         conditioning on :math:`x_o`).
 
+        When ``input_transform`` is set, this method transforms the input to
+        unconstrained space before calling ``forward()``, so the ODE integrates
+        in the unconstrained space.
+
         Args:
             input: :math:`\theta_t`.
             condition: Conditioning variable :math:`x_o`.
@@ -357,6 +381,8 @@ def ode_fn(self, input: Tensor, condition: Tensor, times: Tensor) -> Tensor:
             Estimated vector field :math:`v(\theta_t, t; x_o)`.
             The shape is the same as the input.
         """
+        if self._input_transform is not None:
+            input = self._input_transform(input)
         return self.forward(input, condition, times)
 
     def score(self, input: Tensor, condition: Tensor, t: Tensor) -> Tensor:
@@ -372,6 +398,12 @@ def score(self, input: Tensor, condition: Tensor, t: Tensor) -> Tensor:
         :math:`\nabla_{\theta_t} \log p(\theta_t | x_o) =
             (- (1 - t) v(\theta_t, t; x_o) - \theta_0 ) / (t + \sigma_{min})`.
 
+        Note:
+            This method does **not** apply ``input_transform`` internally. When
+            ``input_transform`` is set, this method expects input in unconstrained
+            (z) space. Use ``ode_fn()`` for automatic transform handling during
+            ODE/SDE sampling.
+
         Args:
             input: variable whose distribution is estimated.
             condition: Conditioning variable.

sbi/neural_nets/estimators/mixture_density_estimator.py

@@ -17,6 +17,7 @@
 import numpy as np
 import torch
 from torch import Tensor, nn
+from torch.distributions import Transform as TorchTransform
 from torch.nn import functional as F
 
 from sbi.neural_nets.estimators.base import ConditionalDensityEstimator
@@ -337,6 +338,7 @@ def __init__(
         condition_shape: torch.Size,
         embedding_net: Optional[nn.Module] = None,
         transform_input: Optional[Tensor] = None,
+        input_transform: Optional[TorchTransform] = None,
     ) -> None:
         """Initialize the mixture density estimator.
 
@@ -354,12 +356,19 @@ def __init__(
                 evaluation, and samples are inverse transformed as:
                 x = z * scale + shift.
                 This is used for z-scoring inputs to improve numerical stability.
+            input_transform: Optional PyTorch Transform for general bijective
+                transformation of inputs. Mutually exclusive with transform_input.
         """
         super().__init__(net, input_shape, condition_shape)
         self._embedding_net = (
             embedding_net if embedding_net is not None else nn.Identity()
         )
 
+        if transform_input is not None and input_transform is not None:
+            raise ValueError(
+                "Only one of transform_input and input_transform can be provided."
+            )
+
         # Validate that embedding_net output matches MDN input if embedding is provided
         if embedding_net is not None:
             with torch.no_grad():
@@ -372,6 +381,8 @@ def __init__(
                     f"MDN context_features ({net.context_features})"
                 )
 
+        self._input_transform = input_transform
+
         # Store z-score transform parameters as buffers (not trained, moved with model)
         if transform_input is not None:
             if transform_input.shape[0] != 2:
@@ -399,49 +410,39 @@ def embedding_net(self) -> nn.Module:
 
     @property
     def has_input_transform(self) -> bool:
-        """Whether input z-score transform is enabled."""
-        return self._transform_shift is not None
+        """Whether input transform is enabled (z-score or general transform)."""
+        return self._input_transform is not None or self._transform_shift is not None
 
     def _transform_input(self, input: Tensor) -> Tensor:
-        """Apply z-score transform to input: z = (x - shift) / scale."""
-        if not self.has_input_transform:
+        """Apply input transform: z = transform(x) or z = (x - shift) / scale."""
+        if self._input_transform is not None:
+            return self._input_transform(input)
+        if self._transform_shift is None:
             return input
         return (input - self._transform_shift) / self._transform_scale
 
     def _inverse_transform_input(self, z: Tensor) -> Tensor:
-        """Apply inverse z-score transform: x = z * scale + shift."""
-        if not self.has_input_transform:
+        """Apply inverse input transform: x = transform.inv(z) or x = z * scale + shift."""
+        if self._input_transform is not None:
+            return self._input_transform.inv(z)
+        if self._transform_shift is None:
             return z
         return z * self._transform_scale + self._transform_shift
 
-    def _log_det_jacobian_forward(self, input: Tensor) -> Tensor:
-        """Compute log determinant of Jacobian for the forward z-score transform.
-
-        For the forward affine transform z = (x - shift) / scale:
-            - The Jacobian matrix is: dz/dx = diag(1/scale)
-            - The determinant is: |det(dz/dx)| = prod(1/scale)
-            - The log determinant is: log|det(dz/dx)| = -sum(log(scale))
-
-        Change of Variables Formula:
-            When we have a density p_z(z) and want p_x(x), we use:
-            p_x(x) = p_z(z(x)) * |det(dz/dx)|
-
-            In log space:
-            log p_x(x) = log p_z(z(x)) + log|det(dz/dx)|
-
-            Since log|det(dz/dx)| = -sum(log(scale)), we have:
-            log p_x(x) = log p_z(z) - sum(log(scale))
+    def _log_det_jacobian_forward(
+        self, input: Tensor, transformed_input: Tensor
+    ) -> Tensor:
+        """Log determinant of the forward input transform Jacobian.
 
-        This method returns log|det(dz/dx)| = -sum(log(scale)), which should
-        be ADDED to log p_z(z) to get log p_x(x).
-
-        Args:
-            input: Input tensor, used to determine device and dtype.
-
-        Returns:
-            Log determinant of forward Jacobian (scalar), on the same device as input.
+        For the affine z-score transform: -sum(log(scale)).
+        For a general transform: transform.log_abs_det_jacobian(x, z).sum(dim=-1).
         """
-        if not self.has_input_transform:
+        if self._input_transform is not None:
+            jac = self._input_transform.log_abs_det_jacobian(
+                input, transformed_input
+            )
+            return jac.sum(dim=-1)
+        if self._transform_shift is None:
             return torch.zeros(1, device=input.device, dtype=input.dtype).squeeze()
         return -torch.log(self._transform_scale).sum()
 
@@ -476,9 +477,10 @@ def log_prob(self, input: Tensor, condition: Tensor, **kwargs) -> Tensor:
 
         # MoG.log_prob handles (sample_dim, batch_dim, dim) input
         # Change of variables: log p(x) = log p(z) + log|det(dz/dx)|
-        # where z = (x - shift) / scale and log|det(dz/dx)| = -sum(log(scale))
         log_probs = mog.log_prob(transformed_input)  # (sample_dim, batch_dim)
-        log_probs = log_probs + self._log_det_jacobian_forward(input)
+        log_probs = log_probs + self._log_det_jacobian_forward(
+            input, transformed_input
+        )
 
         if not has_sample_dim:
             log_probs = log_probs.squeeze(0)