openstat-cli 1.0.0__py3-none-any.whl

openstat/stats/nonparametric.py

@@ -0,0 +1,130 @@
+"""Nonparametric hypothesis tests and rank-based statistics."""
+
+from __future__ import annotations
+
+import numpy as np
+import polars as pl
+from scipy import stats as sp_stats
+
+
+# ── Spearman rank correlation ──────────────────────────────────────────────
+
+def spearman_corr(df: pl.DataFrame, cols: list[str]) -> dict:
+    """Spearman rank correlation matrix."""
+    X = df.select(cols).to_numpy().astype(float)
+    n = X.shape[1]
+    rho = np.eye(n)
+    pvals = np.zeros((n, n))
+    for i in range(n):
+        for j in range(n):
+            if i == j:
+                pvals[i, j] = 0.0
+                continue
+            mask = ~(np.isnan(X[:, i]) | np.isnan(X[:, j]))
+            r, p = sp_stats.spearmanr(X[mask, i], X[mask, j])
+            rho[i, j] = r
+            pvals[i, j] = p
+    return {"rho": rho.tolist(), "pvalues": pvals.tolist(), "cols": cols}
+
+
+# ── Mann-Whitney / Wilcoxon rank-sum ──────────────────────────────────────
+
+def ranksum_test(
+    df: pl.DataFrame,
+    var: str,
+    group: str,
+    *,
+    alternative: str = "two-sided",
+) -> dict:
+    """
+    Wilcoxon rank-sum test (Mann-Whitney U) for two independent groups.
+
+    alternative: 'two-sided', 'less', 'greater'
+    """
+    grp_vals = df[group].drop_nulls().unique().to_list()
+    if len(grp_vals) != 2:
+        raise ValueError(f"'{group}' must have exactly 2 groups, found {len(grp_vals)}")
+
+    g1 = df.filter(pl.col(group) == grp_vals[0])[var].drop_nulls().to_numpy().astype(float)
+    g2 = df.filter(pl.col(group) == grp_vals[1])[var].drop_nulls().to_numpy().astype(float)
+
+    stat, p = sp_stats.mannwhitneyu(g1, g2, alternative=alternative)
+    z = (stat - len(g1) * len(g2) / 2) / np.sqrt(
+        len(g1) * len(g2) * (len(g1) + len(g2) + 1) / 12
+    )
+    return {
+        "test": "Wilcoxon rank-sum (Mann-Whitney U)",
+        "var": var,
+        "group": group,
+        "groups": grp_vals,
+        "n1": len(g1),
+        "n2": len(g2),
+        "U_statistic": float(stat),
+        "z_statistic": float(z),
+        "p_value": float(p),
+        "alternative": alternative,
+    }
+
+
+# ── Wilcoxon signed-rank ───────────────────────────────────────────────────
+
+def signrank_test(
+    df: pl.DataFrame,
+    var1: str,
+    var2: str | None = None,
+    *,
+    mu: float = 0.0,
+    alternative: str = "two-sided",
+) -> dict:
+    """
+    Wilcoxon signed-rank test.
+
+    One-sample: var2=None, tests median of var1 == mu.
+    Paired: tests median of (var1 - var2) == 0.
+    """
+    x = df[var1].drop_nulls().to_numpy().astype(float)
+    if var2 is not None:
+        y = df[var2].drop_nulls().to_numpy().astype(float)
+        diff = x[: len(y)] - y[: len(x)]
+    else:
+        diff = x - mu
+
+    stat, p = sp_stats.wilcoxon(diff, alternative=alternative)
+    return {
+        "test": "Wilcoxon signed-rank",
+        "var1": var1,
+        "var2": var2,
+        "mu": mu,
+        "n": len(diff),
+        "W_statistic": float(stat),
+        "p_value": float(p),
+        "alternative": alternative,
+    }
+
+
+# ── Kruskal-Wallis ─────────────────────────────────────────────────────────
+
+def kruskal_wallis_test(
+    df: pl.DataFrame,
+    var: str,
+    group: str,
+) -> dict:
+    """Kruskal-Wallis H test for k independent groups."""
+    groups = df[group].drop_nulls().unique().to_list()
+    samples = [
+        df.filter(pl.col(group) == g)[var].drop_nulls().to_numpy().astype(float)
+        for g in groups
+    ]
+    stat, p = sp_stats.kruskal(*samples)
+    df_stat = len(groups) - 1
+    return {
+        "test": "Kruskal-Wallis H",
+        "var": var,
+        "group": group,
+        "k_groups": len(groups),
+        "H_statistic": float(stat),
+        "df": df_stat,
+        "p_value": float(p),
+        "groups": groups,
+        "n_per_group": [len(s) for s in samples],
+    }

openstat/stats/panel.py

@@ -0,0 +1,179 @@
+"""Panel data models: Fixed Effects, Random Effects, Between, Hausman test."""
+
+from __future__ import annotations
+
+import io
+from dataclasses import dataclass, field
+
+import numpy as np
+import polars as pl
+from rich.table import Table
+from rich.console import Console
+
+from openstat.stats.models import FitResult
+
+
+def _try_import_linearmodels():
+    try:
+        import linearmodels  # noqa: F401
+    except ImportError:
+        raise ImportError(
+            "Panel data models require linearmodels. "
+            "Install it with: pip install openstat[panel]"
+        )
+
+
+def _panel_to_fit_result(result, model_type: str, dep: str, indeps: list[str]) -> FitResult:
+    """Convert a linearmodels PanelResults to FitResult."""
+    params = {name: float(val) for name, val in result.params.items()}
+    std_errors = {name: float(val) for name, val in result.std_errors.items()}
+    t_values = {name: float(val) for name, val in result.tstats.items()}
+    p_values = {name: float(val) for name, val in result.pvalues.items()}
+    ci = result.conf_int()
+    conf_low = {name: float(ci.loc[name, "lower"]) for name in params}
+    conf_high = {name: float(ci.loc[name, "upper"]) for name in params}
+
+    warnings_list: list[str] = []
+
+    return FitResult(
+        model_type=model_type,
+        formula=f"{dep} ~ {' + '.join(indeps)}",
+        dep_var=dep,
+        indep_vars=indeps,
+        n_obs=int(result.nobs),
+        params=params,
+        std_errors=std_errors,
+        t_values=t_values,
+        p_values=p_values,
+        conf_int_low=conf_low,
+        conf_int_high=conf_high,
+        r_squared=float(result.rsquared) if hasattr(result, "rsquared") else None,
+        f_statistic=float(result.f_statistic.stat) if hasattr(result, "f_statistic") and result.f_statistic is not None else None,
+        f_pvalue=float(result.f_statistic.pval) if hasattr(result, "f_statistic") and result.f_statistic is not None else None,
+        warnings=warnings_list,
+    )
+
+
+def fit_panel_fe(
+    df: pl.DataFrame,
+    dep: str,
+    indeps: list[str],
+    entity_col: str,
+    time_col: str,
+    robust: bool = False,
+    cluster: str | None = None,
+) -> tuple[FitResult, object]:
+    """Fit a Fixed Effects panel model."""
+    _try_import_linearmodels()
+    from linearmodels.panel import PanelOLS
+
+    pdf = df.select([entity_col, time_col, dep] + indeps).to_pandas().dropna()
+    pdf = pdf.set_index([entity_col, time_col])
+
+    import statsmodels.api as sm
+    y = pdf[dep]
+    X = sm.add_constant(pdf[indeps])
+
+    model = PanelOLS(y, X, entity_effects=True)
+
+    cov_type = "unadjusted"
+    cov_kwds: dict = {}
+    if cluster:
+        cov_type = "clustered"
+        cov_kwds["cluster_entity"] = True
+    elif robust:
+        cov_type = "robust"
+
+    result = model.fit(cov_type=cov_type, **cov_kwds)
+    fit = _panel_to_fit_result(result, "Panel FE", dep, ["const"] + indeps)
+    return fit, result
+
+
+def fit_panel_re(
+    df: pl.DataFrame,
+    dep: str,
+    indeps: list[str],
+    entity_col: str,
+    time_col: str,
+    robust: bool = False,
+) -> tuple[FitResult, object]:
+    """Fit a Random Effects panel model."""
+    _try_import_linearmodels()
+    from linearmodels.panel import RandomEffects
+
+    pdf = df.select([entity_col, time_col, dep] + indeps).to_pandas().dropna()
+    pdf = pdf.set_index([entity_col, time_col])
+
+    import statsmodels.api as sm
+    y = pdf[dep]
+    X = sm.add_constant(pdf[indeps])
+
+    model = RandomEffects(y, X)
+    cov_type = "robust" if robust else "unadjusted"
+    result = model.fit(cov_type=cov_type)
+    fit = _panel_to_fit_result(result, "Panel RE", dep, ["const"] + indeps)
+    return fit, result
+
+
+def fit_panel_be(
+    df: pl.DataFrame,
+    dep: str,
+    indeps: list[str],
+    entity_col: str,
+    time_col: str,
+) -> tuple[FitResult, object]:
+    """Fit a Between Effects panel model."""
+    _try_import_linearmodels()
+    from linearmodels.panel import BetweenOLS
+
+    pdf = df.select([entity_col, time_col, dep] + indeps).to_pandas().dropna()
+    pdf = pdf.set_index([entity_col, time_col])
+
+    import statsmodels.api as sm
+    y = pdf[dep]
+    X = sm.add_constant(pdf[indeps])
+
+    model = BetweenOLS(y, X)
+    result = model.fit()
+    fit = _panel_to_fit_result(result, "Panel BE", dep, ["const"] + indeps)
+    return fit, result
+
+
+def hausman_test(fe_result, re_result) -> str:
+    """Perform the Hausman test for FE vs RE.
+
+    H0: RE is consistent and efficient (prefer RE).
+    H1: RE is inconsistent (prefer FE).
+    """
+    b_fe = fe_result.params
+    b_re = re_result.params
+
+    # Use common coefficients (exclude const)
+    common = [k for k in b_fe.index if k in b_re.index and k != "const"]
+    if not common:
+        return "No common coefficients for Hausman test."
+
+    b_diff = np.array([b_fe[k] - b_re[k] for k in common])
+    cov_fe = fe_result.cov.loc[common, common].values
+    cov_re = re_result.cov.loc[common, common].values
+    cov_diff = cov_fe - cov_re
+
+    try:
+        chi2_stat = float(b_diff @ np.linalg.inv(cov_diff) @ b_diff)
+    except np.linalg.LinAlgError:
+        chi2_stat = float(b_diff @ np.linalg.pinv(cov_diff) @ b_diff)
+
+    from scipy import stats as sp_stats
+    df = len(common)
+    p_value = float(1 - sp_stats.chi2.cdf(chi2_stat, df))
+
+    recommendation = "Use Fixed Effects (FE)" if p_value < 0.05 else "Use Random Effects (RE)"
+
+    lines = [
+        "Hausman Test (FE vs RE)",
+        f"  H0: Random Effects model is consistent",
+        f"  chi2({df}) = {chi2_stat:.4f}",
+        f"  p-value = {p_value:.4f}",
+        f"  Recommendation: {recommendation}",
+    ]
+    return "\n".join(lines)

openstat/stats/power.py

@@ -0,0 +1,295 @@
+"""Power analysis — one/two-sample means, proportions, OLS."""
+
+from __future__ import annotations
+
+import math
+
+from scipy import stats as sp_stats
+
+
+# ── Helpers ────────────────────────────────────────────────────────────────
+
+def _solve(fn, lo=1, hi=1_000_000, tol=1e-6):
+    """Bisection solver: find x in [lo, hi] where fn(x) ≈ 0."""
+    for _ in range(60):
+        mid = (lo + hi) / 2
+        if fn(mid) < 0:
+            lo = mid
+        else:
+            hi = mid
+        if hi - lo < tol:
+            break
+    return (lo + hi) / 2
+
+
+# ── One-sample mean ────────────────────────────────────────────────────────
+
+def power_onemean(
+    effect_size: float | None = None,
+    alpha: float = 0.05,
+    n: int | None = None,
+    power: float | None = None,
+    sd: float = 1.0,
+    delta: float | None = None,
+    two_sided: bool = True,
+) -> dict:
+    """
+    Power analysis for one-sample t-test.
+
+    Provide exactly two of: effect_size (or delta/sd), n, power.
+    """
+    if delta is not None and effect_size is None:
+        effect_size = delta / sd
+
+    sides = 2 if two_sided else 1
+    za2 = sp_stats.norm.ppf(1 - alpha / sides)
+
+    def _power_from_n(n_val):
+        zb = abs(effect_size) * math.sqrt(n_val) - za2
+        return sp_stats.norm.cdf(zb)
+
+    if n is None and power is not None and effect_size is not None:
+        # Solve for n
+        n_val = _solve(lambda x: _power_from_n(x) - power)
+        n = math.ceil(n_val)
+        achieved = _power_from_n(n)
+    elif power is None and n is not None and effect_size is not None:
+        achieved = _power_from_n(n)
+        power = achieved
+    elif effect_size is None and n is not None and power is not None:
+        # Solve for detectable effect size
+        zb = sp_stats.norm.ppf(power)
+        effect_size = (za2 + zb) / math.sqrt(n)
+        achieved = power
+    else:
+        raise ValueError("Provide exactly two of: effect_size, n, power")
+
+    return {
+        "test": "One-sample t-test",
+        "effect_size": round(effect_size, 6),
+        "alpha": alpha,
+        "n": n,
+        "power": round(power, 6),
+        "two_sided": two_sided,
+    }
+
+
+# ── Two-sample means ───────────────────────────────────────────────────────
+
+def power_twomeans(
+    effect_size: float | None = None,
+    alpha: float = 0.05,
+    n: int | None = None,
+    power: float | None = None,
+    ratio: float = 1.0,
+    sd: float = 1.0,
+    delta: float | None = None,
+    two_sided: bool = True,
+) -> dict:
+    """Power analysis for two-sample independent t-test."""
+    if delta is not None and effect_size is None:
+        effect_size = delta / sd
+
+    sides = 2 if two_sided else 1
+    za2 = sp_stats.norm.ppf(1 - alpha / sides)
+
+    def _power_from_n(n1):
+        n2 = n1 * ratio
+        se = math.sqrt(1 / n1 + 1 / n2)
+        zb = abs(effect_size) / se - za2
+        return sp_stats.norm.cdf(zb)
+
+    if n is None and power is not None and effect_size is not None:
+        n_val = _solve(lambda x: _power_from_n(x) - power)
+        n = math.ceil(n_val)
+        achieved = _power_from_n(n)
+        power = achieved
+    elif power is None and n is not None and effect_size is not None:
+        power = _power_from_n(n)
+    elif effect_size is None and n is not None and power is not None:
+        zb = sp_stats.norm.ppf(power)
+        n2 = n * ratio
+        se = math.sqrt(1 / n + 1 / n2)
+        effect_size = (za2 + zb) * se
+        achieved = power
+    else:
+        raise ValueError("Provide exactly two of: effect_size, n, power")
+
+    return {
+        "test": "Two-sample t-test",
+        "effect_size": round(effect_size, 6),
+        "alpha": alpha,
+        "n1": n,
+        "n2": math.ceil(n * ratio),
+        "power": round(power, 6),
+        "ratio": ratio,
+        "two_sided": two_sided,
+    }
+
+
+# ── One proportion ─────────────────────────────────────────────────────────
+
+def power_oneproportion(
+    p0: float,
+    pa: float,
+    alpha: float = 0.05,
+    n: int | None = None,
+    power: float | None = None,
+    two_sided: bool = True,
+) -> dict:
+    """Power analysis for one-sample proportion z-test."""
+    sides = 2 if two_sided else 1
+    za2 = sp_stats.norm.ppf(1 - alpha / sides)
+    effect_size = abs(pa - p0) / math.sqrt(p0 * (1 - p0))
+
+    def _power_from_n(n_val):
+        se_null = math.sqrt(p0 * (1 - p0) / n_val)
+        se_alt = math.sqrt(pa * (1 - pa) / n_val)
+        z = (abs(pa - p0) - za2 * se_null) / se_alt
+        return sp_stats.norm.cdf(z)
+
+    if n is None and power is not None:
+        n_val = _solve(lambda x: _power_from_n(x) - power)
+        n = math.ceil(n_val)
+        power = _power_from_n(n)
+    elif power is None and n is not None:
+        power = _power_from_n(n)
+    else:
+        raise ValueError("Provide exactly one of: n, power")
+
+    return {
+        "test": "One-sample proportion z-test",
+        "p0": p0,
+        "pa": pa,
+        "effect_size": round(effect_size, 6),
+        "alpha": alpha,
+        "n": n,
+        "power": round(power, 6),
+        "two_sided": two_sided,
+    }
+
+
+# ── Two proportions ────────────────────────────────────────────────────────
+
+def power_twoproportions(
+    p1: float,
+    p2: float,
+    alpha: float = 0.05,
+    n: int | None = None,
+    power: float | None = None,
+    two_sided: bool = True,
+) -> dict:
+    """Power analysis for two-sample proportion z-test."""
+    sides = 2 if two_sided else 1
+    za2 = sp_stats.norm.ppf(1 - alpha / sides)
+    p_avg = (p1 + p2) / 2
+    effect_size = abs(p2 - p1) / math.sqrt(p_avg * (1 - p_avg))
+
+    def _power_from_n(n_val):
+        se_null = math.sqrt(2 * p_avg * (1 - p_avg) / n_val)
+        se_alt = math.sqrt((p1 * (1 - p1) + p2 * (1 - p2)) / n_val)
+        z = (abs(p2 - p1) - za2 * se_null) / se_alt
+        return sp_stats.norm.cdf(z)
+
+    if n is None and power is not None:
+        n_val = _solve(lambda x: _power_from_n(x) - power)
+        n = math.ceil(n_val)
+        power = _power_from_n(n)
+    elif power is None and n is not None:
+        power = _power_from_n(n)
+    else:
+        raise ValueError("Provide exactly one of: n, power")
+
+    return {
+        "test": "Two-sample proportion z-test",
+        "p1": p1,
+        "p2": p2,
+        "effect_size": round(effect_size, 6),
+        "alpha": alpha,
+        "n": n,
+        "power": round(power, 6),
+        "two_sided": two_sided,
+    }
+
+
+# ── OLS / multiple regression ──────────────────────────────────────────────
+
+def power_ols(
+    f2: float | None = None,
+    alpha: float = 0.05,
+    n: int | None = None,
+    power: float | None = None,
+    k: int = 1,
+) -> dict:
+    """
+    Power analysis for OLS / multiple regression (Cohen's f²).
+
+    f2 = R² / (1 - R²)
+    k  = number of predictors
+    """
+
+    def _power_from_n(n_val):
+        df1 = k
+        df2 = n_val - k - 1
+        if df2 <= 0:
+            return 0.0
+        nc = f2 * n_val
+        return 1 - sp_stats.f.cdf(
+            sp_stats.f.ppf(1 - alpha, df1, df2), df1, df2, nc
+        )
+
+    if n is None and power is not None and f2 is not None:
+        n_val = _solve(lambda x: _power_from_n(x) - power, lo=k + 2)
+        n = math.ceil(n_val)
+        power = _power_from_n(n)
+    elif power is None and n is not None and f2 is not None:
+        power = _power_from_n(n)
+    elif f2 is None and n is not None and power is not None:
+        f2 = _solve(
+            lambda f: _power_from_n_f2(n, f, alpha, k) - power,  # type: ignore[arg-type]
+            lo=0.0001,
+            hi=10,
+        )
+        power = _power_from_n(n)
+    else:
+        raise ValueError("Provide exactly two of: f2, n, power")
+
+    return {
+        "test": "OLS / Multiple Regression",
+        "f2": round(f2, 6),
+        "alpha": alpha,
+        "n": n,
+        "k": k,
+        "power": round(power, 6),
+    }
+
+
+def _power_from_n_f2(n, f2, alpha, k):
+    df1 = k
+    df2 = n - k - 1
+    if df2 <= 0:
+        return 0.0
+    nc = f2 * n
+    return 1 - sp_stats.f.cdf(
+        sp_stats.f.ppf(1 - alpha, df1, df2), df1, df2, nc
+    )
+
+
+# ── sampsi (Stata-style) ───────────────────────────────────────────────────
+
+def sampsi(
+    mu1: float,
+    mu2: float,
+    sd: float = 1.0,
+    alpha: float = 0.05,
+    power: float = 0.80,
+    two_sided: bool = True,
+) -> dict:
+    """Compute required sample size for two-sample t-test (Stata sampsi style)."""
+    effect_size = abs(mu2 - mu1) / sd
+    return power_twomeans(
+        effect_size=effect_size,
+        alpha=alpha,
+        power=power,
+        two_sided=two_sided,
+    )