refactor mixed density estimation

[janfb]

What does this implement/fix? Explain your changes

A couple of fixes and improvements around MNLE:

• log prob for iid data is now much faster on the discrete data because we can actually just pass the iid data as sample dim to CategoricalNet. There is no need to do tricks with the repetitions in the categorical data (i.e., log_prob_iid can be removed (it was not used anyway)).
• allow all kinds of flows, not just nsf for MNLE
• unify z-scoring API for build_categoricalmassestimator
• fix embedding net handling for MNLE. It did not allow for theta embeddings before. Now it does. Importantly, one has to use a "mixed embedding" for the conditioning of the flow because the condition contains embedded theta and "raw" discrete data.

Does this close any currently open issues?

Fixes #1134
Fixes #1136
Fixes #1172

Any other comments?

the first commit is for avoiding circular imports.

[codecov]

Codecov Report

Attention: Patch coverage is 98.07692% with 1 line in your changes missing coverage. Please review.

Project coverage is 75.97%. Comparing base (ba19688) to head (5aa7c2a).
Report is 10 commits behind head on main.

Files	Patch %	Lines
sbi/neural_nets/categorial.py	91.66%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1203      +/-   ##
==========================================
- Coverage   84.55%   75.97%   -8.59%     
==========================================
  Files          96       97       +1     
  Lines        7603     7668      +65     
==========================================
- Hits         6429     5826     -603     
- Misses       1174     1842     +668     

Flag	Coverage Δ
unittests	75.97% <98.07%> (-8.59%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files	Coverage Δ
.../neural_nets/density_estimators/categorical_net.py	97.82% <100.00%> (-0.22%)	⬇️
...nets/density_estimators/mixed_density_estimator.py	98.11% <100.00%> (+28.54%)	⬆️
sbi/neural_nets/density_estimators/zuko_flow.py	64.44% <ø> (ø)
sbi/neural_nets/flow.py	93.12% <100.00%> (ø)
sbi/neural_nets/mnle.py	100.00% <100.00%> (ø)
sbi/utils/__init__.py	100.00% <ø> (ø)
sbi/neural_nets/categorial.py	94.73% <91.66%> (+4.73%)	⬆️

... and 40 files with indirect coverage changes

[michaeldeistler]

Thanks a ton, this is awesome!

For now, this can only handle 1D discrete dimensions, right?

In the long run, the ultimate solution would be to have an Autoregressive abstraction which just concatenates a list of Estimators (can be DensityEstimator or MassEstimator) in an autoregressive way. No need to do this now ofc.

[michaeldeistler]

we do not have a sample_dim for the condition, but we have a sample_dim for the input. I think the docstring is correct

[michaeldeistler]

One more comment for the PartialEmbedding: I think we could also have the embedding_net as an attribute of MixedDensityEstimator instead of as an attribute of the DensityEstimator and the MassEstimator. Not fully sure if this would work, but it feels a bit wasteful to have the embedding_net twice.

[janfb]

Thanks a ton, this is awesome!

For now, this can only handle 1D discrete dimensions, right?

In the long run, the ultimate solution would be to have an Autoregressive abstraction which just concatenates a list of Estimators (can be DensityEstimator or MassEstimator) in an autoregressive way. No need to do this now ofc.

Yes, the CategoricalNet can only handle 1D because it's not straight forward to extend this. Maybe it would work better with @coschroeder 's Grassmann distribution? Although that one is only binary.

Yes, the autoregressive approach of just concatenating conditionals and using one Categorical Net for each dimension would be a really nice feature.

[janfb]

One more comment for the PartialEmbedding: I think we could also have the embedding_net as an attribute of MixedDensityEstimator instead of as an attribute of the DensityEstimator and the MassEstimator. Not fully sure if this would work, but it feels a bit wasteful to have the embedding_net twice.

I think both estimator need an embedding net because these are different embeddings. For the MassEstimator, it's the y / theta embedding that is sticked together with the potential standardizing net in build_categoricalmassestimator. For the DensityEstimator, it's the PartialEmbedding that contains both the discrete and the y / theta data.
But I am not sure whether that's what you mean?

Looking at this now, I think there is a problem: the density estimation build function, e.g., build_nsf will build a standardizing net for the PartialEmbedding using the entire batch of concatenated discrete and continuous data. So effectively, it will z-score the discrete data as well. E.g., when y is an image so that the y_batch passed to the embedding has shape (batch, 32, 32), including the expanded and repeated discrete x values, then the discrete x values will influence the z-scoring of the image, right?

[janfb]

Update: The PartialEmbedding and the expanding and repeating of discrete data is not needed anymore.

We now build the y-embedding inside of build_mnle so that we can pass the concatenation of the embedded y with the discrete data into the continuous density estimator. There, we also pass a combined_embedding_net that combines that two conditions.

This also enables us to handle sample_shape>1 by just repeating the embedded continuous condition accordingly, and concatenating it with the discrete input that has a sample_shape>1.

[michaeldeistler]

This is awesome!!!!!

A few comments below, but good to go then!