Add mixed data task and MNLE mode to mini SBIBM - #2034
Conversation
|
Navigate logical layers of code changes, visualize relationships, and explore their blast radius. 🧰 Additional context used📚 Code guidelines (2)No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configuration
📒 Files selected for processing (5)
💤 Files with no reviewable changes (1)
🚧 Files skipped from review as they are similar to previous changes (1)
Included review availability: This review used your included allowance. Your plan provides up to 4 included reviews per hour; 3 remain after this review. 📝 WalkthroughWalkthroughThe benchmark suite adds a mixed-data task for MNLE evaluation. A pytest option sets the number of independent trials. Contributor instructions document the default count and option. ChangesMNLE mixed-data benchmark
Priority: ⬇️ Low Estimated code review effort: 3 (Moderate) | ~20 minutes Change: Feature Sequence Diagram(s)sequenceDiagram
participant Pytest
participant TaskFixture
participant BenchmarkTest
participant EvaluationFunction
Pytest->>TaskFixture: Provide task parameter
TaskFixture->>BenchmarkTest: Supply constructed Task
BenchmarkTest->>EvaluationFunction: Pass Task to evaluation
Suggested reviewers: Merge Risk: ⚪ Minimal · up to The configurable-trial MNLE benchmark appears ready to merge after normal checks. Security Architecture ReviewSecurity architecture risk: ⚪ Minimal · up to The benchmark contract changes, but the inspected execution paths remain test-local and use synthetic data and predefined task constructors. No material security risk introduced or worsened by this change was identified. Retained concerns Security review detailsSecurity Blast Radius
Trust Boundaries and Controls
Resilience and Maintainability Implications
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2034 +/- ##
==========================================
- Coverage 89.65% 89.42% -0.24%
==========================================
Files 142 142
Lines 14736 15164 +428
==========================================
+ Hits 13212 13560 +348
- Misses 1524 1604 +80
Flags with carried forward coverage won't be shown. Click here to find out more. |
janfb
left a comment
janfb
left a comment
There was a problem hiding this comment.
@SahilKumar75 thanks for implementing this! I ran the new mode with the default 2000 simulations and it works well: C2ST is about 0.6 with 10 trials and the case takes about 40 seconds.
I have a few comments below, mostly small. The one that matters most is that runs with different trial counts end up with the same name in the results.
| """ | ||
| torch.manual_seed(SEED) | ||
| task = get_task(task_name) | ||
| task = _get_benchmark_task(task_name, benchmark_num_iid_trials) |
There was a problem hiding this comment.
Runs with a different --bm-num-iid-trials get the same test id and the same row in the results (MNLE{} on mixed_data), and the results file does not store the trial count. So a run with 3 trials silently replaces a run with 10 trials, although they evaluate different posteriors. Please make the trial count part of the task name in the results, e.g. mixed_data-10trials, so these rows stay apart.
| "snpe": [NPE_C], # NPE_B not implemented, NPE_A need Gaussian prior | ||
| "snle": [NLE], | ||
| "snre": [NRE_A, NRE_B, NRE_C, BNRE], | ||
| "mnle": [MNLE], |
There was a problem hiding this comment.
--bm-mode mnle --bm-estimators mnle currently fails during collection with "requires a benchmark mode whose methods use the same estimator argument", because the new mode has no entry in the estimator argument map. MNLE takes its estimator as density_estimator, so adding "mnle": "density_estimator" to that map should be enough.
| return "-".join(str(value) for value in parameters.values()) or "default" | ||
|
|
||
|
|
||
| def _get_benchmark_task(task_name: str, num_iid_trials: int) -> Task: |
There was a problem hiding this comment.
The trial count is only used by the mixed data task, but it is now passed through the sequential path too, which never runs this task, and the runner checks for the task name as a string. I think it would be cleaner if the mnle mode carried the trial count to its task, so that the rest of the runner stays as it is.
| choices = observation[:, 1:] | ||
|
|
||
| rate_posterior = Gamma( | ||
| torch.tensor([1.0 + self.stimulus_condition * self.num_trials]), |
There was a problem hiding this comment.
The posterior update takes the number of trials from the task setting and repeats the prior parameters as numbers. Both are right today, but if someone passes an observation with a different number of trials, or changes the prior, the reference becomes wrong without any error. Please take the number of trials from the observation and the parameters from the prior.
| from .mixed_data import MixedData | ||
|
|
||
|
|
||
| def test_mixed_data_task_shapes(): |
There was a problem hiding this comment.
The likelihood test below covers what matters, and it would fail if the reference were wrong. The shape test, the parameter test and the test for zero trials check things that either fail loudly anyway or are already covered by the likelihood test, so please remove them and keep only the likelihood test.
| rate, choice = task._get_reference_posterior(observation) | ||
| reference_log_prob = rate.log_prob(theta[:, 0]) + choice.log_prob(theta[:, 1]) | ||
|
|
||
| likelihood = InverseGamma(2.0, theta[:, :1]).log_prob(observation[:, 0]) |
There was a problem hiding this comment.
Small one: the likelihood here uses a fixed 2.0 for the reaction time concentration. Please take it from the task, so the test stays correct if the stimulus condition changes.
|
Thanks @janfb, addressed all six points in 4217c20. Trial counts now appear in both test IDs and result task names. MNLE estimator selection works, and the mode configures its task without passing trial counts through the sequential runner. The reference posterior now reads the observation length and prior parameters. I kept the likelihood test, added differing observation lengths and prior parameters to it, and removed the three redundant tests. The six likelihood cases pass. Benchmark runs with 20 simulations and 3 then 10 trials both pass and remain separate in the same results CSV. These are execution checks, not accuracy claims. Default collection remains 24 cases, and Ruff, formatting, pre commit checks and Pyright pass. |
janfb
left a comment
janfb
left a comment
There was a problem hiding this comment.
@SahilKumar75 thanks, this looks good now, and the stronger likelihood test is neat. Merging.
Bring in the mnle mode and the mixed data task from PR #2034. The mode now gets run labels, seeds and the mean and std errors like the other modes. Its docs move to the mini-sbibm developer page.
2 participants
What does this PR do?
Adds a mixed data task and an
mnlemode to Mini SBIBM, as discussed in #1396. The task shares the simulator and prior used in the MNLE tests. Its reference posterior is available analytically as independent Gamma and Beta distributions.Training uses individual trials. Evaluation observations contain 10 IID trials by default, configurable through
--bm-num-iid-trials. Existing default benchmark collection remains at 24 cases; the new mode collects one mixed data case.Does this close any issues?
Related to #1396. This implements the mixed data task only; the broader issue remains open.
Anything else we should know?
This follows the merged #2024 and is based on current main.
Validation performed:
uv run --no-sync pytest tests/mini_sbibm/mixed_data_test.py -q: six cases pass after rebasing, including agreement between the analytic reference density and simulator likelihood for 1, 4, and 10 trials.uv run --no-sync pyright sbipass after rebasing. Validation used the existing environment because downloads for the updated lockfile timed out.Benchmark execution and reporting pass with a fresh results file (20 simulations, three trials, C2ST 0.851). A previous collection run had created a CSV containing only the
test_idheader; appending execution results to it caused the existing summary error. The old file was preserved. Reporting changes remain outside this task, as maintainers are already working on that area.Implementation and tests were developed with AI assistance and checked locally.
Checklist
uv run pytest -n auto -m "not slow and not gpu"passes.uv run pre-commit run --all-filespasses (ruff and formatting).uv run pyright sbipasses.