fix prior scores standardize

Author: arrjon

bayesflow/approximators/continuous_approximator.py

@@ -14,7 +14,6 @@
     squeeze_inner_estimates_dict,
     concatenate_valid,
     concatenate_valid_shapes,
-    expand_right_as,
 )
 from bayesflow.utils.serialization import serialize, deserialize, serializable
 
@@ -695,25 +694,52 @@ def compositional_sample(
         # Prepare prior scores to handle adapter
         def compute_prior_score_pre(_samples: Tensor) -> Tensor:
             if "inference_variables" in self.standardize:
-                _samples, log_det_jac_standardize = self.standardize_layers["inference_variables"](
-                    _samples, forward=False, log_det_jac=True
-                )
-            else:
-                log_det_jac_standardize = keras.ops.cast(0.0, dtype="float32")
+                _samples = self.standardize_layers["inference_variables"](_samples, forward=False)
             _samples = keras.tree.map_structure(keras.ops.convert_to_numpy, {"inference_variables": _samples})
             adapted_samples, log_det_jac = self.adapter(
                 _samples, inverse=True, strict=False, log_det_jac=True, **kwargs
             )
+
+            if len(log_det_jac) > 0:
+                problematic_keys = [key for key in log_det_jac if log_det_jac[key] != 0.0]
+                raise NotImplementedError(
+                    f"Cannot use compositional sampling with adapters "
+                    f"that have non-zero log_det_jac. Problematic keys: {problematic_keys}"
+                )
+
             prior_score = compute_prior_score(adapted_samples)
             for key in adapted_samples:
-                if isinstance(prior_score[key], np.ndarray):
-                    prior_score[key] = prior_score[key].astype("float32")
-                if len(log_det_jac) > 0 and key in log_det_jac:
-                    prior_score[key] -= expand_right_as(log_det_jac[key], prior_score[key])
+                prior_score[key] = prior_score[key].astype(np.float32)
 
             prior_score = keras.tree.map_structure(keras.ops.convert_to_tensor, prior_score)
             out = keras.ops.concatenate([prior_score[key] for key in adapted_samples], axis=-1)
-            return out - keras.ops.expand_dims(log_det_jac_standardize, axis=-1)
+
+            if "inference_variables" in self.standardize:
+                # Apply jacobian correction from standardization
+                # For standardization T^{-1}(z) = z * std + mean, the jacobian is diagonal with std on diagonal
+                # The gradient of log|det(J)| w.r.t. z is 0 since log|det(J)| = sum(log(std)) is constant w.r.t. z
+                # But we need to transform the score: score_z = score_x * std where x = T^{-1}(z)
+                standardize_layer = self.standardize_layers["inference_variables"]
+
+                # Compute the correct standard deviation for all components
+                std_components = []
+                for idx in range(len(standardize_layer.moving_mean)):
+                    std_val = standardize_layer.moving_std(idx)
+                    std_components.append(std_val)
+
+                # Concatenate std components to match the shape of out
+                if len(std_components) == 1:
+                    std = std_components[0]
+                else:
+                    std = keras.ops.concatenate(std_components, axis=-1)
+
+                # Expand std to match batch dimension of out
+                std_expanded = keras.ops.expand_dims(std, (0, 1))  # Add batch, sample dimensions
+                std_expanded = keras.ops.tile(std_expanded, [n_datasets, num_samples, 1])
+
+                # Apply the jacobian: score_z = score_x * std
+                out = out * std_expanded
+            return out
 
         # Test prior score function, useful for debugging
         test = self.inference_network.base_distribution.sample((n_datasets, num_samples))

tests/test_approximators/conftest.py

@@ -249,6 +249,21 @@ def sample(self, batch_shape: Shape) -> dict[str, Tensor]:
     return SimpleSimulator()
 
 
+@pytest.fixture
+def identity_adapter():
+    """Create an adapter that applies no transformation to the parameters."""
+    from bayesflow.adapters import Adapter
+
+    adapter = Adapter()
+    adapter.to_array()
+    adapter.convert_dtype("float64", "float32")
+
+    adapter.concatenate(["loc"], into="inference_variables")
+    adapter.concatenate(["conditions"], into="inference_conditions")
+    adapter.keep(["inference_variables", "inference_conditions"])
+    return adapter
+
+
 @pytest.fixture
 def transforming_adapter():
     """Create an adapter that applies log transformation to scale parameters."""

tests/test_approximators/test_compositional_prior_score.py

@@ -1,110 +1,43 @@
 """Tests for compositional sampling and prior score computation with adapters."""
 
 import numpy as np
-import keras
 
 from bayesflow import ContinuousApproximator
-from bayesflow.utils import expand_right_as
 
 
 def mock_prior_score_original_space(data_dict):
-    """Mock prior score function that expects data in original (loc, scale) space."""
-    # The function receives data in the same format the compute_prior_score_pre creates
-    # after running the inverse adapter
+    """Mock prior score function that expects data in original space."""
     loc = data_dict["loc"]
-    scale = data_dict["scale"]
 
-    # Simple prior: N(0,1) for loc, LogNormal(0,0.5) for scale
+    # Simple prior: N(0,1) for loc
     loc_score = -loc
-    scale_score = -1.0 / scale - np.log(scale) / (0.25 * scale)
+    return {"loc": loc_score}
 
-    return {"loc": loc_score, "scale": scale_score}
 
-
-def test_prior_score_transforming_adapter(simple_log_simulator, transforming_adapter, diffusion_network):
+def test_prior_score_identity_adapter(simple_log_simulator, identity_adapter, diffusion_network):
     """Test that prior scores work correctly with transforming adapter (log transformation)."""
 
     # Create approximator with transforming adapter
     approximator = ContinuousApproximator(
-        adapter=transforming_adapter,
+        adapter=identity_adapter,
         inference_network=diffusion_network,
     )
 
     # Generate test data and adapt it
     data = simple_log_simulator.sample((2,))
-    adapted_data = transforming_adapter(data)
+    adapted_data = identity_adapter(data)
 
     # Build approximator
     approximator.build_from_data(adapted_data)
 
     # Test compositional sampling
     n_datasets, n_compositional = 3, 5
     conditions = {"conditions": np.random.normal(0.0, 1.0, (n_datasets, n_compositional, 3)).astype("float32")}
-
-    # This should work - the compute_prior_score_pre function should handle the inverse transformation
     samples = approximator.compositional_sample(
         num_samples=10,
         conditions=conditions,
         compute_prior_score=mock_prior_score_original_space,
     )
 
     assert "loc" in samples
-    assert "scale" in samples
     assert samples["loc"].shape == (n_datasets, 10, 2)
-    assert samples["scale"].shape == (n_datasets, 10, 2)
-
-
-def test_prior_score_jacobian_correction(simple_log_simulator, transforming_adapter, diffusion_network):
-    """Test that Jacobian correction is applied correctly in compute_prior_score_pre."""
-
-    # Create approximator with transforming adapter
-    approximator = ContinuousApproximator(
-        adapter=transforming_adapter, inference_network=diffusion_network, standardize=[]
-    )
-
-    # Build with dummy data
-    dummy_data_dict = simple_log_simulator.sample((1,))
-    adapted_dummy_data = transforming_adapter(dummy_data_dict)
-    approximator.build_from_data(adapted_dummy_data)
-
-    # Get the internal compute_prior_score_pre function
-    def get_compute_prior_score_pre():
-        def compute_prior_score_pre(_samples):
-            if "inference_variables" in approximator.standardize:
-                _samples, log_det_jac_standardize = approximator.standardize_layers["inference_variables"](
-                    _samples, forward=False, log_det_jac=True
-                )
-            else:
-                log_det_jac_standardize = keras.ops.cast(0.0, dtype="float32")
-
-            _samples = keras.tree.map_structure(keras.ops.convert_to_numpy, {"inference_variables": _samples})
-            adapted_samples, log_det_jac = approximator.adapter(_samples, inverse=True, strict=False, log_det_jac=True)
-
-            prior_score = mock_prior_score_original_space(adapted_samples)
-            for key in adapted_samples:
-                if isinstance(prior_score[key], np.ndarray):
-                    prior_score[key] = prior_score[key].astype("float32")
-                if len(log_det_jac) > 0 and key in log_det_jac:
-                    prior_score[key] -= expand_right_as(log_det_jac[key], prior_score[key])
-
-            prior_score = keras.tree.map_structure(keras.ops.convert_to_tensor, prior_score)
-            out = keras.ops.concatenate([prior_score[key] for key in adapted_samples], axis=-1)
-            return out - keras.ops.expand_dims(log_det_jac_standardize, axis=-1)
-
-        return compute_prior_score_pre
-
-    compute_prior_score_pre = get_compute_prior_score_pre()
-
-    # Test with a known transformation
-    y_samples = adapted_dummy_data["inference_variables"]
-    scores = compute_prior_score_pre(y_samples)
-    scores_np = keras.ops.convert_to_numpy(scores)[0]  # Remove batch dimension
-
-    # With Jacobian correction: score_transformed = score_original - log|J|
-    old_scores = mock_prior_score_original_space(dummy_data_dict)
-    # order of parameters is flipped due to concatenation in adapter
-    det_jac_scale = y_samples[0, :2].sum()
-    expected_scores = np.array([old_scores["scale"][0] - det_jac_scale, old_scores["loc"][0]]).flatten()
-
-    # Check that scores are reasonably close
-    np.testing.assert_allclose(scores_np, expected_scores, rtol=1e-5, atol=1e-6)