Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 13 additions & 1 deletion .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -10,4 +10,16 @@ target
.hypothesis/
.pytest_cache/
.ruff_cache/
*.proptest-regressions
__pycache__/
*.py[cod]
*.proptest-regressions

# R build/check artifacts
*.Rcheck/
*.tar.gz
.Rhistory
.RData
.Rproj.user/

# Serena MCP
.serena/
3 changes: 3 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -11,10 +11,13 @@ and this project follows [Semantic Versioning](https://semver.org/).

- **Locality sort:** `Design` construction reorders observations by the highest-cardinality factor when unsorted, copying them once into an internal sorted store (the caller's store is never mutated; the `Store` trait stays read-only). Transparent — results return in caller row order.
- Coalesced scatter for large sorted factors: one atomic add per equal-level run per chunk instead of one per row.
- `withinr` preconditioner handles now expose `variant` and `build_time_seconds` metadata for cache compatibility checks and diagnostics.
- `withinr` manual tests now cover diagonal/additive/off correctness parity, prebuilt preconditioner reuse, and sorted-vs-unsorted caller-order equivalence.

### Changed

- Category views are borrowed only when the dominant factor is already sorted; otherwise the columns are copied once for the locality sort. The reorder changes summation order, so unsorted-input results match 0.2.0 within solver tolerance, not bitwise.
- `withinr` development builds patch `within` to the workspace crate, while offline builds patch it to the vendored local crate sources.

## [0.2.0] - 2026-06-04

Expand Down
39 changes: 39 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ members = [
"crates/schwarz-precond",
"crates/within",
"crates/within-py",
"crates/within-r",
]
resolver = "2"

Expand Down
92 changes: 57 additions & 35 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -43,69 +43,89 @@ result = solve(fe, y, weights=np.ones(n))
result = solve(fe, y, preconditioner=PreconditionerConfig.Diagonal)
```

### FWL regression example
## R quickstart

```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]
Requires R and a Rust toolchain (`cargo` on `PATH`).

From the repository root, use `devtools` to install R dependencies and build the
package:

```r
install.packages("devtools")
Sys.setenv(NOT_CRAN = "true")
devtools::install_deps("withinr/", dependencies = TRUE)
devtools::load_all("withinr/")
```

## Python API
Example (FWL with two-way fixed effects):

### High-level functions
```r
set.seed(42)
n <- 1000
n_firms <- 50L
n_years <- 20L

| Function | Description |
|---|---|
| `solve(categories, y, options?, weights?, preconditioner?)` | Solve a single right-hand side. Returns `SolveResult`. |
| `solve_batch(categories, Y, options?, weights?, preconditioner?)` | Solve multiple RHS vectors in parallel. `Y` has shape `(n_obs, k)`. Returns `BatchSolveResult`. |
# 1-based fixed-effect ids in R
firm <- rep(seq_len(n_firms), each = n_years)
year <- rep(seq_len(n_years), times = n_firms)
categories <- cbind(firm, year)

`categories` is a 2-D `uint32` array of shape `(n_obs, n_factors)`. A `UserWarning` is emitted when a C-contiguous array is passed — use `np.asfortranarray(categories)` for best performance.
beta <- 1.5
firm_fe <- rnorm(n_firms, sd = 3)[firm]
year_fe <- rnorm(n_years, sd = 1)[year]
x <- rnorm(n) + 0.3 * firm_fe
y <- beta * x + firm_fe + year_fe + rnorm(n, sd = 0.5)

### Persistent solver
res <- withinr::solve_batch(categories, cbind(y, x))
y_tilde <- res$demeaned[, 1]
x_tilde <- res$demeaned[, 2]
beta_hat <- sum(x_tilde * y_tilde) / sum(x_tilde^2)

For repeated solves with the same design matrix, `Solver` builds the preconditioner once and reuses it.
print(beta_hat)
print(res$converged)
```

```python
from within import Solver
| Function | Description |
|---|---|
| `solve(categories, y, options?, weights?, preconditioner?)` | Solve a single right-hand side. Returns a list shaped like `SolveResult`. |
| `solve_batch(categories, Y, options?, weights?, preconditioner?)` | Solve multiple RHS vectors in parallel. `Y` has shape `(n_obs, k)`. |

For repeated solves with the same design matrix, `Solver` builds the preconditioner once and reuses it. In R, the solver is an environment with methods.

solver = Solver(fe)
r = solver.solve(y) # reuses preconditioner
r = solver.solve_batch(np.column_stack([y, X]))
```r
solver <- withinr::Solver(categories)
r <- solver$solve(y)
r <- solver$solve_batch(cbind(y, x))

precond = solver.preconditioner # picklable property
solver2 = Solver(fe, preconditioner=precond) # skip re-factorization
precond <- solver$preconditioner()
payload <- precond$serialize()
solver2 <- withinr::Solver(categories, preconditioner = withinr::Preconditioner(payload))
```

| Property / Method | Description |
|---|---|
| `Solver(categories, 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(fe, preconditioner=p)`. |
| `$solve(y, options?)` | Solve a single RHS with the given LSMR tuning. |
| `$solve_batch(Y, options?)` | Solve multiple RHS columns in parallel. |
| `$preconditioner()` | Return the built `Preconditioner`, or `NULL`. Reuse via `Solver(categories, preconditioner=p)`. |


### Solver configuration

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

### Preconditioner (5-form Union)

The `preconditioner` argument accepts any of:

| Form | Meaning |
|---|---|
| `None` (default) | Library default — Additive Schwarz with sensible defaults. |
| `PreconditionerConfig.Off` | Explicit identity — solve unpreconditioned. |
| `PreconditionerConfig.Additive` | Additive Schwarz shortcut, equivalent to `None`. |
| `PreconditionerConfig.Diagonal` | Diagonal/Jacobi preconditioner using `diag(D^T W D)^{-1}`. |
| `None` / `NULL` (default) | Library default — Additive Schwarz with sensible defaults. |
| `PreconditionerConfig.Off` / `PreconditionerConfig$Off` | Explicit identity — solve unpreconditioned. |
| `PreconditionerConfig.Additive` / `PreconditionerConfig$Additive` | Additive Schwarz shortcut, equivalent to the default. |
| `PreconditionerConfig.Diagonal` / `PreconditionerConfig$Diagonal` | Diagonal/Jacobi preconditioner using `diag(D^T W D)^{-1}`. |
| `AdditiveSchwarz(local_solver?, reduction?)` | Tuned Schwarz config — import from `within.config`. |
| `Preconditioner` instance | Reuse a previously-built preconditioner across solvers. |

Expand All @@ -114,7 +134,7 @@ The `preconditioner` argument accepts any of:
| Class | Description |
|---|---|
| `LocalSolverConfig(approx_chol?, approx_schur?, dense_threshold=24)` | Schur reduction + approximate Cholesky. Omit `approx_schur` for the library-default approximate variant; pass `approx_schur=None` to request an exact Schur (slower, used for validation). |
| `ApproxCholConfig(seed=0, split=1)` | Approximate Cholesky parameters. |
| `ApproxCholConfig(seed=0, split_merge=None)` | Approximate Cholesky parameters. |
| `ApproxSchurConfig(seed=0, split=1)` | Approximate Schur complement sampling parameters. |
| `ReductionStrategy` enum | `Auto` (default), `AtomicScatter`, `ParallelReduction`. |

Expand Down Expand Up @@ -194,7 +214,9 @@ crates/
schwarz-precond/ Generic domain decomposition library (traits, solvers, Schwarz preconditioners)
within/ Core fixed-effects solver (observation stores, domains, operators, orchestration)
within-py/ PyO3 bridge (cdylib → within._within)
within-r/ Workspace mirror for the extendr bridge
python/within/ Python package re-exporting the Rust extension
withinr/ R package using the published within 0.2.0 crate
benchmarks/ Python benchmark framework
```

Expand Down
13 changes: 13 additions & 0 deletions crates/within-r/Cargo.toml
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
[package]
name = "within-r"
version = "0.1.0"
edition = "2021"

[lib]
crate-type = ["staticlib"]

[dependencies]
within = { path = "../within" }
extendr-api = "0.8"
ndarray = "0.16"
postcard = { version = "1.1", features = ["use-std"] }
1 change: 1 addition & 0 deletions crates/within-r/src/lib.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
include!("../../../withinr/src/rust/src/lib.rs");
11 changes: 11 additions & 0 deletions withinr/.Rbuildignore
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
^.*\.Rproj$
^\.Rproj\.user$
^src/rust/target
^src/rust/vendor$
^src/vendor$
^src/\.cargo
^src/.*\.o$
^src/.*\.so$
^src/.*\.dll$
^test_bindings\.R$
^benchmarks$
7 changes: 7 additions & 0 deletions withinr/.gitignore
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
src/.cargo/
src/rust/target/
src/vendor/
src/*.o
src/*.so
src/*.dll
src/withinr.lib
19 changes: 19 additions & 0 deletions withinr/DESCRIPTION
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
Package: withinr
Title: High-Performance Fixed Effects Solver
Version: 0.1.0
Authors@R: c(
person("Alexander", "Fischer", role = c("aut", "cre"),
email = "alexander-fischer1801@t-online.de"),
person("Kristof", "Schroeder", role = "aut"))
Author: Alexander Fischer [aut, cre], Kristof Schroeder [aut]
Maintainer: Alexander Fischer <alexander-fischer1801@t-online.de>
Description: Fast modified LSMR solvers with Schwarz and diagonal
preconditioners for absorbing high-dimensional fixed effects in panel
data regressions. The computational core is written in Rust via version
0.2.0 of the 'within' crate and accessed through 'extendr'.
License: MIT + file LICENSE
Encoding: UTF-8
Roxygen: list(markdown = TRUE)
RoxygenNote: 7.3.3
SystemRequirements: Rust tool chain w/ cargo, rustc
Config/rextendr/version: 0.5.0
2 changes: 2 additions & 0 deletions withinr/LICENSE
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
YEAR: 2025
COPYRIGHT HOLDER: Alexander Fischer, Kristof Schröder
17 changes: 17 additions & 0 deletions withinr/NAMESPACE
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
# Generated by roxygen2: do not edit by hand

S3method(print,within_preconditioner)
S3method(print,within_solver)
export(AdditiveSchwarz)
export(ApproxCholConfig)
export(ApproxSchurConfig)
export(LocalSolverConfig)
export(LsmrOptions)
export(Preconditioner)
export(PreconditionerConfig)
export(ReductionStrategy)
export(Solver)
export(lsmr_options)
export(solve)
export(solve_batch)
useDynLib(withinr, .registration = TRUE)
83 changes: 83 additions & 0 deletions withinr/R/extendr-wrappers.R
Original file line number Diff line number Diff line change
@@ -0,0 +1,83 @@
# Generated by extendr: Do not edit by hand
#
# This file was created with the following call:
# .Call("wrap__make_withinr_wrappers", use_symbols = TRUE, package_name = "withinr")

#' @usage NULL
#' @useDynLib withinr, .registration = TRUE
NULL

# Internal .Call wrapper.
solve_impl <- function(categories, y, weights, tol, maxiter, local_size, preconditioner) {
.Call("wrap__solve_impl", categories, y, weights, tol, maxiter, local_size, preconditioner)
}

# Internal .Call wrapper.
solve_batch_impl <- function(categories, y_matrix, weights, tol, maxiter, local_size, preconditioner) {
.Call("wrap__solve_batch_impl", categories, y_matrix, weights, tol, maxiter, local_size, preconditioner)
}

# Internal .Call wrapper.
solver_new_impl <- function(categories, weights, preconditioner) {
.Call("wrap__solver_new_impl", categories, weights, preconditioner)
}

# Internal .Call wrapper.
solver_solve_impl <- function(solver, y, tol, maxiter, local_size) {
.Call("wrap__solver_solve_impl", solver, y, tol, maxiter, local_size)
}

# Internal .Call wrapper.
solver_solve_batch_impl <- function(solver, y_matrix, tol, maxiter, local_size) {
.Call("wrap__solver_solve_batch_impl", solver, y_matrix, tol, maxiter, local_size)
}

# Internal .Call wrapper.
solver_preconditioner_impl <- function(solver) {
.Call("wrap__solver_preconditioner_impl", solver)
}

# Internal .Call wrapper.
solver_n_dofs_impl <- function(solver) {
.Call("wrap__solver_n_dofs_impl", solver)
}

# Internal .Call wrapper.
solver_n_obs_impl <- function(solver) {
.Call("wrap__solver_n_obs_impl", solver)
}

# Internal .Call wrapper.
preconditioner_apply_impl <- function(preconditioner, x) {
.Call("wrap__preconditioner_apply_impl", preconditioner, x)
}

# Internal .Call wrapper.
preconditioner_nrows_impl <- function(preconditioner) {
.Call("wrap__preconditioner_nrows_impl", preconditioner)
}

# Internal .Call wrapper.
preconditioner_ncols_impl <- function(preconditioner) {
.Call("wrap__preconditioner_ncols_impl", preconditioner)
}

# Internal .Call wrapper.
preconditioner_variant_impl <- function(preconditioner) {
.Call("wrap__preconditioner_variant_impl", preconditioner)
}

# Internal .Call wrapper.
preconditioner_build_time_seconds_impl <- function(preconditioner) {
.Call("wrap__preconditioner_build_time_seconds_impl", preconditioner)
}

# Internal .Call wrapper.
preconditioner_serialize_impl <- function(preconditioner) {
.Call("wrap__preconditioner_serialize_impl", preconditioner)
}

# Internal .Call wrapper.
preconditioner_deserialize_impl <- function(data) {
.Call("wrap__preconditioner_deserialize_impl", data)
}
Loading
Loading