add size invariant iid embedding nets, tests.

[janfb]

goal: have an embedding net for NPE that can handle varying number of iid trials, i.e., learn how the posterior changes as we change the number of trials.

problem: our training procedure assume that x is a tensor with fixed dimensions, e.g., the trial dimension must be the same for all training data points.

solution: given a training data set with a varying number of trials where the maximum number of trials is max_num_trials, pad all xs with smaller number of trials with NaNs such that x.shape = (num_thetas, max_num_trials, dim_x). Adapt the PermutationInvariantEmbeddingNet such that it detects the NaNs and applies the TrialEmbedding only to the valid entries (using a loop over the batch).

functional tests: seems to work fine for a Gaussian example

Questions:

• any ideas how to change our training procedure to be more flexible, e.g., allows x to be a list? treat x a torch.DataSet from the very start in append_simulations?

[michaeldeistler]

Thanks! A few points below. I'd like to think a bit about this, so let's maybe not merge it today yet?

[coschroeder]

let's discuss again if we can avoid the loop for varying trial numbers.

[codecov-commenter]

Codecov Report

Merging #808 (0ef585a) into main (2a114f4) will increase coverage by 0.05%.
The diff coverage is 87.50%.

📣 This organization is not using Codecov’s GitHub App Integration. We recommend you install it so Codecov can continue to function properly for your repositories. Learn more

@@            Coverage Diff             @@
##             main     #808      +/-   ##
==========================================
+ Coverage   74.76%   74.81%   +0.05%     
==========================================
  Files          80       80              
  Lines        6190     6191       +1     
==========================================
+ Hits         4628     4632       +4     
+ Misses       1562     1559       -3     

Flag	Coverage Δ
unittests	74.81% <87.50%> (+0.05%)	⬆️

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

Impacted Files	Coverage Δ
sbi/analysis/plot.py	56.43% <ø> (ø)
sbi/analysis/tensorboard_output.py	86.25% <ø> (ø)
sbi/inference/abc/smcabc.py	12.26% <0.00%> (ø)
sbi/inference/posteriors/base_posterior.py	80.82% <ø> (ø)
sbi/inference/posteriors/direct_posterior.py	98.21% <ø> (ø)
...inference/potentials/likelihood_based_potential.py	100.00% <ø> (ø)
sbi/inference/snle/snle_base.py	91.58% <ø> (ø)
sbi/inference/snpe/snpe_base.py	89.17% <ø> (ø)
sbi/inference/snre/bnre.py	44.00% <ø> (ø)
sbi/inference/snre/snre_base.py	94.82% <ø> (ø)
... and 10 more

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

[janfb]

this is working accurately for a Gaussian iid example with up 100 trials, trained with varying number of trials.

thanks to @manuelgloeckler's input, the forward pass is performed batched as well by masking the NaNs before the forward pass and setting them zero afterwards.

[manuelgloeckler]

Hey,

Looks good. One minor comment: I like the option to pass a custom aggregation function; currently, it would perfectly work for torch.sum . Yet for others like torch.max, torch.min or torch.median it would compute a slightly different value then expected as we substitute the invalid outputs with zero. We may should add this to the docstring, such that the user can adjust for this explicitly.

[michaeldeistler]

This is getting really awesome, I like the way it is implemented now!

As far as I understand, this does not work if x has a NaN value (which is not just because of missing trials). We should spell this out clearly in the docstring of this embedding net. We could even add a check which tests whether there are x for which only some summary stats are NaN and raise an error in this case (I'm also happy to leave this as TODO in the code for now)

Regarding my comments: I'd be happy to have a quick call if anything is unclear

[janfb]

I added asserts to catch NaNs in standardizing nets. This will inform the user when they want to use NaNs to encode varying number of trials that they have to turn of z-scoring and set exclude_invalid_x=False.