Skip to content

Proposal for Caching the Within-Preconditioner#1264

Open
s3alfisc wants to merge 10 commits intomasterfrom
cache-within
Open

Proposal for Caching the Within-Preconditioner#1264
s3alfisc wants to merge 10 commits intomasterfrom
cache-within

Conversation

@s3alfisc
Copy link
Copy Markdown
Member

@s3alfisc s3alfisc commented Mar 28, 2026

Summary

This PR introduces a typed fixed-effects demeaning API, passes it through the code base, and adds reusable within preconditioners to speed up repeated multi-way FE estimation.

The main user-facing additions are:

  • We introduce typed demeaner= objects:
    • MapDemeaner(...)
    • WithinDemeaner(...)
    • LsmrDemeaner(...)
  • demeaner now takes precedence over demeaner_backend, fixef_tol, and fixef_maxiter
  • feols now exposes fit.preconditioner_ for WithinDemeaner
  • users can reuse that object via WithinDemeaner(preconditioner=...)
  • WithinPreconditioner can be pickled across sessions, following the upstream within-py pattern

What Changed

1. Typed demeaner API

We now support a typed demeaner= interface instead of relying only on string backend selection via demeaner_backend. This makes the demeaner backend configuration explicit, extensible, and easier to validate.

The typed demeaner objects control actual runtime execution.

  • WithinDemeaner options flowing into the within backend
  • LsmrDemeaner options flowing into the SciPy/CuPy solvers
  • MapDemeaner options flowing into the MAP solvers

2. Reusable within preconditioners

For multi-way FE WithinDemeaner fits, we now build and cache a reusable WithinPreconditioner.

That preconditioner is:

  • exposed on feols as fit.preconditioner_
  • reusable across later estimations via WithinDemeaner(preconditioner=...)
  • reused internally for multiple estimation syntax
  • reused internally across IWLS iterations in feglm and fepois

3. IWLS reuse

For feglm and fepois, we now reuse a fit-local within preconditioner across IWLS iterations.

Current policy is intentionally simple:

  • build once
  • reuse by default
  • refresh once when inner FE tolerance tightens

@schroedk - For generalized linear models, we need to demean repeatedly with "different weights". Here we currently say it is better to run on "stale" preconditions for a while than to update the preconditioner in every iteration. This can of course be refined later.

4. Persistence

WithinPreconditioner is pickleable across sessions.

This mirrors within-py closely:

  • the Rust FePreconditioner is serialized with postcard
  • Python pickle support is implemented through __reduce__

Key Design Choices

Why expose preconditioners, but not solvers?

We intentionally expose preconditioners as reusable objects, but not solvers.

Because reusing a within preconditioner is where most of the practical speedup comes from. It is "ok" if the preconditioner becomes "a bit stale" - this is a feature, not a bug imo. Plus, we do not need to check that the fixed effects and the sample and weights are identical, which is what we would have needed to do if we allowed users to pass solvers.

For iterated weighted least squares as used for GLMs, the preconditioner is the better choice as in each iteration, a stale preconditioner will be "good enough".

Stay close to within-py

Tried to stay close to the implementation in within-py.

In particular:

  • preconditioners are built through the same conceptual flow as upstream
  • reused preconditioners are passed back into the solver through the same upstream abstractions
  • persistence uses the same postcard + pickle pattern as within-py

Workflow

Typical usage now looks like:

fit1 = pf.feols(
    "Y ~ X1 | f1 + f2",
    data=data,
    demeaner=pf.WithinDemeaner(),
)

# use the cached preconditioner for the next fit
fit2 = pf.feols(
    "Y ~ X2 | f1 + f2",
    data=data,
    demeaner=pf.WithinDemeaner(preconditioner=fit1.preconditioner_),
)

For persistent reuse, you can do:

import pickle

payload = pickle.dumps(fit1.preconditioner_)
restored = pickle.loads(payload)

And then fit the regression model.

@codecov
Copy link
Copy Markdown

codecov bot commented Mar 28, 2026

Codecov Report

❌ Patch coverage is 86.88946% with 51 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
pyfixest/core/demean.py 81.60% 16 Missing ⚠️
pyfixest/demeaners.py 84.37% 15 Missing ⚠️
pyfixest/estimation/internals/demeaner_options.py 73.17% 11 Missing ⚠️
pyfixest/estimation/internals/demean_.py 92.30% 5 Missing ⚠️
pyfixest/core/collinear.py 87.50% 1 Missing ⚠️
pyfixest/estimation/cupy/demean_cupy_.py 80.00% 1 Missing ⚠️
pyfixest/estimation/models/feglm_.py 94.44% 1 Missing ⚠️
pyfixest/estimation/models/fepois_.py 94.44% 1 Missing ⚠️
Flag Coverage Δ
core-tests 71.74% <84.02%> (+0.71%) ⬆️
test-r-core 53.85% <80.71%> (+1.26%) ⬆️
test-r-extended 19.60% <38.81%> (+0.84%) ⬆️
test-r-fixest 41.96% <80.46%> (+2.11%) ⬆️
tests-extended ?

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines Coverage Δ
pyfixest/__init__.py 81.81% <100.00%> (ø)
pyfixest/estimation/FixestMulti_.py 82.19% <100.00%> (+0.18%) ⬆️
pyfixest/estimation/api/feglm.py 91.17% <100.00%> (+1.52%) ⬆️
pyfixest/estimation/api/feols.py 100.00% <100.00%> (ø)
pyfixest/estimation/api/fepois.py 96.66% <100.00%> (+0.66%) ⬆️
pyfixest/estimation/internals/literals.py 87.50% <ø> (ø)
pyfixest/estimation/models/fegaussian_.py 87.50% <100.00%> (+0.40%) ⬆️
pyfixest/estimation/models/feiv_.py 87.27% <100.00%> (ø)
pyfixest/estimation/models/felogit_.py 88.88% <100.00%> (+0.31%) ⬆️
pyfixest/estimation/models/feols_.py 87.45% <100.00%> (-4.76%) ⬇️
... and 12 more

... and 6 files with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

This comment was marked as outdated.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants