openstat-cli 1.0.0__py3-none-any.whl

openstat/stats/discrete.py

@@ -0,0 +1,311 @@
+"""Discrete / censored models: Tobit, Multinomial Logit, Ordered Logit/Probit."""
+
+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 scipy.optimize import minimize
+
+from openstat.stats.models import FitResult, _prepare_data, _cov_args, _model_type_suffix
+
+
+# ── Tobit (censored regression) ─────────────────────────────────────
+
+def _tobit_loglike(params, y, X, lower, upper):
+    """Tobit log-likelihood for scipy.optimize."""
+    beta = params[:-1]
+    log_sigma = params[-1]
+    sigma = np.exp(log_sigma)
+
+    Xb = X @ beta
+    resid = (y - Xb) / sigma
+
+    ll = 0.0
+    for i in range(len(y)):
+        if lower is not None and y[i] <= lower:
+            # Left-censored
+            cdf_val = sp_stats.norm.cdf((lower - Xb[i]) / sigma)
+            ll += np.log(max(cdf_val, 1e-300))
+        elif upper is not None and y[i] >= upper:
+            # Right-censored
+            cdf_val = sp_stats.norm.sf((upper - Xb[i]) / sigma)
+            ll += np.log(max(cdf_val, 1e-300))
+        else:
+            # Uncensored
+            ll += sp_stats.norm.logpdf(resid[i]) - log_sigma
+
+    return -ll  # minimize negative log-likelihood
+
+
+def fit_tobit(
+    df: pl.DataFrame,
+    dep: str,
+    indeps: list[str],
+    *,
+    lower_limit: float | None = None,
+    upper_limit: float | None = None,
+    robust: bool = False,
+    cluster_col: str | None = None,
+) -> tuple[FitResult, object]:
+    """Fit a Tobit (censored) regression model via MLE."""
+    y, X, warnings_list, var_names, groups = _prepare_data(
+        df, dep, indeps, cluster_col=cluster_col,
+    )
+
+    if lower_limit is None and upper_limit is None:
+        warnings_list.append("Note: No censoring limits specified. Results are equivalent to OLS.")
+
+    n_censored_low = int(np.sum(y <= lower_limit)) if lower_limit is not None else 0
+    n_censored_high = int(np.sum(y >= upper_limit)) if upper_limit is not None else 0
+    n_uncensored = len(y) - n_censored_low - n_censored_high
+
+    if n_uncensored < len(var_names) + 2:
+        raise ValueError(
+            f"Too few uncensored observations ({n_uncensored}) for {len(var_names)} parameters."
+        )
+
+    # Initial values: OLS estimates + log(sigma)
+    ols = sm.OLS(y, X).fit()
+    init_beta = ols.params
+    init_log_sigma = np.log(np.std(ols.resid))
+    init_params = np.append(init_beta, init_log_sigma)
+
+    result = minimize(
+        _tobit_loglike, init_params, args=(y, X, lower_limit, upper_limit),
+        method="BFGS",
+    )
+
+    if not result.success:
+        warnings_list.append(f"Warning: Optimization did not fully converge: {result.message}")
+
+    beta = result.x[:-1]
+    log_sigma = result.x[-1]
+    sigma = np.exp(log_sigma)
+
+    # Standard errors from inverse Hessian
+    try:
+        hess_inv = result.hess_inv
+        if hasattr(hess_inv, 'todense'):
+            hess_inv = hess_inv.todense()
+        se_all = np.sqrt(np.diag(np.abs(hess_inv)))
+    except Exception:
+        # Fallback: numerical Hessian
+        from scipy.optimize import approx_fprime
+        eps = 1e-5
+        n_p = len(result.x)
+        hess = np.zeros((n_p, n_p))
+        for i in range(n_p):
+            def grad_i(p):
+                g = approx_fprime(p, _tobit_loglike, eps, y, X, lower_limit, upper_limit)
+                return g[i]
+            hess[i, :] = approx_fprime(result.x, grad_i, eps)
+        try:
+            se_all = np.sqrt(np.diag(np.linalg.inv(hess)))
+        except np.linalg.LinAlgError:
+            se_all = np.full(n_p, np.nan)
+
+    se_beta = se_all[:-1]
+    se_sigma = se_all[-1]
+
+    # Build coefficient results
+    t_vals = beta / se_beta
+    p_vals = 2 * (1 - sp_stats.norm.cdf(np.abs(t_vals)))
+    ci_low = beta - 1.96 * se_beta
+    ci_high = beta + 1.96 * se_beta
+
+    # Add sigma as extra parameter
+    all_var_names = var_names + ["sigma"]
+    all_params = np.append(beta, sigma)
+    all_se = np.append(se_beta, se_sigma)
+    sigma_t = sigma / se_sigma if se_sigma > 0 else np.nan
+    sigma_p = 2 * (1 - sp_stats.norm.cdf(np.abs(sigma_t)))
+    all_t = np.append(t_vals, sigma_t)
+    all_p = np.append(p_vals, sigma_p)
+    all_ci_low = np.append(ci_low, sigma - 1.96 * se_sigma)
+    all_ci_high = np.append(ci_high, sigma + 1.96 * se_sigma)
+
+    suffix = _model_type_suffix(robust, groups is not None)
+    ll_val = -result.fun
+
+    censor_info = []
+    if lower_limit is not None:
+        censor_info.append(f"Left-censored at {lower_limit}: {n_censored_low} obs")
+    if upper_limit is not None:
+        censor_info.append(f"Right-censored at {upper_limit}: {n_censored_high} obs")
+    censor_info.append(f"Uncensored: {n_uncensored} obs")
+    warnings_list.extend(censor_info)
+
+    fit = FitResult(
+        model_type="Tobit" + suffix,
+        formula=f"{dep} ~ {' + '.join(indeps)}",
+        dep_var=dep,
+        indep_vars=indeps,
+        n_obs=len(y),
+        params=dict(zip(all_var_names, all_params)),
+        std_errors=dict(zip(all_var_names, all_se)),
+        t_values=dict(zip(all_var_names, all_t)),
+        p_values=dict(zip(all_var_names, all_p)),
+        conf_int_low=dict(zip(all_var_names, all_ci_low)),
+        conf_int_high=dict(zip(all_var_names, all_ci_high)),
+        log_likelihood=ll_val,
+        aic=-2 * ll_val + 2 * len(result.x),
+        bic=-2 * ll_val + np.log(len(y)) * len(result.x),
+        warnings=warnings_list,
+    )
+    return fit, result
+
+
+# ── Multinomial Logit ────────────────────────────────────────────────
+
+def fit_mlogit(
+    df: pl.DataFrame,
+    dep: str,
+    indeps: list[str],
+    *,
+    robust: bool = False,
+    cluster_col: str | None = None,
+) -> tuple[FitResult, object]:
+    """Fit a Multinomial Logit model."""
+    y, X, warnings_list, var_names, groups = _prepare_data(
+        df, dep, indeps, cluster_col=cluster_col,
+    )
+
+    cov_type, cov_kwds = _cov_args(robust, groups)
+    model = sm.MNLogit(y, X).fit(disp=0, cov_type=cov_type, cov_kwds=cov_kwds)
+
+    # Get unique categories
+    categories = sorted(np.unique(y))
+    base_cat = categories[0]
+    other_cats = categories[1:]
+
+    # Flatten per-category coefficients
+    params_dict = {}
+    se_dict = {}
+    t_dict = {}
+    p_dict = {}
+    ci_low_dict = {}
+    ci_high_dict = {}
+
+    ci = model.conf_int()  # shape: (n_cats-1, n_vars, 2)
+
+    for j, cat in enumerate(other_cats):
+        cat_label = f"y={int(cat)}" if cat == int(cat) else f"y={cat}"
+        for i, var in enumerate(var_names):
+            key = f"{var} ({cat_label})"
+            params_dict[key] = float(model.params[i, j])
+            se_dict[key] = float(model.bse[i, j])
+            t_dict[key] = float(model.tvalues[i, j])
+            p_dict[key] = float(model.pvalues[i, j])
+            ci_low_dict[key] = float(ci[j, i, 0])
+            ci_high_dict[key] = float(ci[j, i, 1])
+
+    suffix = _model_type_suffix(robust, groups is not None)
+    warnings_list.append(f"Base category: {int(base_cat) if base_cat == int(base_cat) else base_cat}")
+
+    fit = FitResult(
+        model_type="MNLogit" + suffix,
+        formula=f"{dep} ~ {' + '.join(indeps)}",
+        dep_var=dep,
+        indep_vars=indeps,
+        n_obs=int(model.nobs),
+        params=params_dict,
+        std_errors=se_dict,
+        t_values=t_dict,
+        p_values=p_dict,
+        conf_int_low=ci_low_dict,
+        conf_int_high=ci_high_dict,
+        pseudo_r2=float(model.prsquared),
+        log_likelihood=float(model.llf),
+        aic=float(model.aic),
+        bic=float(model.bic),
+        warnings=warnings_list,
+    )
+    return fit, model
+
+
+# ── Ordered Logit / Probit ───────────────────────────────────────────
+
+def fit_ordered(
+    df: pl.DataFrame,
+    dep: str,
+    indeps: list[str],
+    *,
+    link: str = "logit",
+    robust: bool = False,
+    cluster_col: str | None = None,
+) -> tuple[FitResult, object]:
+    """Fit an Ordered Logit or Ordered Probit model."""
+    from statsmodels.miscmodels.ordinal_model import OrderedModel
+
+    y, X, warnings_list, var_names, groups = _prepare_data(
+        df, dep, indeps, cluster_col=cluster_col,
+    )
+
+    # OrderedModel does not want a constant — it estimates thresholds instead
+    # Remove the constant column (first column from sm.add_constant)
+    X_no_const = X[:, 1:]
+    var_names_no_const = var_names[1:]  # remove "_cons"
+
+    distr = "logit" if link == "logit" else "probit"
+
+    model = OrderedModel(y, X_no_const, distr=distr)
+    result = model.fit(disp=0)
+
+    # params is a numpy array; exog_names gives labels
+    n_coefs = len(var_names_no_const)
+    n_total = len(result.params)
+    ci = result.conf_int()  # shape: (n_total, 2)
+
+    # Coefficients
+    params_dict = {}
+    se_dict = {}
+    t_dict = {}
+    p_dict = {}
+    ci_low_dict = {}
+    ci_high_dict = {}
+
+    for i, var in enumerate(var_names_no_const):
+        params_dict[var] = float(result.params[i])
+        se_dict[var] = float(result.bse[i])
+        t_dict[var] = float(result.tvalues[i])
+        p_dict[var] = float(result.pvalues[i])
+        ci_low_dict[var] = float(ci[i, 0])
+        ci_high_dict[var] = float(ci[i, 1])
+
+    # Threshold (cut-point) parameters
+    for i in range(n_coefs, n_total):
+        cut_label = f"cut{i - n_coefs + 1}"
+        params_dict[cut_label] = float(result.params[i])
+        se_dict[cut_label] = float(result.bse[i])
+        t_dict[cut_label] = float(result.tvalues[i])
+        p_dict[cut_label] = float(result.pvalues[i])
+        ci_low_dict[cut_label] = float(ci[i, 0])
+        ci_high_dict[cut_label] = float(ci[i, 1])
+
+    suffix = _model_type_suffix(robust, groups is not None)
+    model_name = f"O{link.capitalize()}" + suffix  # OLogit or OProbit
+
+    categories = sorted(np.unique(y))
+    warnings_list.append(f"Ordered categories: {[int(c) if c == int(c) else c for c in categories]}")
+
+    fit = FitResult(
+        model_type=model_name,
+        formula=f"{dep} ~ {' + '.join(indeps)}",
+        dep_var=dep,
+        indep_vars=indeps,
+        n_obs=int(result.nobs),
+        params=params_dict,
+        std_errors=se_dict,
+        t_values=t_dict,
+        p_values=p_dict,
+        conf_int_low=ci_low_dict,
+        conf_int_high=ci_high_dict,
+        pseudo_r2=float(result.prsquared) if hasattr(result, "prsquared") else None,
+        log_likelihood=float(result.llf),
+        aic=float(result.aic),
+        bic=float(result.bic),
+        warnings=warnings_list,
+    )
+    return fit, result

openstat/stats/epidemiology.py

@@ -0,0 +1,119 @@
+"""Epidemiology functions: risk ratios, odds ratios, incidence rates, NNT."""
+
+from __future__ import annotations
+
+import math
+
+import numpy as np
+import polars as pl
+
+
+def _ci_log(est: float, n: int, p: float, alpha: float = 0.05) -> tuple[float, float]:
+    """Approximate log-based CI for ratio estimates."""
+    from scipy.stats import norm
+    z = float(norm.ppf(1 - alpha / 2))
+    if p <= 0 or p >= 1 or n == 0:
+        return float("nan"), float("nan")
+    se_log = math.sqrt((1 - p) / (n * p))
+    return math.exp(math.log(est) - z * se_log), math.exp(math.log(est) + z * se_log)
+
+
+def cohort_study(df: pl.DataFrame, outcome: str, exposure: str) -> dict:
+    """
+    Cohort study: compute RR, ARR, NNT from a 2×2 table.
+    outcome and exposure must be binary (0/1).
+    """
+    from scipy.stats import chi2_contingency, fisher_exact
+    sub = df.select([outcome, exposure]).drop_nulls()
+    exp = sub[exposure].to_numpy().astype(int)
+    out = sub[outcome].to_numpy().astype(int)
+
+    a = int(((exp == 1) & (out == 1)).sum())   # exposed, outcome
+    b = int(((exp == 1) & (out == 0)).sum())   # exposed, no outcome
+    c = int(((exp == 0) & (out == 1)).sum())   # unexposed, outcome
+    d = int(((exp == 0) & (out == 0)).sum())   # unexposed, no outcome
+
+    n_exp = a + b
+    n_unexp = c + d
+    r_exp = a / n_exp if n_exp > 0 else float("nan")
+    r_unexp = c / n_unexp if n_unexp > 0 else float("nan")
+
+    rr = r_exp / r_unexp if r_unexp > 0 else float("nan")
+    arr = r_exp - r_unexp
+    nnt = 1 / abs(arr) if arr != 0 else float("nan")
+
+    table = [[a, b], [c, d]]
+    chi2, p_chi2, _, _ = chi2_contingency(table)
+    _, p_fisher = fisher_exact(table)
+
+    rr_lo, rr_hi = _ci_log(rr, n_exp, r_exp) if rr == rr else (float("nan"), float("nan"))
+
+    return {
+        "test": "Cohort Study (RR)",
+        "exposure": exposure, "outcome": outcome,
+        "table_2x2": {"a": a, "b": b, "c": c, "d": d},
+        "n_exposed": n_exp, "n_unexposed": n_unexp,
+        "risk_exposed": r_exp, "risk_unexposed": r_unexp,
+        "risk_ratio": rr, "rr_ci_95_lo": rr_lo, "rr_ci_95_hi": rr_hi,
+        "arr": arr, "nnt": nnt,
+        "chi2": float(chi2), "p_chi2": float(p_chi2),
+        "p_fisher": float(p_fisher),
+    }
+
+
+def case_control(df: pl.DataFrame, outcome: str, exposure: str) -> dict:
+    """
+    Case-control study: compute OR with 95% CI (Woolf method).
+    """
+    from scipy.stats import chi2_contingency, fisher_exact
+    sub = df.select([outcome, exposure]).drop_nulls()
+    exp = sub[exposure].to_numpy().astype(int)
+    out = sub[outcome].to_numpy().astype(int)
+
+    a = int(((exp == 1) & (out == 1)).sum())
+    b = int(((exp == 0) & (out == 1)).sum())
+    c = int(((exp == 1) & (out == 0)).sum())
+    d = int(((exp == 0) & (out == 0)).sum())
+
+    or_ = (a * d) / (b * c) if b * c > 0 else float("nan")
+    # Woolf 95% CI
+    if or_ > 0 and or_ == or_:
+        from scipy.stats import norm
+        z = float(norm.ppf(0.975))
+        se_log_or = math.sqrt(1/max(a, 1) + 1/max(b, 1) + 1/max(c, 1) + 1/max(d, 1))
+        or_lo = math.exp(math.log(or_) - z * se_log_or)
+        or_hi = math.exp(math.log(or_) + z * se_log_or)
+    else:
+        or_lo = or_hi = float("nan")
+
+    table = [[a, b], [c, d]]
+    chi2, p_chi2, _, _ = chi2_contingency(table)
+    _, p_fisher = fisher_exact(table)
+
+    return {
+        "test": "Case-Control (OR)",
+        "exposure": exposure, "outcome": outcome,
+        "table_2x2": {"a": a, "b": b, "c": c, "d": d},
+        "odds_ratio": or_, "or_ci_95_lo": or_lo, "or_ci_95_hi": or_hi,
+        "chi2": float(chi2), "p_chi2": float(p_chi2),
+        "p_fisher": float(p_fisher),
+    }
+
+
+def incidence_rate(df: pl.DataFrame, outcome: str, person_time: str) -> dict:
+    """Compute incidence rate = cases / total person-time."""
+    sub = df.select([outcome, person_time]).drop_nulls()
+    cases = int(sub[outcome].sum())
+    pt = float(sub[person_time].sum())
+    ir = cases / pt if pt > 0 else float("nan")
+    # Exact Poisson CI (Byar's approximation)
+    from scipy.stats import chi2
+    lo = 0.5 * float(chi2.ppf(0.025, 2 * cases)) / pt if cases > 0 else 0.0
+    hi = 0.5 * float(chi2.ppf(0.975, 2 * (cases + 1))) / pt
+    return {
+        "test": "Incidence Rate",
+        "outcome": outcome, "person_time_col": person_time,
+        "cases": cases, "person_time": pt,
+        "incidence_rate": ir,
+        "ir_ci_95_lo": lo, "ir_ci_95_hi": hi,
+    }

openstat/stats/equiv_tobit.py

@@ -0,0 +1,163 @@
+"""Equivalence tests (TOST) and Tobit/Heckman regression."""
+
+from __future__ import annotations
+
+import numpy as np
+import polars as pl
+from scipy import stats as sp_stats
+
+
+def tost_onemean(
+    df: pl.DataFrame,
+    col: str,
+    mu: float = 0.0,
+    delta: float = 0.5,
+    alpha: float = 0.05,
+) -> dict:
+    """
+    Two One-Sided Tests (TOST) for equivalence: one-sample.
+    H0: |mean - mu| >= delta  vs  H1: |mean - mu| < delta
+    """
+    data = df[col].drop_nulls().to_numpy().astype(float)
+    n = len(data)
+    x_bar = data.mean()
+    se = data.std(ddof=1) / np.sqrt(n)
+
+    # Lower test: H0: mean <= mu - delta
+    t_lo = (x_bar - (mu - delta)) / se
+    p_lo = float(sp_stats.t.sf(t_lo, df=n - 1))  # one-sided upper
+
+    # Upper test: H0: mean >= mu + delta
+    t_hi = (x_bar - (mu + delta)) / se
+    p_hi = float(sp_stats.t.cdf(t_hi, df=n - 1))  # one-sided lower
+
+    p_tost = max(p_lo, p_hi)
+    equivalent = p_tost < alpha
+
+    return {
+        "test": "TOST Equivalence (one-sample)",
+        "col": col,
+        "n_obs": n,
+        "mean": float(x_bar),
+        "mu": mu,
+        "delta": delta,
+        "alpha": alpha,
+        "t_lower": float(t_lo),
+        "t_upper": float(t_hi),
+        "p_lower": p_lo,
+        "p_upper": p_hi,
+        "p_tost": p_tost,
+        "equivalent_at_alpha": equivalent,
+    }
+
+
+def tost_twomeans(
+    df: pl.DataFrame,
+    col: str,
+    by: str,
+    delta: float = 0.5,
+    alpha: float = 0.05,
+) -> dict:
+    """TOST for equivalence of two independent group means."""
+    groups = df[by].drop_nulls().unique().sort().to_list()
+    if len(groups) != 2:
+        raise ValueError(f"tost_twomeans 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)
+
+    diff = float(g1.mean() - g2.mean())
+    se = float(np.sqrt(g1.var(ddof=1) / len(g1) + g2.var(ddof=1) / len(g2)))
+    df_welch = int((g1.var(ddof=1) / len(g1) + g2.var(ddof=1) / len(g2))**2 /
+                   ((g1.var(ddof=1) / len(g1))**2 / (len(g1) - 1) +
+                    (g2.var(ddof=1) / len(g2))**2 / (len(g2) - 1)))
+
+    t_lo = (diff - (-delta)) / se
+    t_hi = (diff - delta) / se
+    p_lo = float(sp_stats.t.sf(t_lo, df=df_welch))
+    p_hi = float(sp_stats.t.cdf(t_hi, df=df_welch))
+    p_tost = max(p_lo, p_hi)
+
+    return {
+        "test": "TOST Equivalence (two-sample)",
+        "col": col, "by": by,
+        "groups": [str(g) for g in groups],
+        "mean_diff": diff,
+        "delta": delta,
+        "alpha": alpha,
+        "p_tost": p_tost,
+        "equivalent_at_alpha": p_tost < alpha,
+        "p_lower": p_lo,
+        "p_upper": p_hi,
+    }
+
+
+def fit_tobit(
+    df: pl.DataFrame,
+    dep: str,
+    indeps: list[str],
+    left: float | None = 0.0,
+    right: float | None = None,
+) -> dict:
+    """
+    Tobit regression for censored outcomes via MLE (scipy optimize).
+    Handles left-censoring (default at 0), right-censoring, or both.
+    """
+    from scipy.optimize import minimize
+    from scipy.stats import norm
+
+    sub = df.select([dep] + indeps).drop_nulls()
+    y = sub[dep].to_numpy().astype(float)
+    X_raw = sub.select(indeps).to_numpy().astype(float)
+    n, k = X_raw.shape
+    X = np.column_stack([np.ones(n), X_raw])
+    kp = k + 1
+
+    def neg_ll(params):
+        beta = params[:kp]
+        log_sigma = params[kp]
+        sigma = np.exp(log_sigma)
+        xb = X @ beta
+        ll = np.zeros(n)
+
+        for i in range(n):
+            if left is not None and y[i] <= left:
+                ll[i] = norm.logcdf((left - xb[i]) / sigma)
+            elif right is not None and y[i] >= right:
+                ll[i] = norm.logsf((right - xb[i]) / sigma)
+            else:
+                ll[i] = norm.logpdf(y[i], loc=xb[i], scale=sigma)
+        return -ll.sum()
+
+    # OLS start
+    beta0 = np.linalg.lstsq(X, y, rcond=None)[0]
+    resid0 = y - X @ beta0
+    log_sigma0 = np.log(max(resid0.std(), 1e-4))
+    x0 = np.concatenate([beta0, [log_sigma0]])
+
+    try:
+        res = minimize(neg_ll, x0, method="L-BFGS-B", options={"maxiter": 500})
+        beta_hat = res.x[:kp]
+        sigma_hat = float(np.exp(res.x[kp]))
+        llf = -res.fun
+        aic = 2 * (kp + 1) - 2 * llf
+        bic = (kp + 1) * np.log(n) - 2 * llf
+
+        param_names = ["_cons"] + indeps
+        params = {nm: float(v) for nm, v in zip(param_names, beta_hat)}
+
+        return {
+            "method": "Tobit",
+            "dep": dep, "indeps": indeps,
+            "left_censoring": left,
+            "right_censoring": right,
+            "params": params,
+            "sigma": sigma_hat,
+            "log_likelihood": float(llf),
+            "aic": float(aic),
+            "bic": float(bic),
+            "n_obs": n,
+            "n_censored_left": int((y <= left).sum()) if left is not None else 0,
+            "n_censored_right": int((y >= right).sum()) if right is not None else 0,
+        }
+    except Exception as exc:
+        raise RuntimeError(f"Tobit failed: {exc}") from exc