Add survey_estimation: design-based standard errors for clustered surveys

README.md

@@ -34,6 +34,7 @@ This is the **data pipeline layer** of [OpenStacks for Change](https://openstack
 | `modelling/` | Multicollinearity checks (VIF) | Ready |
 | `visualisation/` | Annotated bar charts, district-level choropleth maps | Ready |
 | `social_sector/` | Public health access index | Ready |
+| `survey_estimation/` | Design-based proportions and means for stratified, clustered surveys, with a worked NFHS-5 example | Ready |
 
 ### Notebooks
 
@@ -47,7 +48,7 @@ This is the **data pipeline layer** of [OpenStacks for Change](https://openstack
 | Directory | What It Contains |
 |-----------|-----------------|
 | `sample_data/` | Gender sample and time-use sample datasets |
-| `tests/` | 10 pytest test files covering all core modules |
+| `tests/` | 11 pytest test files covering all core modules |
 | `scripts/` | Standalone export utilities |
 
 ## Getting Started

requirements.txt

@@ -9,4 +9,5 @@ openpyxl
 xlsxwriter
 pyreadstat
 pydantic
-pandas-profiling
+pandas-profiling
+scipy

survey_estimation/README.md

@@ -0,0 +1,81 @@
+# Survey estimation
+
+Design-based estimates from complex survey data: weighted proportions and means
+with standard errors that account for stratification and clustering.
+
+Weights are the easy half and the half everyone remembers. A national household
+survey is also stratified and clustered, so a standard error computed as though
+the sample were independent is too small, often by a factor of two or more, and
+every confidence interval and test built on it is wrong in the direction that
+flatters the finding.
+
+`design_based_estimates.py` implements the Taylor linearisation (ultimate
+cluster) variance estimator, the same one behind Stata's `svy:` prefix and R's
+`survey` package. It is not tied to any particular survey.
+
+```python
+from survey_estimation import svy_prop_by
+svy_prop_by(df, "stunted", by="wealth_quintile")
+```
+
+The DataFrame needs a weight, a PSU and a stratum column. Everything else is
+optional.
+
+## Two things it does that a hand-rolled version usually does not
+
+**Subgroups are estimated as domains, not subsets.** Filtering the data before
+estimating a subgroup throws away the PSUs that contain none of its members.
+Those PSUs are still part of the design and still count towards the stratum's
+cluster total, so dropping them understates the standard error. Pass `domain=`
+or use `svy_prop_by`, and they are kept.
+
+**Proportions get a logit interval.** Near zero or one, a linear interval runs
+outside [0, 1] and reports something that is not a proportion. In the test
+suite, a 2.5 percent outcome concentrated in one cluster produces a linear
+lower bound of -0.024 and a logit lower bound of 0.003.
+
+**The interval says which distribution produced it.** Degrees of freedom are
+clusters minus strata, and with a few dozen clusters the t quantile is
+noticeably larger than 1.96. Every result carries a `ci_dist` field. Without
+`scipy` the module falls back to the normal quantile, which makes intervals
+slightly too narrow, and it warns rather than doing so quietly: on a 40-cluster,
+8-stratum example the lowest quintile's interval is [17.1, 32.8] under the
+normal and [16.9, 33.2] under t on 32 degrees of freedom.
+
+## Worked example
+
+`dhs_stunting.py` reproduces India's published NFHS-5 stunting table from the
+raw children's recode, by wealth quintile, with design-based intervals. It is
+the second half of a chain that starts in
+[InsightStack](https://github.com/Varnasr/InsightStack)'s
+`data_starters/dhs-south-asia/`:
+
+```
+# in InsightStack
+python load_dhs.py IAKR7EFL.DTA --vars v190 v025 hw70 b5 --anthro --out children.csv
+
+# here
+python -m survey_estimation.dhs_stunting children.csv
+```
+
+The two repositories are coupled through a CSV rather than an import, so
+neither needs the other installed.
+
+The comparison step is the point. DHS published NFHS-5 stunting at 35.5 percent
+nationally, 46.1 in the poorest wealth quintile and 22.9 in the richest. A run
+that does not land within a few tenths of that has a fault upstream, and the
+script names the three that account for almost all of them: an unscaled weight,
+a subgroup filtered before estimation, an anthropometry flag divided instead of
+dropped.
+
+## Tests
+
+```
+python -m pytest tests/test_survey_estimation.py
+```
+
+Twelve checks, each pinned to a case with a known closed-form answer rather than
+to the estimator's own output: one unit per cluster must reduce to `s / sqrt(n)`,
+multiplying every weight by a million must leave the standard error unchanged,
+and a sample where every unit inside a cluster is identical must produce a
+design effect of exactly `(N - 1) / (n - 1)`.

survey_estimation/__init__.py

@@ -0,0 +1,9 @@
+"""Design-based estimation for complex survey data."""
+
+from survey_estimation.design_based_estimates import (
+    svy_prop,
+    svy_prop_by,
+    compare_to_published,
+)
+
+__all__ = ["svy_prop", "svy_prop_by", "compare_to_published"]

survey_estimation/design_based_estimates.py

@@ -0,0 +1,217 @@
+"""
+Design-based estimation for complex survey data: weighted proportions and means
+with standard errors that account for stratification and clustering.
+
+Most Python survey work stops at a weighted mean. The weights are the easy half.
+A national household survey is stratified and clustered, so an unadjusted
+standard error is too small, often by a factor of two or more, and every
+confidence interval and significance test built on it is wrong in the direction
+that flatters the finding.
+
+This module implements the Taylor linearisation (ultimate cluster) variance
+estimator, which is what Stata's `svy:` prefix and R's `survey` package use.
+
+    from survey_estimation import svy_prop_by
+    svy_prop_by(df, "stunted", by="wealth_quintile")
+
+Works with any survey that carries a weight, a PSU and a stratum. It is not
+specific to DHS; see dhs_stunting_example.py for one worked application.
+"""
+
+from __future__ import annotations
+
+import warnings
+from statistics import NormalDist
+
+import numpy as np
+import pandas as pd
+
+try:
+    from scipy import stats as _st
+    _HAVE_SCIPY = True
+except ImportError:  # pragma: no cover - depends on the environment
+    _HAVE_SCIPY = False
+
+_SCIPY_WARNED = False
+
+
+def _t_quantile(conf: float, df: float) -> tuple[float, str]:
+    """Two-sided t quantile, and the name of the distribution actually used.
+
+    A survey design has finite degrees of freedom, clusters minus strata, and
+    with a few dozen clusters the t quantile is visibly larger than 1.96. Where
+    scipy is unavailable this falls back to the normal, which makes every
+    interval slightly too narrow, so the fallback says so out loud and the
+    returned distribution name records which one produced the number.
+    """
+    global _SCIPY_WARNED
+    alpha = 1 - conf
+    if _HAVE_SCIPY and df > 0:
+        return float(_st.t.ppf(1 - alpha / 2, df)), f"t({df:g})"
+    if not _SCIPY_WARNED:
+        warnings.warn(
+            "scipy is not installed, so confidence intervals use the normal "
+            "quantile rather than t on the design's degrees of freedom. Intervals "
+            "will be slightly too narrow. Install scipy to fix this.",
+            RuntimeWarning, stacklevel=3)
+        _SCIPY_WARNED = True
+    return float(NormalDist().inv_cdf(1 - alpha / 2)), "normal (scipy absent)"
+
+
+def _linearised_variance(u: pd.Series, psu: pd.Series, strata: pd.Series,
+                         singleunit: str = "centered") -> float:
+    """Ultimate cluster variance of a total whose linearised values are `u`.
+
+    The sum of squared deviations is taken between PSU totals within each
+    stratum, which is what makes this a *cluster* standard error rather than an
+    independence one.
+
+    A stratum containing a single PSU has no within-stratum variation to
+    measure. `singleunit="centered"` centres its contribution on the grand mean
+    of the PSU totals, matching Stata's `singleunit(centered)`;
+    `"certainty"` treats it as a certainty unit contributing nothing.
+    """
+    frame = pd.DataFrame({"u": u.to_numpy(), "psu": psu.to_numpy(),
+                          "strata": strata.to_numpy()})
+    psu_totals = frame.groupby(["strata", "psu"], sort=False)["u"].sum().reset_index()
+    grand_mean = psu_totals["u"].mean()
+
+    variance = 0.0
+    for _stratum, block in psu_totals.groupby("strata", sort=False):
+        n_h = len(block)
+        if n_h > 1:
+            centre = block["u"].mean()
+            variance += (n_h / (n_h - 1)) * float(((block["u"] - centre) ** 2).sum())
+        elif singleunit == "centered":
+            variance += float(((block["u"] - grand_mean) ** 2).sum())
+        # "certainty" adds nothing: a stratum with one PSU is taken as selected
+        # with certainty and contributes no sampling variance.
+    return variance
+
+
+def svy_prop(df: pd.DataFrame, outcome: str, weight: str = "weight",
+             psu: str = "psu", strata: str = "strata", domain=None,
+             conf: float = 0.95, ci: str = "logit",
+             singleunit: str = "centered") -> dict:
+    """Design-based estimate of a proportion or mean, with its standard error.
+
+    Parameters
+    ----------
+    df : DataFrame
+        The full sample. Do not subset it to estimate a subgroup; pass `domain`
+        instead, for the reason given below.
+    outcome : str
+        Column to average. Binary 0/1 for a proportion, numeric for a mean.
+        Rows where it is missing are dropped from the numerator and denominator
+        but their PSUs still count towards the degrees of freedom.
+    weight, psu, strata : str
+        Survey design columns.
+    domain : array-like of bool, optional
+        Subgroup indicator. Estimating a subgroup by filtering the DataFrame
+        first understates the standard error, because PSUs that contain no
+        members of the subgroup still belong to the design and still count in
+        the stratum's PSU total. Passing `domain` keeps them.
+    ci : {"logit", "linear"}
+        Logit keeps a proportion's interval inside [0, 1], which matters when
+        the estimate is near either bound. Ignored for a non-binary outcome.
+
+    Returns
+    -------
+    dict with estimate, se, ci_low, ci_high, n (unweighted rows used),
+    n_clusters, n_strata, df (degrees of freedom) and deff for binary outcomes.
+    """
+    for col in (outcome, weight, psu, strata):
+        if col not in df.columns:
+            raise KeyError(f"column {col!r} is not in the DataFrame")
+
+    d = df.copy()
+    d["_in"] = np.ones(len(d)) if domain is None else np.asarray(domain, dtype=float)
+    y = pd.to_numeric(d[outcome], errors="coerce")
+    # A missing outcome leaves the estimate but not the design: the PSU stays.
+    d["_in"] = d["_in"].where(y.notna(), 0.0)
+    d["_y"] = y.fillna(0.0)
+    w = pd.to_numeric(d[weight], errors="coerce").fillna(0.0)
+
+    denom = float((w * d["_in"]).sum())
+    if denom <= 0:
+        raise ValueError("The domain has no weighted observations.")
+    est = float((w * d["_in"] * d["_y"]).sum()) / denom
+
+    # Linearised value of the ratio estimator: the residual, weighted, scaled by
+    # the estimated domain size.
+    u = w * d["_in"] * (d["_y"] - est) / denom
+    var = _linearised_variance(u, d[psu], d[strata], singleunit=singleunit)
+    se = float(np.sqrt(max(var, 0.0)))
+
+    n_clusters = int(d.groupby([strata, psu], sort=False).ngroups)
+    n_strata = int(d[strata].nunique())
+    dof = max(n_clusters - n_strata, 1)
+    tq, ci_dist = _t_quantile(conf, dof)
+
+    n_used = int(d["_in"].sum())
+    binary = bool(np.isin(d.loc[d["_in"] > 0, "_y"].dropna().unique(), [0, 1]).all())
+
+    if binary and ci == "logit" and 0 < est < 1 and se > 0:
+        # Delta-method SE on the logit scale, back-transformed. Keeps the
+        # interval inside [0, 1] instead of reporting a negative lower bound.
+        logit = np.log(est / (1 - est))
+        se_logit = se / (est * (1 - est))
+        lo, hi = (1 / (1 + np.exp(-(logit + s * tq * se_logit))) for s in (-1, 1))
+    else:
+        lo, hi = est - tq * se, est + tq * se
+
+    out = {"estimate": est, "se": se, "ci_low": float(lo), "ci_high": float(hi),
+           "n": n_used, "n_clusters": n_clusters, "n_strata": n_strata, "df": dof,
+           "ci_dist": ci_dist}
+
+    if binary and n_used > 1:
+        var_srs = est * (1 - est) / (n_used - 1)
+        out["deff"] = float(var / var_srs) if var_srs > 0 else float("nan")
+    return out
+
+
+def svy_prop_by(df: pd.DataFrame, outcome: str, by: str, weight: str = "weight",
+                psu: str = "psu", strata: str = "strata", conf: float = 0.95,
+                ci: str = "logit", singleunit: str = "centered",
+                include_total: bool = True) -> pd.DataFrame:
+    """`svy_prop` across the levels of `by`, one row per level.
+
+    Each level is estimated as a domain of the full sample rather than by
+    subsetting, so the standard errors are right.
+    """
+    if by not in df.columns:
+        raise KeyError(f"column {by!r} is not in the DataFrame")
+    rows = []
+    for level in sorted(df[by].dropna().unique()):
+        res = svy_prop(df, outcome, weight=weight, psu=psu, strata=strata,
+                       domain=(df[by] == level), conf=conf, ci=ci,
+                       singleunit=singleunit)
+        rows.append({by: level, **res})
+    if include_total:
+        res = svy_prop(df, outcome, weight=weight, psu=psu, strata=strata,
+                       conf=conf, ci=ci, singleunit=singleunit)
+        rows.append({by: "Total", **res})
+    return pd.DataFrame(rows)
+
+
+def compare_to_published(estimates: pd.DataFrame, published: pd.DataFrame,
+                         on: str, est_col: str = "estimate",
+                         pub_col: str = "published", scale: float = 100.0,
+                         tolerance: float = 1.0) -> pd.DataFrame:
+    """Line your estimates up against a published table and flag the gaps.
+
+    Reproducing the published figures is the only cheap check that a survey
+    pipeline is correct end to end. A weight left unscaled, a domain filtered
+    too early, an anthropometry flag kept as data: each of these produces a
+    number that looks reasonable on its own and visibly wrong beside the
+    report the survey agency published.
+
+    `tolerance` is in the same units as `published` (percentage points by
+    default). A difference inside it is consistent with rounding in the
+    published table; outside it, something in the pipeline needs finding.
+    """
+    merged = estimates.merge(published, on=on, how="outer", suffixes=("", "_pub"))
+    merged["estimate_pct"] = merged[est_col] * scale
+    merged["difference"] = merged["estimate_pct"] - merged[pub_col]
+    merged["within_tolerance"] = merged["difference"].abs() <= tolerance
+    return merged[[on, "estimate_pct", pub_col, "difference", "within_tolerance"]]