pysofra 0.1.0a1__py3-none-any.whl

pysofra/summary/calibrate.py

@@ -0,0 +1,214 @@
+"""Survey-weight calibration — post-stratification and raking.
+
+Both algorithms scale the design weights so that the weighted marginal
+totals over selected variables match supplied population targets.
+
+* :func:`post_stratify` solves the case where the calibration variables
+  partition the sample into a single cross-classification (e.g. age × sex
+  cells). The new weight for each row is the design weight multiplied
+  by the cell-level ratio of population total to weighted sample total.
+
+* :func:`rake` (a.k.a. iterative proportional fitting) handles the more
+  common case where targets are given for several variables marginally
+  but not for their joint cross-classification. Weights are scaled one
+  variable at a time, repeatedly, until convergence.
+
+Both functions return a new pandas Series of calibrated weights with the
+same index as the input. Use the result as the ``weights`` column of a
+:class:`SurveyDesign`.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+
+import numpy as np
+import pandas as pd
+
+
+def post_stratify(
+    data: pd.DataFrame,
+    base_weights: pd.Series | str,
+    *,
+    strata_cols: list[str] | tuple[str, ...],
+    targets: Mapping[tuple[object, ...], float] | pd.Series,
+) -> pd.Series:
+    """Post-stratification calibration over a complete cross-classification.
+
+    Parameters
+    ----------
+    data
+        Source dataframe.
+    base_weights
+        Either the column name of design weights in ``data`` or a Series
+        aligned to ``data.index``.
+    strata_cols
+        One or more columns whose Cartesian product defines the
+        post-strata.
+    targets
+        Population totals for each stratum. Accepts either:
+
+        * a ``dict``-like keyed by tuples whose length equals
+          ``len(strata_cols)`` (e.g. ``{('M', '<50'): 1200, ...}``), or
+        * a ``pandas.Series`` indexed by those tuples
+          (a ``MultiIndex.Series``).
+
+    Returns
+    -------
+    pandas.Series
+        Calibrated weights, aligned to ``data.index``.
+
+    Raises
+    ------
+    KeyError
+        When a stratum present in the data is missing from ``targets``.
+    """
+    if isinstance(base_weights, str):
+        bw = pd.to_numeric(data[base_weights], errors="coerce").astype(float)
+    else:
+        bw = pd.to_numeric(base_weights, errors="coerce").astype(float)
+
+    strata_cols = list(strata_cols)
+    if not strata_cols:
+        raise ValueError("post_stratify requires at least one strata column.")
+    key = data[strata_cols].apply(
+        lambda row: tuple(row.tolist()) if len(strata_cols) > 1 else row.iloc[0],
+        axis=1,
+    )
+    weighted_totals = bw.groupby(key).sum()
+
+    targets_dict = (
+        targets.to_dict()
+        if isinstance(targets, pd.Series)
+        else dict(targets)
+    )
+    # Allow scalar keys when there's only one strata column.
+    if len(strata_cols) == 1:
+        targets_dict = {(k if isinstance(k, tuple) else (k,)): v
+                        for k, v in targets_dict.items()}
+        key = key.map(lambda x: x if isinstance(x, tuple) else (x,))
+        weighted_totals.index = [
+            (i if isinstance(i, tuple) else (i,))
+            for i in weighted_totals.index
+        ]
+
+    missing = [
+        k for k in weighted_totals.index
+        if (k if isinstance(k, tuple) else (k,)) not in targets_dict
+    ]
+    if missing:
+        raise KeyError(f"post_stratify: targets missing strata {missing[:5]}...")
+
+    scale_map = {
+        stratum: float(targets_dict[stratum if isinstance(stratum, tuple) else (stratum,)])
+        / float(total)
+        if total > 0 else 0.0
+        for stratum, total in weighted_totals.items()
+    }
+    return bw * key.map(scale_map).astype(float)
+
+
+def rake(
+    data: pd.DataFrame,
+    base_weights: pd.Series | str,
+    *,
+    margins: Mapping[str, Mapping[object, float]],
+    max_iter: int = 50,
+    tol: float = 1e-6,
+) -> pd.Series:
+    """Raking (iterative proportional fitting) over marginal targets.
+
+    Parameters
+    ----------
+    data
+        Source dataframe.
+    base_weights
+        Either the column name in ``data`` or an aligned Series.
+    margins
+        Mapping of *variable → {level: target_total}*. Each variable's
+        targets are summed during one iteration; the algorithm cycles
+        through the variables until the weights stabilise.
+    max_iter
+        Maximum number of full sweeps over ``margins``.
+    tol
+        Convergence threshold on the largest relative change in any
+        weight between iterations.
+
+    Returns
+    -------
+    pandas.Series
+        Calibrated weights aligned to ``data.index``.
+    """
+    if isinstance(base_weights, str):
+        w = pd.to_numeric(data[base_weights], errors="coerce").astype(float)
+    else:
+        w = pd.to_numeric(base_weights, errors="coerce").astype(float)
+    w = w.copy()
+
+    if not margins:
+        return w
+
+    # Validate columns / targets up front.
+    for col, target_levels in margins.items():
+        if col not in data.columns:
+            raise KeyError(f"rake: column {col!r} not in data")
+        present = set(data[col].dropna().unique())
+        missing_targets = present - set(target_levels)
+        if missing_targets:
+            raise KeyError(
+                f"rake: column {col!r} has levels with no target: {sorted(missing_targets)[:5]}"
+            )
+
+    for _ in range(max_iter):
+        max_rel = 0.0
+        for col, target_levels in margins.items():
+            for lvl, target in target_levels.items():
+                mask = data[col] == lvl
+                total = float(w[mask].sum())
+                if total <= 0:
+                    continue
+                factor = float(target) / total
+                old = w[mask].copy()
+                w.loc[mask] = old * factor
+                # Track the worst relative change for convergence.
+                if old.sum() > 0:
+                    rel = abs((w[mask].sum() - old.sum()) / old.sum())
+                    max_rel = max(max_rel, rel)
+        if max_rel < tol:
+            break
+
+    return w
+
+
+def design_effect(weights: pd.Series) -> float:
+    """Kish's design-effect estimate: ``DEFF ≈ n · Σw² / (Σw)²``.
+
+    A quick QC check after calibration — large DEFF (≫ 1) means the
+    weights are highly variable and effective sample size is low.
+
+    Negative weights are not meaningful in a design context (they would
+    flip the contribution of a row), so they are excluded from the
+    computation. If any are present, a ``UserWarning`` flags how many
+    rows were dropped — matching the same behaviour as ``tbl_one(...,
+    weights=...)``. Returns ``nan`` when no positive weights remain.
+    """
+    w_raw = pd.to_numeric(weights, errors="coerce").dropna()
+    n_negative = int((w_raw < 0).sum())
+    if n_negative:
+        import warnings
+        warnings.warn(
+            f"design_effect: weights column contains {n_negative} negative "
+            "value(s); rows with negative weight are excluded from the "
+            "design-effect estimate.",
+            UserWarning,
+            stacklevel=2,
+        )
+    w = w_raw[w_raw > 0]
+    if w.empty:
+        return float("nan")
+    n = len(w)
+    return float(n * (w ** 2).sum() / (w.sum() ** 2))
+
+
+# silence unused-import lint
+_ = np

pysofra/summary/design.py

@@ -0,0 +1,246 @@
+"""Survey design object for variance estimation under complex sampling.
+
+The :class:`SurveyDesign` dataclass mirrors the headline fields of R's
+``survey::svydesign``:
+
+* ``weights`` — column carrying the sampling weight for each row.
+* ``strata`` — optional stratification variable. Within each stratum,
+  PSUs are assumed independent; variance is summed across strata.
+* ``cluster`` — optional primary-sampling-unit (PSU) variable. When
+  given, the variance of any estimator is computed across cluster
+  totals rather than individual observations (Taylor linearization
+  for the mean).
+* ``fpc`` — optional finite-population-correction column (population
+  size in each stratum, or per-cluster if no strata). Used to scale
+  the variance by ``(1 - n/N)`` per stratum.
+
+This is a *first-order* implementation: it covers what the vast
+majority of survey-weighted clinical / epidemiology pipelines need
+(stratified single-stage and clustered single-stage designs with FPC).
+Multi-stage designs and post-stratification calibration remain on the
+roadmap.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+import pandas as pd
+
+
+@dataclass(frozen=True)
+class SurveyDesign:
+    """Column-name bundle describing a survey-design structure.
+
+    ``cluster`` accepts either a single column name (single-stage
+    cluster sampling) or a tuple of names (multi-stage). For multi-stage
+    designs PySofra currently uses the outermost PSU for variance
+    estimation and a footnote will name the second-stage column as
+    "nested within" — full multi-stage Taylor linearisation is planned.
+
+    ``replicate_weights`` and ``replicate_scale`` enable the jackknife
+    family of variance estimators: every replicate column carries
+    weights with one PSU dropped, and the variance is computed as
+    ``replicate_scale * Σ (θ̂_r − θ̂)²``.  The ``"jk1"`` default sets
+    ``replicate_scale`` to ``(n − 1)/n`` automatically.
+    """
+
+    weights: str
+    strata: str | None = None
+    cluster: str | tuple[str, ...] | None = None
+    fpc: str | None = None
+    replicate_weights: tuple[str, ...] | None = None
+    replicate_type: str = "jk1"   # 'jk1' | 'jkn' | 'bootstrap'
+
+    @property
+    def primary_cluster(self) -> str | None:
+        if self.cluster is None:
+            return None
+        if isinstance(self.cluster, tuple):
+            return self.cluster[0] if self.cluster else None
+        return self.cluster
+
+    def validate(self, data: pd.DataFrame) -> None:
+        for name, col in (("weights", self.weights),
+                          ("strata", self.strata),
+                          ("fpc", self.fpc)):
+            if col is not None and col not in data.columns:
+                raise KeyError(f"{name} column {col!r} not in data")
+        if self.cluster is not None:
+            cluster_cols = (
+                self.cluster if isinstance(self.cluster, tuple)
+                else (self.cluster,)
+            )
+            for c in cluster_cols:
+                if c not in data.columns:
+                    raise KeyError(f"cluster column {c!r} not in data")
+        if self.replicate_weights is not None:
+            missing = [c for c in self.replicate_weights if c not in data.columns]
+            if missing:
+                raise KeyError(f"replicate_weights columns not in data: {missing}")
+            if self.replicate_type not in ("jk1", "jkn", "bootstrap"):
+                raise ValueError(
+                    f"replicate_type must be 'jk1', 'jkn', or 'bootstrap'; "
+                    f"got {self.replicate_type!r}."
+                )
+
+
+# ----------------------------------------------------------------------
+# Variance estimators
+# ----------------------------------------------------------------------
+
+
+def design_mean_var(
+    values: pd.Series,
+    weights: pd.Series,
+    *,
+    strata: pd.Series | None = None,
+    cluster: pd.Series | None = None,
+    fpc: pd.Series | None = None,
+) -> tuple[float, float, float]:
+    """Estimate the survey-weighted mean and its design-based variance.
+
+    Returns ``(mean, variance, n_eff)``.
+
+    For a simple stratified design, this implements the Taylor-series
+    linearisation:
+
+        Var(ŷ) = Σ_h (1 - f_h) · (n_h / (n_h - 1)) · Σ_{i in h} (w_i (y_i - ŷ))²
+                                                       /  (Σ w_i)²
+
+    For a clustered design (no strata, single-stage), the variance is
+    computed across PSU totals.
+
+    When both strata and clusters are given, the formula nests cluster
+    variance within strata.
+    """
+    v = pd.to_numeric(values, errors="coerce").astype(float)
+    w = pd.to_numeric(weights, errors="coerce").astype(float)
+    mask = ~(v.isna() | w.isna()) & (w > 0)
+    v = v[mask]
+    w = w[mask]
+    if strata is not None:
+        strata = strata[mask].reset_index(drop=True)
+    if cluster is not None:
+        cluster = cluster[mask].reset_index(drop=True)
+    if fpc is not None:
+        fpc = pd.to_numeric(fpc, errors="coerce")[mask].reset_index(drop=True)
+    v = v.reset_index(drop=True)
+    w = w.reset_index(drop=True)
+
+    total_w = float(w.sum())
+    if total_w <= 0 or v.size == 0:
+        return float("nan"), float("nan"), 0.0
+
+    mean = float((w * v).sum() / total_w)
+    n = int(v.size)
+
+    # Residuals for the mean estimator.
+    e = w.to_numpy() * (v.to_numpy() - mean)
+
+    if strata is None and cluster is None:
+        var_num = float(np.sum(e ** 2)) * (n / max(n - 1, 1))
+    elif strata is None:
+        assert cluster is not None
+        # Single-stage clusters; sum residuals within each cluster, take
+        # variance across cluster totals.
+        s_per_cluster = pd.Series(e).groupby(cluster).sum().to_numpy()
+        n_clust = int(s_per_cluster.size)
+        if n_clust > 1:
+            mean_cluster = float(s_per_cluster.mean())
+            var_num = float(np.sum((s_per_cluster - mean_cluster) ** 2)) \
+                * (n_clust / (n_clust - 1))
+        else:
+            var_num = 0.0
+    else:
+        # Stratified, possibly with clusters within strata.
+        var_num = 0.0
+        s = pd.Series(e)
+        strata_series = pd.Series(strata)
+        for _stratum, idx in strata_series.groupby(strata_series).indices.items():
+            idx_arr = np.asarray(idx)
+            e_h = s.iloc[idx_arr].to_numpy()
+            if cluster is not None:
+                c_h = cluster.iloc[idx_arr].to_numpy()
+                psu_totals = pd.Series(e_h).groupby(c_h).sum().to_numpy()
+                n_h = int(psu_totals.size)
+                if n_h > 1:
+                    mean_h = float(psu_totals.mean())
+                    contrib = float(np.sum((psu_totals - mean_h) ** 2)) \
+                        * (n_h / (n_h - 1))
+                else:
+                    contrib = 0.0
+            else:
+                n_h = int(e_h.size)
+                contrib = (
+                    float(np.sum(e_h ** 2)) * (n_h / (n_h - 1))
+                    if n_h > 1
+                    else 0.0
+                )
+
+            if fpc is not None and idx_arr.size > 0:
+                fpc_h = float(fpc.iloc[idx_arr].iloc[0])
+                # FPC = 1 - n/N
+                f_h = min(1.0, idx_arr.size / max(fpc_h, 1.0))
+                contrib *= 1.0 - f_h
+
+            var_num += contrib
+
+    variance = var_num / (total_w ** 2)
+    return mean, variance, total_w
+
+
+# ----------------------------------------------------------------------
+# Replicate-weight variance
+# ----------------------------------------------------------------------
+
+
+def replicate_mean_var(
+    values: pd.Series,
+    base_weights: pd.Series,
+    replicate_weights: list[pd.Series] | tuple[pd.Series, ...],
+    *,
+    replicate_type: str = "jk1",
+) -> tuple[float, float, float]:
+    """Variance of a weighted mean from replicate weights.
+
+    The full-sample estimator uses ``base_weights``; each replicate gives
+    a perturbed estimate, and the variance is
+
+        Var(θ̂) = c · Σ_r (θ̂_r − θ̂)²
+
+    where ``c`` is the replicate-type scaling: ``(R-1)/R`` for ``jk1``,
+    ``1/R`` for ``bootstrap``. ``jkn`` (BRR/stratified jackknife) uses
+    ``(R-1)/R`` too; users who need a different scale should pass the
+    appropriate replicate weights and use ``bootstrap`` for the
+    unscaled form.
+    """
+    v = pd.to_numeric(values, errors="coerce").astype(float)
+    bw = pd.to_numeric(base_weights, errors="coerce").astype(float)
+    mask = ~(v.isna() | bw.isna()) & (bw > 0)
+    v = v[mask].reset_index(drop=True)
+    bw = bw[mask].reset_index(drop=True)
+
+    total_w = float(bw.sum())
+    if total_w <= 0 or v.empty:
+        return float("nan"), float("nan"), 0.0
+    theta_hat = float((v * bw).sum() / total_w)
+
+    R = len(replicate_weights)
+    if R == 0:
+        return theta_hat, 0.0, total_w
+
+    sq_dev = 0.0
+    for rw in replicate_weights:
+        rw_arr = pd.to_numeric(rw, errors="coerce").astype(float)
+        rw_arr = rw_arr[mask].reset_index(drop=True)
+        w_pos = rw_arr.where(rw_arr > 0, 0.0)
+        denom = float(w_pos.sum())
+        if denom <= 0:
+            continue
+        theta_r = float((v * w_pos).sum() / denom)
+        sq_dev += (theta_r - theta_hat) ** 2
+
+    scale = 1.0 / R if replicate_type == "bootstrap" else (R - 1.0) / R
+    return theta_hat, scale * sq_dev, total_w

pysofra/summary/effect_size.py

@@ -0,0 +1,187 @@
+"""Effect-size helpers — Cohen's d, Hedges' g, Cramér's V, eta-squared.
+
+Companion functions to the inferential tests in
+:mod:`pysofra.summary.tests`. Effect sizes describe the *magnitude* of
+a difference / association independently of sample size, and are
+frequently requested alongside the p-values in clinical reports.
+
+All functions accept aligned :class:`pandas.Series` for ``values`` /
+``groups``; missing values are dropped pairwise. They return floats
+(or ``None`` for degenerate input).
+"""
+
+from __future__ import annotations
+
+import math
+
+import numpy as np
+import pandas as pd
+
+# ----------------------------------------------------------------------
+# Continuous
+# ----------------------------------------------------------------------
+
+def cohen_d(a: pd.Series | np.ndarray, b: pd.Series | np.ndarray) -> float | None:
+    """Cohen's d using the pooled standard deviation.
+
+    ``d = (μ₁ − μ₂) / s_pool``, where the pooled SD weights the two
+    samples by their degrees of freedom.
+    """
+    a_arr = pd.to_numeric(pd.Series(a), errors="coerce").dropna().to_numpy(dtype=float)
+    b_arr = pd.to_numeric(pd.Series(b), errors="coerce").dropna().to_numpy(dtype=float)
+    n_a, n_b = a_arr.size, b_arr.size
+    if n_a < 2 or n_b < 2:
+        return None
+    v_a = float(np.var(a_arr, ddof=1))
+    v_b = float(np.var(b_arr, ddof=1))
+    s_pool = math.sqrt(((n_a - 1) * v_a + (n_b - 1) * v_b) / (n_a + n_b - 2))
+    if s_pool == 0:
+        return 0.0 if a_arr.mean() == b_arr.mean() else float("inf")
+    return (float(a_arr.mean()) - float(b_arr.mean())) / s_pool
+
+
+def hedges_g(a: pd.Series | np.ndarray, b: pd.Series | np.ndarray) -> float | None:
+    """Hedges' g — Cohen's d with the small-sample bias correction.
+
+    ``g = d · J``, where ``J ≈ 1 − 3/(4(n_a+n_b) − 9)`` (Hedges 1981).
+    """
+    d = cohen_d(a, b)
+    if d is None or math.isinf(d):
+        return d
+    n_a = int(pd.to_numeric(pd.Series(a), errors="coerce").dropna().size)
+    n_b = int(pd.to_numeric(pd.Series(b), errors="coerce").dropna().size)
+    denom = 4 * (n_a + n_b) - 9
+    if denom <= 0:  # pragma: no cover — unreachable given cohen_d's n>=2 guard
+        return d
+    j = 1.0 - 3.0 / denom
+    return d * j
+
+
+def eta_squared(values: pd.Series, groups: pd.Series) -> float | None:
+    """One-way ANOVA effect size: between-group / total sum-of-squares.
+
+    Ranges ``[0, 1]``. Small ≈ 0.01, medium ≈ 0.06, large ≈ 0.14
+    (Cohen 1988).
+    """
+    df = pd.DataFrame({"v": pd.to_numeric(values, errors="coerce"),
+                       "g": groups}).dropna()
+    if df.empty:
+        return None
+    grand = float(df["v"].mean())
+    ss_between = float((df.groupby("g")["v"]
+                         .apply(lambda x: x.size * (x.mean() - grand) ** 2))
+                        .sum())
+    ss_total = float(((df["v"] - grand) ** 2).sum())
+    if ss_total <= 0:
+        return 0.0
+    return ss_between / ss_total
+
+
+def omega_squared(values: pd.Series, groups: pd.Series) -> float | None:
+    """Less-biased counterpart to ``eta_squared`` (Hays 1973)."""
+    df = pd.DataFrame({"v": pd.to_numeric(values, errors="coerce"),
+                       "g": groups}).dropna()
+    if df.empty:
+        return None
+    k = int(df["g"].nunique())
+    n = int(df.shape[0])
+    if n - k <= 0 or k <= 1:
+        return None
+    grand = float(df["v"].mean())
+    ss_between = float((df.groupby("g")["v"]
+                         .apply(lambda x: x.size * (x.mean() - grand) ** 2))
+                        .sum())
+    ss_total = float(((df["v"] - grand) ** 2).sum())
+    if ss_total <= 0:
+        return 0.0
+    ms_within = (ss_total - ss_between) / (n - k)
+    omega = (ss_between - (k - 1) * ms_within) / (ss_total + ms_within)
+    return float(max(0.0, omega))
+
+
+# ----------------------------------------------------------------------
+# Categorical
+# ----------------------------------------------------------------------
+
+def cramers_v(values: pd.Series, groups: pd.Series) -> float | None:
+    """Cramér's V — chi-square effect size normalised to ``[0, 1]``.
+
+    ``V = √(χ² / (N · (min(R, C) − 1)))``.
+    """
+    import warnings as _w
+
+    from scipy import stats as sp_stats
+    df = pd.DataFrame({"v": values, "g": groups}).dropna()
+    if df.empty:
+        return None
+    ctab = pd.crosstab(df["v"], df["g"])
+    if ctab.shape[0] < 2 or ctab.shape[1] < 2:
+        return None
+    with np.errstate(invalid="ignore", over="ignore", divide="ignore"), \
+            _w.catch_warnings():
+        _w.simplefilter("ignore", RuntimeWarning)
+        chi2, _, _, _ = sp_stats.chi2_contingency(ctab.to_numpy(), correction=False)
+    n = float(ctab.values.sum())
+    if n <= 0:  # pragma: no cover — guarded by the shape >= 2x2 check above
+        return None
+    min_dim = min(ctab.shape) - 1
+    if min_dim <= 0:  # pragma: no cover — shape >= 2x2 guarantees min_dim >= 1
+        return None
+    return float(math.sqrt(chi2 / (n * min_dim)))
+
+
+def phi_coefficient(values: pd.Series, groups: pd.Series) -> float | None:
+    """Phi — Cramér's V special case for 2×2 tables; ``φ = √(χ²/N)``."""
+    import warnings as _w
+
+    from scipy import stats as sp_stats
+    df = pd.DataFrame({"v": values, "g": groups}).dropna()
+    if df.empty:
+        return None
+    ctab = pd.crosstab(df["v"], df["g"])
+    if ctab.shape != (2, 2):
+        return None
+    with np.errstate(invalid="ignore", over="ignore", divide="ignore"), \
+            _w.catch_warnings():
+        _w.simplefilter("ignore", RuntimeWarning)
+        chi2, _, _, _ = sp_stats.chi2_contingency(ctab.to_numpy(), correction=False)
+    n = float(ctab.values.sum())
+    if n <= 0:  # pragma: no cover — guarded by the (2,2) shape check above
+        return None
+    return float(math.sqrt(chi2 / n))
+
+
+# ----------------------------------------------------------------------
+# Auto dispatch (mirrors auto-test selection)
+# ----------------------------------------------------------------------
+
+
+def auto_effect_size(values: pd.Series, groups: pd.Series) -> tuple[str, float | None]:
+    """Pick a sensible effect size for the variable kind / number of groups.
+
+    Returns ``(name, value)`` so callers can both display the metric and
+    label it in a footnote.
+    """
+    g_unique = pd.Series(groups).dropna().unique()
+    n_groups = len(g_unique)
+
+    # Continuous-looking?
+    try:
+        pd.to_numeric(values, errors="raise")
+        continuous = True
+    except (ValueError, TypeError):
+        continuous = False
+
+    if continuous and n_groups == 2:
+        return "Cohen's d", cohen_d(
+            values[groups == g_unique[0]],
+            values[groups == g_unique[1]],
+        )
+    if continuous and n_groups >= 3:
+        return "η²", eta_squared(values, groups)
+    if not continuous and n_groups >= 2:
+        ctab = pd.crosstab(values, groups)
+        if ctab.shape == (2, 2):
+            return "φ", phi_coefficient(values, groups)
+        return "Cramér's V", cramers_v(values, groups)
+    return "—", None