Automatically scale synth control sigma prior to the data

[szego]

This PR adds a feature to SyntheticControl that scales the prior on sigma to match the scale of the data.

Before, it always applied the prior sigma ~ HalfNormal(1), which works well for unit-scale data but will generally be too tight for data on world-scale (in units of counts, dollars, etc.). The user could make the model work for their data by doing the scaling/unscaling manually - say, by standardizing the data within cells, though that does change the model - or they would need to change the prior on that sigma.

We go the latter route and set a prior sigma ~ Exponential(2/s), where s is the pre-treatment standard deviation of the treated cells. The prior mean of sigma is then s/2, which seems reasonable as a weak prior on the size of the residuals. We want the prior on sigma to be wide enough not to get in the way but not, like, orders of magnitude too wide.

Motivation

We ran into this when developing this simulation study: https://research.getrecast.com/geolift-sim-study

We were using data in world-scale units and noticed that the confidence intervals from CausalPy were way too narrow. Here's a comparison between some CausalPy CIs to some Google Matched Markets CIs:

We realized we were probably just using CausalPy wrong, and a close read of its code lead us to discovering the sigma ~ HalfNormal(1) prior as the probable root cause. Putting a wider prior on sigma (roughly like we do in this PR) produced CIs more like we expected:

In the final version of the sim study we ended up going the standardize-the-data route and didn't change the HalfNormal(1) prior on sigma just to keep things simple. That also produced reasonable CIs since the HalfNormal(1) prior is fine in that regime.

Changes to the notebooks

Generally the changes to the results in the notebooks are minor since the data isn't too far from unit scale in them.

One thing we did change was deleting this line from docs/source/notebooks/geolift1.ipynb:

We can see that our sampling went well. PyMC returns no sampling warnings, and we have no divergences.

It's true that there were no divergences during sampling in the original example, but that was only by accident. When I tried different seeds in that notebook (on main, NOT on this PR's branch), most of them produced divergences there.

On this PR's branch we do get some divergences for that seed and other seeds, but it didn't seem much better or worse than the behavior on main. Similar can be said for the other notebooks.

Alternative proposal

If you'd rather not add this feature, we've also prepared a branch that only includes changes to the docs:

main...getrecast:CausalPy:docs/synth-control-scaling-reco

It recommends either applying this prior manually or standardizing the data. Happy to open a PR for that if you'd prefer.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[github-actions]

👋 Welcome to CausalPy, @szego!

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:

• ✅ Make sure all CI checks pass (tests, linting, type checking)
• 📝 Run prek run --all-files locally before pushing
• 📖 Check our Contributing Guide for more details

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.

[read-the-docs-community]

Documentation build overview

📚 causalpy | 🛠️ Build #33345896 | 📁 Comparing 5e8858e against latest (d513de3)

  🔍 Preview build  

151 files changed · ± 151 modified

± Modified

• 404.html
• notebooks/geolift1.html
• notebooks/sc_pymc.html
• notebooks/sc_pymc_brexit.html
• api/generated/causalpy.checks.base.CheckResult.html
• api/generated/causalpy.checks.placebo_in_time.AssuranceResult.html
• api/generated/causalpy.checks.placebo_in_time.PlaceboInTime.bayesian_rope_decision.html
• api/generated/causalpy.data.simulate_data.generate_ancova_data.html
• api/generated/causalpy.data.simulate_data.generate_did.html
• api/generated/causalpy.data.simulate_data.generate_geolift_data.html
• and 141 more...

[drbenvincent]

Amazing, thanks for this!

Very initial thought - sometimes you can get numerical issues with the exponential prior on a sigma because of so much mass on zero. So that might be a reason for the divergences.

We also recently added the SoftmaxWeightedSumFitter model which uses a different approach for the prior over weights constrained to sum to 1. Sometimes that offers more robust sampling.

But I'll look forward to reviewing this properly soon.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 95.35%. Comparing base (d513de3) to head (5e8858e).

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #976      +/-   ##
==========================================
+ Coverage   95.32%   95.35%   +0.02%     
==========================================
  Files          96       97       +1     
  Lines       15221    15302      +81     
  Branches      878      881       +3     
==========================================
+ Hits        14510    14591      +81     
  Misses        502      502              
  Partials      209      209              

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[szego]

I don't think the exponential prior leads to significantly more divergences than the current HalfNormal(1) prior.

I had claude write a messy script that sweeps across seeds 1:100 for the four models in the three notebooks, comparing the results on main to the results on this PR's branch:
drive_sweep.sh
run_sweep.py
(It took about 13.5 hours to run on my little laptop)

Here are the resulting seed x model x branch divergence counts:
divergence_sweep.csv

And here are the totals across seeds:

model	main branch	PR branch
sc	21	2
geolift1	170	147
brexit	88	89
brexit sparse	10785	12410

It did produce about 2% more divergences in the sparse-Dirichlet brexit model, but when you have that many divergences you have bigger problems already.

[daimon-pymclabs]

Review

Overall: the change fixes a real, well-diagnosed correctness problem — the default sigma ~ HalfNormal(1) is far too tight on world-scale data and produces artificially narrow CIs. The motivation, the simulation evidence, and the math all hold up: Exponential(lam=2/s) has mean s/2, matching the stated weak prior on residual size. The mechanism is also well-placed — the override sets model.priors["y_hat"] before fit(), and since WeightedSumFitter.priors_from_data only returns beta, the merge {**priors_from_data(X, y), **self.priors} preserves the override with correct precedence. _user_priors is left untouched, so _clone() still reflects the user's actual intent, and the explicit-y_hat guard works.

A few things I'd like to see addressed before merge:

1. No tests (blocking)

The PR adds a new public parameter and a new branch in algorithm() with zero dedicated tests. Codecov's "all lines covered" is incidental — the default path just happens to execute the branch, but nothing asserts behavior. I'd ask for at least:

• (a) after fit, the y_hat sigma prior is Exponential with lam ≈ 2/s
• (b) auto_scale_sigma=False preserves HalfNormal(1)
• (c) a user-supplied y_hat prior is respected (guard not triggered)
• (d) the multi-treated-unit case

2. Single scalar scale across multiple treated units

y_std = float(self.pre_design["treated"].std(dim="obs_ind", ddof=1).mean())

This collapses to one scalar, then applies it to every treated unit via the dims=["treated_units"] broadcast. If treated units differ in scale by orders of magnitude, this reintroduces exactly the mis-scaling the PR fixes — for the smaller-scale units. A per-unit rate vector (lam shaped to treated_units) would be strictly more correct and costs nothing, since the Prior already carries the treated_units dim.

3. Edge cases on s

2/y_std blows up if y_std == 0 (constant pre-treatment treated series) → infinite rate / degenerate prior, and ddof=1 yields NaN with a single pre-treatment observation. Worth a guard or at least a clear error message.

4. SoftmaxWeightedSumFitter is silently excluded

It's a sibling of WeightedSumFitter (both subclass PyMCModel, not each other), so the isinstance(self.model, WeightedSumFitter) check skips it — yet it has the same sigma-scale problem. This interacts with the suggestion to use the softmax fitter for more robust sampling: a user who follows that advice silently loses auto-scaling. Either extend the check or document the limitation explicitly. Relatedly, passing auto_scale_sigma=True with an OLS or softmax model does nothing, with no warning.

5. Minor: persistent mutation

algorithm() permanently mutates self.model.priors. Low risk since models are typically built per-experiment, and _clone() is protected — but reusing a model instance across experiments would carry the derived prior over. A docstring note, or computing on a clone, would be cleaner.

On the divergences

The instinct that Exponential's mass at zero could cause divergences is reasonable, but the 100-seed sweep is fairly convincing that it's not worse than the status quo here. If you'd rather be conservative, a Gamma/HalfNormal with matched mean and less mass near zero would sidestep the concern while keeping the scale fix — but I don't think that's necessary to block on given the evidence.

Recommendation: approve in principle. I'd ask for tests (1) and the per-unit scale fix (2) before merge, with (3)–(4) addressed or explicitly deferred.

[juanitorduz]

Thank you @szego ! I just triggered our internal Bayesian bot to check the PR :)

In addition, I will look at it this week as well :)

[szego]

I'm working on implementing the cell-specific priors as recommended by the bot but I'm getting totally different data when I run multi_cell_geolift.ipynb compared to what's in the current version. I think that's because that notebook was last executed before #794 which changed the dataset generation behavior.

I reran the notebook on main and it produces similar results to the cell-specific exponential prior, but it'll be worth comparing the results of rerunning that notebook on the current main with the results in this branch when reviewing this. (Neither version actually fits very well.)

I'll leave that notebook as-is for now but I'm happy to commit the newly-run version if you'd prefer.

[szego]

Ok, changed it so that it scales the prior separately to each treated unit, erroring if any have zero or undefined standard deviation. This was a good recommendation from the bot!

Also added support for SoftmaxWeightedSumFitter which was actually really straightforward (kudos on your design there).