stratindex 0.1.0__py3-none-any.whl

stratindex/__init__.py

@@ -0,0 +1,13 @@
+"""stratindex: a nonparametric index of stratification (Zhou 2012).
+
+Python port of the R package 'strat' by Xiang Zhou
+(https://cran.r-project.org/package=strat).
+"""
+
+from .core import srank, strat
+from .datasets import load_cpsmarch2015
+from .results import SrankResult, StratResult
+
+__version__ = "0.1.0"
+
+__all__ = ["srank", "strat", "SrankResult", "StratResult", "load_cpsmarch2015", "__version__"]

stratindex/_kernel.py

@@ -0,0 +1,95 @@
+"""Pairwise-comparison kernels.
+
+Blocked NumPy ports of ``strat_cpp`` and ``strat_cpp_by`` from
+``src/strat_cpp.cpp`` in the R package. Inputs must be sorted ascending by
+``r``. A pair (i, j), i < j, contributes ``w_i * w_j`` to the denominator and
+``sign(y_j - y_i) * w_i * w_j`` to the numerator, unless ``r_j == r_i`` (same
+stratum position) or ``y_j == y_i`` (tied outcome), in which case it is
+skipped.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+_BLOCK = 2048
+
+
+def pair_sums(
+    y: np.ndarray, r: np.ndarray, w: np.ndarray, block: int = _BLOCK
+) -> tuple[float, float]:
+    """Return ``(deno, nume)`` summed over all valid pairs."""
+    n = y.size
+    deno = 0.0
+    nume = 0.0
+    for i0 in range(0, n, block):
+        i1 = min(i0 + block, n)
+        yi = y[i0:i1, None]
+        ri = r[i0:i1, None]
+        wi = w[i0:i1, None]
+        for j0 in range(i0, n, block):
+            j1 = min(j0 + block, n)
+            yj = y[None, j0:j1]
+            rj = r[None, j0:j1]
+            valid = (rj > ri) & (yj != yi)
+            if j0 == i0:
+                valid &= np.arange(j0, j1)[None, :] > np.arange(i0, i1)[:, None]
+            wij = wi * w[None, j0:j1] * valid
+            deno += wij.sum()
+            nume += (np.sign(yj - yi) * wij).sum()
+    return deno, nume
+
+
+def pair_sums_by(
+    y: np.ndarray,
+    r: np.ndarray,
+    w: np.ndarray,
+    c: np.ndarray,
+    n_groups: int,
+    block: int = _BLOCK,
+) -> dict[str, np.ndarray | float]:
+    """Group decomposition of the pairwise sums.
+
+    Pairs within the same group accumulate into per-group ``deno_by`` /
+    ``nume_by`` (indexed by the group code of both members); pairs across
+    groups accumulate into the between-group sums.
+    """
+    n = y.size
+    deno_by = np.zeros(n_groups)
+    nume_by = np.zeros(n_groups)
+    deno_between = 0.0
+    nume_between = 0.0
+    for i0 in range(0, n, block):
+        i1 = min(i0 + block, n)
+        yi = y[i0:i1, None]
+        ri = r[i0:i1, None]
+        wi = w[i0:i1, None]
+        ci = c[i0:i1]
+        for j0 in range(i0, n, block):
+            j1 = min(j0 + block, n)
+            yj = y[None, j0:j1]
+            rj = r[None, j0:j1]
+            valid = (rj > ri) & (yj != yi)
+            if j0 == i0:
+                valid &= np.arange(j0, j1)[None, :] > np.arange(i0, i1)[:, None]
+            wij = wi * w[None, j0:j1] * valid
+            s = np.sign(yj - yi) * wij
+            same = c[None, j0:j1] == ci[:, None]
+            w_same = np.where(same, wij, 0.0)
+            s_same = np.where(same, s, 0.0)
+            deno_by += np.bincount(ci, weights=w_same.sum(axis=1), minlength=n_groups)
+            nume_by += np.bincount(ci, weights=s_same.sum(axis=1), minlength=n_groups)
+            # element-wise differences are exact where same-group entries cancel
+            deno_between += (wij - w_same).sum()
+            nume_between += (s - s_same).sum()
+
+    deno_within = deno_by.sum()
+    nume_within = nume_by.sum()
+    return {
+        "deno_by": deno_by,
+        "nume_by": nume_by,
+        "deno_within": deno_within,
+        "nume_within": nume_within,
+        "deno_between": deno_between,
+        "nume_between": nume_between,
+    }

stratindex/_utils.py

@@ -0,0 +1,128 @@
+"""Input cleaning and weighted percentile ranks.
+
+Ports the internal helpers of the R package ``strat`` (``R/utils.R``):
+``wtd_rank`` (a thin wrapper around ``Hmisc::wtd.rank``) and ``clean``.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+
+import numpy as np
+
+
+def wtd_rank(x: np.ndarray, weights: np.ndarray | None = None) -> np.ndarray:
+    """Weighted midranks of ``x``.
+
+    Equivalent to ``Hmisc::wtd.rank(x, weights, normwt = TRUE)``: weights are
+    normalized to sum to ``len(x)``, tied values share the midrank of their
+    weight block.
+    """
+    x = np.asarray(x, dtype=float)
+    if weights is None:
+        w = np.ones_like(x)
+    else:
+        w = np.asarray(weights, dtype=float)
+    w = w / w.sum() * x.size  # normwt = TRUE
+    _, inverse = np.unique(x, return_inverse=True)
+    freqs = np.bincount(inverse, weights=w)
+    ranks = np.cumsum(freqs) - 0.5 * (freqs - 1.0)
+    return ranks[inverse]
+
+
+def _is_na(arr: np.ndarray) -> np.ndarray:
+    if arr.dtype.kind == "f":
+        return np.isnan(arr)
+    if arr.dtype.kind == "O":
+        return np.array(
+            [v is None or (isinstance(v, float) and math.isnan(v)) for v in arr],
+            dtype=bool,
+        )
+    return np.zeros(arr.shape, dtype=bool)
+
+
+@dataclass
+class CleanData:
+    """Complete cases of all inputs, ready for the pairwise kernel.
+
+    Mirrors the data frame returned by ``clean()`` in the R package: ``prank``
+    is the weighted percentile rank in (0, 1], ``weights`` are normalized to
+    sum to ``n``, strata and group are stored as integer codes plus levels
+    (sorted unique values, as ``factor()`` does).
+    """
+
+    prank: np.ndarray
+    strata_codes: np.ndarray
+    strata_levels: np.ndarray
+    weights: np.ndarray
+    group_codes: np.ndarray | None
+    group_levels: np.ndarray | None
+    n: int
+
+
+def clean(
+    outcome,
+    strata,
+    weights=None,
+    group=None,
+) -> CleanData:
+    """Validate inputs, drop incomplete cases, compute weighted percentile ranks."""
+    outcome = np.asarray(outcome)
+    if (
+        outcome.ndim != 1
+        or outcome.size == 0
+        or not np.issubdtype(outcome.dtype, np.number)
+        or outcome.dtype.kind == "b"
+    ):
+        raise ValueError("outcome has to be a non-empty numeric 1-d array")
+    outcome = outcome.astype(float)
+
+    strata = np.asarray(strata)
+    if strata.ndim != 1 or len(strata) != len(outcome):
+        raise ValueError("outcome and strata have to be of equal length")
+
+    if weights is None:
+        weights_arr = np.ones(outcome.size)
+    else:
+        weights_arr = np.asarray(weights)
+        if weights_arr.ndim != 1 or not np.issubdtype(weights_arr.dtype, np.number):
+            raise ValueError("weights has to be a numeric 1-d array")
+        if len(weights_arr) != len(outcome):
+            raise ValueError("outcome and weights have to be of equal length")
+        weights_arr = weights_arr.astype(float)
+
+    group_arr: np.ndarray | None = None
+    if group is not None:
+        group_arr = np.asarray(group)
+        if group_arr.ndim != 1 or len(group_arr) != len(outcome):
+            raise ValueError("outcome and group have to be of equal length")
+        # R's strat() rejects missing values in group up front
+        if _is_na(group_arr).any():
+            raise ValueError("group contains missing values")
+
+    ok = ~(_is_na(outcome) | _is_na(strata) | _is_na(weights_arr))
+    n = int(ok.sum())
+    if n == 0:
+        raise ValueError("no complete cases!")
+
+    outcome = outcome[ok]
+    weights_arr = weights_arr[ok]
+    weights_arr = weights_arr / weights_arr.sum() * n
+    strata_levels, strata_codes = np.unique(strata[ok], return_inverse=True)
+
+    group_codes = group_levels = None
+    if group_arr is not None:
+        group_levels, group_codes = np.unique(group_arr[ok], return_inverse=True)
+
+    prank = wtd_rank(outcome, weights_arr) / n
+
+    return CleanData(
+        prank=prank,
+        strata_codes=strata_codes,
+        strata_levels=strata_levels,
+        weights=weights_arr,
+        group_codes=group_codes,
+        group_levels=group_levels,
+        n=n,
+    )

stratindex/core.py

@@ -0,0 +1,161 @@
+"""Public API: :func:`srank` and :func:`strat`.
+
+Port of ``R/srank.R`` and ``R/strat.R`` from the R package ``strat``
+(Xiang Zhou, https://cran.r-project.org/package=strat).
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from ._kernel import pair_sums, pair_sums_by
+from ._utils import CleanData, clean
+from .results import SrankResult, StratResult
+
+__all__ = ["srank", "strat"]
+
+
+def _summarize(cd: CleanData) -> dict[str, np.ndarray]:
+    """Per-stratum population share and average percentile rank."""
+    k = len(cd.strata_levels)
+    w_by_stratum = np.bincount(cd.strata_codes, weights=cd.weights, minlength=k)
+    share = w_by_stratum / cd.weights.sum()
+    s_prank = (
+        np.bincount(cd.strata_codes, weights=cd.weights * cd.prank, minlength=k) / w_by_stratum
+    )
+    return {"strata": cd.strata_levels, "share": share, "s_prank": s_prank}
+
+
+def _raw(cd: CleanData) -> dict[str, np.ndarray]:
+    raw = {
+        "prank": cd.prank,
+        "strata": cd.strata_levels[cd.strata_codes],
+        "weights": cd.weights,
+    }
+    if cd.group_codes is not None:
+        raw["group"] = cd.group_levels[cd.group_codes]
+    return raw
+
+
+def srank(outcome, strata, weights=None, group=None) -> SrankResult:
+    """Rank strata by the average percentile rank of their members.
+
+    Parameters
+    ----------
+    outcome:
+        Numeric array of outcomes.
+    strata:
+        Array of the same length indicating strata membership.
+    weights:
+        Optional numeric array of sampling weights.
+    group:
+        Optional grouping factor, carried through to ``raw``.
+
+    Returns
+    -------
+    SrankResult
+        ``raw`` (complete cases with percentile ranks) and ``summary``
+        (per-stratum share and average percentile rank).
+    """
+    cd = clean(outcome, strata, weights=weights, group=group)
+    return SrankResult(raw=_raw(cd), summary=_summarize(cd))
+
+
+def strat(
+    outcome,
+    strata,
+    weights=None,
+    ordered: bool = False,
+    group=None,
+    group_name: str = "group",
+) -> StratResult:
+    """Compute the stratification index proposed in Zhou (2012).
+
+    Parameters
+    ----------
+    outcome:
+        Numeric array of outcomes.
+    strata:
+        Array of the same length indicating strata membership.
+    weights:
+        Optional numeric array of sampling weights.
+    ordered:
+        If True, strata are taken as pre-ordered ascendingly (by their sorted
+        unique values); otherwise they are ordered by average percentile rank.
+    group:
+        Optional grouping factor. If supplied (with more than one level), the
+        result includes a between-/within-group decomposition of the overall
+        stratification.
+    group_name:
+        Label used for the group in printed output (R derives it from the
+        expression passed as ``group``; Python cannot, so pass it explicitly).
+
+    Returns
+    -------
+    StratResult
+        Overall index with approximate standard error (Goodman & Kruskal
+        1963), per-stratum information, and the group decomposition when a
+        group is supplied.
+
+    References
+    ----------
+    Zhou, Xiang. 2012. "A Nonparametric Index of Stratification."
+    Sociological Methodology, 42(1): 365-389.
+    """
+    if not isinstance(ordered, bool):
+        raise ValueError("ordered has to be a valid logical scalar")
+
+    cd = clean(outcome, strata, weights=weights, group=group)
+    strata_info = _summarize(cd)
+
+    row_s_prank = strata_info["s_prank"][cd.strata_codes]
+    if ordered:
+        sort_by = cd.strata_codes.astype(float)
+    else:
+        sort_by = row_s_prank
+    order = np.argsort(sort_by, kind="stable")
+    y = cd.prank[order]
+    r = sort_by[order]
+    w = cd.weights[order]
+
+    decomposition = None
+    within_group = None
+    if cd.group_codes is None or len(cd.group_levels) == 1:
+        deno, nume = pair_sums(y, r, w)
+        with np.errstate(invalid="ignore"):
+            index = nume / deno
+    else:
+        c = cd.group_codes[order]
+        sums = pair_sums_by(y, r, w, c, len(cd.group_levels))
+        deno = sums["deno_within"] + sums["deno_between"]
+        with np.errstate(invalid="ignore", divide="ignore"):
+            index = (sums["nume_within"] + sums["nume_between"]) / deno
+            decomposition = {
+                "within": {
+                    "weight": sums["deno_within"] / deno,
+                    "strat": sums["nume_within"] / sums["deno_within"],
+                },
+                "between": {
+                    "weight": sums["deno_between"] / deno,
+                    "strat": sums["nume_between"] / sums["deno_between"],
+                },
+            }
+            within_group = {
+                group_name: cd.group_levels,
+                "weight": sums["deno_by"] / sums["deno_within"],
+                "strat": sums["nume_by"] / sums["deno_by"],
+            }
+
+    # approximate standard error (Goodman & Kruskal 1963)
+    with np.errstate(invalid="ignore", divide="ignore"):
+        arg = deno / (1.0 - index**2) / cd.n
+        std_error = 1.0 / np.sqrt(arg)
+
+    return StratResult(
+        strat=float(index),
+        std_error=float(std_error),
+        strata_info=strata_info,
+        decomposition=decomposition,
+        within_group=within_group,
+        group_name=group_name,
+    )

stratindex/datasets.py

@@ -0,0 +1,56 @@
+"""Bundled example dataset.
+
+``cpsmarch2015`` is the dataset shipped with the R package ``strat``: income,
+big class, microclass, education and sampling weight for 14,358 male
+respondents aged 35-64 from the March 2015 Current Population Survey.
+"""
+
+from __future__ import annotations
+
+import csv
+import gzip
+from importlib.resources import files
+
+import numpy as np
+
+__all__ = ["load_cpsmarch2015"]
+
+_COLUMNS = ("income", "big_class", "micro_class", "education", "weight")
+
+
+def load_cpsmarch2015(as_pandas: bool = False):
+    """Load the cpsmarch2015 dataset.
+
+    Parameters
+    ----------
+    as_pandas:
+        If True, return a pandas DataFrame (requires pandas); otherwise a
+        dict of NumPy arrays keyed by column name.
+
+    Returns
+    -------
+    dict[str, numpy.ndarray] | pandas.DataFrame
+        Columns: ``income`` (float, personal market income in US dollars),
+        ``big_class`` (str), ``micro_class`` (int), ``education`` (str),
+        ``weight`` (float, CPS sampling weight).
+    """
+    resource = files("stratindex.data").joinpath("cpsmarch2015.csv.gz")
+    with resource.open("rb") as raw, gzip.open(raw, "rt", newline="") as fh:
+        reader = csv.reader(fh)
+        header = tuple(next(reader))
+        assert header == _COLUMNS
+        rows = list(reader)
+
+    columns = list(zip(*rows, strict=True))
+    data = {
+        "income": np.array(columns[0], dtype=float),
+        "big_class": np.array(columns[1]),
+        "micro_class": np.array(columns[2], dtype=int),
+        "education": np.array(columns[3]),
+        "weight": np.array(columns[4], dtype=float),
+    }
+    if as_pandas:
+        import pandas as pd
+
+        return pd.DataFrame(data)
+    return data

stratindex/results.py

@@ -0,0 +1,129 @@
+"""Result containers for :func:`stratindex.strat` and :func:`stratindex.srank`.
+
+The ``__str__`` output mirrors the ``print.strat`` / ``print.srank`` methods
+of the R package.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import numpy as np
+
+
+def _fmt(value: float, digits: int) -> str:
+    if isinstance(value, float) and np.isnan(value):
+        return "nan"
+    return f"{value:.{digits}g}"
+
+
+def _table(columns: dict[str, list | np.ndarray], digits: int) -> str:
+    """Render a dict of equal-length columns as an aligned text table."""
+    cells: dict[str, list[str]] = {}
+    for name, col in columns.items():
+        cells[name] = [
+            _fmt(float(v), digits) if isinstance(v, (int, float, np.floating)) else str(v)
+            for v in col
+        ]
+    widths = {name: max(len(name), *(len(c) for c in col)) for name, col in cells.items()}
+    lines = ["  ".join(name.rjust(widths[name]) for name in cells)]
+    n_rows = len(next(iter(cells.values())))
+    for i in range(n_rows):
+        lines.append("  ".join(cells[name][i].rjust(widths[name]) for name in cells))
+    return "\n".join(lines)
+
+
+@dataclass
+class SrankResult:
+    """Stratum-specific information: population share and average percentile rank.
+
+    Attributes
+    ----------
+    raw:
+        Complete cases of all inputs — dict with ``prank``, ``strata``,
+        ``weights`` (and ``group`` when supplied), each an ndarray of length n.
+    summary:
+        Per-stratum table — dict with ``strata`` (level labels), ``share``
+        and ``s_prank`` arrays.
+    """
+
+    raw: dict[str, np.ndarray] = field(repr=False)
+    summary: dict[str, np.ndarray]
+
+    def to_pandas(self):
+        """Return ``(raw, summary)`` as pandas DataFrames (requires pandas)."""
+        import pandas as pd
+
+        return pd.DataFrame(self.raw), pd.DataFrame(self.summary)
+
+    def format(self, digits: int = 3) -> str:
+        return _table(self.summary, digits)
+
+    def __str__(self) -> str:
+        return self.format()
+
+
+@dataclass
+class StratResult:
+    """The stratification index and its approximate standard error.
+
+    Attributes
+    ----------
+    strat:
+        The overall stratification index.
+    std_error:
+        Approximate standard error (Goodman & Kruskal 1963).
+    strata_info:
+        Per-stratum table — dict with ``strata``, ``share`` and ``s_prank``.
+    decomposition:
+        ``None`` unless a group was supplied; otherwise a dict with rows
+        ``within`` / ``between``, each a dict with ``weight`` and ``strat``.
+    within_group:
+        ``None`` unless a group was supplied; otherwise a dict with the group
+        levels (keyed by ``group_name``) and per-group ``weight`` / ``strat``
+        arrays.
+    """
+
+    strat: float
+    std_error: float
+    strata_info: dict[str, np.ndarray]
+    decomposition: dict[str, dict[str, float]] | None = None
+    within_group: dict[str, np.ndarray] | None = None
+    group_name: str = "group"
+
+    @property
+    def overall(self) -> dict[str, float]:
+        return {"strat": self.strat, "std_error": self.std_error}
+
+    def to_pandas(self):
+        """Return ``strata_info`` (and ``within_group`` if any) as DataFrames."""
+        import pandas as pd
+
+        strata_info = pd.DataFrame(self.strata_info)
+        if self.within_group is None:
+            return strata_info
+        return strata_info, pd.DataFrame(self.within_group)
+
+    def format(self, digits: int = 3) -> str:
+        lines = [
+            "overall stratification:",
+            "",
+            _table({"strat": [self.strat], "std_error": [self.std_error]}, digits),
+        ]
+        if self.decomposition is not None:
+            rows = {
+                "": [f"within {self.group_name}", f"between {self.group_name}"],
+                "weight": [
+                    self.decomposition["within"]["weight"],
+                    self.decomposition["between"]["weight"],
+                ],
+                "strat": [
+                    self.decomposition["within"]["strat"],
+                    self.decomposition["between"]["strat"],
+                ],
+            }
+            lines += ["", f"decomposition by {self.group_name}:", "", _table(rows, digits)]
+        return "\n".join(lines)
+
+    def __str__(self) -> str:
+        return self.format()