Combining Embedding net output with input parameters that are continuous and discrete

[btetaren]

Hello,

I have a simulator where my theta is a set of 5 parameters (2 continuous and 3 discrete), producing x, which is a set of 10 time-series (each of length (50,) ).

I see that the MNLE tutorial shows how to deal with a combination continuous/discrete x.

Would it be possible to do something similar, but for a combination continuous/discrete theta, all while being able to use an embedding net for x?

If so, a basic example showing the workflow of how to implement something like this in sbi would be very much appreciated.

thank you.

[janfb]

Hi @btetaren ,

yes, we recently implemented MNPE, the analog to MNLE but for the parameter space.

You can use it like this:

from sbi.inference import MNPE

theta = prior.sample((num_simulations,))
x = your_simulator(theta)

trainer = MNPE(prior=prior)
trainer.append_simulations(theta, x).train()

posterior = trainer.build_posterior()
samples = posterior.sample(
    sample_shape=(num_samples,),
    x=x_o,
)

Does this work for your case?

[michaeldeistler]

(and with an embedding net):

import torch
from sbi.inference import MNPE
from sbi.neural_nets.embedding_nets import FCEmbedding
from sbi.neural_nets import posterior_nn

dim_theta_discrete = 3
dim_theta_continuous = 2
dim_theta = 5
dim_x = 50

num_sims = 100

discrete_theta = torch.randint(low=0, high=2, size=(100, dim_theta_discrete))
continuous_theta = torch.randn((num_sims, dim_theta_discrete))
theta = torch.cat([continuous_theta, discrete_theta], dim=1)
x = torch.randn((num_sims, dim_x))

embedding_net = FCEmbedding(dim_x)
net_builder = posterior_nn("mnpe", embedding_net=embedding_net)

inference = MNPE(density_estimator=net_builder)
_ = inference.append_simulations(theta, x).train()

posterior = inference.build_posterior()

x_o = torch.randn((1, dim_x))
samples = posterior.sample((100,), x=x_o)

[btetaren]

Yes this is what I need. THanks.

Follow-up question about the FCembedding:
I have only ever used CNNs where the input dimension is a tuple. In my case, my x dimensions are (10,50). Do I have to flatten my x's as FCEmbedding only takes an integer as input dimension? or am I misunderstanding here?

I end up getting a RunTimeError: mat1 and mat2 shapes cannot be multiplied with .train(). So I am assuming my x's are the wrong shape here.

Additionally, how would I go about building a prior? I tried using MultipleIndependent, but end up with an error when trying to build the posterior, "mismatch between transformed and untransformed space. Float did not match double."

thanks.

[btetaren]

If I flatten the x's to (500,), I no longer have the issues above. However, a new problem occurs. I get a RuntimeError (similar to the traceback below). The traceback below is what happens when I re-create the problem with your basic example above, building discrete_theta in multiple pieces i.e., a situation where each discrete theta dimension may lie in a different range (like my real training data). For example, for dim_theta_discrete = 3, say the first two dimensions range between (0,2) and the last one between (10,20).

It seems the only way I can get this to train is if discrete_theta is defined as one tensor like above. If I, for example, concatenate my own continuous_theta (read from disc) with a discrete_theta = torch.randint(low=0, high=2, size=(num_sims, dim_theta_discrete)), and use that as theta with my x's (read from disc), it will train.

My only guess is that it is not properly inferring which theta's are continuous and which are discrete?? I do not know why this is happening. Any thoughts on how to fix this?

RuntimeError                              Traceback (most recent call last)
File ~/anaconda3/envs/sbi_env2/lib/python3.10/site-packages/sbi/utils/user_input_checks.py:795, in test_posterior_net_for_multi_d_x(net, theta, x)
    792     if hasattr(net, "log_prob"):
    793         # This only is checked for density estimators, not for classifiers and
    794         # others
--> 795         net.log_prob(theta[:, :2], condition=x[:2])
    797 except RuntimeError as rte:

File ~/anaconda3/envs/sbi_env2/lib/python3.10/site-packages/sbi/neural_nets/estimators/mixed_density_estimator.py:168, in MixedDensityEstimator.log_prob(self, input, condition)
    167 # Get log probs from discrete data. Pass with singleton sample dim.
--> 168 disc_log_prob = self.discrete_net.log_prob(
    169     input=disc_input, condition=condition
    170 )
    172 cont_input_reshaped = cont_input.reshape((1, combined_batch_size, -1))

File ~/anaconda3/envs/sbi_env2/lib/python3.10/site-packages/sbi/neural_nets/estimators/categorical_net.py:207, in CategoricalMassEstimator.log_prob(self, input, condition, **kwargs)
    205 # The CatetoricalNet can actually handle torch shape conventions and
    206 # just returns log-probabilities of shape `(sample_dim, batch_dim)`.
--> 207 return self.net.log_prob(input, condition, **kwargs)

File ~/anaconda3/envs/sbi_env2/lib/python3.10/site-packages/sbi/neural_nets/estimators/categorical_net.py:116, in CategoricalMADE.log_prob(self, input, condition)
    115 outputs = outputs.reshape(*input.shape, self.num_categories)
--> 116 log_prob = Categorical(logits=outputs).log_prob(input).sum(dim=-1)
    118 return log_prob

File ~/anaconda3/envs/sbi_env2/lib/python3.10/site-packages/torch/distributions/categorical.py:143, in Categorical.log_prob(self, value)
    142 value = value[..., :1]
--> 143 return log_pmf.gather(-1, value).squeeze(-1)

RuntimeError: index 12 is out of bounds for dimension 3 with size 10

The above exception was the direct cause of the following exception:

RuntimeError                              Traceback (most recent call last)
Cell In[637], line 2
      1 #MNPE stuff
----> 2 density_estimator= inference.append_simulations(theta,x).train(show_train_summary=True,
      3                                                             stop_after_epochs = 20)

File ~/anaconda3/envs/sbi_env2/lib/python3.10/site-packages/sbi/inference/trainers/npe/mnpe.py:89, in MNPE.train(self, training_batch_size, learning_rate, validation_fraction, stop_after_epochs, max_num_epochs, clip_max_norm, resume_training, discard_prior_samples, retrain_from_scratch, show_train_summary, dataloader_kwargs)
     75 def train(
     76     self,
     77     training_batch_size: int = 200,
   (...)
     87     dataloader_kwargs: Optional[Dict] = None,
     88 ) -> MixedDensityEstimator:
---> 89     density_estimator = super().train(
     90         **del_entries(locals(), entries=("self", "__class__"))
     91     )
     92     assert isinstance(
     93         density_estimator, MixedDensityEstimator
     94     ), f"""Internal net must be of type
     95         MixedDensityEstimator but is {type(density_estimator)}."""
     96     return density_estimator

File ~/anaconda3/envs/sbi_env2/lib/python3.10/site-packages/sbi/inference/trainers/npe/npe_c.py:203, in NPE_C.train(self, num_atoms, training_batch_size, learning_rate, validation_fraction, stop_after_epochs, max_num_epochs, clip_max_norm, calibration_kernel, resume_training, force_first_round_loss, discard_prior_samples, use_combined_loss, retrain_from_scratch, show_train_summary, dataloader_kwargs)
    199     if self.use_non_atomic_loss:
    200         # Take care of z-scoring, pre-compute and store prior terms.
    201         self._set_state_for_mog_proposal()
--> 203 return super().train(**kwargs)

File ~/anaconda3/envs/sbi_env2/lib/python3.10/site-packages/sbi/inference/trainers/npe/npe_base.py:355, in PosteriorEstimatorTrainer.train(self, training_batch_size, learning_rate, validation_fraction, stop_after_epochs, max_num_epochs, clip_max_norm, calibration_kernel, resume_training, force_first_round_loss, discard_prior_samples, retrain_from_scratch, show_train_summary, dataloader_kwargs)
    351     theta = reshape_to_sample_batch_event(
    352         theta.to("cpu"), self._neural_net.input_shape
    353     )
    354     x = reshape_to_batch_event(x.to("cpu"), self._neural_net.condition_shape)
--> 355     test_posterior_net_for_multi_d_x(self._neural_net, theta, x)
    357     del theta, x
    359 # Move entire net to device for training.

File ~/anaconda3/envs/sbi_env2/lib/python3.10/site-packages/sbi/utils/user_input_checks.py:808, in test_posterior_net_for_multi_d_x(net, theta, x)
    805 else:
    806     message = ""
--> 808 raise RuntimeError(message) from rte

RuntimeError: ```

[michaeldeistler]

Hi,

does the minimal example I posted above work for you? If yes, do you think you could provide a minimal example of the bug you are facing?

Thank you!
Michael

[btetaren]

Apologies that my previous comment was unclear. Yes, the minimal example works fine. It is when I tried to insert my own training data (theta,x) where I get a traceback like above. In trying to figure out what was causing the error, I was able to reproduce it with an edit to your minimal example as follows:

import torch
from sbi.inference import MNPE
from sbi.neural_nets.embedding_nets import FCEmbedding
from sbi.neural_nets import posterior_nn

dim_theta_discrete = 3
dim_theta_continuous = 2
dim_theta = 5
dim_x = 50

num_sims = 100

dt1 = torch.randint(low=0, high=2, size=(100, 2))
dt2= torch.randint(low=10, high=20, size=(100, 1))
discrete_theta = torch.cat([dt1, dt2], dim=1)

continuous_theta = torch.randn((num_sims, dim_theta_discrete))
theta = torch.cat([continuous_theta, discrete_theta], dim=1)
x = torch.randn((num_sims, dim_x))

embedding_net = FCEmbedding(dim_x)
net_builder = posterior_nn("mnpe", embedding_net=embedding_net)

inference = MNPE(density_estimator=net_builder)
_ = inference.append_simulations(theta, x).train()

which emulates what my training data looks like i.e., my individual discrete theta dimensions lie in different numerical ranges.

[michaeldeistler]

This seems like a bug, I have opened an issue here.