Metadata-Version: 2.4
Name: svy-sae
Version: 0.4.0
Summary: Scalable Small Area Estimation with JAX
Author-email: Samplics LLC <msdiallo@samplics.org>
Keywords: sampling,sample,weighting,estimation,survey
Classifier: Programming Language :: Python :: Implementation :: CPython
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Development Status :: 4 - Beta
Classifier: Operating System :: OS Independent
Classifier: Topic :: Scientific/Engineering
Requires-Python: <3.15,>=3.12
Description-Content-Type: text/markdown
Requires-Dist: jax>=0.11.0
Requires-Dist: jaxlib>=0.11.0
Requires-Dist: scipy>=1.18.0
Requires-Dist: svy[report]<0.27.0,>=0.26.0
Requires-Dist: tqdm>=4.70.0

# svy-sae

**Small area estimation at census scale, in Python.**

[![PyPI](https://img.shields.io/pypi/v/svy-sae?color=blue)](https://pypi.org/project/svy-sae/)
[![Python](https://img.shields.io/pypi/pyversions/svy-sae)](https://pypi.org/project/svy-sae/)

svy-sae produces estimates for domains too small for direct survey estimation — districts, municipalities, enumeration areas — by borrowing strength from a model linking a survey to auxiliary population data. It implements the standard area-level and unit-level estimators, validated against R's `sae` package, on compiled JAX kernels that scale to a multi-million-household census.

🌐 [svylab.com](https://svylab.com) · 📦 [PyPI](https://pypi.org/project/svy-sae/) · 🧭 Built on [svy](https://pypi.org/project/svy/)

---

## Why svy-sae?

- **The three estimators practitioners actually use.** Fay–Herriot for area-level data, the Battese–Harter–Fuller nested-error model for unit-level, and Molina–Rao Empirical Best Prediction for non-linear indicators like poverty rates and Gini.
- **Uncertainty you can report.** Analytical MSE (Prasad–Rao for Fay–Herriot) where a closed form exists, and a parametric bootstrap where it does not — because a small area estimate without a defensible MSE is not usable for policy.
- **Non-linear indicators, not just means.** Poverty headcount, poverty gap, poverty severity, Gini and the quintile share ratio, computed per simulated population rather than back-transformed from a mean, plus any custom indicator you pass as a callable.
- **Census scale.** Memory is bounded independently of the Monte Carlo count, so a 2.5M-household census at 200 draws runs in about 16 seconds within roughly 2.5 GB.
- **Cheap to import.** `import svy_sae` costs about 5 ms; JAX and the model machinery load on first use, so CLI tools, test collection, and doc builds do not pay for a fit they never run.

## Installation

```bash
pip install svy-sae      # or: uv add svy-sae
```

Requires Python 3.12–3.14. Wheels are published for macOS (Apple silicon), Linux (x86_64 and aarch64), and Windows.

## Area level: Fay–Herriot

When you have direct estimates per area and their sampling variances:

```python
import polars as pl
import svy
from svy_sae.models import AreaLevel

milk = svy.datasets.load("milk").with_columns(variance=pl.col("SD") ** 2)

result = AreaLevel(milk).fh(
    y="yi",
    x=svy.Cat("MajorArea", ref=1),
    variance="variance",
    area="SmallArea",
    method="reml",
)

print(result.stats.sigma_u**2)          # 0.018550  (R sae::eblupFH: 0.0186)
for pred in result.prediction[:3]:
    print(pred.area, round(pred.pred, 4), round(pred.cv, 4))
```

## Unit level: nested-error EBLUP

When you have unit records in the survey and area-level means from the census:

```python
import svy
from svy_sae.models import UnitLevel

sample = svy.datasets.load("cornsoybean")
census = svy.datasets.load("cornsoybeanmeans")

result = UnitLevel(sample).eblup(
    y="CornHec",
    x=["CornPix", "SoyBeansPix"],
    area="County",
    pop_data=census,
    pop_size="PopnSegments",
    mse="analytical",          # or "bootstrap" for a parametric bootstrap MSE
)

print(result.stats.sigma_u**2)          # 63.3145   (R: 63.3149)
print(result.stats.sigma_e**2)          # 297.7133  (R: 297.7128)
```

## Non-linear indicators: Molina–Rao EBP

Poverty and inequality measures are not linear in the outcome, so they cannot be obtained by back-transforming a predicted mean. EBP simulates populations from the fitted model and computes the indicator on each:

```python
from svy_sae.core.enumerations import Indicator

result = UnitLevel(sample).ebp(
    y="CornHec",
    x=["CornPix", "SoyBeansPix"],
    area="County",
    pop_data=census,
    transformation="log",       # "boxcox" (default), "log", or "none"
    indicators=[Indicator.MEAN, Indicator.POVERTY_HEADCOUNT],
    threshold=100.0,            # absolute; pass a callable for e.g. 0.6 x median
    n_mc=200,
    rstate=42,
)

for pred in result.prediction[:4]:
    print(pred.area, pred.indicator, round(pred.pred, 4))
```

Pass `n_reps` to add a parametric bootstrap MSE for each indicator.

## Capabilities

| | |
|---|---|
| **Area level** | Fay–Herriot (REML, ML), Prasad–Rao and parametric bootstrap MSE |
| **Unit level** | Battese–Harter–Fuller EBLUP (REML, ML), analytical and bootstrap MSE |
| **Non-linear** | Molina–Rao EBP with Box–Cox, log, or no transformation |
| **Indicators** | Mean, poverty headcount, poverty gap, poverty severity, Gini, quintile share ratio, plus custom callables |
| **Inputs** | Polars frames or a `svy.Sample`, with the survey design carried through |
| **Reproducibility** | Explicit `rstate` seeds; typed, serializable result objects |

## First-run performance

Model kernels compile on first use and are cached to disk (`~/.cache/svy_sae`),
so only the first fit of a given data shape on a machine pays compilation — a
few extra seconds. To pay that cost deliberately instead of during real work
(a class demo, a production window), warm the cache first:

```python
import svy_sae

svy_sae.warmup()  # every estimator path, tiny synthetic data

# Or warm the exact shapes you will fit — e.g. once before a workshop:
svy_sae.warmup(course_sample, y="pc_exp", x=["hhsize", "rooms"],
               area="geo2", pop_data=course_census)
```

The cache is keyed by input shapes: warming with your own data covers your
later fits of that data exactly; the no-argument form initialises JAX and
covers only the built-in example shapes.

## Validation

Estimates are checked against R's [`sae`](https://cran.r-project.org/package=sae) package on its own benchmark datasets, which ship with `svy` and make the comparisons above reproducible:

| | svy-sae | R `sae` |
|---|---|---|
| `cornsoybean` σ²_e | 297.7133 | 297.7128 |
| `cornsoybean` σ²_u | 63.3145 | 63.3149 |
| `milk` σ²_u | 0.018550 | 0.018600 |

Bootstrap MSE is compared against `sae::pbmseBHF`, where mean agreement across the twelve counties is within 0.4% at 4,000 replicates — the residual being Monte Carlo error rather than bias.

## Ecosystem

svy-sae builds on [svy](https://pypi.org/project/svy/), which handles the survey side: design metadata, weighting, direct estimation, and the datasets used above. Pass a `svy.Sample` and its design travels with your data into the model.

## Status

svy-sae is in alpha and pre-1.0. The API may change between minor versions; each release documents what moved in [CHANGELOG.md](CHANGELOG.md), including any change that alters reported numbers.

## Feedback

Issues and discussions: [github.com/samplics-org/svy-sae](https://github.com/samplics-org/svy-sae)

---

**Small area estimates carry policy weight. svy-sae aims to make the uncertainty around them as defensible as the estimates themselves.**
