feat: add power_analysis() with sigmoid fitting strategy (#820)#989
feat: add power_analysis() with sigmoid fitting strategy (#820)#989genrichez wants to merge 5 commits into
Conversation
Add efficient power curve estimation for geo-experiment design. Two strategies are supported: - 'grid': brute-force Monte Carlo at every user-supplied effect size - 'sigmoid': evaluate at ~5 points, fit a two-parameter logistic, and extract the MDE analytically (74% computation reduction) The function operates on a completed PlaceboInTime CheckResult, leveraging the learned null distribution to simulate detection probability at hypothetical effect sizes using the same ROPE-based Bayesian decision rule. New public API: - power_analysis() function - PowerCurveResult dataclass (with .plot() method) - LogisticFit dataclass (with .predict() and .mde() methods) Validation (5-point sigmoid vs 20-point grid): - Max absolute error: 1.4% - Mean absolute error: 0.6% - Speedup: 3.8x (74% computation reduction) Closes pymc-labs#820
|
👋 Welcome to CausalPy, @genrichez! Thank you for opening your first pull request! We're excited to have you contribute to the project. 🎉 Here are a few tips to help your PR get merged smoothly:
A maintainer will review your changes soon. Thanks for helping make CausalPy better! 🚀 💼 LinkedIn Shoutout: Once your PR is merged, we'd love to give you a shoutout on LinkedIn to thank you for your contribution! If you're interested, just drop your LinkedIn profile URL in a comment below. |
Documentation build overview
57 files changed ·
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #989 +/- ##
==========================================
+ Coverage 95.41% 95.48% +0.07%
==========================================
Files 97 99 +2
Lines 15667 15967 +300
Branches 920 936 +16
==========================================
+ Hits 14948 15246 +298
- Misses 507 509 +2
Partials 212 212 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
… branch coverage tests - Moved test_power_analysis.py to causalpy/tests/ (where CI discovers tests) - Added pragma no cover to plot() method (standard for visualization code) - Added 5 new tests covering: fold_results fallback, empty fold_results error, sigmoid with >2 effect sizes, MDE outside range warning, fitting failure - Patch coverage now 98% (was 24.59% due to test file in wrong location)
daimon-pymclabs
left a comment
daimon-pymclabs
left a comment
There was a problem hiding this comment.
Thanks for this — the sigmoid idea is well-motivated and the module is clean, well-documented, and comes with a genuinely thorough test suite (nice that the tests are pure-numpy against a hand-built CheckResult, so CI stays fast with no MCMC). I verified the metadata contract: PlaceboInTime.run() does emit null_samples, fold_sds, rope_half_width, and threshold, so the function wires up correctly against real output. A few things worth discussing before merge, one design-level and a handful of code-level.
Design: is this the API we want?
The biggest question is scope alignment. Issue #820 framed this as adding a strategy="sigmoid" parameter to the existing SyntheticControl.power_analysis() method introduced in #819, and #819 is still open. This PR instead ships a standalone free function cp.checks.power_analysis(pit_result, ...) that operates on a completed PlaceboInTime result rather than a SyntheticControl experiment. So we'd end up with two parallel power-analysis surfaces that don't share code or signature:
SyntheticControl.power_analysis(...)(method, from #819)causalpy.checks.power_analysis(pit_result, ...)(function, this PR)
That's worth reconciling explicitly. Options I see: (a) land the sigmoid strategy inside the #819 method as the issue proposed; (b) keep the free function but factor the shared sigmoid-fitting + LogisticFit + PowerCurveResult into one place both call; (c) decide the two are genuinely different tools (SC dress-rehearsal vs PlaceboInTime learned-null) and document why. Right now #820's acceptance criteria aren't literally met, so I'd want @juanitorduz to confirm which direction is intended before investing more in either.
Statistical: null-component sampling (shared with _compute_assurance)
The detection simulation draws the null component as null_samples[i % len(null_samples)] (deterministic prefix), so with the typical n_simulations=200 against ~4000 theta_new draws you only ever use the first 200 samples. Those draws come from posterior_predictive["theta_new"].stack(sample=("chain","draw")), which orders chain-slowest — so the first 200 are all from chain 0. random_seed doesn't change which null draws are used (it only affects the sigma pick and the normal noise). This is a real (mild) bias and it's shared with the existing PlaceboInTime._compute_assurance, which uses the identical i % len idiom. Suggest fixing it centrally with rng.choice(null_samples, size=n_simulations, replace=True) and reusing that in both places. See inline note.
Correctness note (in the PR's favor)
The analytic MDE in code is x0 + (1/k)·ln(tau/(1-tau)), which is the correct inversion (MDE > x0 for tau>0.5). Note the issue text wrote ln((1-tau)/tau), which has the sign flipped — the code is right, the issue was wrong, good catch by whoever implemented it.
Smaller items (details inline)
- Decision rule is re-implemented inline instead of reusing
PlaceboInTime.bayesian_rope_decision; they can drift. except RuntimeErroraroundcurve_fitis too narrow (scipy raisesValueError/OptimizeWarningon degenerate data too), andLogisticFit.mdedivides bykwith no guard fork → 0.- Docstring defaults for
effect_sizesdon't match the code (4*tauvsmax(4*tau, 3*rope)). - The
fold_results→fr.fold_sdfallback is effectively dead against real output (whenevernull_samplesexists,fold_sdsdoes too), only reachable via hand-built metadata. - No end-to-end smoke test running an actual tiny
PlaceboInTime.run()intopower_analysis— all tests use a syntheticCheckResult, so metadata-shape drift wouldn't be caught.
None of the code-level items are blockers; the design question is the one I'd resolve first.
| detections = 0 | ||
| for i in range(n_simulations): | ||
| # Draw null component (structural noise) | ||
| null_component = float(null_samples[i % len(null_samples)]) |
There was a problem hiding this comment.
Null draw is a deterministic prefix, not seeded. null_samples[i % len(null_samples)] walks the first n_simulations entries in order. With n_simulations=200 and ~4000 theta_new draws, only the first 200 are ever used, and because theta_new is stack(sample=("chain","draw")) (chain-slowest) those 200 all come from chain 0. random_seed never touches this draw. Prefer rng.choice(null_samples, size=n_simulations, replace=True).
This exact idiom also lives in PlaceboInTime._compute_assurance, so ideally fix once and share (a small module-level _draw_null_components(rng, null_samples, n) helper both call) rather than duplicating the biased version here.
| loc=true_effect, scale=sigma, size=n_posterior_samples | ||
| ) | ||
| # Apply ROPE decision | ||
| prob_positive = float((simulated_posterior > rope_half_width).mean()) |
There was a problem hiding this comment.
The ROPE decision is re-implemented inline (prob_positive >= threshold) instead of calling the existing PlaceboInTime.bayesian_rope_decision(...) == "positive". It's numerically equivalent to the one-sided "positive" branch today, but if the shared decision rule ever changes (e.g. ROPE semantics), this copy silently drifts. Reusing the static method keeps power analysis and assurance consistent by construction.
| fold_results = meta.get("fold_results", []) | ||
| if not fold_results: | ||
| raise ValueError("pit_result has no fold_sds or fold_results in metadata.") | ||
| fold_sds_raw = [fr.fold_sd for fr in fold_results] |
There was a problem hiding this comment.
This fold_results → fr.fold_sd fallback is effectively unreachable against real PlaceboInTime output: whenever null_samples is present (required a few lines up), the run also stored fold_sds. The 0-folds-completed path that lacks fold_sds also lacks null_samples, so it fails the earlier check first. It's only exercised by hand-built metadata in the tests. Not harmful, but worth a comment noting it's a defensive path, or drop it.
| # Set up effect sizes based on strategy | ||
| if strategy == "grid": | ||
| if effect_sizes is None: | ||
| effect_sizes_arr = np.linspace(0, max(4.0 * tau, 3.0 * rope_half_width), 8) |
There was a problem hiding this comment.
Docstring drift: the docstring says the default is np.linspace(0, 4 * tau, 8) (grid) and [0, 4 * tau] (sigmoid), but the code uses max(4.0 * tau, 3.0 * rope_half_width) as the upper bound in both. Sync the docstring so users know the ROPE term participates in the default range.
| f"Consider widening the effect_sizes range.", | ||
| stacklevel=2, | ||
| ) | ||
| except RuntimeError as e: |
There was a problem hiding this comment.
except RuntimeError is too narrow. scipy.optimize.curve_fit raises ValueError on degenerate inputs (NaNs, inconsistent bounds/p0) and can emit OptimizeWarning while returning a junk fit rather than raising at all. The test_sigmoid_fitting_failure_warns case happens to land on a RuntimeError, but flat/constant detection rates can surface as ValueError. Suggest except (RuntimeError, ValueError).
Related: LogisticFit.mde computes (1.0 / self.k) * ... with k bounded [0, inf]. If the fit returns k == 0, this yields inf (numpy) or raises (python float); either way mde is meaningless. Consider guarding k <= 0 and emitting the same 'fitting failed' warning path.
juanitorduz
left a comment
juanitorduz
left a comment
There was a problem hiding this comment.
Hi @genrichez ! Thank you vary much for your contribution. I left some innitial comments on the code. I will take a detailed look into the logic next week :)
| # --------------------------------------------------------------------------- | ||
| # Tests for LogisticFit | ||
| # --------------------------------------------------------------------------- |
| def plot( # pragma: no cover | ||
| self, | ||
| power_threshold: float = 0.80, | ||
| ax: Any = None, |
There was a problem hiding this comment.
Can we have better type hints? Something like pllt.Axes | None = None
|
by the way daimon-pymclabs in an internal agent reviewer ;) |
Got it :) |
daimon-pymclabs
left a comment
daimon-pymclabs
left a comment
There was a problem hiding this comment.
Re-review after the latest push (a600dfd). Thanks @genrichez — the split into _simulate_detection_rate / _extract_null_distribution / _build_effect_sizes / _fit_sigmoid reads much better. Two of my earlier points are confirmed resolved:
- ✅ Null draw is now
rng.choice(null_samples)(was the deterministic first-nprefix walk) — properly seeded and it samples the whole null distribution rather than a chain-slow prefix. - ✅
_fit_sigmoidnow catches(RuntimeError, ValueError), so degeneratecurve_fitinputs (NaNs, inconsistent bounds/p0) no longer escape.
One thing still outstanding, plus two optional nits:
- Docstring drift (please fix). The
effect_sizesdocstring still says the default isnp.linspace(0, 4 * tau, 8)(grid) /[0, 4 * tau](sigmoid), but_build_effect_sizesusesmax(4.0 * tau, 3.0 * rope_half_width)as the upper bound in both branches. Worth syncing the docstring so users know the ROPE half-width participates in the default range. - (optional) The ROPE decision is still inlined as
(simulated_posterior > rope_half_width).mean() >= thresholdrather than reusingPlaceboInTime.bayesian_rope_decision. Numerically equivalent to the one-sided"positive"branch today; only matters if the shared rule ever diverges. - (optional) The
fold_sds→fold_results/fr.fold_sdfallback in_extract_null_distributionis still effectively unreachable against realPlaceboInTimeoutput — a result carryingnull_samplesalways carriesfold_sdstoo. Harmless as defensive code.
Nothing blocking from my side beyond the docstring. Deferring to @juanitorduz on the detailed logic pass.
Automated triageRecommendation: Why:
Review focus:
Confidence: high |
drbenvincent
added
the
review:high
4 participants
Summary
Implements efficient power curve estimation for geo-experiment design, as requested in #820.
What's New
power_analysis()functionComputes detection probability as a function of true effect size, using the learned null distribution from a completed
PlaceboInTimecheck. Two strategies:grid— evaluate at every user-supplied effect size (brute-force Monte Carlo)sigmoid— evaluate at ~5 points, fit a two-parameter logistic, extract MDE analyticallyNew public API
causalpy.checks.power_analysis(pit_result, ...)— main entry pointPowerCurveResult— dataclass with.plot()methodLogisticFit— fitted logistic with.predict()and.mde()methodsUsage
Validation
Compared 5-point sigmoid vs 20-point grid (ground truth):
Robust across different null parameter regimes (null_std 20–200, rope 10–50).
Tests
42 tests total (27 pytest + 15 edge case/business logic), covering:
Design Notes
PlaceboInTimeCheckResult, leveraging the learned null distribution (null_samples,fold_sds) to simulate detection at hypothetical effect sizes using the same ROPE-based Bayesian decision rule.P(posterior > rope) >= threshold), consistent with howPlaceboInTimeassurance works. Users wanting two-sided power should use|effect_size|.P(detect | x) = 1 / (1 + exp(-k*(x - x0))). MDE is extracted analytically:x0 + (1/k) * ln(tau / (1 - tau)).Closes #820