multi poi support - #69
Conversation
…c supplied per generate call
… criteria and toy key threading
…lculator owns the fixed distribution; ToyCalculator requires its generator and key
… the Phi(-N) quantile of q_alt)
📝 WalkthroughWalkthroughThe PR refactors hypothesis testing around mapping-based POI inputs, adds ChangesHypothesis-testing API overhaul
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related issues
Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
⚔️ Resolve 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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 8
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
docs/extending-hypotest.md (1)
159-182: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winImport
jax.numpyin the custom sampling example.The changed code calls
jnp.sqrtandjnp.maximum, but imports onlyjax, so copying the example raisesNameError.Proposed fix
import jax +import jax.numpy as jnp🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/extending-hypotest.md` around lines 159 - 182, Update the custom sampling example imports to include jax.numpy under the jnp alias used by sample_fn, so jnp.sqrt and jnp.maximum resolve when the example is copied.docs/tips.md (1)
12-26: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winImport
AsymptoticCalculatorin the JIT example.The updated snippet constructs
AsymptoticCalculatorbut never imports it, so it cannot run as written.Proposed fix
import jax +from everwillow.hypotest.calculators import AsymptoticCalculator from everwillow.hypotest.limit_solvers import RootFindingLimitSolver🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/tips.md` around lines 12 - 26, Add the missing AsymptoticCalculator import to the JIT example so the compute_limit function can construct it successfully, while preserving the existing RootFindingLimitSolver import and calculation flow.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/extending-hypotest.md`:
- Around line 205-214: Update the ToyGenerator examples to define a minimal
predict function before the first use of predict, or clearly state that it must
come from earlier user code. Ensure both map_fn variants continue using the same
defined prediction function.
In `@docs/multi-poi.md`:
- Around line 166-173: Update the “Toy-based regions” example to define toy_calc
before calling confidence_region, either by adding a concise ToyCalculator
construction with the required import or by explicitly stating that toy_calc is
an already-configured instance. Ensure the example is self-contained enough to
avoid an undefined-name error.
- Around line 83-84: Add the missing imports for BisectionLimitSolver and
Matplotlib’s pyplot as plt before the examples in docs/multi-poi.md, ensuring
both references in the upper_limit examples are defined.
In `@docs/quickstart.md`:
- Around line 231-232: Update the quickstart documentation around the
toy-hypothesis descriptions to distinguish null toys, which use the tested POI
of 1.0, from alternative toys, which use poi_alt (0.0 here). Apply the same
correction to the corresponding explanation near the additionally referenced
section, while preserving the existing tail-counting and distribution_factory
guidance.
In `@src/everwillow/_src/inference/hypotest/calculators.py`:
- Around line 10-12: Update the module docstring entry for AsymptoticCalculator
to replace the stale mu_asimov field name with poi_asimov, while preserving the
existing predict_fn and asimov_observation references and description.
In `@src/everwillow/_src/inference/hypotest/limit_solvers.py`:
- Around line 283-291: Update the bisect iteration logic around the converged
flag so lo and hi remain unchanged when the evaluated mid satisfies the
tolerance; only narrow the bracket when not converged. Preserve the existing
directional updates for non-converged iterations, ensuring the post-loop
final_mid equals the converged mid.
In `@src/everwillow/_src/inference/hypotest/results.py`:
- Line 132: Update the docstring Attributes entry in the relevant results class
to use the field name distribution instead of the stale dist name, keeping the
description unchanged.
In `@src/everwillow/hypotest/limit_solvers.py`:
- Around line 20-26: Reorder the entries in the __all__ declaration
alphabetically to satisfy Ruff RUF022: place BisectionLimitSolver,
GridScanLimitSolver, LimitSolver, RootFindingLimitSolver, and
StochasticLimitSolver in that order.
---
Outside diff comments:
In `@docs/extending-hypotest.md`:
- Around line 159-182: Update the custom sampling example imports to include
jax.numpy under the jnp alias used by sample_fn, so jnp.sqrt and jnp.maximum
resolve when the example is copied.
In `@docs/tips.md`:
- Around line 12-26: Add the missing AsymptoticCalculator import to the JIT
example so the compute_limit function can construct it successfully, while
preserving the existing RootFindingLimitSolver import and calculation flow.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: e832f783-de5c-4557-827d-bba471830251
📒 Files selected for processing (29)
docs/api/inference/hypotest.mddocs/asimov-datasets.mddocs/extending-hypotest.mddocs/index.mddocs/multi-poi.mddocs/quickstart.mddocs/tips.mdpyproject.tomlsrc/everwillow/_src/inference/hypotest/calculators.pysrc/everwillow/_src/inference/hypotest/distributions.pysrc/everwillow/_src/inference/hypotest/limit_solvers.pysrc/everwillow/_src/inference/hypotest/results.pysrc/everwillow/_src/inference/hypotest/test_statistics.pysrc/everwillow/_src/inference/hypotest/toys.pysrc/everwillow/_src/inference/hypotest/upper_limit.pysrc/everwillow/_src/inference/hypotest/utils.pysrc/everwillow/hypotest/__init__.pysrc/everwillow/hypotest/calculators.pysrc/everwillow/hypotest/limit_solvers.pysrc/everwillow/hypotest/results.pysrc/everwillow/hypotest/upper_limit.pytests/inference/hypotest/_counting_model.pytests/inference/hypotest/test_calculators.pytests/inference/hypotest/test_distributions.pytests/inference/hypotest/test_limit_solvers.pytests/inference/hypotest/test_test_statistics.pytests/inference/hypotest/test_toys.pytests/inference/hypotest/test_upper_limit.pytests/inference/hypotest/test_utils.py
💤 Files with no reviewable changes (3)
- src/everwillow/hypotest/upper_limit.py
- tests/inference/hypotest/test_upper_limit.py
- src/everwillow/_src/inference/hypotest/upper_limit.py
| ToyGenerator( | ||
| test_statistic=QTilde(), | ||
| predict_fn=predict, | ||
| map_fn=lambda fn: partial(jax.lax.map, fn, batch_size=8), | ||
| ) | ||
|
|
||
| # Python loop (no JIT, useful for step-through debugging) | ||
| ToyGenerator( | ||
| test_statistic=QTilde(), | ||
| predict_fn=predict, | ||
| map_fn=lambda fn: lambda keys: jnp.stack([fn(k) for k in keys]), | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Define predict before using it.
The new ToyGenerator(predict_fn=predict, ...) example references no predict defined in this page. Add a minimal prediction function or explicitly state that it must be supplied from earlier user code.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/extending-hypotest.md` around lines 205 - 214, Update the ToyGenerator
examples to define a minimal predict function before the first use of predict,
or clearly state that it must come from earlier user code. Ensure both map_fn
variants continue using the same defined prediction function.
| calc.upper_limit(BisectionLimitSolver(bounds=(0.0, 5.0)), poi_key="mu_a") | ||
| calc.upper_limit(BisectionLimitSolver(bounds=(0.0, 12.0)), poi_key="mu_b") |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Add imports for symbols used by the examples.
BisectionLimitSolver and plt are referenced but never imported in this guide. Add the corresponding limit-solver and Matplotlib imports before the examples.
Also applies to: 146-147
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/multi-poi.md` around lines 83 - 84, Add the missing imports for
BisectionLimitSolver and Matplotlib’s pyplot as plt before the examples in
docs/multi-poi.md, ensuring both references in the upper_limit examples are
defined.
| ## Toy-based regions | ||
|
|
||
| `ToyCalculator.confidence_region` regenerates the toy ensembles at every | ||
| point, threading an independent subkey into each. The whole scan is | ||
| reproducible from the calculator's key: | ||
|
|
||
| ```python | ||
| region = toy_calc.confidence_region(points, criterion=lambda r: r.pnull) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Define or explicitly assume toy_calc.
This page never constructs a ToyCalculator or imports it, so copying the toy-region example raises NameError. Add a short calculator construction, or state that toy_calc is an already-configured instance.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/multi-poi.md` around lines 166 - 173, Update the “Toy-based regions”
example to define toy_calc before calling confidence_region, either by adding a
concise ToyCalculator construction with the required import or by explicitly
stating that toy_calc is an already-configured instance. Ensure the example is
self-contained enough to avoid an undefined-name error.
| # Toys are thrown at poi=1.0 under both hypotheses; p-values come from tail | ||
| # counting by default (swap distribution_factory for e.g. smoothed variants). |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Distinguish null and alternative toy hypotheses.
With poi_alt=0.0, null toys use the tested POI (1.0 here), while alternative toys use poi_alt. The current wording says both ensembles match the tested POI, which is inaccurate for CLs. (raw.githubusercontent.com)
Also applies to: 294-297
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/quickstart.md` around lines 231 - 232, Update the quickstart
documentation around the toy-hypothesis descriptions to distinguish null toys,
which use the tested POI of 1.0, from alternative toys, which use poi_alt (0.0
here). Apply the same correction to the corresponding explanation near the
additionally referenced section, while preserving the existing tail-counting and
distribution_factory guidance.
| - ``AsymptoticCalculator``: Extends the base with Asimov dataset config | ||
| (``predict_fn``/``mu_asimov`` or ``asimov_observation``) for Cowan et al. | ||
| asymptotic workflows. |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Stale field name in module docstring. AsymptoticCalculator no longer exposes mu_asimov; the Asimov config field is poi_asimov (line 450). Update the module docstring to match.
📝 Proposed doc fix
- ``AsymptoticCalculator``: Extends the base with Asimov dataset config
- (``predict_fn``/``mu_asimov`` or ``asimov_observation``) for Cowan et al.
+ (``predict_fn``/``poi_asimov`` or ``asimov_observation``) for Cowan et al.
asymptotic workflows.📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| - ``AsymptoticCalculator``: Extends the base with Asimov dataset config | |
| (``predict_fn``/``mu_asimov`` or ``asimov_observation``) for Cowan et al. | |
| asymptotic workflows. | |
| ``AsymptoticCalculator``: Extends the base with Asimov dataset config | |
| (``predict_fn``/``poi_asimov`` or ``asimov_observation``) for Cowan et al. | |
| asymptotic workflows. |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/everwillow/_src/inference/hypotest/calculators.py` around lines 10 - 12,
Update the module docstring entry for AsymptoticCalculator to replace the stale
mu_asimov field name with poi_asimov, while preserving the existing predict_fn
and asimov_observation references and description.
| converged = jnp.abs(mid_value - level) < self.tol | ||
|
|
||
| # Update the search bracket based on the middle value | ||
| # narrow it: if mid_value is below level, lo moves to mid; if above, hi moves to mid | ||
| # This assumes the objective is montonically decreasing with increasing POIs | ||
| new_lo = jnp.where(mid_value > level, mid, lo) | ||
| new_hi = jnp.where(mid_value <= level, mid, hi) | ||
|
|
||
| return iteration + 1, new_lo, new_hi, converged |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Bisection early-exit returns a POI offset from the converged point.
bisect always moves lo/hi to mid, including on the iteration where converged becomes True. After the loop exits, final_mid = (lo + hi) / 2 is therefore (mid + old_bound) / 2, not mid — an offset of up to a quarter of the current bracket. This contradicts the docstring pseudocode, which stops before updating lo/hi so that (lo + hi) / 2 == mid. The error is largest for flat objectives, where the value tolerance is met while the POI bracket is still wide, so a wrong limit is returned silently.
Freeze the bracket on the converged step so final_mid equals the evaluated mid:
🐛 Proposed fix
- new_lo = jnp.where(mid_value > level, mid, lo)
- new_hi = jnp.where(mid_value <= level, mid, hi)
+ new_lo = jnp.where(converged, lo, jnp.where(mid_value > level, mid, lo))
+ new_hi = jnp.where(converged, hi, jnp.where(mid_value <= level, mid, hi))📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| converged = jnp.abs(mid_value - level) < self.tol | |
| # Update the search bracket based on the middle value | |
| # narrow it: if mid_value is below level, lo moves to mid; if above, hi moves to mid | |
| # This assumes the objective is montonically decreasing with increasing POIs | |
| new_lo = jnp.where(mid_value > level, mid, lo) | |
| new_hi = jnp.where(mid_value <= level, mid, hi) | |
| return iteration + 1, new_lo, new_hi, converged | |
| converged = jnp.abs(mid_value - level) < self.tol | |
| # Update the search bracket based on the middle value | |
| # narrow it: if mid_value is below level, lo moves to mid; if above, hi moves to mid | |
| # This assumes the objective is montonically decreasing with increasing POIs | |
| new_lo = jnp.where(converged, lo, jnp.where(mid_value > level, mid, lo)) | |
| new_hi = jnp.where(converged, hi, jnp.where(mid_value <= level, mid, hi)) | |
| return iteration + 1, new_lo, new_hi, converged |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/everwillow/_src/inference/hypotest/limit_solvers.py` around lines 283 -
291, Update the bisect iteration logic around the converged flag so lo and hi
remain unchanged when the evaluated mid satisfies the tolerance; only narrow the
bracket when not converged. Preserve the existing directional updates for
non-converged iterations, ensuring the post-loop final_mid equals the converged
mid.
| pnull: p-value under the tested hypothesis (poi_test). | ||
| palt: p-value under the alternative hypothesis (poi_alt / background-only). | ||
| test_stat_result: Full test statistic result with fit information. | ||
| dist: The distribution of the test statistic which was used to compute p-values. |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Docstring attribute name is stale. The Attributes entry is named dist, but the field is distribution. Rename for accuracy.
📝 Proposed doc fix
- dist: The distribution of the test statistic which was used to compute p-values.
+ distribution: The distribution of the test statistic which was used to compute p-values.📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| dist: The distribution of the test statistic which was used to compute p-values. | |
| distribution: The distribution of the test statistic which was used to compute p-values. |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/everwillow/_src/inference/hypotest/results.py` at line 132, Update the
docstring Attributes entry in the relevant results class to use the field name
distribution instead of the stale dist name, keeping the description unchanged.
| __all__ = [ | ||
| "LimitSolver", | ||
| "StochasticLimitSolver", | ||
| "RootFindingLimitSolver", | ||
| "GridScanLimitSolver", | ||
| "BisectionLimitSolver", | ||
| ] |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Sort __all__ to satisfy Ruff RUF022.
Use isort-style alphabetical ordering, for example BisectionLimitSolver, GridScanLimitSolver, LimitSolver, RootFindingLimitSolver, StochasticLimitSolver.
🧰 Tools
🪛 Ruff (0.15.21)
[warning] 20-26: __all__ is not sorted
Apply an isort-style sorting to __all__
(RUF022)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/everwillow/hypotest/limit_solvers.py` around lines 20 - 26, Reorder the
entries in the __all__ declaration alphabetically to satisfy Ruff RUF022: place
BisectionLimitSolver, GridScanLimitSolver, LimitSolver, RootFindingLimitSolver,
and StochasticLimitSolver in that order.
Source: Linters/SAST tools
1 participant
Summary by CodeRabbit
New Features
Documentation
Tests