fix: nflows inverse transform and retarget lc2st (#1865)

* tests: add seed to lc2st session fixtures

* ruff

* fix: typo in inverse transform, fix usage

Author: janfb

docs/advanced_tutorials/13_diagnostics_lc2st.ipynb

@@ -604,7 +604,7 @@
    "source": [
     "from sbi.diagnostics.lc2st import LC2ST_NF\n",
     "\n",
-    "flow_inverse_transform = lambda theta, x: npe.net._transform(theta, context=x)[0]\n",
+    "flow_inverse_transform = npe.inverse_transform\n",
     "flow_base_dist = torch.distributions.MultivariateNormal(\n",
     "    torch.zeros(2), torch.eye(2)\n",
     ") # same as npe.net._distribution\n",

sbi/neural_nets/estimators/nflows_flow.py

@@ -70,8 +70,8 @@ def inverse_transform(self, input: Tensor, condition: Tensor) -> Tensor:
         input = input.reshape(-1, input.shape[-1])
         condition = condition.reshape(-1, *self.condition_shape)
 
-        noise, _ = self.net._transorm(input, context=condition)
-        noise = noise.reshape(batch_shape)
+        noise, _ = self.net._transform(input, context=condition)
+        noise = noise.reshape(batch_shape + (noise.shape[-1],))
         return noise
 
     def log_prob(self, input: Tensor, condition: Tensor) -> Tensor:

tests/lc2st_test.py

@@ -17,6 +17,7 @@
     gaussian_mixture,
     uniform_prior_gaussian_mixture,
 )
+from sbi.utils.sbiutils import seed_all_backends
 
 # =============================================================================
 # Fixtures
@@ -55,6 +56,8 @@ def sim_setup() -> SimulatorSetup:
 @pytest.fixture(scope="session")
 def badly_trained_npe(sim_setup):
     """A poorly trained NPE for testing LC2ST detection of bad posteriors."""
+    # seed explicitly to keep session scope deterministic.
+    seed_all_backends(1)
     theta_train = sim_setup.prior.sample((50,))
     x_train = sim_setup.simulator(theta_train)
 
@@ -66,6 +69,8 @@ def badly_trained_npe(sim_setup):
 @pytest.fixture(scope="session")
 def well_trained_npe(sim_setup):
     """A well-trained NPE for testing LC2ST false positive rate."""
+    # seed explicitly to keep session scope deterministic.
+    seed_all_backends(1)
     theta_train = sim_setup.prior.sample((5_000,))
     x_train = sim_setup.simulator(theta_train)
 
@@ -77,6 +82,8 @@ def well_trained_npe(sim_setup):
 @pytest.fixture(scope="session")
 def cal_data(sim_setup, badly_trained_npe) -> CalibrationData:
     """Calibration data for LC2ST tests."""
+    # seed explicitly to keep session scope deterministic.
+    seed_all_backends(1)
     num_cal = 100
     thetas = sim_setup.prior.sample((num_cal,))
     xs = sim_setup.simulator(thetas)
@@ -129,7 +136,7 @@ def test_lc2st_methods(method, cal_data, badly_trained_npe, theta_o, x_o):
     else:
         npe = badly_trained_npe
         kwargs_init = {
-            "flow_inverse_transform": lambda t, x: npe.net._transform(t, context=x)[0],
+            "flow_inverse_transform": npe.inverse_transform,
             "flow_base_dist": torch.distributions.MultivariateNormal(
                 torch.zeros(2), torch.eye(2)
             ),
@@ -180,7 +187,7 @@ def test_lc2st_parameter_combinations(
     else:
         npe = badly_trained_npe
         kwargs_init = {
-            "flow_inverse_transform": lambda t, x: npe.net._transform(t, context=x)[0],
+            "flow_inverse_transform": npe.inverse_transform,
             "flow_base_dist": torch.distributions.MultivariateNormal(
                 torch.zeros(2), torch.eye(2)
             ),
@@ -294,7 +301,7 @@ def test_lc2st_nf_with_pretrained_null_is_ready_after_observed_training(
     """
     npe = badly_trained_npe
     kwargs_init = {
-        "flow_inverse_transform": lambda t, x: npe.net._transform(t, context=x)[0],
+        "flow_inverse_transform": npe.inverse_transform,
         "flow_base_dist": torch.distributions.MultivariateNormal(
             torch.zeros(2), torch.eye(2)
         ),
@@ -612,7 +619,7 @@ def test_lc2st_true_positive_rate(method, sim_setup, badly_trained_npe):
     else:
         npe = badly_trained_npe
         kwargs_init = {
-            "flow_inverse_transform": lambda t, x: npe.net._transform(t, context=x)[0],
+            "flow_inverse_transform": npe.inverse_transform,
             "flow_base_dist": torch.distributions.MultivariateNormal(
                 torch.zeros(2), torch.eye(2)
             ),
@@ -665,7 +672,7 @@ def test_lc2st_false_positive_rate(method, sim_setup, well_trained_npe, set_seed
     else:
         npe = well_trained_npe
         kwargs_init = {
-            "flow_inverse_transform": lambda t, x: npe.net._transform(t, context=x)[0],
+            "flow_inverse_transform": npe.inverse_transform,
             "flow_base_dist": torch.distributions.MultivariateNormal(
                 torch.zeros(2), torch.eye(2)
             ),