Swap out patsy for formulaic

[ksolarski]

Solving issue #386

Starting with DiD, will continue with other methods if you with general design @drbenvincent

Seems like the key practical difference between formulae and patsy is lack of build_design_matrices method in formulae. User has to then provide formula again.

Edit: After discussion, it was decided that formulaic suits us best.

---

📚 Documentation preview 📚: https://causalpy--463.org.readthedocs.build/en/463/

[drbenvincent]

Cool. Thanks @ksolarski, just a quick reply from my phone...

Don't do this for the synthetic control because I have an in progress PR that will change it. It won't have a formula input.

But can I just get some clarification... does this change the API? Can we get the exact same functionality? If not, let's think again.

Will try to look at the code properly when I can 👍🏻

[drbenvincent]

I can't find where I saw it in the patsy docs at this point. But I think one of the things that build_design_matrices did was to ensure that predictions on new/out of sample data are correct. For example, you could get a situation where you don't have all levels of a categorical variable in one predictor for out of sample data. So I think if you to it naively, you can get silent errors.

I'm not 100% sure that this is a problem, and apologies I can't find the relevant part in the docs. But does my concern make sense?

[ksolarski]

You're right, Patsy has the power of preserving the transformation / encoding of variables through build_design_matrices method. There's no equivalent way in formulae so it's certainly not straightforward to copy paste the current behaviour with formulae.

However, Patsy repo suggests migration to https://github.com/matthewwardrop/formulaic instead, which is capable of "reusing the encoding choices made during conversion of one data-set on other datasets." (see https://matthewwardrop.github.io/formulaic/latest/). There's also a migration guide from Patsy to Formulaic to switch would be easy. It also supports many operators: https://matthewwardrop.github.io/formulaic/latest/guides/grammar/

Did you check out this library before? What do you think about using this instead of formulae?

[ksolarski]

@drbenvincent any strong opinions about using formulaic instead of formulae package?

[drbenvincent]

Sorry for the delayed response @ksolarski. So as far as I understand, formulaic is considered the successor to patsy. Though it is only formulae which allows hierarchical modeling?

Right now there are no use-cases for hierarchical modelling. That might change in the future, though I don't have any specific use cases in mind.

So I guess the only choice at the moment is formulaic. But I'll just ask @tomicapretto if there's any plan on how it handles out of sample prediction when not all levels of categorical variables are in the out of sample data (see short discussion above).

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 94.45%. Comparing base (1cf3387) to head (7ad2d5b).

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #463      +/-   ##
==========================================
+ Coverage   94.44%   94.45%   +0.01%     
==========================================
  Files          80       80              
  Lines       12707    12701       -6     
  Branches      770      770              
==========================================
- Hits        12001    11997       -4     
+ Misses        498      497       -1     
+ Partials      208      207       -1     

☔ View full report in Codecov by Sentry.
📢 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.

[ksolarski]

@drbenvincent Yes, from the docs it seems that no hierarchical models are allowed in the formulaic. I think it properly handles the transformation of categorical variables in out of sample data. Here's some small example that I generated:

import pandas as pd
from formulaic import model_matrix
import formulaic

# Create a training dataset
train_data = pd.DataFrame(
    {
        "feature1": ["A", "B", "C", "D"],
        "target": [0, 1, 0, 1],
    }
)

# Create a test dataset
test_data = pd.DataFrame(
    {
        "feature1": [
            "A",  # In training
            "D",  # In training
            "E",  # Not in training
        ],
        "target": [0, 1, 0],
    }
)

# Generate the model matrix for the training data
train_matrix = model_matrix("target ~ 0 + feature1", train_data)

# Print the training matrix and spec
print("Training Matrix:")
print(train_matrix)

# Use the same spec to transform the test data
test_matrix = model_matrix(spec=train_matrix.model_spec, data=test_data)

# Print the test matrix - see that columns are properly aligned from the training data transformation
print("\nTest Matrix:")
print(test_matrix)

Is that the problem you had in mind or something else?

[drbenvincent]

Is that the problem you had in mind or something else?

@ksolarski Yes that's pretty much it. Turns out the phrase I was looking for was "stateful transforms" which you pretty much said here.

I'm actually wondering - if we don't get the additional functionality of hierarchical modeling, is there much benefit from moving from patsy to formulaic? Not saying it's not worth it, but we should be clear about the rationale. Is it because patsy is no longer maintained and formulaic might see more features, or does it have a richer formula API?

Sorry for the extended conversation on this by the way - but given the formula aspect is a core part of the API it's worth thinking it through :)

[tomicapretto]

Hi @ksolarski and @drbenvincent, sorry for the delay in my response. I can add a few pieces of information.

formulae:

• supports stateful transformations
• allows to evaluate new data, see https://bambinos.github.io/formulae/notebooks/getting_started.html#Evaluating-new-data
• supports generation of design matrices for unseen categories, it was added in Allow evaluations on new data with new categories bambinos/formulae#96. Unfortunately, there's no proper documentation for it, but here you have the tests https://github.com/bambinos/formulae/blob/d4ea54f87652a064659df29273b12e872a904b3c/tests/test_eval_new_data.py#L315-L402

With that said, unless you need hierarchical models supported via the | operator, I would consider formulaic as it's more popular and better maintained I think. There's a PR they have to support mixed-effects models, but it has not been merged yet. I'm not sure if there's anything fundamental that prevents that. matthewwardrop/formulaic#34

[ksolarski]

@drbenvincent I went with formulaic as suggested. Quite a straightforward switch it was. I also noticed that the library mentions that there's some speed benefit: https://github.com/matthewwardrop/formulaic?tab=readme-ov-file#benchmarks

[drbenvincent]

Sorry for being slow on this @ksolarski ! After #641 is merged, we should revisit this. That will see increased modularity, so all the action can happen in the new _build_design_matrices method

[drbenvincent]

This comment was generated with LLM assistance and may have been edited by the commenter.

---

Issue Evaluation

Summary

This PR aims to replace patsy with formulaic as the formula parsing library in CausalPy. The original issue (#386) proposed switching to formulae, but after discussion, formulaic was chosen as the better alternative because:

1. It's the official successor to patsy (recommended by the patsy repo itself)
2. It supports stateful transforms for proper out-of-sample prediction handling
3. It offers speed improvements according to benchmarks
4. While it doesn't support hierarchical modeling (like formulae does), CausalPy doesn't currently have that requirement

Current Relevance

This issue is still relevant and valid. The codebase continues to use patsy:

• Files using patsy: 9 experiment files import from patsy
• Key imports: dmatrices, build_design_matrices
• Dependency: Listed in pyproject.toml

Code references checked:

• causalpy/experiments/diff_in_diff.py - uses dmatrices, build_design_matrices
• causalpy/experiments/interrupted_time_series.py - uses dmatrices, build_design_matrices
• All other experiment files in causalpy/experiments/

Related PRs:

• Refactor Experiment Classes: Extract Methods from __init__ #641 (OPEN): "Refactor Experiment Classes: Extract Methods from __init__" - This PR extracts design matrix logic into dedicated _build_design_matrices() methods, which would make the patsy→formulaic migration much cleaner.

Discussion Summary

1. Apr 21, 2025: PR opened targeting formulae library
2. Apr 21-22: @drbenvincent raised concern about stateful transforms for out-of-sample data (ensuring categorical levels are handled correctly)
3. Apr 22: @ksolarski discovered formulaic supports stateful transforms via model_spec and is the patsy-recommended migration path
4. May 15: @drbenvincent confirmed formulaic is the right choice; asked @tomicapretto about formulae capabilities
5. May 26: @tomicapretto confirmed both libraries support the needed functionality, but recommended formulaic as it's better maintained
6. May 28: @ksolarski updated PR to use formulaic
7. Jan 11, 2026 (today): @drbenvincent commented that PR Refactor Experiment Classes: Extract Methods from __init__ #641 should be merged first, as the new _build_design_matrices method will centralize the migration work

Recommendation

• Status: Blocked by Refactor Experiment Classes: Extract Methods from __init__ #641
• Category: The path forward is clear; this is a straightforward dependency ordering issue.

Proposed next steps:

1. Wait for Refactor Experiment Classes: Extract Methods from __init__ #641 to merge - This PR extracts all patsy logic into dedicated _build_design_matrices() methods, making the migration much cleaner
2. Rebase this PR on the updated main once Refactor Experiment Classes: Extract Methods from __init__ #641 is merged
3. Focus the migration work on the new _build_design_matrices() methods in each experiment class
4. Update pyproject.toml to replace patsy with formulaic in dependencies

Migration Scope (for reference)

Once #641 is merged, the migration will involve:

File	Changes Needed
causalpy/experiments/diff_in_diff.py	Update _build_design_matrices()
causalpy/experiments/prepostnegd.py	Update _build_design_matrices()
causalpy/experiments/regression_discontinuity.py	Update _build_design_matrices()
causalpy/experiments/regression_kink.py	Update _build_design_matrices()
causalpy/experiments/interrupted_time_series.py	Update _build_design_matrices()
causalpy/experiments/instrumental_variable.py	Update _build_design_matrices()
causalpy/experiments/inverse_propensity_weighting.py	Update _build_design_matrices()
causalpy/experiments/staggered_did.py	Update _build_design_matrices()
causalpy/pymc_models.py	Update import only
pyproject.toml	Replace patsy with formulaic

The migration pattern (from patsy to formulaic) would be:

# Before (patsy)
from patsy import dmatrices, build_design_matrices
y, X = dmatrices(formula, data)

# After (formulaic)
from formulaic import model_matrix
mm = model_matrix(formula, data)
# Then use mm.model_spec for new data

---

cc @ksolarski - Heads up that this is blocked on #641. Once that merges, this PR should be straightforward to complete!

[ksolarski]

Sounds good, I see it's merged already so I'll continue here

[cursor]

PR Summary

Medium Risk
Touches core data-to-design-matrix plumbing used by many estimators, so subtle column-order/encoding differences could change fitted coefficients or predictions despite mostly mechanical changes.

Overview
Replaces patsy with formulaic for formula parsing and design-matrix generation across multiple experiment classes (DiD, ITS, IV, IPW, PrePostNEGD, RD, Regression Kink, and Staggered DiD), storing model_spec so prediction-time matrices can be rebuilt consistently without build_design_matrices.

Updates the PyMC models layer to use formulaic for spline basis construction and tweaks pm.set_data usage to pass raw X.data plus an explicitly-typed dummy y. Dependency/docs metadata are updated to drop patsy and add formulaic, with minor wording updates in tests/docs.

^{Written by Cursor Bugbot for commit 2f7e57e. This will update automatically on new commits. Configure here.}

[ksolarski]

@drbenvincent now I believe this is ready to be looked at. I removed Patsy from all the places except for notebooks. Let me know if you think I should also regenerate those then.

[read-the-docs-community]

Documentation build overview

📚 causalpy | 🛠️ Build #32278975 | 📁 Comparing 7ad2d5b against latest (1cf3387)

  🔍 Preview build  

13 files changed · ± 13 modified

± Modified

• 404.html
• notebooks/rd_pymc_drinking.html
• _modules/causalpy/pymc_models.html
• api/generated/causalpy.experiments.interrupted_time_series.InterruptedTimeSeries.html
• _modules/causalpy/experiments/diff_in_diff.html
• _modules/causalpy/experiments/instrumental_variable.html
• _modules/causalpy/experiments/interrupted_time_series.html
• _modules/causalpy/experiments/inverse_propensity_weighting.html
• _modules/causalpy/experiments/panel_regression.html
• _modules/causalpy/experiments/prepostnegd.html
• and 3 more...

[drbenvincent]

Thanks @ksolarski - I'm hoping to take a few looks at this because this is a major change :)

At the moment I'll just flag up a few issues:

• inv_prop_latent.ipynb contains patsy code, so that will break once patsy is removed as a dependency
• patsy is also mentioned (just text, no code) in rkink_pymc.ipynb
• in pymc_models.py can you do some kind of manual one-off check that the new way. B = model_matrix("bs(...) - 1", data=ps_df, context={"knots": ...}) results in the same B-spline matrices to patsy, B = dmatrix("bs(...) - 1", {"ps": p, "knots": ...})
• I can also look into this, but can you say something about the added dtype in np.zeros((new_no_of_observations, n_treated_units), dtype=np.int32) in pymc_models.py
• apparently the column naming for categorical variables (e.g. C(group)) might be slighly different. So if we've got any internal logic which looks out for categorical variable names we need to test for that. I'm not 100% sure on this - but would presumably be picked up by existing tests or from manually running the notebooks. (It's not there yet, but Add full notebook re-execution command #672 will hopefully do automated notebook execution.)

[drbenvincent]

PS. I'm fine if you don't re-run or update notebooks in this PR. But given that it's a big change I think I would want to do that and push changes all in this one PR, so that this is done in one decisive PR. Reason is because I'd like to be able to be able to create a new release at any point without having a PR which breaks things until some future hypothetical PR fixes the notebooks :)

PPS. I updated this branch from main by the way, so you'll need to pull changes to your local machine @ksolarski

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

[cursor]

Prediction data converted to buffer object

High Severity

_data_setter() now passes X.data into pm.set_data. When predict() is called with a NumPy array (common in several experiments), X.data is a raw buffer, not the feature matrix. This can corrupt X shape/content during posterior prediction and break out-of-sample predictions.

 

[ksolarski]

in pymc_models.py can you do some kind of manual one-off check that the new way. B = model_matrix("bs(...) - 1", data=ps_df, context={"knots": ...}) results in the same B-spline matrices to patsy, B = dmatrix("bs(...) - 1", {"ps": p, "knots": ...})

They match, see the screenshots.

I can also look into this, but can you say something about the added dtype in np.zeros((new_no_of_observations, n_treated_units), dtype=np.int32) in pymc_models.py

That came from an error in the test and it was a suggested fix. I looked there deeper now and the reason is small discrepancy between dmatrices() from Patsy and dm.lhs.to_numpy() from formulaic. Patsy makes the matrix always a float while formulaic can keep it as an integer. Then down the line, np.zeros() is a float and we try to use it inside pm.set_data which expects that y is an integer because the initial matrix had integer dtype. Let me know if you have some strong opinion how this issue should be handled.

Regarding notebooks, I'll rerun them once you're happy with all other things, should be easy

[drbenvincent]

Apologies about letting this go stale @ksolarski. I'm going to get this to green and then I'll recap any remaining issues to be resolved.

[drbenvincent]

PR #463 status update — sync with main + bug fixes

@ksolarski I've merged main into the branch and fixed a few issues. Here's a summary.

What was done

1. Merged main into branch — resolved 2 conflicts:

   • regression_discontinuity.py: kept main's improved warning logic with formulaic's model_matrix
   • environment.yml: kept formulaic + main's new dependencies (prek, pymc-bart, pymc-extras)

2. Fixed _data_setter dtype mismatch (pymc_models.py) — reverted the X.data change (X is always xr.DataArray here, so this was unnecessary and could be fragile). For the dtype issue you identified, instead of hardcoding dtype=np.int32, the fix now dynamically queries the existing y data node's dtype via self["y"].type.dtype. This correctly handles both integer and float outcomes without assumptions.

3. Fixed IPW outcome_variable_name collision (inverse_propensity_weighting.py) — for IPW the formula LHS is the treatment variable (trt), not the outcome. Storing it as outcome_variable_name broke the __maketables_depvar__ hook (which uses outcome_variable_name with higher precedence than outcome_variable). Renamed to treatment_variable_name.

What passes now

• prek run --all-files: all hooks pass
• pytest causalpy/tests/ (excluding test_piecewise_its.py): 729 passed, 3 skipped, 0 failures
• pre-commit.ci: passed
• ReadTheDocs: build in progress

Review items from @drbenvincent (Jan 26) — status

#	Item	Status
1	inv_prop_latent.ipynb contains patsy code	Open — needs updating before merge
2	rkink_pymc.ipynb mentions patsy in text	Open — minor text fix needed
3	Verify B-spline equivalence (pymc_models.py)	Resolved — @ksolarski confirmed matrices match (Feb 14 screenshots)
4	Explain / fix dtype=np.int32 in _data_setter	Resolved — now uses dynamic self["y"].type.dtype instead of hardcoding
5	Check categorical column naming (C(group))	Resolved — I verified formulaic and patsy produce identical column names (e.g. C(group)[T.1], C(group)[T.treatment]) for both integer and string categoricals
6	Notebooks should be re-run in this PR	Open — @ksolarski offered to do this once code changes are settled

Remaining blocker: piecewise_its.py + transforms.py

Three files added to main after this PR was opened still use patsy:

File	patsy usage
causalpy/transforms.py	import patsy + patsy.stateful_transform(StepTransform/RampTransform)
causalpy/experiments/piecewise_its.py	from patsy import dmatrices
causalpy/tests/test_piecewise_its.py	from patsy import dmatrix, PatsyError

Since the PR removes patsy from pyproject.toml, these files will break at import time if left unmigrated.

The good news: the step and ramp transforms are only truly stateful for datetime inputs (they remember the datetime origin). For numeric time — the common case — they're pure functions. formulaic can handle them as plain callables via its context parameter:

dm = model_matrix(formula, data, context={"step": step, "ramp": ramp})

The patsy.stateful_transform() wrapper can be replaced with simple functions that instantiate StepTransform/RampTransform and call transform() directly.

Suggested approach:

1. In transforms.py: replace patsy.stateful_transform(StepTransform) with a plain callable wrapper that instantiates the class and calls its transform() method
2. In piecewise_its.py: switch dmatrices to model_matrix and pass step/ramp via context
3. In test_piecewise_its.py: replace dmatrix and PatsyError imports with formulaic equivalents

Let me know if you'd prefer me to tackle this, or if you'd rather do it yourself!

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[drbenvincent]

Synced branch with main and fixed failing checks:

• Merge conflict in inverse_propensity_weighting.py: resolved by keeping the formulaic implementation and incorporating the improved docstring from main.
• Notebook CI failure (rd_pymc_drinking.ipynb): formulaic treats boolean columns as numeric, so the coefficient name is treated rather than patsy's treated[T.True]. Updated all sel(coeffs=...) and set(title=...) calls in the notebook accordingly.
• mypy failure in panel_regression.py (from main): corrected type: ignore error codes (attr-defined -> union-attr).

All local checks pass (prek run --all-files). The GitHub Actions CI workflows need maintainer approval to run (fork PR).