openstat-cli 1.0.0__py3-none-any.whl

openstat/stats/resampling.py

@@ -0,0 +1,203 @@
+"""Bootstrap and permutation test statistics."""
+
+from __future__ import annotations
+
+import numpy as np
+import polars as pl
+
+
+def bootstrap_ci(
+    df: pl.DataFrame,
+    col: str,
+    stat: str = "mean",
+    n_boot: int = 2000,
+    ci: float = 0.95,
+    seed: int = 42,
+) -> dict:
+    """Bootstrap confidence interval for a statistic."""
+    rng = np.random.default_rng(seed)
+    data = df[col].drop_nulls().to_numpy().astype(float)
+    n = len(data)
+
+    stat_fns = {
+        "mean": np.mean,
+        "median": np.median,
+        "std": np.std,
+        "var": np.var,
+        "min": np.min,
+        "max": np.max,
+    }
+    if stat not in stat_fns:
+        raise ValueError(f"Unknown statistic: {stat}. Use: {', '.join(stat_fns)}")
+
+    fn = stat_fns[stat]
+    observed = float(fn(data))
+
+    boot_stats = np.array([
+        fn(rng.choice(data, size=n, replace=True))
+        for _ in range(n_boot)
+    ])
+
+    alpha = 1 - ci
+    lo = float(np.quantile(boot_stats, alpha / 2))
+    hi = float(np.quantile(boot_stats, 1 - alpha / 2))
+    se = float(boot_stats.std())
+    bias = float(boot_stats.mean() - observed)
+
+    return {
+        "test": f"Bootstrap CI ({stat})",
+        "col": col,
+        "stat": stat,
+        "observed": observed,
+        "n_obs": n,
+        "n_boot": n_boot,
+        "ci_level": ci,
+        "ci_lo": lo,
+        "ci_hi": hi,
+        "se_boot": se,
+        "bias": bias,
+    }
+
+
+def bootstrap_diff(
+    df: pl.DataFrame,
+    col: str,
+    by: str,
+    stat: str = "mean",
+    n_boot: int = 2000,
+    ci: float = 0.95,
+    seed: int = 42,
+) -> dict:
+    """Bootstrap CI for the difference in a statistic between two groups."""
+    rng = np.random.default_rng(seed)
+    groups = df[by].drop_nulls().unique().sort().to_list()
+    if len(groups) != 2:
+        raise ValueError(f"bootstrap_diff requires exactly 2 groups, got {len(groups)}")
+
+    g1 = df.filter(pl.col(by) == groups[0])[col].drop_nulls().to_numpy().astype(float)
+    g2 = df.filter(pl.col(by) == groups[1])[col].drop_nulls().to_numpy().astype(float)
+
+    stat_fns = {
+        "mean": np.mean, "median": np.median, "std": np.std,
+        "var": np.var, "min": np.min, "max": np.max,
+    }
+    fn = stat_fns.get(stat, np.mean)
+    observed_diff = float(fn(g1) - fn(g2))
+
+    boot_diffs = np.array([
+        fn(rng.choice(g1, size=len(g1), replace=True)) -
+        fn(rng.choice(g2, size=len(g2), replace=True))
+        for _ in range(n_boot)
+    ])
+
+    alpha = 1 - ci
+    lo = float(np.quantile(boot_diffs, alpha / 2))
+    hi = float(np.quantile(boot_diffs, 1 - alpha / 2))
+    # Shift bootstrap distribution to null (mean=0) for p-value
+    boot_centered = boot_diffs - boot_diffs.mean()
+    p_value = float((np.abs(boot_centered) >= np.abs(observed_diff)).mean())
+
+    return {
+        "test": f"Bootstrap Difference ({stat})",
+        "col": col, "by": by,
+        "groups": [str(g) for g in groups],
+        "observed_diff": observed_diff,
+        "n_boot": n_boot,
+        "ci_level": ci,
+        "ci_lo": lo,
+        "ci_hi": hi,
+        "se_boot": float(boot_diffs.std()),
+        "p_value": p_value,
+    }
+
+
+def permutation_test(
+    df: pl.DataFrame,
+    col: str,
+    by: str,
+    stat: str = "mean",
+    n_perm: int = 2000,
+    alternative: str = "two-sided",
+    seed: int = 42,
+) -> dict:
+    """Permutation test for difference between two groups."""
+    rng = np.random.default_rng(seed)
+    groups = df[by].drop_nulls().unique().sort().to_list()
+    if len(groups) != 2:
+        raise ValueError(f"permutation_test requires exactly 2 groups, got {len(groups)}")
+
+    g1 = df.filter(pl.col(by) == groups[0])[col].drop_nulls().to_numpy().astype(float)
+    g2 = df.filter(pl.col(by) == groups[1])[col].drop_nulls().to_numpy().astype(float)
+
+    stat_fns = {"mean": np.mean, "median": np.median, "std": np.std}
+    fn = stat_fns.get(stat, np.mean)
+
+    observed = float(fn(g1) - fn(g2))
+    combined = np.concatenate([g1, g2])
+    n1 = len(g1)
+
+    perm_stats = np.array([
+        fn(perm := rng.permutation(combined), ) - fn(perm[n1:])  # noqa: confusing but valid
+        for _ in range(n_perm)
+    ])
+    # Fix: proper permutation
+    perm_stats = np.zeros(n_perm)
+    for i in range(n_perm):
+        perm = rng.permutation(combined)
+        perm_stats[i] = fn(perm[:n1]) - fn(perm[n1:])
+
+    if alternative == "two-sided":
+        p_value = float((np.abs(perm_stats) >= np.abs(observed)).mean())
+    elif alternative == "greater":
+        p_value = float((perm_stats >= observed).mean())
+    else:
+        p_value = float((perm_stats <= observed).mean())
+
+    return {
+        "test": "Permutation Test",
+        "col": col, "by": by,
+        "stat": stat,
+        "groups": [str(g) for g in groups],
+        "observed_diff": observed,
+        "n_perm": n_perm,
+        "alternative": alternative,
+        "p_value": p_value,
+        "reject_5pct": p_value < 0.05,
+    }
+
+
+def jackknife_ci(
+    df: pl.DataFrame,
+    col: str,
+    stat: str = "mean",
+) -> dict:
+    """Jackknife (leave-one-out) bias and standard error estimate."""
+    data = df[col].drop_nulls().to_numpy().astype(float)
+    n = len(data)
+    stat_fns = {
+        "mean": np.mean, "median": np.median, "std": np.std,
+        "var": np.var, "min": np.min, "max": np.max,
+    }
+    if stat not in stat_fns:
+        raise ValueError(f"Unknown statistic: {stat}")
+    fn = stat_fns[stat]
+    observed = float(fn(data))
+
+    jack_stats = np.array([
+        fn(np.delete(data, i))
+        for i in range(n)
+    ])
+    jack_mean = jack_stats.mean()
+    bias = float((n - 1) * (jack_mean - observed))
+    se = float(np.sqrt((n - 1) / n * np.sum((jack_stats - jack_mean) ** 2)))
+
+    return {
+        "test": f"Jackknife ({stat})",
+        "col": col,
+        "stat": stat,
+        "observed": observed,
+        "n_obs": n,
+        "bias": bias,
+        "se_jackknife": se,
+        "bias_corrected": observed - bias,
+    }

openstat/stats/survey.py

@@ -0,0 +1,213 @@
+"""Survey-weighted estimation: weighted means, WLS, Taylor linearization, DEFF."""
+
+from __future__ import annotations
+
+import numpy as np
+import polars as pl
+import statsmodels.api as sm
+from scipy import stats as sp_stats
+
+from openstat.stats.models import FitResult
+
+
+def weighted_summary(df: pl.DataFrame, cols: list[str], weight_var: str) -> str:
+    """Compute weighted summary statistics."""
+    pdf = df.to_pandas()
+    weights = pdf[weight_var].values
+
+    lines = ["Survey-Weighted Summary Statistics:"]
+    lines.append(f"{'Variable':>15} {'Wt.Mean':>12} {'Wt.SE':>12} {'N':>8}")
+    lines.append("-" * 50)
+
+    for col in cols:
+        vals = pdf[col].values
+        mask = ~np.isnan(vals) & ~np.isnan(weights)
+        v = vals[mask]
+        w = weights[mask]
+        n = len(v)
+        if n == 0:
+            lines.append(f"{col:>15} {'—':>12} {'—':>12} {0:>8}")
+            continue
+
+        # Weighted mean
+        wt_mean = np.average(v, weights=w)
+        # Weighted variance (for SE estimation)
+        w_sum = w.sum()
+        w_sum2 = (w ** 2).sum()
+        wt_var = np.average((v - wt_mean) ** 2, weights=w) * w_sum ** 2 / (w_sum ** 2 - w_sum2)
+        wt_se = np.sqrt(wt_var / n)
+
+        lines.append(f"{col:>15} {wt_mean:>12.4f} {wt_se:>12.4f} {n:>8}")
+
+    return "\n".join(lines)
+
+
+def fit_weighted_ols(
+    df: pl.DataFrame,
+    dep: str,
+    indeps: list[str],
+    weight_var: str,
+    strata_var: str | None = None,
+    psu_var: str | None = None,
+) -> tuple[FitResult, object]:
+    """Fit weighted OLS with optional Taylor linearization for complex survey SE."""
+    all_cols = [dep] + indeps + [weight_var]
+    if strata_var:
+        all_cols.append(strata_var)
+    if psu_var:
+        all_cols.append(psu_var)
+
+    pdf = df.select(all_cols).to_pandas().dropna()
+    weights = pdf[weight_var].values
+    y = pdf[dep].values
+    X = sm.add_constant(pdf[indeps].values)
+    var_names = ["const"] + indeps
+
+    # Fit WLS
+    model = sm.WLS(y, X, weights=weights)
+    result = model.fit()
+
+    # Taylor linearization for SE if PSU/strata provided
+    if psu_var and strata_var:
+        vcov = _taylor_linearization(pdf, y, X, weights, result.params,
+                                     psu_var, strata_var)
+        se = np.sqrt(np.diag(vcov))
+    else:
+        se = result.bse
+
+    params = {name: float(result.params[i]) for i, name in enumerate(var_names)}
+    std_errors = {name: float(se[i]) for i, name in enumerate(var_names)}
+    t_vals = {name: float(result.params[i] / se[i]) if se[i] > 0 else 0.0 for i, name in enumerate(var_names)}
+    p_vals = {name: float(2 * (1 - sp_stats.t.cdf(abs(t_vals[name]), result.df_resid))) for name in var_names}
+    ci_low = {name: params[name] - 1.96 * std_errors[name] for name in var_names}
+    ci_high = {name: params[name] + 1.96 * std_errors[name] for name in var_names}
+
+    warnings_list = [f"Weight variable: {weight_var}"]
+    if strata_var:
+        warnings_list.append(f"Strata: {strata_var}")
+    if psu_var:
+        warnings_list.append(f"PSU: {psu_var}")
+
+    fit = FitResult(
+        model_type="Svy: OLS",
+        formula=f"{dep} ~ {' + '.join(indeps)}",
+        dep_var=dep,
+        indep_vars=var_names,
+        n_obs=int(len(pdf)),
+        params=params,
+        std_errors=std_errors,
+        t_values=t_vals,
+        p_values=p_vals,
+        conf_int_low=ci_low,
+        conf_int_high=ci_high,
+        r_squared=float(result.rsquared),
+        warnings=warnings_list,
+    )
+
+    return fit, result
+
+
+def fit_weighted_logit(
+    df: pl.DataFrame,
+    dep: str,
+    indeps: list[str],
+    weight_var: str,
+) -> tuple[FitResult, object]:
+    """Fit weighted logistic regression."""
+    all_cols = [dep] + indeps + [weight_var]
+    pdf = df.select(all_cols).to_pandas().dropna()
+    weights = pdf[weight_var].values
+    y = pdf[dep].values
+    X = sm.add_constant(pdf[indeps].values)
+    var_names = ["const"] + indeps
+
+    model = sm.Logit(y, X)
+    result = model.fit(disp=0, freq_weights=weights)
+
+    params = {name: float(result.params[i]) for i, name in enumerate(var_names)}
+    std_errors = {name: float(result.bse[i]) for i, name in enumerate(var_names)}
+    t_vals = {name: float(result.tvalues[i]) for i, name in enumerate(var_names)}
+    p_vals = {name: float(result.pvalues[i]) for i, name in enumerate(var_names)}
+    ci = result.conf_int()
+    ci_low = {name: float(ci[i, 0]) for i, name in enumerate(var_names)}
+    ci_high = {name: float(ci[i, 1]) for i, name in enumerate(var_names)}
+
+    fit = FitResult(
+        model_type="Svy: Logit",
+        formula=f"{dep} ~ {' + '.join(indeps)}",
+        dep_var=dep,
+        indep_vars=var_names,
+        n_obs=int(len(pdf)),
+        params=params,
+        std_errors=std_errors,
+        t_values=t_vals,
+        p_values=p_vals,
+        conf_int_low=ci_low,
+        conf_int_high=ci_high,
+        pseudo_r2=float(result.prsquared),
+        log_likelihood=float(result.llf),
+        warnings=[f"Weight variable: {weight_var}"],
+    )
+
+    return fit, result
+
+
+def _taylor_linearization(pdf, y, X, weights, beta, psu_var, strata_var):
+    """Taylor series (sandwich) variance estimation for complex survey designs."""
+    import pandas as pd
+
+    resid = y - X @ beta
+    score = X * (resid * weights)[:, np.newaxis]
+
+    strata = pdf[strata_var].values
+    psu = pdf[psu_var].values
+
+    unique_strata = np.unique(strata)
+    n_h = len(unique_strata)
+    k = X.shape[1]
+    meat = np.zeros((k, k))
+
+    for h in unique_strata:
+        mask_h = strata == h
+        psus_in_h = np.unique(psu[mask_h])
+        m_h = len(psus_in_h)
+        if m_h < 2:
+            continue
+
+        # Sum of scores within each PSU
+        score_psu = np.zeros((m_h, k))
+        for j, p in enumerate(psus_in_h):
+            mask_p = mask_h & (psu == p)
+            score_psu[j] = score[mask_p].sum(axis=0)
+
+        score_mean = score_psu.mean(axis=0)
+        for j in range(m_h):
+            diff = score_psu[j] - score_mean
+            meat += np.outer(diff, diff) * m_h / (m_h - 1)
+
+    bread = np.linalg.inv(X.T @ np.diag(weights) @ X)
+    vcov = bread @ meat @ bread
+    return vcov
+
+
+def compute_deff(df: pl.DataFrame, col: str, weight_var: str,
+                 psu_var: str | None, strata_var: str | None) -> float:
+    """Compute design effect (DEFF) for a variable.
+
+    DEFF = var(complex design) / var(SRS of same size)
+    """
+    pdf = df.select([col, weight_var] + ([psu_var] if psu_var else []) + ([strata_var] if strata_var else [])).to_pandas().dropna()
+    vals = pdf[col].values
+    weights = pdf[weight_var].values
+    n = len(vals)
+
+    # SRS variance
+    var_srs = np.var(vals, ddof=1) / n
+
+    # Weighted variance
+    wt_mean = np.average(vals, weights=weights)
+    w_sum = weights.sum()
+    w_sum2 = (weights ** 2).sum()
+    # Kish's approximation of design effect
+    deff = 1 + (w_sum2 / w_sum ** 2 - 1 / n) * n
+    return max(deff, 1.0)

openstat/stats/survival.py

@@ -0,0 +1,196 @@
+"""Survival analysis: Cox PH, Kaplan-Meier, log-rank test."""
+
+from __future__ import annotations
+
+import numpy as np
+import polars as pl
+
+from openstat.stats.models import FitResult
+
+
+def _try_import_lifelines():
+    try:
+        import lifelines  # noqa: F401
+    except ImportError:
+        raise ImportError(
+            "Survival analysis requires lifelines. "
+            "Install it with: pip install openstat[survival]"
+        )
+
+
+def fit_cox_ph(
+    df: pl.DataFrame,
+    time_var: str,
+    event_var: str,
+    covariates: list[str],
+    robust: bool = False,
+) -> tuple[FitResult, object]:
+    """Fit a Cox Proportional Hazards model."""
+    _try_import_lifelines()
+    from lifelines import CoxPHFitter
+
+    cols = [time_var, event_var] + covariates
+    pdf = df.select(cols).to_pandas().dropna()
+
+    cph = CoxPHFitter()
+    cph.fit(pdf, duration_col=time_var, event_col=event_var, robust=robust)
+
+    summary = cph.summary
+    params = {name: float(summary.loc[name, "coef"]) for name in covariates}
+    std_errors = {name: float(summary.loc[name, "se(coef)"]) for name in covariates}
+    z_values = {name: float(summary.loc[name, "z"]) for name in covariates}
+    p_values = {name: float(summary.loc[name, "p"]) for name in covariates}
+
+    ci_cols = [c for c in summary.columns if "lower" in c.lower()]
+    ci_low_col = ci_cols[0] if ci_cols else None
+    ci_high_cols = [c for c in summary.columns if "upper" in c.lower()]
+    ci_high_col = ci_high_cols[0] if ci_high_cols else None
+
+    conf_low = {}
+    conf_high = {}
+    for name in covariates:
+        conf_low[name] = float(summary.loc[name, ci_low_col]) if ci_low_col else 0.0
+        conf_high[name] = float(summary.loc[name, ci_high_col]) if ci_high_col else 0.0
+
+    warnings_list = [
+        f"Concordance: {cph.concordance_index_:.4f}",
+        f"Partial log-likelihood: {cph.log_likelihood_:.2f}",
+    ]
+
+    # Hazard ratios
+    hr_lines = ["Hazard Ratios:"]
+    for name in covariates:
+        hr = float(summary.loc[name, "exp(coef)"])
+        hr_lines.append(f"  {name}: {hr:.4f}")
+    warnings_list.append("\n".join(hr_lines))
+
+    fit = FitResult(
+        model_type="Cox PH",
+        formula=f"h(t) ~ {' + '.join(covariates)}",
+        dep_var=f"{time_var} (event: {event_var})",
+        indep_vars=covariates,
+        n_obs=int(len(pdf)),
+        params=params,
+        std_errors=std_errors,
+        t_values=z_values,
+        p_values=p_values,
+        conf_int_low=conf_low,
+        conf_int_high=conf_high,
+        log_likelihood=float(cph.log_likelihood_),
+        warnings=warnings_list,
+    )
+
+    return fit, cph
+
+
+def kaplan_meier(
+    df: pl.DataFrame,
+    time_var: str,
+    event_var: str,
+    group_var: str | None = None,
+) -> tuple[str, object | list]:
+    """Fit Kaplan-Meier survival curves.
+
+    Returns summary string and fitted KMF object(s).
+    """
+    _try_import_lifelines()
+    from lifelines import KaplanMeierFitter
+
+    pdf = df.select([time_var, event_var] + ([group_var] if group_var else [])).to_pandas().dropna()
+
+    if group_var is None:
+        kmf = KaplanMeierFitter()
+        kmf.fit(pdf[time_var], event_observed=pdf[event_var])
+        median = kmf.median_survival_time_
+        lines = [
+            f"Kaplan-Meier Estimate (N={len(pdf)})",
+            f"  Median survival time: {median:.2f}" if np.isfinite(median) else "  Median survival time: not reached",
+            f"  Events: {int(pdf[event_var].sum())}",
+            f"  Censored: {int(len(pdf) - pdf[event_var].sum())}",
+        ]
+        return "\n".join(lines), kmf
+    else:
+        groups = sorted(pdf[group_var].unique())
+        kmfs = []
+        lines = [f"Kaplan-Meier Estimates by {group_var}:"]
+        for g in groups:
+            mask = pdf[group_var] == g
+            sub = pdf[mask]
+            kmf = KaplanMeierFitter()
+            kmf.fit(sub[time_var], event_observed=sub[event_var], label=str(g))
+            kmfs.append(kmf)
+            median = kmf.median_survival_time_
+            lines.append(
+                f"\n  Group {g} (N={len(sub)}):"
+                f"\n    Median survival: {median:.2f}" if np.isfinite(median)
+                else f"\n  Group {g} (N={len(sub)}):\n    Median survival: not reached"
+            )
+            lines.append(f"    Events: {int(sub[event_var].sum())}")
+        return "\n".join(lines), kmfs
+
+
+def log_rank_test(
+    df: pl.DataFrame,
+    time_var: str,
+    event_var: str,
+    group_var: str,
+) -> str:
+    """Log-rank test comparing survival between groups."""
+    _try_import_lifelines()
+    from lifelines.statistics import logrank_test
+
+    pdf = df.select([time_var, event_var, group_var]).to_pandas().dropna()
+    groups = sorted(pdf[group_var].unique())
+
+    if len(groups) < 2:
+        return "Log-rank test requires at least 2 groups."
+
+    if len(groups) == 2:
+        g1 = pdf[pdf[group_var] == groups[0]]
+        g2 = pdf[pdf[group_var] == groups[1]]
+        result = logrank_test(
+            g1[time_var], g2[time_var],
+            event_observed_A=g1[event_var],
+            event_observed_B=g2[event_var],
+        )
+        lines = [
+            f"Log-Rank Test: {group_var}",
+            f"  Groups: {groups[0]} vs {groups[1]}",
+            f"  Test statistic: {result.test_statistic:.4f}",
+            f"  p-value: {result.p_value:.4f}",
+        ]
+        if result.p_value < 0.05:
+            lines.append("  ⚠ Significant difference in survival between groups")
+        else:
+            lines.append("  ✓ No significant difference in survival")
+        return "\n".join(lines)
+    else:
+        # Pairwise for >2 groups
+        lines = [f"Log-Rank Tests (pairwise): {group_var}"]
+        for i, g1 in enumerate(groups):
+            for g2 in groups[i + 1:]:
+                d1 = pdf[pdf[group_var] == g1]
+                d2 = pdf[pdf[group_var] == g2]
+                result = logrank_test(
+                    d1[time_var], d2[time_var],
+                    event_observed_A=d1[event_var],
+                    event_observed_B=d2[event_var],
+                )
+                sig = "*" if result.p_value < 0.05 else ""
+                lines.append(f"  {g1} vs {g2}: chi2={result.test_statistic:.3f}, p={result.p_value:.4f}{sig}")
+        return "\n".join(lines)
+
+
+def schoenfeld_test(cph_result) -> str:
+    """Test proportional hazards assumption via Schoenfeld residuals."""
+    try:
+        ph_test = cph_result.check_assumptions(show_plots=False, p_value_threshold=1.0)
+        lines = ["Proportional Hazards Test (Schoenfeld Residuals):"]
+        # check_assumptions returns summary or prints it
+        if ph_test is not None:
+            lines.append(str(ph_test))
+        else:
+            lines.append("  PH assumption appears satisfied for all covariates.")
+        return "\n".join(lines)
+    except Exception as e:
+        return f"PH test: {e}"