MNPE class similar to MNLE #1362

dgedon · 2025-01-10T15:57:48Z

Implementation of mixed NPE where we have some continuous parameters theta followed by one (or multiple with this PR #1269) discrete parameters. The observation space is fully continuous.

Deprecated mnle.py in net_builders und unified mnle/mnpe as mixed_nets.py.

dgedon · 2025-01-17T10:22:21Z

Update:

MNPE and tests are implemented
for test with Bernoulli prior, I had to change mcmc_transforms to handle discrete distirbution. As default we just compute mean/std for discrete distributions
currently MNPE with embedding nets does not work yet. Gives some backwards inplace operations error that I couldn't solve yet.

janfb · 2025-02-25T18:59:53Z

@dgedon #1269 is now merged 🙌

…ODO: MNPE class + test

…dent; embedding net in mnpe not working yet

…this PR)

…not allow gpu handling yet though

dgedon · 2025-03-18T14:30:59Z

Updates:

bug fix so everything works now. Essentially need to handle normalization with care when switching from mnle to mnpe
remove unnecessary GPU handling (hackathon task). This limits MultipleIndependent as not allowing device argument yet

janfb

looks good overall, except one central question about the call signature of MNPE.

janfb · 2025-03-18T13:54:44Z

sbi/neural_nets/estimators/mixed_density_estimator.py

@@ -11,10 +11,12 @@


 class MixedDensityEstimator(ConditionalDensityEstimator):
-    """Class performing Mixed Neural Likelihood Estimation.
+    """Class performing Mixed Neural Density Estimation.


janfb · 2025-03-18T13:55:30Z

sbi/inference/trainers/npe/mnpe.py

+            assert isinstance(
+                density_estimator, MixedDensityEstimator
+            ), f"""net must be of type MixedDensityEstimator but is {
+                type(density_estimator)
+            }."""


we could also change the type above to be MixedDensityEstimator to have a static check.

Good point. I'll add this additionally to the assertion

janfb · 2025-03-18T13:59:21Z

sbi/neural_nets/estimators/mixed_density_estimator.py

+    This estimator combines a Categorical net and a neural density estimator to model
+    data with mixed types (discrete and continuous), e.g., as they occur in
+    decision-making models. It can be used for both likelihood and posterior estimation
+    of mixed data.


Suggested change

This estimator combines a Categorical net and a neural density estimator to model

data with mixed types (discrete and continuous), e.g., as they occur in

decision-making models. It can be used for both likelihood and posterior estimation

of mixed data.

This estimator combines a categorical mass estimator and a density estimator to model

variables with mixed types (discrete and continuous). It can be used for both likelihood

estimation (e.g., for discrete decisions and continuous reaction times in decision-making

models) or posterior estimation (e.g., for models that have both discrete and continuous

parameters).

janfb · 2025-03-18T13:59:49Z

sbi/neural_nets/estimators/mixed_density_estimator.py

-            """The forward method is not implemented for MNLE, use '.sample(...)' to
-            generate samples though a forward pass."""
+            """The forward method is not implemented for mixed neural density
+            estimation,use '.sample(...)' to generate samples though a forward


Suggested change

estimation,use '.sample(...)' to generate samples though a forward

estimation, use '.sample(...)' to generate samples though a forward

janfb · 2025-03-18T22:20:20Z

sbi/neural_nets/net_builders/mixed_nets.py

+def build_mnle(
+    batch_x: Tensor,
+    batch_y: Tensor,
+    **kwargs,
+) -> MixedDensityEstimator:
+    """Returns a mixed neural likelihood estimator.
+
+    This estimator models p(x|theta) where x contains both continuous and discrete data.
+
+    Args:
+        batch_x: Batch of xs (data), used to infer dimensionality.
+        batch_y: Batch of ys (parameters), used to infer dimensionality.
+        **kwargs: Additional arguments passed to _build_mixed_density_estimator.
+
+    Returns:
+        MixedDensityEstimator for MNLE.
+    """
+    return _build_mixed_density_estimator(
+        batch_x=batch_x, batch_y=batch_y, mode="mnle", **kwargs
+    )
+
+
+def build_mnpe(
+    batch_x: Tensor,
+    batch_y: Tensor,
+    **kwargs,
+) -> MixedDensityEstimator:
+    """Returns a mixed neural posterior estimator.
+
+    This estimator models p(theta|x) where x contains both continuous and discrete data.


I am confused by these two call functions. maybe I am missing something, but couldn't we just call _build_mixed_density_estimator with batch_x and batch_y swapped for MNPE and MNLE?

To me it seems this swapping is not happening, i.e., we need to make sure that in MNPE we are only embedding x and not theta, and v.v. in MNLE. Let's discuss tomorrow.

True, we could remove both functions and just use _build_mixed_density_estimator.

The swapping of x/theta is happening because the function is once called as likelihood_nn and once as posterior_nn.

…onflict with tutorial

dgedon · 2025-03-19T09:14:12Z

Update:

simplify _build_mixed_density_estimator by not having mode='mnpe'/'mnle'
add default log_transform_x as kwarg to build_mnle and build_mnpe

janfb

some more comments on the tests and the refactoring.

I am suggestion a toy example with ground-truth posterior for the MNPE scenario to test the accuracy.

sbi/neural_nets/net_builders/mnle.py

sbi/utils/sbiutils.py

janfb · 2025-03-19T08:58:47Z

sbi/utils/sbiutils.py

+                try:
+                    prior_mean = prior.mean.to(device)
+                    prior_std = prior.stddev.to(device)
+                except (NotImplementedError, AttributeError):
+                    warnings.warn(
+                        "The passed discrete prior has no mean or stddev attribute, "
+                        "estimating them from samples to build affine standardizing "
+                        "transform.",
+                        stacklevel=2,
+                    )
+                    theta = prior.sample(torch.Size((num_prior_samples_for_zscoring,)))
+                    prior_mean = theta.mean(dim=0).to(device)
+                    prior_std = theta.std(dim=0).to(device)


move this into a small function to avoid code duplication?

janfb · 2025-03-19T09:05:09Z

tests/mnpe_test.py

+    x = mixed_param_simulator(theta)
+
+    # Build estimator manually
+    theta_embedding = FCEmbedding(1, 1)  # simple embedding net, 1 continuous parameter


this should be an x_embedding to avoid confusion

janfb · 2025-03-19T09:05:19Z

tests/mnpe_test.py

+        flow_model=flow_model,
+        z_score_theta=z_score_theta,
+        embedding_net=theta_embedding if use_embed_net else torch.nn.Identity(),
+        log_transform_x=False,


janfb · 2025-03-19T09:05:52Z

tests/mnpe_test.py

+        log_transform_x=False,
+    )
+    trainer = MNPE(density_estimator=density_estimator)
+    trainer.append_simulations(theta, x).train(max_num_epochs=5)


max_num_epochs=1

to speed up tests.

janfb · 2025-03-19T09:07:16Z

tutorials/Example_01_DecisionMakingModel.ipynb

please remove all diffs here. we probably want to have some kind of tutorial or how-to-guide with MNPE, but let's wait for the new documentation setup.

janfb · 2025-03-19T09:15:07Z

tests/mnpe_test.py

+def mixed_param_simulator(theta: Tensor) -> Tensor:
+    """Simulator for continuous data with mixed parameters.
+
+    Args:
+        theta: Parameters with mixed types - continuous and discrete.
+    Returns:
+        x: Continuous observation.
+    """
+    device = theta.device
+
+    # Extract parameters
+    a, b = theta[:, 0], theta[:, 1]
+    noise = 0.05 * torch.randn(a.shape, device=device).reshape(-1, 1)
+    return (a + 2 * b).reshape(-1, 1) + noise


can we come up with a simulator for which we have ground truth posterior to make a test on the accuracy?

e.g., a mixture of Gaussians where the discrete params just selects the component:

z ~ Categorical
x ~ N(mu_z, 1)

and Gaussian priors on mu_z with different prior means, e.g., -1 and 1.

then for a fixed z, the posterior will just be a Gaussian composed from the likelihood and prior Gaussian just like in our standard linear Gaussian example.

janfb mentioned this pull request Mar 13, 2025

Discrete (mixed) density estimators #907

Open

dgedon added 8 commits March 18, 2025 12:20

wip: build_mnpe integrated, build_mnle refactored, mnle_test fails; T…

a567fc8

…ODO: MNPE class + test

wip: added MNPE class and test case for it

b9a5b92

wip: added MNPE class and test case for it, not working yet

fe274c9

fix: tests, embned+maf not working

6016f07

wip: fixed discrete data issue in mcmc_transform with MultipleIndepen…

736e3ac

…dent; embedding net in mnpe not working yet

wip: remove unnecessary helper function (introduced while working on …

0934e86

…this PR)

bug fix with normalization when using embedding nets

6412964

revert unnecessary gpu handling things. Now MultipleIndependent does …

48f0e63

…not allow gpu handling yet though

dgedon force-pushed the mnpe branch from eb7ad8b to 48f0e63 Compare March 18, 2025 11:29

dgedon requested review from janfb and michaeldeistler March 18, 2025 11:30

janfb reviewed Mar 18, 2025

View reviewed changes

dgedon added 2 commits March 19, 2025 09:39

review changes: comments, missing import, static type check

4f9524f

simplify mixed nets (default logtransform set for mnle/mnpe), merge c…

7b140b4

…onflict with tutorial

janfb reviewed Mar 19, 2025

View reviewed changes

dgedon added 4 commits March 19, 2025 11:45

remove legacy mnle.py that was not interacted with by users

b475446

refactor prior transform function (code duplication)

3d7f870

cleanup cosmetics

07783d2

cleanup cosmetics (again)

8b94843

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

MNPE class similar to MNLE #1362

MNPE class similar to MNLE #1362

dgedon commented Jan 10, 2025

dgedon commented Jan 17, 2025

janfb commented Feb 25, 2025

dgedon commented Mar 18, 2025

janfb left a comment

janfb Mar 18, 2025

janfb Mar 18, 2025

dgedon Mar 19, 2025

janfb Mar 18, 2025

janfb Mar 18, 2025

janfb Mar 18, 2025

dgedon Mar 19, 2025

dgedon commented Mar 19, 2025

janfb left a comment

janfb Mar 19, 2025

janfb Mar 19, 2025

janfb Mar 19, 2025

janfb Mar 19, 2025

janfb Mar 19, 2025

janfb Mar 19, 2025

-    This estimator combines a Categorical net and a neural density estimator to model
-    data with mixed types (discrete and continuous), e.g., as they occur in
-    decision-making models. It can be used for both likelihood and posterior estimation
-    of mixed data.
+    This estimator combines a categorical mass estimator and a density estimator to model
+    variables with mixed types (discrete and continuous). It can be used for both likelihood
+    estimation (e.g., for discrete decisions and continuous reaction times in decision-making
+    models) or posterior estimation (e.g., for models that have both discrete and continuous
+    parameters).

	estimation,use '.sample(...)' to generate samples though a forward
	estimation, use '.sample(...)' to generate samples though a forward

MNPE class similar to MNLE #1362

Are you sure you want to change the base?

MNPE class similar to MNLE #1362

Conversation

dgedon commented Jan 10, 2025

dgedon commented Jan 17, 2025

janfb commented Feb 25, 2025

dgedon commented Mar 18, 2025

janfb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

dgedon commented Mar 19, 2025

janfb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment