openstat-cli 1.0.0__py3-none-any.whl

openstat/commands/meta_cmds.py

@@ -0,0 +1,284 @@
+"""Meta-analysis commands: meta, forest plot, funnel plot."""
+
+from __future__ import annotations
+
+import math
+from pathlib import Path
+
+import numpy as np
+from scipy import stats
+
+from openstat.commands.base import command
+from openstat.session import Session
+
+
+def _meta_fe(effects: np.ndarray, variances: np.ndarray) -> dict:
+    """Fixed-effects meta-analysis (inverse-variance weighting)."""
+    weights = 1.0 / variances
+    pooled = np.sum(weights * effects) / np.sum(weights)
+    se = math.sqrt(1.0 / np.sum(weights))
+    z = pooled / se
+    p = 2 * (1 - stats.norm.cdf(abs(z)))
+    ci_low = pooled - 1.96 * se
+    ci_high = pooled + 1.96 * se
+
+    # Cochran's Q
+    q = np.sum(weights * (effects - pooled) ** 2)
+    df = len(effects) - 1
+    q_pval = 1 - stats.chi2.cdf(q, df) if df > 0 else float("nan")
+    i2 = max(0.0, (q - df) / q * 100) if q > df else 0.0
+
+    return {
+        "pooled": pooled, "se": se, "z": z, "p": p,
+        "ci_low": ci_low, "ci_high": ci_high,
+        "Q": q, "Q_df": df, "Q_p": q_pval, "I2": i2,
+    }
+
+
+def _meta_re(effects: np.ndarray, variances: np.ndarray) -> dict:
+    """Random-effects meta-analysis (DerSimonian-Laird)."""
+    weights_fe = 1.0 / variances
+    pooled_fe = np.sum(weights_fe * effects) / np.sum(weights_fe)
+    q = np.sum(weights_fe * (effects - pooled_fe) ** 2)
+    df = len(effects) - 1
+
+    # Between-study variance (tau^2) via DerSimonian-Laird
+    c = np.sum(weights_fe) - np.sum(weights_fe ** 2) / np.sum(weights_fe)
+    tau2 = max(0.0, (q - df) / c)
+
+    # RE weights
+    weights_re = 1.0 / (variances + tau2)
+    pooled_re = np.sum(weights_re * effects) / np.sum(weights_re)
+    se_re = math.sqrt(1.0 / np.sum(weights_re))
+    z = pooled_re / se_re
+    p = 2 * (1 - stats.norm.cdf(abs(z)))
+    ci_low = pooled_re - 1.96 * se_re
+    ci_high = pooled_re + 1.96 * se_re
+
+    q_pval = 1 - stats.chi2.cdf(q, df) if df > 0 else float("nan")
+    i2 = max(0.0, (q - df) / q * 100) if q > df else 0.0
+
+    return {
+        "pooled": pooled_re, "se": se_re, "z": z, "p": p,
+        "ci_low": ci_low, "ci_high": ci_high,
+        "tau2": tau2, "tau": math.sqrt(tau2),
+        "Q": q, "Q_df": df, "Q_p": q_pval, "I2": i2,
+        "model": "Random-effects (DerSimonian-Laird)",
+    }
+
+
+def _sig(p: float) -> str:
+    if p < 0.001: return "***"
+    if p < 0.01: return "**"
+    if p < 0.05: return "*"
+    return ""
+
+
+def _plot_forest(
+    study_labels: list[str],
+    effects: np.ndarray,
+    ci_low: np.ndarray,
+    ci_high: np.ndarray,
+    pooled: float,
+    pooled_ci_low: float,
+    pooled_ci_high: float,
+    output_dir: Path,
+    title: str = "Forest Plot",
+) -> Path:
+    import matplotlib
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    n = len(study_labels)
+    fig_h = max(4, n * 0.45 + 2.5)
+    fig, ax = plt.subplots(figsize=(10, fig_h))
+
+    y_pos = np.arange(n, 0, -1, dtype=float)
+
+    # Individual studies
+    ax.errorbar(
+        effects, y_pos,
+        xerr=[effects - ci_low, ci_high - effects],
+        fmt="s", color="#4C72B0", ecolor="#4C72B0",
+        capsize=3, markersize=5, linewidth=1.2, label="Studies",
+    )
+
+    # Pooled diamond
+    diamond_x = [pooled_ci_low, pooled, pooled_ci_high, pooled, pooled_ci_low]
+    diamond_y = [0.0, 0.25, 0.0, -0.25, 0.0]
+    ax.fill(diamond_x, diamond_y, color="#E66100", zorder=5, label="Pooled")
+    ax.plot([pooled, pooled], [-0.3, n + 0.3], color="#E66100",
+            linestyle="--", linewidth=1, alpha=0.6)
+
+    ax.axvline(0, color="gray", linestyle="--", linewidth=1, alpha=0.7)
+
+    ax.set_yticks(y_pos)
+    ax.set_yticklabels(study_labels)
+    ax.set_xlabel("Effect Size")
+    ax.set_title(title)
+    ax.legend(loc="upper right", fontsize=9)
+    fig.tight_layout()
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+    from openstat.plots.plotter import _unique_path
+    path = _unique_path(output_dir, "forest_plot")
+    fig.savefig(path, dpi=150)
+    plt.close(fig)
+    return path
+
+
+def _plot_funnel(
+    effects: np.ndarray,
+    se: np.ndarray,
+    pooled: float,
+    output_dir: Path,
+) -> Path:
+    import matplotlib
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    max_se = se.max()
+    se_range = np.linspace(0, max_se * 1.1, 100)
+    ci_l = pooled - 1.96 * se_range
+    ci_u = pooled + 1.96 * se_range
+
+    fig, ax = plt.subplots(figsize=(7, 6))
+    ax.scatter(effects, se, alpha=0.7, s=40, color="#4C72B0")
+    ax.plot(ci_l, se_range, "r--", linewidth=1, label="95% CI")
+    ax.plot(ci_u, se_range, "r--", linewidth=1)
+    ax.axvline(pooled, color="gray", linestyle="--", linewidth=1, label=f"Pooled={pooled:.3f}")
+    ax.invert_yaxis()
+    ax.set_xlabel("Effect Size")
+    ax.set_ylabel("Standard Error")
+    ax.set_title("Funnel Plot")
+    ax.legend(fontsize=9)
+    fig.tight_layout()
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+    from openstat.plots.plotter import _unique_path
+    path = _unique_path(output_dir, "funnel_plot")
+    fig.savefig(path, dpi=150)
+    plt.close(fig)
+    return path
+
+
+@command("meta", usage="meta <effect_col> <se_col> [study=<label_col>] [--re|--fe] [--forest] [--funnel]")
+def cmd_meta(session: Session, args: str) -> str:
+    """Meta-analysis: fixed-effects or random-effects pooling with forest/funnel plots.
+
+    Requires columns for effect sizes and their standard errors.
+    Uses DerSimonian-Laird for random-effects (default).
+
+    Examples:
+      meta es se
+      meta es se study=author --re
+      meta es se study=author --forest --funnel
+      meta logOR se_logOR study=trial --fe
+    """
+    import re
+    import polars as pl
+
+    df = session.require_data()
+    args = args.strip()
+
+    # Parse options
+    fe = "--fe" in args
+    re_flag = "--re" in args or not fe  # default: random-effects
+    forest = "--forest" in args
+    funnel = "--funnel" in args
+    args_clean = re.sub(r"--\w+", "", args).strip()
+
+    # study=col
+    study_col = None
+    m = re.search(r"study[= ](\w+)", args_clean)
+    if m:
+        study_col = m.group(1)
+        args_clean = args_clean[:m.start()] + args_clean[m.end():]
+
+    tokens = args_clean.split()
+    if len(tokens) < 2:
+        return (
+            "Usage: meta <effect_col> <se_col> [study=<label_col>] [--re|--fe] [--forest] [--funnel]\n"
+            "Effect sizes must be pre-computed (e.g., Cohen's d, log OR, standardized mean diff)."
+        )
+
+    eff_col, se_col = tokens[0], tokens[1]
+    for c in [eff_col, se_col]:
+        if c not in df.columns:
+            return f"Column not found: {c}"
+
+    sub = df.select([c for c in [eff_col, se_col, study_col] if c]).drop_nulls()
+    effects = sub[eff_col].to_numpy().astype(float)
+    ses = sub[se_col].to_numpy().astype(float)
+    variances = ses ** 2
+    k = len(effects)
+
+    if k < 2:
+        return "Need at least 2 studies for meta-analysis."
+
+    labels = (
+        [str(v) for v in sub[study_col].to_list()]
+        if study_col
+        else [f"Study {i+1}" for i in range(k)]
+    )
+
+    # Run analysis
+    method = "FE" if fe else "RE"
+    res = _meta_fe(effects, variances) if fe else _meta_re(effects, variances)
+
+    lines = [
+        f"Studies (k): {k}",
+        f"Method: {'Fixed-effects (IV)' if fe else 'Random-effects (DerSimonian-Laird)'}",
+        "",
+        "Individual Studies:",
+        f"  {'Study':<20} {'Effect':>8} {'SE':>7} {'95% CI':>18}",
+        "  " + "-" * 55,
+    ]
+    ci_low_arr = effects - 1.96 * ses
+    ci_high_arr = effects + 1.96 * ses
+    for lbl, eff, se_i, lo, hi in zip(labels, effects, ses, ci_low_arr, ci_high_arr):
+        lines.append(f"  {lbl:<20} {eff:>8.4f} {se_i:>7.4f} [{lo:>7.4f}, {hi:>7.4f}]")
+
+    lines += [
+        "",
+        "Pooled Estimate:",
+        f"  Effect = {res['pooled']:.4f}  (95% CI: [{res['ci_low']:.4f}, {res['ci_high']:.4f}])",
+        f"  SE = {res['se']:.4f}   z = {res['z']:.3f}   p = {res['p']:.4f} {_sig(res['p'])}",
+        "",
+        "Heterogeneity:",
+        f"  Q({res['Q_df']}) = {res['Q']:.3f}  p = {res['Q_p']:.4f}",
+        f"  I² = {res['I2']:.1f}%",
+    ]
+    if not fe:
+        lines.append(f"  τ² = {res['tau2']:.4f}   τ = {res['tau']:.4f}")
+
+    out = "\n" + "=" * 60 + "\nMeta-Analysis Results\n" + "=" * 60 + "\n"
+    out += "\n".join(lines) + "\n" + "=" * 60
+
+    # Plots
+    plot_msgs = []
+    if forest:
+        try:
+            p = _plot_forest(
+                labels, effects, ci_low_arr, ci_high_arr,
+                res["pooled"], res["ci_low"], res["ci_high"],
+                session.output_dir,
+                title=f"Forest Plot (pooled={res['pooled']:.3f})",
+            )
+            session.plot_paths.append(str(p))
+            plot_msgs.append(f"Forest plot saved: {p}")
+        except Exception as exc:
+            plot_msgs.append(f"Forest plot error: {exc}")
+
+    if funnel:
+        try:
+            p = _plot_funnel(effects, ses, res["pooled"], session.output_dir)
+            session.plot_paths.append(str(p))
+            plot_msgs.append(f"Funnel plot saved: {p}")
+        except Exception as exc:
+            plot_msgs.append(f"Funnel plot error: {exc}")
+
+    if plot_msgs:
+        out += "\n" + "\n".join(plot_msgs)
+
+    return out

openstat/commands/mi_cmds.py

@@ -0,0 +1,228 @@
+"""Multiple imputation commands: mi impute, mi estimate, mi describe."""
+
+from __future__ import annotations
+
+import re
+import io
+
+import polars as pl
+from rich.console import Console
+from rich.table import Table
+
+from openstat.session import Session, ModelResult
+from openstat.commands.base import command, CommandArgs, rich_to_str, friendly_error
+from openstat.dsl.parser import parse_formula, ParseError
+
+
+def _parse_impute_specs(args: str) -> tuple[list[tuple[str, str]], dict]:
+    """Parse imputation specification.
+
+    Formats:
+        chained (regress) x1 (logit) x2, add(5)
+        pmm x1 x2, add(10) knn(5)
+
+    Returns (specs, options).
+    """
+    # Split on comma for options
+    if ',' in args:
+        spec_part, opt_part = args.split(',', 1)
+    else:
+        spec_part, opt_part = args, ""
+
+    # Parse options
+    options: dict = {}
+    m_add = re.search(r'add\((\d+)\)', opt_part)
+    if m_add:
+        options["m"] = int(m_add.group(1))
+    m_knn = re.search(r'knn\((\d+)\)', opt_part)
+    if m_knn:
+        options["knn"] = int(m_knn.group(1))
+
+    spec_part = spec_part.strip()
+    specs: list[tuple[str, str]] = []
+
+    if spec_part.startswith("chained"):
+        # Parse (method) var pairs
+        spec_str = spec_part[len("chained"):].strip()
+        pattern = r'\((\w+)\)\s+(\w+)'
+        for match in re.finditer(pattern, spec_str):
+            method = match.group(1)
+            col = match.group(2)
+            specs.append((method, col))
+    elif spec_part.startswith("pmm"):
+        # PMM for all listed variables
+        cols = spec_part[len("pmm"):].strip().split()
+        for col in cols:
+            specs.append(("pmm", col))
+    else:
+        # Default: regress for all variables
+        cols = spec_part.split()
+        for col in cols:
+            specs.append(("regress", col))
+
+    return specs, options
+
+
+@command("mi", usage="mi impute|estimate|describe ...")
+def cmd_mi(session: Session, args: str) -> str:
+    """Multiple imputation: impute missing data, estimate models, describe imputations."""
+    df = session.require_data()
+
+    parts = args.strip().split(None, 1)
+    subcmd = parts[0].lower() if parts else ""
+    rest = parts[1] if len(parts) > 1 else ""
+
+    if subcmd == "impute":
+        return _mi_impute(session, df, rest)
+    elif subcmd.startswith("estimate"):
+        # Handle "estimate:" or "estimate :"
+        model_cmd = rest
+        if model_cmd.startswith(":"):
+            model_cmd = model_cmd[1:].strip()
+        elif ":" in args:
+            model_cmd = args.split(":", 1)[1].strip()
+        return _mi_estimate(session, model_cmd)
+    elif subcmd == "describe":
+        return _mi_describe(session)
+    else:
+        return "Usage: mi impute|estimate|describe\n  mi impute chained (regress) x1 (logit) x2, add(5)\n  mi estimate: ols y ~ x1 + x2\n  mi describe"
+
+def _mi_impute(session: Session, df: pl.DataFrame, args: str) -> str:
+    """Run MICE imputation."""
+    specs, options = _parse_impute_specs(args)
+    if not specs:
+        return "No variables specified for imputation."
+
+    m = options.get("m", 5)
+    knn = options.get("knn", 5)
+
+    # Validate columns
+    for method, col in specs:
+        if col not in df.columns:
+            return f"Column not found: {col}"
+
+    # Check for missing values (both null and NaN)
+    has_missing = False
+    for _, col in specs:
+        if df[col].null_count() > 0:
+            has_missing = True
+            break
+        if df[col].dtype.is_float() and df[col].is_nan().sum() > 0:
+            has_missing = True
+            break
+    if not has_missing:
+        return "No missing values found in specified columns."
+
+    try:
+        from openstat.stats.imputation import mice_impute
+
+        datasets = mice_impute(df, specs, m=m)
+        session._imputed_datasets = datasets
+        session._mi_m = m
+
+        lines = [f"Created {m} imputed datasets."]
+        for method, col in specs:
+            n_missing = df[col].null_count()
+            lines.append(f"  {col}: {n_missing} missing values imputed ({method})")
+        lines.append(f"\nUse 'mi estimate: <model>' to run analysis across imputed datasets.")
+        return "\n".join(lines)
+    except Exception as e:
+        return friendly_error(e, "mi impute")
+
+
+def _mi_estimate(session: Session, model_cmd: str) -> str:
+    """Run model on each imputed dataset and combine with Rubin's rules."""
+    if not session._imputed_datasets:
+        return "No imputed datasets. Run 'mi impute' first."
+
+    model_cmd = model_cmd.strip()
+    if not model_cmd:
+        return "Usage: mi estimate: ols y ~ x1 + x2"
+
+    # Parse model type and formula
+    parts = model_cmd.split(None, 1)
+    model_type = parts[0].lower()
+    formula_str = parts[1] if len(parts) > 1 else ""
+
+    try:
+        dep, indeps = parse_formula(formula_str)
+    except ParseError as e:
+        return f"Formula error: {e}"
+
+    estimates: list[dict[str, float]] = []
+    std_errors: list[dict[str, float]] = []
+
+    for i, imp_df in enumerate(session._imputed_datasets):
+        try:
+            if model_type == "ols":
+                from openstat.stats.models import fit_ols
+                result, _ = fit_ols(imp_df, dep, indeps)
+            elif model_type == "logit":
+                from openstat.stats.models import fit_logit
+                result, _ = fit_logit(imp_df, dep, indeps)
+            else:
+                return f"Unsupported model for MI: {model_type}. Use ols or logit."
+
+            estimates.append(dict(result.params))
+            std_errors.append(dict(result.std_errors))
+        except Exception as e:
+            return f"Error fitting imputation {i + 1}: {e}"
+
+    # Combine with Rubin's rules
+    from openstat.stats.imputation import rubins_rules
+    n_obs = session._imputed_datasets[0].height
+    mi_result = rubins_rules(estimates, std_errors, n_obs)
+    mi_result.model_type = f"MI ({mi_result.m} imputations): {model_type.upper()}"
+    mi_result.formula = f"{dep} ~ {' + '.join(indeps)}"
+
+    # Display as table
+    def render(console: Console) -> None:
+        table = Table(title=mi_result.model_type)
+        table.add_column("Variable", style="cyan")
+        table.add_column("Coef", justify="right")
+        table.add_column("MI Std.Err", justify="right")
+        table.add_column("t", justify="right")
+        table.add_column("P>|t|", justify="right")
+        table.add_column("[95% CI Low]", justify="right")
+        table.add_column("[95% CI High]", justify="right")
+        table.add_column("FMI", justify="right")
+
+        for var in mi_result.params:
+            sig = ""
+            pv = mi_result.p_values[var]
+            if pv < 0.001:
+                sig = " ***"
+            elif pv < 0.01:
+                sig = " **"
+            elif pv < 0.05:
+                sig = " *"
+
+            table.add_row(
+                var,
+                f"{mi_result.params[var]:.4f}",
+                f"{mi_result.std_errors[var]:.4f}",
+                f"{mi_result.t_values[var]:.3f}",
+                f"{pv:.4f}{sig}",
+                f"{mi_result.conf_int_low[var]:.4f}",
+                f"{mi_result.conf_int_high[var]:.4f}",
+                f"{mi_result.fmi[var]:.3f}",
+            )
+        console.print(table)
+
+    output = rich_to_str(render)
+    output += f"\nN = {mi_result.n_obs}  |  Imputations = {mi_result.m}"
+    return output
+
+
+def _mi_describe(session: Session) -> str:
+    """Describe imputed datasets."""
+    if not session._imputed_datasets:
+        return "No imputed datasets. Run 'mi impute' first."
+
+    lines = [
+        f"Multiple Imputation Summary:",
+        f"  Number of imputations (m): {session._mi_m}",
+        f"  Rows per dataset: {session._imputed_datasets[0].height}",
+        f"  Columns: {session._imputed_datasets[0].width}",
+    ]
+    return "\n".join(lines)

openstat/commands/mixed_cmds.py

@@ -0,0 +1,79 @@
+"""Mixed / hierarchical model commands: mixed, estat icc."""
+
+from __future__ import annotations
+
+from openstat.session import Session, ModelResult
+from openstat.commands.base import command, CommandArgs, friendly_error
+from openstat.dsl.parser import parse_formula, ParseError
+
+
+def _parse_mixed_formula(raw: str) -> tuple[str, list[str], str, list[str]]:
+    """Parse mixed model formula: y ~ x1 || group: x1.
+
+    Returns (dep, fixed_effects, group_var, random_effects).
+    """
+    if '||' not in raw:
+        raise ParseError("Mixed model requires '||' to specify grouping: y ~ x1 || group: [re_vars]")
+
+    fixed_part, random_part = raw.split('||', 1)
+    dep, fixed = parse_formula(fixed_part.strip())
+
+    random_part = random_part.strip()
+    if ':' in random_part:
+        group_str, re_str = random_part.split(':', 1)
+        group_var = group_str.strip()
+        re_vars = re_str.strip().split() if re_str.strip() else []
+    else:
+        group_var = random_part.strip()
+        re_vars = []
+
+    if not group_var:
+        raise ParseError("No grouping variable specified after '||'")
+
+    return dep, fixed, group_var, re_vars
+
+
+@command("mixed", usage="mixed y ~ x1 || group: [re_vars]")
+def cmd_mixed(session: Session, args: str) -> str:
+    """Fit a mixed/hierarchical linear model with random intercepts and/or slopes."""
+    df = session.require_data()
+
+    try:
+        dep, fixed, group_var, re_vars = _parse_mixed_formula(args.strip())
+    except ParseError as e:
+        return f"Formula error: {e}"
+
+    # Validate columns
+    all_cols = [dep] + fixed + [group_var] + re_vars
+    missing = [c for c in all_cols if c not in df.columns]
+    if missing:
+        return f"Columns not found: {', '.join(missing)}"
+
+    try:
+        from openstat.stats.mixed import fit_mixed
+
+        result, raw = fit_mixed(df, dep, fixed, group_var, re_vars or None)
+
+        session._last_model = raw
+        session._last_model_vars = (dep, fixed)
+        session._last_fit_result = result
+        session._last_fit_kwargs = {"group_var": group_var, "re_vars": re_vars}
+
+        md = result.to_markdown() if hasattr(result, "to_markdown") else ""
+        session.results.append(ModelResult(
+            name="Mixed LM", formula=result.formula,
+            table=md, details={
+                "n_obs": result.n_obs,
+                "params": dict(result.params),
+                "aic": result.aic,
+                "bic": result.bic,
+                "log_likelihood": result.log_likelihood,
+            },
+        ))
+
+        output = result.summary_table()
+        if result.warnings:
+            output += "\n" + "\n".join(result.warnings)
+        return output
+    except Exception as e:
+        return friendly_error(e, "mixed")