Fix MixtureSameFamily log probability computations as described in #3188

[justjhong]

Fixes #3188.

This PR fixes the issue in both the get_aggregated_posterior function as well as the Mixture of Gaussians prior.

Tested on tutorial: integration remains the same, differential abundance results now take a much larger range as expected (see issue). To mitigate spikiness in favor of the sample of origin's covariate, we added the default omit_original_sample=True to avoid incorporating the sample of origin's log prob into the calculation.

CC @rastogiruchir

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 86.11%. Comparing base (c3926eb) to head (798eee6).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3189      +/-   ##
==========================================
- Coverage   89.27%   86.11%   -3.16%     
==========================================
  Files         185      185              
  Lines       16265    16270       +5     
==========================================
- Hits        14520    14011     -509     
- Misses       1745     2259     +514     

Files with missing lines	Coverage Δ
src/scvi/external/mrvi/_model.py	89.17% <100.00%> (+0.14%)	⬆️
src/scvi/external/mrvi/_module.py	96.01% <100.00%> (-0.04%)	⬇️

... and 19 files with indirect coverage changes

[PierreBoyeau]

LGTM

[PierreBoyeau]

snippet to double-check it works in loss (it works):

        pu = generative_outputs["pu"]
        
        # approach 1: below is the reference implementation of the log_prob of the mixture distribution
        log_pu = generative_outputs["pu"].log_prob(inference_outputs["u"])
        
        # approach 2: manual computation of the log_prob of the mixture distribution
        mixing_distribution = pu._mixing_distribution
        component_distribution = pu._component_distribution
        pk = mixing_distribution.probs
        mus = component_distribution.base_dist.loc
        stds = component_distribution.base_dist.scale
        
        pu_k = dist.Normal(mus, stds)
        log_pu_k = pu_k.log_prob(inference_outputs["u"][:, None, :]).sum(-1)
        log_pk = jnp.log(pk)
        log_puk = log_pk + log_pu_k
        # above are joint probas
        # shape (n_cells, n_mixtures)
        
        log_prob_u = jax.scipy.special.logsumexp(log_puk, axis=1)  # shape (n_cells,)
        # (log_prob_u == log_pu) ?

[PierreBoyeau]

Looks good to me, see comments!

[PierreBoyeau]

Quick observations:

• These changes do not seem to affect the performed analyses conducted in the paper. Here are the log ratios (COVID vs healthy) from the complete Haniffa dataset, showing a similar behavior as in the MrVI manuscript.
  
• That being said, these log ratios seem slightly different when computed on the drastically subsampled Haniffa dataset. If similar issues are reported on small datasets, we might want to explore options to smooth log ratios then.