Metadata-Version: 2.4
Name: within-py
Version: 0.4.0
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Programming Language :: Python :: Free Threading :: 2 - Beta
Classifier: Programming Language :: Python :: Implementation :: CPython
Classifier: Programming Language :: Rust
Classifier: Topic :: Scientific/Engineering :: Mathematics
Requires-Dist: numpy>=1.24
License-File: LICENSE
Summary: High-performance fixed-effects solver for econometric panel data
Keywords: fixed-effects,econometrics,panel-data,least-squares,solver
Author-email: Alexander Fischer <alexander-fischer1801@t-online.de>, Kristof Schröder <k.schroeder@appliedai-institute.de>
License-Expression: MIT
Requires-Python: >=3.9
Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM
Project-URL: Changelog, https://github.com/py-econometrics/within/blob/main/CHANGELOG.md
Project-URL: Homepage, https://github.com/py-econometrics/within
Project-URL: Repository, https://github.com/py-econometrics/within

# within

`within` provides high-performance solvers for projecting out high-dimensional fixed effects from regression problems.

By the Frisch-Waugh-Lovell theorem, estimating a regression of the form *y = Xβ + Dα + ε* reduces to a sequence of least-squares projections, one for y and one for each column of X, followed by a cheap regression fit on the resulting residuals. The projection step of solving the normal equations *D'Dx = D'z* is the computational bottleneck, which is the problem `within` is designed to solve.

`within`'s solvers are tailored to the structure of fixed effects problems, which can be represented as a graph (as first noted by Correia, 2016). Concretely, `within` uses modified LSMR with a domain decomposition (Schwarz) preconditioner, backed by approximate Cholesky local solvers (Gao et al, 2025).

## Scope

`within` is a low-level fixed-effects kernel. Callers pass categorical labels as `uint32` F-order (column-major) arrays. Labels need not be contiguous: observed labels are compacted internally, while result layouts retain the caller-visible labels. Formula-level convenience — DataFrames, string/object categoricals, `pandas.factorize`, and formula parsing — is intentionally out of scope and belongs to a frontend layer built on top. The pyfixest-style workflow is served by such a frontend calling `within` underneath.

## Installation

You can install Python bindings from PyPi by running 

```bash
pip install within_py
```

Wheels ship for CPython 3.9+ (one `abi3` wheel per platform), free-threaded CPython 3.14+ (`cp314t`; PyO3 dropped experimental 3.13t support), and PyPy 3.11 (PyO3's minimum supported PyPy version; earlier PyPy releases are rejected at build time). CPython 3.15's `abi3t` (PEP 803, [PyO3/maturin#3064](https://github.com/PyO3/maturin/issues/3064)) will let the free-threaded build fold back into a single wheel per platform.

musllinux ships no PyPy wheel at all — the `musllinux_1_2` build image bundles no PyPy interpreter.

## Python Quickstart

`within`'s main user-facing function is `solve`. Provide a 2-D `uint32` array of category codes (one column per fixed-effect factor) and a response vector `y`. The solver finds x in the normal equations **D'D x = D'y**, where D is the sparse categorical design matrix.

```python
import numpy as np
from within import solve, solve_batch, LsmrOptions, PreconditionerConfig

np.random.seed(1)
n = 100_000
fe = np.asfortranarray(np.column_stack([
    np.random.randint(0, 500, n).astype(np.uint32),
    np.random.randint(0, 200, n).astype(np.uint32),
]))
y = np.random.randn(n)

# Default: the adaptive diagonal→Schwarz ladder + LSMR
result = solve(fe, y)

# Custom tolerance / iteration cap
result = solve(fe, y, options=LsmrOptions(tol=1e-10, maxiter=2000))

# Weighted solve
result = solve(fe, y, weights=np.ones(n))

# Opt into diagonal/Jacobi preconditioning
result = solve(fe, y, preconditioner=PreconditionerConfig.Diagonal())
```

### FWL regression example

```python
beta_true = np.array([1.0, -2.0, 0.5])
X = np.random.randn(n, 3)
y = X @ beta_true + np.random.randn(n)

result = solve_batch(fe, np.column_stack([y, X]))
y_tilde, X_tilde = result.demeaned[:, 0], result.demeaned[:, 1:]
beta_hat = np.linalg.lstsq(X_tilde, y_tilde, rcond=None)[0]
print(np.round(beta_hat, 4))  # [ 0.9982 -2.006   0.5005]
```

### Varying slopes

Pass a list of `Effect` terms instead of a categories array. Each term is a
factor's level codes plus an optional intercept and zero or more slope
covariates (per-level slopes, as in fixest's `f[z]` notation).

```python
from within import solve, Effect

firm = np.random.randint(0, 500, n).astype(np.uint32)
year = np.random.randint(0, 20, n).astype(np.uint32)
x = np.random.randn(n)  # covariate whose slope varies by firm

result = solve(
    [
        Effect(firm, intercept=True, slopes=[x]),  # firm intercept + firm-specific x slope
        Effect(year, intercept=True),              # year intercept
    ],
    y,
)

# Read firm level 3's x-slope via the layout map (column 0 = intercept, 1 = first slope):
i = result.layout.index(0, 3, 1)
print(result.x[i])
```

## Python API

### High-level functions

| Function | Description |
|---|---|
| `solve(design, y, weights?, options?, preconditioner?)` | Solve a single right-hand side. Returns `SolveResult`. |
| `solve_batch(design, Y, weights?, options?, preconditioner?)` | Solve multiple RHS vectors in parallel. `Y` has shape `(n_obs, k)`. Returns `BatchSolveResult`. |

`design` is a persistent `Design`, a 2-D `uint32` array of shape `(n_obs, n_factors)`, or a list of `Effect` terms (see [Varying slopes](#varying-slopes)). A `UserWarning` is emitted when a C-contiguous categories array is passed — use `np.asfortranarray(design)` for best performance.

### Persistent solver

For repeated solves with the same design matrix, `Design` reuses the preprocessed observation storage and `Solver` builds the weight-dependent state and preconditioner once.

```python
from within import Design, Solver, solve

design = Design(fe)
solver = Solver(design)
r = solver.solve(y)                            # reuses preconditioner
r = solver.solve_batch(np.column_stack([y, X]))

precond = solver.preconditioner                # picklable property
solver2 = Solver(design, preconditioner=precond)  # skip re-factorization

# Reuse design preprocessing while rebuilding weight-dependent solver state.
weighted = solve(design, y, weights=np.ones(n))
```

| Property / Method | Description |
|---|---|
| `Design(design)` | Build an owned native design for reuse across solves and solvers. |
| `Solver(design, weights?, preconditioner?)` | Build solver. Factorizes the preconditioner at construction. |
| `.solve(y, options?)` | Solve a single RHS with the given LSMR tuning. Returns `SolveResult`. |
| `.solve_batch(Y, options?)` | Solve multiple RHS columns in parallel. Returns `BatchSolveResult`. |
| `.preconditioner` | Return the built `Preconditioner` (picklable), or `None`. Reuse via `Solver(design, preconditioner=p)`. |


### Solver configuration

| Class | Description |
|---|---|
| `LsmrOptions(tol=1e-8, maxiter=1000, local_size=None)` | Modified LSMR. `local_size` enables windowed reorthogonalization. |

### Preconditioners

The `preconditioner` argument accepts any of:

| Form | Meaning |
|---|---|
| `None` (default) | Library default — Adaptive: diagonal, escalating to additive Schwarz on a stalled contraction. |
| `PreconditionerConfig.Off()` | Explicit identity — solve unpreconditioned. |
| `PreconditionerConfig.Additive()` | Additive Schwarz shortcut, built up front. |
| `PreconditionerConfig.Diagonal()` | Diagonal/Jacobi preconditioner using `diag(D^T W D)^{-1}`. |
| `PreconditionerConfig.Additive(local_solver?, reduction?)` | Tuned additive Schwarz. Argument types import from `within.config`. |
| `PreconditionerConfig.Adaptive(local_solver?, reduction?, stall?)` | Tuned form of the default; the factorization is built only on escalation. |
| `Preconditioner` instance | Reuse a previously-built preconditioner across solvers. |

`PreconditionerConfig` is a tagged union: each variant is a subclass, so instances
support `match`/`case` and compare by value. A built preconditioner exposes
the configuration used to construct it as `.config` and its original build duration in seconds
as `.build_duration_seconds`; both are preserved across pickling and reuse.

### Local solver configuration (advanced — `within.config`)

| Class | Description |
|---|---|
| `LocalSolverConfig(approx_chol?, schur?, dense_threshold=24, scaling?)` | Schur reduction + approximate Cholesky. Omit `schur` for the library-default approximate variant; pass `schur=Schur.exact()` to request an exact Schur (slower, used for validation). |
| `Schur.approximate(config?)` / `Schur.exact()` | Schur-reduction mode passed as `LocalSolverConfig(schur=...)`. |
| `ApproxCholConfig(seed=0, split_merge=None)` | Approximate Cholesky parameters. |
| `ApproxSchurConfig(seed=0, split=1)` | Approximate Schur complement sampling parameters. |
| `ReductionStrategy` | `Auto` (default), `AtomicScatter`, `ParallelReduction` (class attributes, not an `Enum`). |

### Result types

**`SolveResult`**: `x` (coefficients), `unidentified` (directions the data cannot identify, as `UnidentifiedDirection(term, level, column)` records), `layout` (a `CoefficientLayout` mapping a `(term, level, column)` address to its flat `x` index and back), `demeaned` (residuals), `converged`, `iterations`, `residual`, `time_total`, `time_setup`, `time_solve`.

**`BatchSolveResult`**: Same fields, with `converged`, `iterations`, `residual`, and `time_solve` as lists (one entry per RHS).

Coefficients for unidentified directions are pinned to the **minimal-norm** value `0` (never NaN). This is why `x` can differ from reference tools that instead drop a reference level; the identified fit — `demeaned` — is unaffected by the choice.

## Rust API

```rust
use ndarray::Array2;
use within::{solve, LsmrOptions, PreconditionerConfig};
use within::config::{LocalSolverConfig, ReductionStrategy};

let categories = /* Array2<u32> of shape (n_obs, n_factors) */;
let y: &[f64] = /* response vector */;

// Default: LSMR + the adaptive ladder (None → library default)
let r = solve(categories.view(), &y, None, &LsmrOptions::default(), None)?;
assert!(r.converged);

// Tighter tolerance with an explicit additive preconditioner
let lsmr = LsmrOptions { tol: 1e-10, ..LsmrOptions::default() };
let precond = PreconditionerConfig::Additive {
    local_solver: LocalSolverConfig::default(),
    reduction: ReductionStrategy::default(),
};
let r = solve(categories.view(), &y, None, &lsmr, &precond)?;

// Opt into diagonal/Jacobi preconditioning
let diagonal = PreconditionerConfig::Diagonal;
let r = solve(categories.view(), &y, None, &lsmr, &diagonal)?;
```

Persistent solver — build once, solve many:

```rust
use ndarray::array;
use within::{solve, Channel, CoefficientAddress, Design, Solver};

// Caller labels are deliberately non-contiguous.
let categories = array![
    [10_u32, 100],
    [20, 100],
    [10, 900],
    [20, 900],
];
let y = [1.0, 2.0, 3.0, 4.0];

// Build and compact the design once, then share it across solvers.
let design = Design::from_categories(categories.view())?;
let unweighted = Solver::new(&design, None, None)?;
let weights = [1.0, 2.0, 1.0, 2.0];
let weighted = Solver::new(&design, Some(&weights), None)?;

let result = unweighted.solve(&y, None)?;
let weighted_result = weighted.solve(&y, None)?;
let one_shot_result = solve(&design, &y, None, None, None)?;
assert!(weighted_result.converged && one_shot_result.converged);

// Internal factor positions stay private; result lookup uses caller labels.
let address = CoefficientAddress {
    channel: Channel { term: 1, column: 0 },
    level: 900,
};
let slot = result.layout.index(address).expect("observed label");
assert_eq!(result.layout.address(slot), Some(address));
assert_eq!(result.x.len(), 4); // two observed levels in each of two factors
```

`solve` and `Solver::new` take the preconditioner as `impl Into<PreconditionerInput>`:
`None` (library default), a `&PreconditionerConfig` or owned `PreconditionerConfig`
(e.g. `PreconditionerConfig::Off` for the identity), or an owned/borrowed
`Preconditioner` for reuse. LSMR options are `impl Into<Option<&LsmrOptions>>`, so
`None` accepts the defaults and `&opts` overrides them.

| Type | Variants / Fields |
|---|---|
| `Design` | Reusable immutable design; build with `Design::from_categories` or `Design::new`, then share via `Solver::new(&design, ...)` or `solve(&design, ...)`. |
| `LsmrOptions` | `{ tol: f64, maxiter: usize, local_size: Option<usize> }` |
| `PreconditionerConfig` | `Off` \| `Additive { local_solver: LocalSolverConfig, reduction: ReductionStrategy }` \| `Diagonal` (`#[non_exhaustive]`) |
| `LocalSolverConfig` | `{ approx_chol, schur: SchurMode, dense_threshold, scaling }` |
| `SchurMode` | `Approximate(ApproxSchurConfig)` \| `Exact` |
| `Preconditioner` | Opaque built handle — reuse via `Solver::new(.., precond)` (owned or `&`) |
| `Effect` | `Effect::new(levels: &[u32], intercept: bool, slopes: impl IntoIterator<Item = &[f64]>) -> Result<Self, BuildError>` |
| `CoefficientAddress` | `{ channel: Channel { term, column }, level: u32 }` |
| `CoefficientLayout` | `index(CoefficientAddress) -> Option<usize>`, `address(usize) -> Option<CoefficientAddress>`, `n_dofs()`, `n_terms()`, `n_levels(term)`, `n_columns(term)` |

### Varying slopes

Pass a `Vec<Effect>` instead of a categories array. Each term is a factor's
level codes plus an optional intercept and zero or more slope covariates
(per-level slopes, as in fixest's `f[z]` notation).

```rust
use within::{solve, Channel, CoefficientAddress, Effect};

let firm: &[u32] = /* level codes */;
let year: &[u32] = /* level codes */;
let x: &[f64] = /* covariate whose slope varies by firm */;

let terms = vec![
    Effect::new(firm, true, [x])?,  // firm intercept + firm-specific x slope
    Effect::new(year, true, [])?,   // year intercept
];
let r = solve(terms, &y, None, None, None)?;

// Read firm level 3's x-slope via the layout map (column 0 = intercept, 1 = first slope):
let at = CoefficientAddress { channel: Channel { term: 0, column: 1 }, level: 3 };
println!("{}", r.x[r.layout.index(at).unwrap()]);
```

`Solver::new` takes the same `Vec<Effect>`, so a slope design can be reused
across solves like any other.

### Lower-level access

| Module | Visibility | Key types |
|---|---|---|
| `within::config` | public | `LsmrOptions`, `PreconditionerConfig`, `LocalSolverConfig`, `SchurMode`, `ApproxCholConfig`, `ApproxSchurConfig`, `ScalingConfig`, `ReductionStrategy` |
| `within::error` | public | `WithinError`, `BuildError`, `SolveError` |
| `block_elim` / `channel` / `csr_block` / `domain` / `linalg` / `operator` / `solver` | `pub(crate)` | implementation layers — public items are re-exported at the crate root |

## Project structure

```
crates/
  schwarz-precond/   Generic domain decomposition library (traits, solvers, Schwarz preconditioners)
  within/            Core fixed-effects solver (designs, domains, operators, solver)
  within-py/         PyO3 bridge (cdylib → within._within)
python/within/       Python package re-exporting the Rust extension
benchmarks/          Python benchmark framework
```

## Development

Uses [pixi](https://pixi.sh) as the task runner.

```bash
pixi run develop                         # Build Rust extension (release mode)
pixi run test                            # Rebuild + pytest
cargo test --workspace                   # Rust tests only
cargo bench -p within                    # Criterion benchmarks
pixi run python -m benchmarks run all    # Python benchmarks
```

Rust changes require rebuilding before running Python code (`pixi run develop`).

## License

MIT

## References

- Correia, Sergio. "A feasible estimator for linear models with multi-way fixed effects." *Preprint* at http://scorreia.com/research/hdfe.pdf (2016).
- Gao, Y., Kyng, R. & Spielman, D. A. (2025). AC(k): Robust Solution of Laplacian Equations by Randomized Approximate Cholesky Factorization. *SIAM Journal on Scientific Computing*.
- Toselli & Widlund (2005). *Domain Decomposition Methods — Algorithms and Theory*. Springer.
- Xu, J. (1992). Iterative Methods by Space Decomposition and Subspace Correction. *SIAM Review*, 34(4), 581--613.

