Fix MoG conditioning for non-contiguous dimensions

[vansh-oberoi]

What does this PR do?

Fixes #2038.

The conditioning of a mixture of Gaussians was incorrect when a free dimension came after a fixed dimension.

The issue was caused by constructing the precision submatrices by slicing the precision factor before computing the full precision matrix. This can give incorrect conditional precisions for non-contiguous dimensions.

This change computes the full precision matrix first, extracts the required blocks, and uses the Schur complement when calculating the marginal likelihood used for the mixture weights.

I also added a regression test using an analytically known Gaussian conditional.

Tests

• uv run pytest tests/mog_test.py -q
• uv run pytest tests/sbiutils_test.py -q
• uv run pytest tests/posterior_nn_test.py -q
• uv run ruff check sbi/neural_nets/estimators/mog.py tests/mog_test.py

All passed.

[coderabbitai]

Navigate logical layers of code changes, visualize relationships, and explore their blast radius.

Note

Currently processing new changes in this PR. This may take a few minutes, please wait...

⚙️ Run configuration

• Configuration used: Organization UI
• Review profile: CHILL
• Plan: Advanced
• Run ID: ed4d18bf-e13f-43cc-bdfe-bd0ee5165bd1

📥 Commits

Reviewing files that changed from the base of the PR and between 101e107 and 2ec3800.

📒 Files selected for processing (1)

• tests/mog_test.py

 _________________________________________________________
< Your config has more knobs than a 2004 Honda dashboard. >
 ---------------------------------------------------------
  \
   \   (\__/)
       (•ㅅ•)
       / 　 づ

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Organization UI
• Review profile: CHILL
• Plan: Advanced
• Run ID: 5e3e69c6-a0d4-4872-9fdd-fe437163eddc

📥 Commits

Reviewing files that changed from the base of the PR and between b8afbd0 and 101e107.

📒 Files selected for processing (2)

• sbi/neural_nets/estimators/mog.py
• tests/mog_test.py

Included review availability: This review used your included allowance. Your plan provides up to 4 included reviews per hour; 3 remain after this review.

---

📝 Walkthrough

Walkthrough

MoG.condition now derives conditional and marginal precision blocks from the full precision matrix. It factors the free-dimension precision block and uses the fixed-dimension Schur complement for component likelihoods and mixture weights. A correlated two-dimensional Gaussian test checks the analytic conditional.

Changes

MoG conditioning

Layer / File(s)	Summary
Precision blocks and conditional validation sbi/neural_nets/estimators/mog.py, tests/mog_test.py	MoG.condition derives precision blocks from the full precision matrix and uses the fixed-dimension Schur complement for component likelihoods. The added test checks the conditioned mean and precision against an analytic result.

Priority: ➖ Normal

Estimated code review effort: 3 (Moderate) | ~20 minutes

Change: Bug fix · Severity of issue fixed: Medium

Merge Risk: ⚪ Minimal · up to 101e1

This change corrects MoG conditioning when a free dimension follows a fixed one, and adds an analytic regression test. No merge-blocking risk is evident from the supplied review.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely identifies the fix for MoG conditioning with non-contiguous dimensions.
Description check	✅ Passed	The description explains the problem and solution, links the issue, and lists tests that passed. It does not include the template’s “Anything else” or checklist sections, but the main required informa…
Linked Issues check	✅ Passed	Issue [#2038] requires correct conditional means, precisions, and mixture weights when free dimensions follow fixed dimensions. MoG.condition now forms the full precision matrix, extracts the free a…
Out of Scope Changes check	✅ Passed	The changes are limited to MoG.condition and its regression test in tests/mog_test.py. Both changes directly address issue [#2038].
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 3 functions across 2 files.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

• Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[BHARATH0153]

#2038 already had pr #2039

[vansh-oberoi]

#2038 already had pr #2039
Thanks for pointing this out. I missed #2039 when I initially picked up the issue. I checked it now and it looks like it already addresses the same problem, so I’ll close this PR to avoid duplicating the work.