typestar

sm_regression_report.py in Python

A regression worked properly: fit, diagnose, correct, compare, forecast.

import numpy as np
import pandas as pd
import statsmodels.formula.api as smf
from statsmodels.stats.anova import anova_lm
from statsmodels.stats.diagnostic import het_breuschpagan
from statsmodels.stats.stattools import durbin_watson, jarque_bera

LANG_EFFECT = {"python": 9.0, "rust": 5.0, "sql": 0.0}


def make_frame(rows=240, seed=31):
    """Practice history: sessions, hours, language, editor mode."""
    rng = np.random.default_rng(seed)
    frame = pd.DataFrame({
        "sessions": rng.integers(1, 80, size=rows),
        "hours": rng.uniform(0.25, 3.0, size=rows),
        "lang": rng.choice(list(LANG_EFFECT), size=rows),
        "editor": rng.integers(0, 2, size=rows),
    })
    signal = (72
              + 0.35 * frame["sessions"]
              + 3.0 * frame["hours"]
              + frame["lang"].map(LANG_EFFECT)
              + 4.0 * frame["editor"])
    # noise grows with sessions, so the errors are heteroscedastic
    frame["tpm"] = signal + rng.normal(0, 1 + frame["sessions"] / 25)
    return frame


def diagnose(model, frame):
    print("diagnostics")
    _, bp_p, _, _ = het_breuschpagan(model.resid, model.model.exog)
    print(f"  breusch-pagan p {bp_p:.2e}"
          f" -> {'heteroscedastic' if bp_p < 0.05 else 'constant'}")
    print(f"  durbin-watson   {durbin_watson(model.resid):.3f}")
    print(f"  jarque-bera p   {jarque_bera(model.resid)[1]:.4f}")
    return bp_p < 0.05


def main():
    frame = make_frame()

    simple = smf.ols("tpm ~ sessions + hours", data=frame).fit()
    full = smf.ols("tpm ~ sessions + hours + C(lang) + editor",
                   data=frame).fit()

    print(f"simple r2 {simple.rsquared:.4f}   full r2 {full.rsquared:.4f}")
    f_stat, f_p, _ = full.compare_f_test(simple)
    print(f"adding language and editor: F={f_stat:.2f} p={f_p:.2e}")

    needs_robust = diagnose(full, frame)
    used = full.get_robustcov_results("HC3") if needs_robust else full
    label = "HC3 robust" if needs_robust else "classical"

    print(f"\ncoefficients ({label} standard errors)")
    for name, coef, err in zip(full.params.index, used.params, used.bse):
        print(f"  {name[:34]:<36}{coef:>8.3f} +/- {err:.3f}")

    print("\nanova (type 2)")
    print(anova_lm(full, typ=2).round(3).to_string())

    future = pd.DataFrame({"sessions": [40, 80], "hours": [1.5, 2.5],
                           "lang": ["python", "rust"], "editor": [1, 0]})
    predicted = full.get_prediction(future).summary_frame(alpha=0.05)
    print("\nforecast")
    for row, (mean, low, high) in enumerate(zip(
            predicted["mean"], predicted["obs_ci_lower"],
            predicted["obs_ci_upper"])):
        print(f"  row {row}: {mean:.1f} tpm  (95% obs {low:.1f}-{high:.1f})")


if __name__ == "__main__":
    main()

How it works

  1. The formula API keeps the model readable next to the data.
  2. The diagnostics decide whether robust errors are needed.
  3. A nested model comparison justifies the extra terms.

Keywords and builtins used here

The run, in numbers

Lines
73
Characters to type
2455
Tokens
712
Three-star pace
115 tpm

At the three-star pace of 115 tokens a minute, this run takes about 371 seconds.

Type this snippet

Step 1 of 1 in Encore, step 19 of 19 in Statistics with statsmodels.

← Previous