SuperGLM Editor Demo#
This notebook demonstrates the optional superglm.editor module on a synthetic Tweedie/log-link pricing model. It fits splines and categorical factors, profiles the Tweedie power parameter, opens an editable notebook widget, applies curve edits to a copied model, and compares the original and edited effects.
The editor only supports 1D main effects. Interactions are intentionally excluded.
Setup#
Install the optional editor dependencies before running this notebook:
uv sync --extra dev
In VS Code, select the Python environment from the worktree you opened: .venv/bin/python (for this branch, /home/mhick/python_projects/superglm/.worktrees/editor-refit-timing-debugging/.venv/bin/python). The editor renders as a local iframe app served by the Python kernel, so it does not require anywidget, ipympl, or other custom Jupyter widget frontend modules.
import sys
import time
from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from sklearn.model_selection import KFold, train_test_split
from superglm import (
Categorical,
Numeric,
OrderedCategorical,
Spline,
SuperGLM,
cross_validate,
families,
generate_tweedie_cpg,
)
from superglm.editor import EditorSession
print(sys.executable)
Synthetic Data#
The response is generated from a Tweedie compound Poisson-Gamma process with a log link, exposure weights, a continuous age effect, one numeric effect, and categorical effects. age_band is deliberately included in the model as a coarse duplicate of age, not as a true separate signal; at low spline capacity it can look useful because it patches residual age shape.
rng = np.random.default_rng(20260703)
n = 150_000
TRUE_TWEEDIE_P = 1.45
TRUE_TWEEDIE_P_INIT = 1.35
TRUE_TWEEDIE_PHI = 0.45
age = np.clip(rng.normal(loc=48.0, scale=13.5, size=n), 18, 80)
mileage = rng.normal(0.0, 1.0, n)
exposure = rng.gamma(shape=4.0, scale=0.25, size=n) + 0.05
region = rng.choice(["North", "South", "West"], n, p=[0.40, 0.35, 0.25])
TERRITORY_LEVELS = [f"T{i:02d}" for i in range(1, 11)]
territory_probs = np.array([0.015, 0.25, 0.18, 0.14, 0.12, 0.10, 0.08, 0.025, 0.07, 0.02])
territory_probs = territory_probs / territory_probs.sum()
territory = rng.choice(TERRITORY_LEVELS, n, p=territory_probs)
territory_effect_map = {
"T01": 0.10,
"T02": 0.11,
"T03": -0.16,
"T04": 0.00,
"T05": 0.24,
"T06": -0.10,
"T07": 0.18,
"T08": 0.09,
"T09": -0.22,
"T10": 0.12,
}
territory_effect = np.array([territory_effect_map[level] for level in territory], dtype=float)
age_band = pd.cut(
age,
bins=[18, 25, 35, 50, 65, 81],
labels=["18-24", "25-34", "35-49", "50-64", "65-80"],
right=False,
).astype(str)
age_effect = (
0.24 * np.exp(-0.5 * ((age - 31) / 6.0) ** 2)
- 0.20 * np.exp(-0.5 * ((age - 52) / 7.5) ** 2)
+ 0.18 * np.exp(-0.5 * ((age - 72) / 5.5) ** 2)
+ 0.08 * np.sin((age - 18) / 4.2)
- 0.05 * np.maximum((age - 66) / 14, 0)
)
age_effect -= np.average(age_effect, weights=exposure)
eta = (
0.15
+ age_effect
+ 0.09 * mileage
+ 0.22 * (region == "South")
- 0.16 * (region == "West")
+ territory_effect
)
mu = np.exp(eta)
y = generate_tweedie_cpg(
n,
mu=mu,
phi=TRUE_TWEEDIE_PHI / exposure,
p=TRUE_TWEEDIE_P,
rng=rng,
)
X = pd.DataFrame(
{
"age": age,
"mileage": mileage,
"region": region,
"age_band": age_band,
"territory": territory,
}
)
X_train_val, X_test, y_train_val, y_test, w_train_val, w_test = train_test_split(
X,
y,
exposure,
test_size=0.15,
random_state=20260704,
stratify=X["age_band"],
)
X_train, X_val, y_train, y_val, w_train, w_val = train_test_split(
X_train_val,
y_train_val,
w_train_val,
test_size=0.1765,
random_state=20260705,
stratify=X_train_val["age_band"],
)
print(len(X_train), len(X_val), len(X_test))
X_train.head()
Choose Spline Capacity With A k Sweep#
k is the public spline basis size: it gives the smooth enough capacity to represent shape, while REML decides how much of that capacity is actually used. Treat the first table as an in-sample adequacy diagnostic, not as proof of the optimal k.
The more defensible check is cross-validated validation deviance. The conservative rule below is to choose the smallest k within one standard error of the best CV deviance, provided the curve and term-competition story still make sense.
AGE_BAND_ORDER = ["18-24", "25-34", "35-49", "50-64", "65-80"]
def make_features(age_k: int):
return {
"age": Spline(kind="bs", k=age_k, knot_strategy="quantile_tempered"),
"mileage": Numeric(),
"region": Categorical(base="first"),
"age_band": OrderedCategorical(
order=AGE_BAND_ORDER,
basis=Spline(kind="ps", k=5),
base="first",
),
"territory": Categorical(base="most_exposed"),
}
def make_model(age_k: int) -> SuperGLM:
return SuperGLM(
family=families.tweedie(p=TRUE_TWEEDIE_P_INIT),
selection_penalty=0.0,
spline_penalty=0.15,
features=make_features(age_k),
discrete=True,
n_bins=512,
)
def fit_age_model(
age_k: int,
X_fit: pd.DataFrame = X_train,
y_fit: np.ndarray = y_train,
w_fit: np.ndarray | None = w_train,
) -> SuperGLM:
candidate = make_model(age_k)
candidate.fit_reml(X_fit, y_fit, sample_weight=w_fit, max_reml_iter=8, reml_tol=1e-4)
return candidate
def mean_unit_deviance(
model: SuperGLM,
X_eval: pd.DataFrame,
y_eval: np.ndarray,
sample_weight: np.ndarray | None = None,
) -> float:
mu = np.asarray(model.predict(X_eval), dtype=float)
dev = model._distribution.deviance_unit(np.asarray(y_eval, dtype=float), mu)
if sample_weight is None:
return float(np.mean(dev))
return float(np.average(dev, weights=np.asarray(sample_weight, dtype=float)))
def term_p(summary, term: str) -> float:
row = next((row for row in summary._coef_rows if row.name == term), None)
if row is None:
return np.nan
p_value = row.wald_p if row.wald_p is not None else row.p
return float(p_value) if p_value is not None else np.nan
def cv_diagnostics(model, X_val, y_val, *, sample_weight=None, offset=None):
age_info = model.diagnostics().get("age", {})
return {
"age_edf": float(age_info.get("edf", np.nan)),
"age_band_p": term_p(model.summary(), "age_band"),
}
def format_sweep_table(frame: pd.DataFrame) -> pd.DataFrame:
out = frame.copy()
if "edf_share" in out:
out["edf_share"] = 100 * out["edf_share"]
out = out.rename(columns={"edf_share": "edf_share_pct"})
return out.round(
{
"mean_train_deviance": 5,
"mean_val_deviance": 5,
"se_val_deviance": 5,
"total_edf": 2,
"age_edf": 2,
"edf_share_pct": 1,
"age_p": 4,
"age_band_p": 4,
"median_age_band_p": 4,
"train_improvement": 5,
"cv_improvement": 5,
"age_edf_increase": 2,
}
)
k_grid = list(range(5, 26))
In-Sample k Adequacy#
This table shows whether the fitted model is still basis-limited. If EDF is pressing against capacity, or a coarse duplicate such as age_band only looks important at low k, the continuous age spline probably needs more room.
Use this for diagnosis; use the CV section below for model-selection evidence.
candidate_models = {}
for age_k in k_grid:
start = time.perf_counter()
candidate_models[age_k] = fit_age_model(age_k)
print(f"k={age_k:2d}: fitted in {time.perf_counter() - start:.1f}s")
records = []
for age_k, candidate in candidate_models.items():
diagnostics = candidate.diagnostics()
summary = candidate.summary()
age_info = diagnostics.get("age", {})
age_edf = float(age_info.get("edf", np.nan))
age_capacity = max(float(age_info.get("n_params", age_k)), 1.0)
records.append(
{
"k": age_k,
"mean_train_deviance": mean_unit_deviance(candidate, X_train, y_train, w_train),
"total_edf": candidate.result.effective_df,
"age_edf": age_edf,
"edf_share": age_edf / age_capacity,
"age_p": term_p(summary, "age"),
"age_band_p": term_p(summary, "age_band"),
}
)
k_sweep = pd.DataFrame(records).sort_values("k").reset_index(drop=True)
k_sweep["train_improvement"] = -k_sweep["mean_train_deviance"].diff()
k_sweep["age_edf_increase"] = k_sweep["age_edf"].diff()
k_sweep_display = format_sweep_table(k_sweep)
fig, axes = plt.subplots(2, 1, figsize=(8, 7), sharex=True)
axes[0].plot(k_sweep["k"], k_sweep["mean_train_deviance"], marker="o", label="train")
axes[0].set_ylabel("weighted mean deviance")
axes[0].legend()
axes[1].plot(k_sweep["k"], k_sweep["total_edf"], marker="o", label="total EDF")
axes[1].plot(k_sweep["k"], k_sweep["age_edf"], marker="o", label="age EDF")
axes[1].set_xlabel("k")
axes[1].set_ylabel("EDF")
axes[1].legend()
fig.tight_layout()
k_sweep_display
Cross-Validated k Check#
This is the actual model-selection check. Each k is refit through superglm.cross_validate, then scored on held-out rows. The one-SE rule is shown as a conservative reference, but the editor recommendation also checks whether the duplicate/coarse age_band term has stopped absorbing residual age shape.
The recommendation below chooses the smallest k that is on the CV deviance plateau and passes the leakage diagnostic. This is closer to a spline-capacity adequacy check than a pure parsimony rule.
n_folds = 3
cv = KFold(n_splits=n_folds, shuffle=True, random_state=20260704)
cv_records = []
for age_k in k_grid:
start = time.perf_counter()
result = cross_validate(
make_model(age_k),
X_train,
y_train,
sample_weight=w_train,
cv=cv,
fit_mode="fit_reml",
scoring=("deviance", cv_diagnostics),
error_score="raise",
)
fold_scores = result.fold_scores
cv_records.append(
{
"k": age_k,
"mean_val_deviance": result.mean_scores["deviance"],
"se_val_deviance": float(fold_scores["deviance"].std(ddof=1) / np.sqrt(n_folds)),
"total_edf": float(fold_scores["effective_df"].mean()),
"age_edf": float(fold_scores["age_edf"].mean()),
"median_age_band_p": float(fold_scores["age_band_p"].median()),
}
)
print(
f"k={age_k:2d}: {n_folds}-fold validation deviance "
f"{result.mean_scores['deviance']:.5f} in {time.perf_counter() - start:.1f}s"
)
def recommend_k_from_cv(
sweep: pd.DataFrame,
*,
leakage_p_min: float = 0.5,
plateau_se_fraction: float = 0.25,
) -> tuple[int, int, float]:
best_idx = sweep["mean_val_deviance"].idxmin()
best_deviance = float(sweep.loc[best_idx, "mean_val_deviance"])
best_se = float(sweep.loc[best_idx, "se_val_deviance"])
one_se_limit = best_deviance + best_se
cv_one_se_k = int(sweep.loc[sweep["mean_val_deviance"] <= one_se_limit, "k"].iloc[0])
plateau_limit = best_deviance + plateau_se_fraction * best_se
clean = sweep["median_age_band_p"] >= leakage_p_min
plateau = sweep["mean_val_deviance"] <= plateau_limit
candidates = sweep.loc[clean & plateau]
if candidates.empty:
candidates = sweep.loc[plateau]
if candidates.empty:
candidates = sweep.loc[[best_idx]]
return int(candidates["k"].iloc[0]), cv_one_se_k, plateau_limit
cv_sweep = pd.DataFrame(cv_records).sort_values("k").reset_index(drop=True)
cv_sweep["cv_improvement"] = -cv_sweep["mean_val_deviance"].diff()
cv_recommended_k, cv_one_se_k, plateau_limit = recommend_k_from_cv(cv_sweep)
cv_sweep_display = format_sweep_table(cv_sweep)
fig, ax = plt.subplots(figsize=(8, 4))
ax.errorbar(
cv_sweep["k"],
cv_sweep["mean_val_deviance"],
yerr=cv_sweep["se_val_deviance"],
marker="o",
capsize=3,
label="CV validation",
)
ax.axhline(plateau_limit, color="0.55", linestyle="--", linewidth=1, label="plateau limit")
ax.axvline(cv_one_se_k, color="tab:orange", linestyle=":", linewidth=2, label=f"one-SE k={cv_one_se_k}")
ax.axvline(cv_recommended_k, color="tab:green", linestyle="-.", linewidth=2, label=f"recommended k={cv_recommended_k}")
ax.set_xlabel("k")
ax.set_ylabel("weighted mean validation deviance")
ax.legend()
fig.tight_layout()
cv_sweep_display
Fit A SuperGLM#
cv_one_se_k is the pure parsimony reference. cv_recommended_k is the capacity-aware recommendation: it stays on the CV deviance plateau while requiring the duplicate age_band term to stop absorbing residual age shape.
The model starts with an intentionally fixed Tweedie power. After fitting the selected spline capacity, the next cell re-profiles p and shows the powers the search evaluated. Override chosen_k if your own curve, exposure support, or business interpretation argues for a different value. The fitted model uses SuperGLM’s discrete path for the 150k generated rows; the editor itself still works on a fixed display grid, so the widget remains light.
chosen_k = cv_recommended_k
model = fit_age_model(chosen_k, X_fit=X_train, y_fit=y_train, w_fit=w_train)
print("Profiling the Tweedie power p...")
start = time.perf_counter()
tweedie_profile = model.estimate_p(
X_train,
y_train,
sample_weight=w_train,
fit_mode="inherit",
xatol=0.005,
)
display(tweedie_profile.evaluations.tail())
print(
f"profiled Tweedie p={tweedie_profile.p_hat:.3f}; "
f"true p={TRUE_TWEEDIE_P:.3f}; phi={tweedie_profile.phi_hat:.4f} "
f"({time.perf_counter() - start:.1f}s)"
)
cv_loss_records = []
for fold, (fit_idx, holdout_idx) in enumerate(cv.split(X_train), start=1):
X_fold_train = X_train.iloc[fit_idx]
y_fold_train = y_train[fit_idx]
w_fold_train = w_train[fit_idx]
X_fold_validation = X_train.iloc[holdout_idx]
y_fold_validation = y_train[holdout_idx]
w_fold_validation = w_train[holdout_idx]
fold_model = fit_age_model(chosen_k, X_fit=X_fold_train, y_fit=y_fold_train, w_fit=w_fold_train)
print(f"fold {fold}: fitted on {len(X_fold_train):,} rows")
fold_train_loss = mean_unit_deviance(fold_model, X_fold_train, y_fold_train, w_fold_train)
fold_validation_loss = mean_unit_deviance(
fold_model,
X_fold_validation,
y_fold_validation,
w_fold_validation,
)
cv_loss_records.append(
{
"fold": fold,
"train_loss": fold_train_loss,
"validation_loss": fold_validation_loss,
"train_rows": len(X_fold_train),
"validation_rows": len(X_fold_validation),
}
)
cv_loss = pd.DataFrame(cv_loss_records)
cv_loss_summary = (
cv_loss[["train_loss", "validation_loss"]]
.agg(["mean", "std"])
.T.rename_axis("loss")
.reset_index()
)
split_loss = pd.DataFrame(
[
{
"split": "train",
"loss": mean_unit_deviance(model, X_train, y_train, w_train),
"n_obs": len(X_train),
},
{
"split": "validation",
"loss": mean_unit_deviance(model, X_val, y_val, w_val),
"n_obs": len(X_val),
},
{
"split": "test",
"loss": mean_unit_deviance(model, X_test, y_test, w_test),
"n_obs": len(X_test),
},
]
)
cv_report = {
"method": "KFold",
"folds": n_folds,
"metric": "weighted mean unit deviance",
"scope": "chosen model on training folds with fixed initial Tweedie p",
"chosen_k": int(chosen_k),
"tweedie_p": float(tweedie_profile.p_hat),
"rows": cv_loss.round(6).to_dict("records"),
"summary": cv_loss_summary.round(6).to_dict("records"),
"split_loss": split_loss.round(6).to_dict("records"),
}
print(f"One-SE k={cv_one_se_k}; capacity-aware recommendation k={chosen_k}")
print("CV fold loss summary")
print(cv_loss_summary.round(6))
print("Train/validation/test loss")
print(split_loss.round(6))
model.summary()
Create An Editor Session#
EditorSession.from_model extracts editable 1D terms from the fitted model. Edits live in the session until you explicitly apply them to a copied model.
session = EditorSession.from_model(
model,
terms=["age", "mileage", "region", "age_band", "territory"],
n_points=260,
train_data=(X_train, y_train, w_train),
validation_data=(X_val, y_val, w_val),
test_data=(X_test, y_test, w_test),
cv_report=cv_report,
)
list(session.terms)
Collapse Sparse Categorical Levels#
The territory term has 10 levels with uneven exposure. This example creates two distinct collapsed groups on the same graph: T01+T03 and T08+T10. The editor marks each group with its own color/marker treatment so grouped levels are visible even when their fitted relativities match.
from superglm.editor.payloads import session_payload
territory_exposure = (
X_train["territory"]
.value_counts(normalize=True)
.rename_axis("territory")
.reset_index(name="train_share")
.sort_values("territory")
)
territory_effect_table = pd.DataFrame(
{
"territory": TERRITORY_LEVELS,
"true_link_effect": [territory_effect_map[level] for level in TERRITORY_LEVELS],
}
)
display(territory_exposure.merge(territory_effect_table, on="territory"))
collapse_session = EditorSession.from_model(
model,
terms=["territory"],
train_data=(X_train, y_train, w_train),
validation_data=(X_val, y_val, w_val),
test_data=(X_test, y_test, w_test),
)
collapse_session.select_levels("territory", ["T01", "T03"])
collapse_session.replace_with_collapsed_levels("territory", method="fit")
collapse_session.select_levels("territory", ["T08", "T10"])
collapsed_model = collapse_session.replace_with_collapsed_levels("territory", method="fit")
grouped_preview = session_payload(collapse_session)["territory"]["level_groups"]
display(pd.DataFrame(grouped_preview))
collapsed_model.summary()
Interactive Editor#
Use Select mode to click a point or press-drag-release a rectangle on the plot to select points in that region. Use Move mode after selecting points to drag a point vertically; if the dragged point is already selected, the whole selected group moves together. Smooth applies to the selected region, or to the whole current term if nothing is selected. Reset restores the selected region, or the whole current term if nothing is selected. Smooth is anchored to neighboring points to reduce edge jumps; Increasing and Decreasing fit the closest monotone curve to the selection, weighted by exposure, and can leave a step at an edge.
editor = session.widget()
editor
Programmatic Edits#
The same operations are available without the widget, which is useful when you want reproducible edit scripts.
# Raise the young-age portion of the spline on the link scale.
session.select_x("age", 18, 28).shift("age", 0.08)
# Smooth a mid-age region.
session.select_x("age", 35, 55).smooth("age", strength=0.65)
# Make the ordered age-band effect non-decreasing.
session.select_levels("age_band", ["18-24", "25-34", "35-49", "50-64", "65-80"])
session.isotonic("age_band", direction="increasing")
session.history[-3:]
Visualize The Edited Curve#
The session stores original and edited link-scale effects. Exponentiating is convenient when you want to inspect relative effects.
term = session.terms["age"]
fig, ax = plt.subplots(figsize=(9, 4.5))
ax.plot(term.x, np.exp(term.original_log_effect), label="original", linewidth=2)
ax.plot(term.x, np.exp(term.edited_log_effect), label="edited", linewidth=2)
ax.axhline(1.0, color="#999999", linestyle="--", linewidth=1)
ax.set_xlabel("age")
ax.set_ylabel("relative effect")
ax.legend(frameon=False)
ax.grid(alpha=0.25)
fig.tight_layout()
Apply Edits To A Copy#
to_model() deep-copies the source model and applies the edited coefficients to the copy. The original fitted model is left untouched.
edited_model = session.to_model()
original_eta = model._predict_eta_exact(X)
edited_eta = edited_model._predict_eta_exact(X)
print(f"original model unchanged: {edited_model is not model}")
print(f"mean prediction delta: {np.mean(edited_eta - original_eta):.4f}")
print(f"max abs prediction delta: {np.max(np.abs(edited_eta - original_eta)):.4f}")
Save And Reload Edits#
Saved edit artifacts are JSON. Reloading requires the same fitted model so the editor can validate term grids and levels.
edit_path = Path("editor_edits.json")
session.save(edit_path)
loaded = EditorSession.load(edit_path, model=model)
loaded.to_model()
edit_path