diff --git a/README.md b/README.md index a236cbf..c313643 100644 --- a/README.md +++ b/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 diff --git a/requirements.txt b/requirements.txt index f42b4c8..d6aef46 100644 --- a/requirements.txt +++ b/requirements.txt @@ -9,4 +9,5 @@ openpyxl xlsxwriter pyreadstat pydantic -pandas-profiling \ No newline at end of file +pandas-profiling +scipy diff --git a/survey_estimation/README.md b/survey_estimation/README.md new file mode 100644 index 0000000..4f68f2b --- /dev/null +++ b/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)`. diff --git a/survey_estimation/__init__.py b/survey_estimation/__init__.py new file mode 100644 index 0000000..add59f1 --- /dev/null +++ b/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"] diff --git a/survey_estimation/design_based_estimates.py b/survey_estimation/design_based_estimates.py new file mode 100644 index 0000000..808fec7 --- /dev/null +++ b/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"]] diff --git a/survey_estimation/dhs_stunting.py b/survey_estimation/dhs_stunting.py new file mode 100644 index 0000000..3cd9f5e --- /dev/null +++ b/survey_estimation/dhs_stunting.py @@ -0,0 +1,131 @@ +""" +Worked example: child stunting by wealth quintile from a DHS children's recode, +with design-based standard errors, checked against the figures DHS published. + +This is the second half of a two-repo chain. InsightStack's +`data_starters/dhs-south-asia/` turns the raw recode into a clean CSV: + + python load_dhs.py IAKR7EFL.DTA --vars v190 v025 hw70 b5 --anthro \ + --out nfhs5_children.csv + +and this script takes it from there: + + python -m survey_estimation.dhs_stunting nfhs5_children.csv + +The two are coupled through a file rather than an import, so neither repo needs +the other installed. + +Why bother reproducing a published table +---------------------------------------- +Because it is the only cheap check that the whole pipeline is right. DHS +published India's NFHS-5 stunting rate as 35.5 percent, and 46.1 percent in the +poorest wealth quintile falling to 22.9 in the richest. If your run does not +land within a few tenths of that, something upstream is wrong, and the usual +suspects are an unscaled weight, a domain filtered before estimation, or an +anthropometry flag kept as though it were a measurement. +""" + +from __future__ import annotations + +import argparse +import sys + +import numpy as np +import pandas as pd + +from survey_estimation.design_based_estimates import svy_prop_by, compare_to_published + +# DHS Program API, indicator CN_NUTS_C_HA2 (children stunted, height-for-age +# below -2 SD of the WHO 2006 median), survey IA2020DHS, retrieved 2026-09-08. +PUBLISHED_NFHS5 = pd.DataFrame({ + "wealth_quintile": ["Lowest", "Second", "Middle", "Fourth", "Highest", "Total"], + "published": [46.1, 39.7, 34.4, 28.1, 22.9, 35.5], +}) + +QUINTILE_LABELS = {1: "Lowest", 2: "Second", 3: "Middle", 4: "Fourth", 5: "Highest"} + + +def prepare(df: pd.DataFrame) -> pd.DataFrame: + """Build the stunting indicator and the wealth quintile labels. + + Two filters matter and neither announces itself if you skip it. The + children's recode covers births in the last five years including children + who have died, so `b5 == 1` is required before any anthropometry. And a + child with no valid height-for-age is not a child who is not stunted; those + rows carry a missing outcome, which the estimator keeps in the design and + out of the numerator. + """ + if "haz" not in df.columns: + raise KeyError( + "No 'haz' column. Run the loader with --anthro so the flags at 9990 " + "and above are dropped before the values are divided by 100." + ) + out = df.copy() + + if "b5" in out.columns: + before = len(out) + out = out[out["b5"] == 1] + print(f" living children: {len(out):,} of {before:,}", file=sys.stderr) + + out["stunted"] = np.where(out["haz"].notna(), (out["haz"] < -2).astype(float), np.nan) + measured = int(out["stunted"].notna().sum()) + print(f" with a valid height-for-age: {measured:,} of {len(out):,}", file=sys.stderr) + + if "v190" in out.columns: + out["wealth_quintile"] = out["v190"].map(QUINTILE_LABELS) + return out + + +def stunting_by_wealth(df: pd.DataFrame) -> pd.DataFrame: + """Stunting prevalence by wealth quintile, with design-based intervals.""" + est = svy_prop_by(df, "stunted", by="wealth_quintile") + order = ["Lowest", "Second", "Middle", "Fourth", "Highest", "Total"] + est["wealth_quintile"] = pd.Categorical(est["wealth_quintile"], order, ordered=True) + return est.sort_values("wealth_quintile").reset_index(drop=True) + + +def report(est: pd.DataFrame, published: pd.DataFrame | None = PUBLISHED_NFHS5, + tolerance: float = 0.5) -> pd.DataFrame: + """Print the estimates, and the comparison when a published table is given.""" + show = est.copy() + for col in ("estimate", "ci_low", "ci_high"): + show[col] = (show[col] * 100).round(1) + show["se"] = (show["se"] * 100).round(2) + print("\nStunting by wealth quintile, percent") + print(show[["wealth_quintile", "estimate", "se", "ci_low", "ci_high", + "n", "n_clusters", "deff"]].to_string(index=False)) + + if published is None: + return show + cmp = compare_to_published(est, published, on="wealth_quintile", tolerance=tolerance) + print(f"\nAgainst the published NFHS-5 table (tolerance {tolerance} points)") + print(cmp.round(2).to_string(index=False)) + off = cmp[~cmp["within_tolerance"].fillna(False)] + if len(off): + print(f"\n{len(off)} row(s) outside tolerance. Check the weight scaling, " + "whether the subgroup was filtered before estimation, and whether " + "the anthropometry flags were dropped before dividing by 100.", + file=sys.stderr) + else: + print("\nEvery row reproduces the published figure. The pipeline is sound.") + return cmp + + +def main(argv=None) -> int: + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("csv", help="Cleaned children's recode, from InsightStack's load_dhs") + p.add_argument("--no-benchmark", action="store_true", + help="Skip the comparison, for a survey other than NFHS-5") + p.add_argument("--tolerance", type=float, default=0.5, + help="Allowed gap in percentage points (default 0.5)") + a = p.parse_args(argv) + + df = prepare(pd.read_csv(a.csv)) + est = stunting_by_wealth(df) + report(est, None if a.no_benchmark else PUBLISHED_NFHS5, a.tolerance) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_survey_estimation.py b/tests/test_survey_estimation.py new file mode 100644 index 0000000..8bad7fc --- /dev/null +++ b/tests/test_survey_estimation.py @@ -0,0 +1,200 @@ +""" +Tests for survey_estimation.design_based_estimates. + +A variance estimator cannot be checked by eye, so each test here pins it to a +case where the right answer is known in closed form. +""" + +import numpy as np +import pandas as pd + +from survey_estimation.design_based_estimates import ( + svy_prop, svy_prop_by, compare_to_published, +) + + +def test_point_estimate_is_the_weighted_mean(): + df = pd.DataFrame({"y": [1, 0, 1, 0], "weight": [1.0, 3.0, 1.0, 5.0], + "psu": [1, 2, 3, 4], "strata": [1, 1, 1, 1]}) + expected = (1 * 1 + 0 * 3 + 1 * 1 + 0 * 5) / (1 + 3 + 1 + 5) + assert abs(svy_prop(df, "y")["estimate"] - expected) < 1e-12 + + +def test_one_unit_per_cluster_reduces_to_the_simple_random_sample_se(): + """With one unit per PSU, one stratum and equal weights, the ultimate + cluster estimator must collapse to s / sqrt(n).""" + rng = np.random.default_rng(11) + y = rng.integers(0, 2, 40).astype(float) + df = pd.DataFrame({"y": y, "weight": np.ones(40), + "psu": np.arange(40), "strata": np.ones(40)}) + expected_se = y.std(ddof=1) / np.sqrt(len(y)) + assert abs(svy_prop(df, "y", ci="linear")["se"] - expected_se) < 1e-12 + + +def test_standard_error_is_invariant_to_the_weight_scale(): + """A ratio estimator does not care whether weights are raw DHS integers or + divided by a million. If the SE moves, the estimator is wrong.""" + rng = np.random.default_rng(7) + n = 200 + df = pd.DataFrame({ + "y": rng.integers(0, 2, n).astype(float), + "weight": rng.uniform(0.4, 2.5, n), + "psu": rng.integers(0, 25, n), + "strata": rng.integers(0, 5, n), + }) + base = svy_prop(df, "y") + scaled = svy_prop(df.assign(weight=df["weight"] * 1_000_000), "y") + assert abs(base["estimate"] - scaled["estimate"]) < 1e-12 + assert abs(base["se"] - scaled["se"]) < 1e-12 + + +def test_perfect_intracluster_correlation_gives_the_known_design_effect(): + """Every unit inside a cluster identical means rho = 1, and the design + effect is then exactly (N - 1) / (n - 1): the sample carries only as much + information as its cluster count.""" + n_clusters, size = 20, 5 + rng = np.random.default_rng(3) + cluster_value = rng.integers(0, 2, n_clusters).astype(float) + df = pd.DataFrame({ + "y": np.repeat(cluster_value, size), + "weight": np.ones(n_clusters * size), + "psu": np.repeat(np.arange(n_clusters), size), + "strata": np.ones(n_clusters * size), + }) + res = svy_prop(df, "y") + n_total = n_clusters * size + expected_deff = (n_total - 1) / (n_clusters - 1) + assert abs(res["deff"] - expected_deff) < 1e-9 + # And the clustered SE must exceed the SE that ignores the design. + naive_se = np.sqrt(res["estimate"] * (1 - res["estimate"]) / (n_total - 1)) + assert res["se"] > naive_se + + +def test_a_subgroup_keeps_the_clusters_that_contain_none_of_it(): + """Filtering before estimating throws away PSUs that are part of the design, + which shrinks the standard error. Passing a domain keeps them.""" + rng = np.random.default_rng(5) + n = 300 + df = pd.DataFrame({ + "y": rng.integers(0, 2, n).astype(float), + "group": rng.integers(0, 5, n), + "weight": rng.uniform(0.5, 2.0, n), + "psu": rng.integers(0, 30, n), + "strata": rng.integers(0, 3, n), + }) + domain = df["group"] == 0 + correct = svy_prop(df, "y", domain=domain) + naive = svy_prop(df[domain].copy(), "y") + + assert correct["estimate"] == naive["estimate"] # same point estimate + assert correct["n_clusters"] > naive["n_clusters"] # more clusters retained + assert correct["df"] > naive["df"] + assert correct["se"] != naive["se"] + + +def test_a_stratum_with_one_cluster_does_not_crash(): + df = pd.DataFrame({ + "y": [1.0, 0.0, 1.0, 1.0, 0.0], + "weight": [1.0] * 5, + "psu": [1, 2, 3, 4, 9], + "strata": [1, 1, 1, 1, 2], # stratum 2 holds a single PSU + }) + centered = svy_prop(df, "y", singleunit="centered") + certainty = svy_prop(df, "y", singleunit="certainty") + assert np.isfinite(centered["se"]) and np.isfinite(certainty["se"]) + assert certainty["se"] < centered["se"] # a certainty unit adds no variance + + +def test_logit_interval_stays_inside_the_unit_interval(): + """Near zero a linear interval goes negative, which is not a proportion.""" + n = 400 + rng = np.random.default_rng(2) + psu = rng.integers(0, 40, n) + # A rare outcome concentrated inside a single cluster: p is 2.5% and the + # design standard error is the same size, so the linear interval crosses zero. + y = np.zeros(n) + y[psu == 0] = 1.0 + df = pd.DataFrame({"y": y, "weight": np.ones(n), "psu": psu, + "strata": np.ones(n)}) + logit = svy_prop(df, "y", ci="logit") + linear = svy_prop(df, "y", ci="linear") + assert logit["ci_low"] > 0 + assert logit["ci_high"] < 1 + assert linear["ci_low"] < 0 # the reason logit is the default + + +def test_missing_outcomes_leave_the_estimate_but_keep_the_design(): + rng = np.random.default_rng(13) + n = 150 + y = rng.integers(0, 2, n).astype(float) + y[::10] = np.nan + df = pd.DataFrame({"y": y, "weight": np.ones(n), + "psu": rng.integers(0, 20, n), + "strata": np.ones(n)}) + res = svy_prop(df, "y") + assert res["n"] == int(np.isfinite(y).sum()) + assert res["n_clusters"] == df.groupby(["strata", "psu"]).ngroups + + +def test_by_returns_one_row_per_level_plus_a_total(): + rng = np.random.default_rng(17) + n = 240 + df = pd.DataFrame({ + "y": rng.integers(0, 2, n).astype(float), + "quintile": rng.integers(1, 6, n), + "weight": rng.uniform(0.5, 2.0, n), + "psu": rng.integers(0, 24, n), + "strata": rng.integers(0, 4, n), + }) + out = svy_prop_by(df, "y", by="quintile") + assert len(out) == 6 # five quintiles and a total + assert out["quintile"].iloc[-1] == "Total" + assert (out["ci_low"] <= out["estimate"]).all() + assert (out["estimate"] <= out["ci_high"]).all() + + +def test_comparison_against_a_published_table_flags_the_gap(): + est = pd.DataFrame({"quintile": ["Lowest", "Highest"], "estimate": [0.461, 0.300]}) + pub = pd.DataFrame({"quintile": ["Lowest", "Highest"], "published": [46.1, 22.9]}) + out = compare_to_published(est, pub, on="quintile", tolerance=1.0) + assert bool(out.loc[out["quintile"] == "Lowest", "within_tolerance"].iloc[0]) + assert not bool(out.loc[out["quintile"] == "Highest", "within_tolerance"].iloc[0]) + assert abs(out.loc[out["quintile"] == "Highest", "difference"].iloc[0] - 7.1) < 1e-9 + + +def test_the_interval_records_which_distribution_produced_it(): + """A survey design has finite degrees of freedom. Where the t quantile is + unavailable the result must say so rather than quietly using 1.96.""" + rng = np.random.default_rng(23) + n = 200 + df = pd.DataFrame({"y": rng.integers(0, 2, n).astype(float), + "weight": np.ones(n), + "psu": rng.integers(0, 20, n), + "strata": np.ones(n)}) + res = svy_prop(df, "y") + assert "ci_dist" in res + assert res["ci_dist"].startswith("t(") or "scipy absent" in res["ci_dist"] + + +def test_a_non_default_confidence_level_is_honoured(): + """The old fallback hardcoded the 95 percent quantile, so a 99 percent + interval came back at 95 percent width with nothing to show for it.""" + rng = np.random.default_rng(29) + n = 300 + df = pd.DataFrame({"y": rng.integers(0, 2, n).astype(float), + "weight": np.ones(n), + "psu": rng.integers(0, 30, n), + "strata": np.ones(n)}) + narrow = svy_prop(df, "y", conf=0.90, ci="linear") + wide = svy_prop(df, "y", conf=0.99, ci="linear") + assert (wide["ci_high"] - wide["ci_low"]) > (narrow["ci_high"] - narrow["ci_low"]) + + +def test_the_documented_top_level_import_works(): + """The README shows `from survey_estimation import svy_prop_by`, so the + package has to re-export it.""" + import survey_estimation + + assert callable(survey_estimation.svy_prop_by) + assert callable(survey_estimation.svy_prop) + assert callable(survey_estimation.compare_to_published)