weightpipe 0.1.0__py3-none-any.whl

weightpipe/__init__.py

@@ -0,0 +1,83 @@
+"""weightpipe: declarative survey weighting with recipe-aware replicates."""
+
+from weightpipe._logging import set_log_level, setup_logging
+from weightpipe.design import Design
+from weightpipe.diagnostics import BalanceReport, balance, design_effect, ess, margins
+from weightpipe.estimate import Estimation, estimate, estimate_glm, point_estimate
+from weightpipe.frame import WeightFrame
+from weightpipe.methods import design_matrix, population_totals
+from weightpipe.pipeline import WeightPipe
+from weightpipe.planning import (
+    allocate_strata,
+    allocation_table,
+    margin_of_error,
+    sample_size,
+    stratified_margin_of_error,
+)
+from weightpipe.recipe import Recipe
+from weightpipe.replicates import (
+    BootstrapResult,
+    JackknifeResult,
+    boot_mean,
+    boot_median,
+    boot_proportion,
+    boot_ratio,
+    boot_total,
+    bootstrap_estimate,
+    bootstrap_weights,
+    jack_mean,
+    jack_median,
+    jack_proportion,
+    jack_ratio,
+    jack_total,
+    jackknife_estimate,
+    jackknife_weights,
+)
+from weightpipe.result import WeightResult, collect_weights, weight_factors
+
+__all__ = [
+    "BalanceReport",
+    "BootstrapResult",
+    "Design",
+    "Estimation",
+    "JackknifeResult",
+    "Recipe",
+    "WeightFrame",
+    "WeightPipe",
+    "WeightResult",
+    "__version__",
+    "allocate_strata",
+    "allocation_table",
+    "balance",
+    "boot_mean",
+    "boot_median",
+    "boot_proportion",
+    "boot_ratio",
+    "boot_total",
+    "bootstrap_estimate",
+    "bootstrap_weights",
+    "collect_weights",
+    "design_effect",
+    "design_matrix",
+    "ess",
+    "estimate",
+    "estimate_glm",
+    "jack_mean",
+    "jack_median",
+    "jack_proportion",
+    "jack_ratio",
+    "jack_total",
+    "jackknife_estimate",
+    "jackknife_weights",
+    "margin_of_error",
+    "margins",
+    "point_estimate",
+    "population_totals",
+    "sample_size",
+    "set_log_level",
+    "setup_logging",
+    "stratified_margin_of_error",
+    "weight_factors",
+]
+
+__version__ = "0.1.0"

weightpipe/_logging.py

@@ -0,0 +1,50 @@
+"""Package logging: quiet by default, opt-in readable output."""
+
+import logging
+import sys
+from typing import TextIO
+
+ROOT_LOGGER_NAME = "weightpipe"
+LOG_FORMAT = "%(asctime)s  %(message)s"
+DATE_FORMAT = "%Y-%m-%d %H:%M:%S"
+
+_HANDLER_FLAG = "_weightpipe_handler"
+
+_root = logging.getLogger(ROOT_LOGGER_NAME)
+# Libraries should not emit anything unless the application asks for it.
+_root.addHandler(logging.NullHandler())
+
+
+def get_logger(name: str) -> logging.Logger:
+    """Logger for a module inside the package (``weightpipe.<module>``)."""
+    return logging.getLogger(name)
+
+
+def setup_logging(
+    level: int | str = "INFO",
+    *,
+    stream: TextIO | None = None,
+    fmt: str = LOG_FORMAT,
+    datefmt: str = DATE_FORMAT,
+) -> logging.Logger:
+    """Send weightpipe messages to ``stream`` with a compact format.
+
+    Calling this repeatedly replaces the handler instead of stacking new ones,
+    and it leaves the application's root logger configuration untouched.
+    """
+    for handler in [h for h in _root.handlers if getattr(h, _HANDLER_FLAG, False)]:
+        _root.removeHandler(handler)
+
+    handler = logging.StreamHandler(stream if stream is not None else sys.stderr)
+    handler.setFormatter(logging.Formatter(fmt=fmt, datefmt=datefmt))
+    setattr(handler, _HANDLER_FLAG, True)
+
+    _root.addHandler(handler)
+    _root.setLevel(level)
+    _root.propagate = False
+    return _root
+
+
+def set_log_level(level: int | str) -> None:
+    """Change the verbosity of weightpipe messages."""
+    _root.setLevel(level)

weightpipe/design.py

@@ -0,0 +1,233 @@
+"""Sampling design: base weights and variance structure (strata / PSU)."""
+
+from collections.abc import Sequence
+from dataclasses import dataclass, field
+from typing import Any
+
+import numpy as np
+import pandas as pd
+
+from weightpipe._logging import get_logger
+
+logger = get_logger(__name__)
+
+
+def _as_col_list(value: str | Sequence[str] | None) -> list[str] | None:
+    if value is None:
+        return None
+    if isinstance(value, str):
+        return [value]
+    cols = [str(c) for c in value]
+    if not cols:
+        raise ValueError("stage column list must be non-empty")
+    return cols
+
+
+def _infer_kind(
+    *,
+    N: Any,
+    N_h: Any,
+    weight: str | None,
+    probabilities: list[str] | None,
+    stage_weights: list[str] | None,
+    strata: str | None,
+    psu: str | None,
+) -> str:
+    if N is not None:
+        return "srs"
+    if N_h is not None:
+        return "stratified"
+    if probabilities is not None or stage_weights is not None:
+        return "stratified_multistage" if strata is not None else "multistage"
+    if psu is not None:
+        return "stratified_cluster" if strata is not None else "cluster"
+    return "custom"
+
+
+def _product_weights(
+    frame: pd.DataFrame,
+    columns: list[str],
+    *,
+    as_probabilities: bool,
+) -> np.ndarray:
+    missing = [c for c in columns if c not in frame.columns]
+    if missing:
+        raise KeyError(f"stage columns not found: {missing}")
+    mat = frame.loc[:, columns].to_numpy(dtype=float)
+    if not np.isfinite(mat).all():
+        raise ValueError("stage columns must be finite")
+    if as_probabilities:
+        if np.any((mat <= 0) | (mat > 1)):
+            raise ValueError("inclusion probabilities must be in (0, 1]")
+        with np.errstate(divide="raise", invalid="raise"):
+            return 1.0 / np.prod(mat, axis=1)
+    if np.any(mat < 0):
+        raise ValueError("stage weights must be non-negative")
+    return np.prod(mat, axis=1)
+
+
+@dataclass(frozen=True, init=False)
+class Design:
+    """Sampling design metadata and base weights.
+
+    Pass the inputs that define the design; ``kind`` is inferred:
+
+    - *(none)* → unit weights ``1`` (with a log message)
+    - ``N=...`` → SRS, weights ``N/n``
+    - ``strata=...``, ``N_h=...`` → stratified SRS, weights ``N_h/n_h``
+    - ``weight=...``, ``psu=...`` → cluster (add ``strata=`` if stratified)
+    - ``probabilities=[...]``, ``psu=...`` → multi-stage, ``w = 1 / ∏ π_k``
+    - ``stage_weights=[...]``, ``psu=...`` → multi-stage, ``w = ∏ w_k``
+    - ``weight=...`` only → use existing weights
+
+    For multi-stage / cluster designs, ``psu`` is the ultimate cluster used for
+    bootstrap/jackknife variance. Stage details are folded into the weight.
+
+    Base weights live in ``data[weight]``.
+    """
+
+    data: pd.DataFrame
+    weight: str
+    strata: str | None = None
+    psu: str | None = None
+    kind: str = "custom"
+    meta: dict[str, Any] = field(default_factory=dict)
+
+    def __init__(
+        self,
+        data: pd.DataFrame,
+        *,
+        weight: str | None = None,
+        strata: str | None = None,
+        psu: str | None = None,
+        N: float | int | None = None,
+        N_h: dict[Any, float] | pd.Series | None = None,
+        probabilities: str | Sequence[str] | None = None,
+        stage_weights: str | Sequence[str] | None = None,
+        copy: bool = True,
+    ) -> None:
+        prob_cols = _as_col_list(probabilities)
+        stage_cols = _as_col_list(stage_weights)
+        sources = (N, N_h, weight, prob_cols, stage_cols)
+        n_sources = sum(x is not None for x in sources)
+        if n_sources > 1:
+            raise ValueError("provide at most one of N=, N_h=, weight=, probabilities=, or stage_weights=")
+        if N is not None and (strata is not None or psu is not None):
+            raise ValueError("N= (SRS) cannot be combined with strata= or psu=")
+        if N_h is not None and strata is None:
+            raise ValueError("N_h= requires strata=")
+        if N_h is not None and psu is not None:
+            raise ValueError("N_h= stratified SRS cannot be combined with psu=")
+        if (prob_cols is not None or stage_cols is not None) and psu is None:
+            raise ValueError("multi-stage designs require psu= (ultimate cluster for variance)")
+
+        frame = data.copy() if copy else data
+        meta: dict[str, Any] = {}
+        weight_col: str
+        kind = _infer_kind(
+            N=N,
+            N_h=N_h,
+            weight=weight,
+            probabilities=prob_cols,
+            stage_weights=stage_cols,
+            strata=strata,
+            psu=psu,
+        )
+
+        if n_sources == 0:
+            weight_col = "base_weight"
+            if len(frame) < 1:
+                raise ValueError("data must have at least one row")
+            if weight_col in frame.columns:
+                raise ValueError(f"column {weight_col!r} already exists")
+            frame = frame.assign(**{weight_col: 1.0})
+            meta = {"unit_weights": True, "note": "no design weight provided; using 1.0"}
+            logger.info("No design weight provided; using base_weight=1.0 for all rows")
+        elif N is not None:
+            weight_col = "base_weight"
+            n = len(frame)
+            if n < 1:
+                raise ValueError("data must have at least one row")
+            if float(N) <= 0:
+                raise ValueError("N must be positive")
+            if weight_col in frame.columns:
+                raise ValueError(f"column {weight_col!r} already exists")
+            frame = frame.assign(**{weight_col: float(N) / n})
+            meta = {"N": float(N), "n": n}
+            strata = None
+            psu = None
+        elif N_h is not None:
+            assert strata is not None
+            weight_col = "base_weight"
+            if strata not in frame.columns:
+                raise KeyError(f"strata column not found: {strata}")
+            if weight_col in frame.columns:
+                raise ValueError(f"column {weight_col!r} already exists")
+            n_h = frame.groupby(strata, observed=True).size()
+            pop = {str(k): float(v) for k, v in dict(N_h).items()}
+            weights = np.empty(len(frame), dtype=float)
+            meta_nh: dict[str, dict[str, float]] = {}
+            labels = frame[strata].astype(str)
+            for label, n in n_h.items():
+                key = str(label)
+                if key not in pop:
+                    raise KeyError(f"N_h missing stratum {label!r}")
+                Nh = pop[key]
+                if n < 1 or Nh <= 0:
+                    raise ValueError(f"invalid N_h/n_h for stratum {label!r}: N={Nh}, n={n}")
+                w = Nh / float(n)
+                weights[labels.to_numpy() == key] = w
+                meta_nh[key] = {"N": Nh, "n": float(n), "weight": w}
+            frame = frame.assign(**{weight_col: weights})
+            meta = {"N_h": meta_nh}
+            psu = None
+        elif prob_cols is not None or stage_cols is not None:
+            weight_col = "base_weight"
+            if weight_col in frame.columns:
+                raise ValueError(f"column {weight_col!r} already exists")
+            if prob_cols is not None:
+                weights = _product_weights(frame, prob_cols, as_probabilities=True)
+                meta = {
+                    "stages": len(prob_cols),
+                    "probabilities": list(prob_cols),
+                    "formula": "1 / prod(probabilities)",
+                }
+            else:
+                assert stage_cols is not None
+                weights = _product_weights(frame, stage_cols, as_probabilities=False)
+                meta = {
+                    "stages": len(stage_cols),
+                    "stage_weights": list(stage_cols),
+                    "formula": "prod(stage_weights)",
+                }
+            frame = frame.assign(**{weight_col: weights})
+        else:
+            assert weight is not None
+            weight_col = weight
+            if weight_col not in frame.columns:
+                raise KeyError(f"weight column not found: {weight_col}")
+
+        if strata is not None and strata not in frame.columns:
+            raise KeyError(f"strata column not found: {strata}")
+        if psu is not None and psu not in frame.columns:
+            raise KeyError(f"psu column not found: {psu}")
+        w = frame[weight_col]
+        if (w < 0).any():
+            raise ValueError("design weights must be non-negative")
+
+        object.__setattr__(self, "data", frame)
+        object.__setattr__(self, "weight", weight_col)
+        object.__setattr__(self, "strata", strata)
+        object.__setattr__(self, "psu", psu)
+        object.__setattr__(self, "kind", kind)
+        object.__setattr__(self, "meta", meta)
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "weight": self.weight,
+            "strata": self.strata,
+            "psu": self.psu,
+            "kind": self.kind,
+            "meta": dict(self.meta),
+            "n": len(self.data),
+        }

weightpipe/diagnostics/__init__.py

@@ -0,0 +1,52 @@
+"""Weight / design-effect / balance diagnostics."""
+
+import numpy as np
+import pandas as pd
+
+from weightpipe.diagnostics.balance import BalanceReport, balance
+from weightpipe.diagnostics.margins import (
+    attach_margin_table,
+    margin_table_from_targets,
+    margins,
+    weighted_category_margins,
+)
+from weightpipe.result import WeightResult
+
+__all__ = [
+    "BalanceReport",
+    "attach_margin_table",
+    "balance",
+    "design_effect",
+    "ess",
+    "margin_table_from_targets",
+    "margins",
+    "weighted_category_margins",
+]
+
+
+def design_effect(weights: pd.Series | WeightResult) -> float:
+    """Kish design effect from unequal weighting: 1 + CV(w)^2 = n * sum(w^2) / sum(w)^2."""
+    w = weights.weights if isinstance(weights, WeightResult) else weights
+    w = w.astype(float)
+    active = w[w > 0]
+    if active.empty:
+        return float("nan")
+    n = float(len(active))
+    s = float(active.sum())
+    if s <= 0:
+        return float("nan")
+    return float(n * float(np.sum(np.square(active.to_numpy()))) / (s * s))
+
+
+def ess(weights: pd.Series | WeightResult) -> float:
+    """Effective sample size under Kish: (sum w)^2 / sum(w^2)."""
+    w = weights.weights if isinstance(weights, WeightResult) else weights
+    w = w.astype(float)
+    active = w[w > 0]
+    if active.empty:
+        return float("nan")
+    s = float(active.sum())
+    ss = float(np.sum(np.square(active.to_numpy())))
+    if ss <= 0:
+        return float("nan")
+    return float((s * s) / ss)