From a raw table to model-ready arrays

import json
import tempfile
import time
from dataclasses import dataclass
from pathlib import Path

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import patsy
import pymc as pm
import xarray as xr

from pymc_challenges import data

RANDOM_SEED = 42
rng = np.random.default_rng(RANDOM_SEED)
az.style.use("arviz-variat")

print(f"PyMC {pm.__version__}, ArviZ {az.__version__}, pandas {pd.__version__}, patsy {patsy.__version__}")

data.describe("penguins_raw")
raw = data.load("penguins_raw")
raw.head(3).T

def audit(df):
    """One row per column: dtype, missing, blank strings, distinct values, an example."""
    is_text = [pd.api.types.is_string_dtype(df[c]) for c in df.columns]
    blank = [int((df[c].str.strip() == "").sum()) if t else 0 for c, t in zip(df.columns, is_text)]
    return pd.DataFrame(
        {
            "dtype": df.dtypes.astype(str),
            "missing": df.isna().sum(),
            "blank": blank,
            "distinct": df.nunique(),
            "example": df.iloc[0].astype(str).str.slice(0, 28),
        }
    )


audit(raw)

key = ["study", "species", "island", "individual_id"]
print("exact duplicate rows:       ", int(peng_all.duplicated().sum()))
print("duplicates of the row key:  ", int(peng_all.duplicated(key).sum()))
print("distinct individual_id:     ", peng_all.individual_id.nunique(), "for", len(peng_all), "rows")
print("(species, id) pairs seen in more than one season:",
      int((peng_all.groupby(["species", "individual_id"]).study.nunique() > 1).sum()))
print("adults per (study, species, island, nest):",
      peng_all.groupby(["study", "species", "island", "nest"]).size().value_counts().to_dict())

data.describe("telco")
telco_raw = data.load("telco")
audit(telco_raw).loc[["customerID", "SeniorCitizen", "tenure", "Contract", "MonthlyCharges", "TotalCharges", "Churn"]]

total = pd.to_numeric(telco_raw.TotalCharges, errors="coerce")
changed = total.isna() & telco_raw.TotalCharges.notna()
print("cells that coercion turned into NaN:", int(changed.sum()),
      "- their raw values:", telco_raw.TotalCharges[changed].unique().tolist())
telco_raw.loc[changed, ["tenure", "Contract", "MonthlyCharges", "TotalCharges", "Churn"]].head()

print("tenure of the blank rows:", telco_raw.tenure[changed].unique(),
      "| rows with tenure 0:", int((telco_raw.tenure == 0).sum()))

def clean_telco(raw):
    df = raw.copy()
    total = pd.to_numeric(df.TotalCharges, errors="coerce")
    assert (total.isna() == (df.tenure == 0)).all(), "blank TotalCharges no longer means 'not billed yet'"
    df["TotalCharges"] = total.fillna(0.0)
    df["churned"] = df.Churn.map({"Yes": 1, "No": 0}).astype(int)
    df["SeniorCitizen"] = df.SeniorCitizen.map({0: "No", 1: "Yes"})  # a 0/1 integer that is really a category
    assert df.customerID.is_unique
    return df.drop(columns="Churn")


telco = clean_telco(telco_raw)
telco[["tenure", "MonthlyCharges", "TotalCharges", "churned"]].describe().round(1)

miss = peng_all[MEASUREMENTS + ["sex", "delta_15n", "delta_13c"]].isna()
print("rows by number of missing fields:", miss.sum(axis=1).value_counts().sort_index().to_dict())
pd.concat(
    [miss.sum().rename("n_missing"), (100 * miss.mean()).round(1).rename("percent"),
     miss.groupby(peng_all.species).sum().T],
    axis=1,
)

# the free-text column earns its keep: it says WHY
peng_all.loc[peng_all.sex.isna(), ["species", "island", "year", "body_mass_g", "comments"]]

sexed = peng_all.dropna(subset=["body_mass_g"]).assign(sex_known=lambda d: d.sex.notna())
sexed.groupby(["species", "sex_known"]).body_mass_g.agg(["size", "mean", "min", "max"]).round(0)

peng = peng_all.dropna(subset=MEASUREMENTS).reset_index(drop=True)      # 342: every measured bird
peng_sexed = peng.dropna(subset=["sex"]).reset_index(drop=True)          # 333: complete cases for models with sex
print(len(peng_all), "->", len(peng), "->", len(peng_sexed))

views = {
    "file order          ": peng,
    "re-sorted by mass   ": peng.sort_values("body_mass_g", ascending=False),
    "Biscoe Island only  ": peng[peng.island == "Biscoe"],
}
for label, frame in views.items():
    print(label, "sort=False:", list(pd.factorize(frame.species)[1]),
          "| sort=True:", list(pd.factorize(frame.species, sort=True)[1]))

SPECIES = ["Adelie", "Chinstrap", "Gentoo"]
SEXES = ["female", "male"]


def encode(values, levels):
    """Integer codes against a declared list of levels; -1 for anything unseen or missing."""
    return pd.Index(levels).get_indexer(pd.Series(values))


new_birds = ["Gentoo", "Adelie", "Emperor", None]
codes = encode(new_birds, SPECIES)
alpha_demo = np.array([3.7, 3.7, 5.1])  # pretend posterior means, kg
print("codes:", codes)
print("alpha[codes]:", alpha_demo[codes], "<- the Emperor and the missing bird were silently treated as Gentoo")

coords_sp = {"species": SPECIES}

with pm.Model(coords=coords_sp) as m_index:
    pm.Normal("group_mean", 4.0, 1.0, dims="species")

with pm.Model(coords={**coords_sp, "non_ref": SPECIES[1:]}) as m_dummy:
    a = pm.Normal("a", 4.0, 1.0)
    b = pm.Normal("b", 0.0, 1.0, dims="non_ref")
    pm.Deterministic("group_mean", a + pm.math.concatenate([[0.0], b]), dims="species")

with pm.Model(coords=coords_sp) as m_zerosum:
    a = pm.Normal("a", 4.0, 1.0)
    delta = pm.ZeroSumNormal("delta", 1.0, dims="species")
    pm.Deterministic("group_mean", a + delta, dims="species")

PAIRS = [("Chinstrap", "Adelie"), ("Gentoo", "Adelie"), ("Gentoo", "Chinstrap")]
rows = {}
for name, m in [("index", m_index), ("dummy (ref = Adelie)", m_dummy), ("sum-to-zero", m_zerosum)]:
    with m:
        gm = pm.sample_prior_predictive(20_000, random_seed=RANDOM_SEED).prior["group_mean"]
    row = {f"sd mean {s}": float(gm.sel(species=s).std()) for s in SPECIES}
    row |= {f"sd {hi} - {lo}": float((gm.sel(species=hi) - gm.sel(species=lo)).std()) for hi, lo in PAIRS}
    rows[name] = row

prior_sd = pd.DataFrame(rows).T.round(2)
prior_sd

fig, axes = plt.subplots(1, 2, figsize=(11, 3.6), sharey=True)
width = 0.26
for ax, cols, title in [
    (axes[0], prior_sd.columns[:3], "prior sd of each species mean (kg)"),
    (axes[1], prior_sd.columns[3:], "prior sd of each difference (kg)"),
]:
    for k, coding in enumerate(prior_sd.index):
        ax.bar(np.arange(3) + (k - 1) * width, prior_sd.loc[coding, cols], width, label=coding)
    ticks = [c.replace("sd mean ", "").replace("sd ", "").replace(" - ", "\n- ") for c in cols]
    ax.set_xticks(np.arange(3), ticks, fontsize=10)
    ax.set_title(title)
    ax.set_ylim(0, 1.9)
axes[0].legend(fontsize=8, loc="upper left", ncols=3);

cells = pd.DataFrame([(s, x) for s in SPECIES for x in SEXES], columns=["species", "sex"])
X_cells = patsy.dmatrix("species * sex", cells)
print(X_cells.design_info.column_names)
cells.assign(
    n_coefficients=np.asarray(X_cells).sum(axis=1).astype(int),
    prior_sd_if_each_is_N01=np.sqrt((np.asarray(X_cells) ** 2).sum(axis=1)).round(2),
)

species_idx = encode(peng_sexed.species, SPECIES)
sex_idx = encode(peng_sexed.sex, SEXES)
assert species_idx.min() >= 0 and sex_idx.min() >= 0

with pm.Model(coords={"species": SPECIES, "sex": SEXES}) as m_cells:
    mu = pm.Normal("mu", 4.0, 1.0, dims=("species", "sex"))
    sigma = pm.HalfNormal("sigma", 1.0)
    pm.Normal("mass", mu[species_idx, sex_idx], sigma, observed=peng_sexed.body_mass_g / 1000)
    idata_cells = pm.sample(random_seed=RANDOM_SEED)

mu_post = idata_cells.posterior["mu"]
dimorphism = mu_post.sel(sex="male") - mu_post.sel(sex="female")
print("divergences:", int(idata_cells.sample_stats["diverging"].sum()))
az.summary(xr.Dataset({"mu": mu_post, "male_minus_female": dimorphism}), ci_kind="hdi", ci_prob=0.94, round_to=2)

fig, axes = plt.subplots(1, 3, figsize=(12, 3.6))
for ax, prep_name in zip(axes, PREPARATIONS):
    draws = az.extract(fits["nutpie", prep_name], num_samples=1000, random_seed=RANDOM_SEED)
    a_draws, b_draws = draws["a"].values, draws["b"].sel(predictor="flipper_length_mm").values
    ax.scatter(a_draws, b_draws, s=5, alpha=0.4)
    corr = np.corrcoef(a_draws, b_draws)[0, 1]
    ax.set(xlabel="intercept a", ylabel="flipper slope", title=f"{prep_name}: corr = {corr:+.2f}")

scale = xr.DataArray(y_sd / x_sd, coords={"predictor": PREDICTORS})
centre = xr.DataArray(x_mean, coords={"predictor": PREDICTORS})

post_z = fits["nutpie", "standardised"].posterior
b_back = post_z["b"] * scale
a_back = y_mean + y_sd * post_z["a"] - (b_back * centre).sum("predictor")
back = xr.Dataset({"a": a_back, "b": b_back, "sigma": y_sd * post_z["sigma"]})

SHOW = ["mean", "sd", "hdi94_lb", "hdi94_ub"]
pd.concat(
    {
        "back-transformed draws": az.summary(back, ci_kind="hdi", ci_prob=0.94, round_to=1)[SHOW],
        "fitted on raw scale": az.summary(fits["nutpie", "raw"], ci_kind="hdi", ci_prob=0.94, round_to=1)[SHOW],
    },
    axis=1,
)

summ_z = az.summary(fits["nutpie", "standardised"], ci_kind="hdi", ci_prob=0.94, round_to=6)
b_lo, b_hi = (summ_z.loc[[f"b[{p}]" for p in PREDICTORS], c].to_numpy() * y_sd / x_sd for c in ["hdi94_lb", "hdi94_ub"])
a_lo, a_hi = (y_mean + y_sd * summ_z.loc["a", c] for c in ["hdi94_lb", "hdi94_ub"])
# "worst case" endpoint arithmetic: what people do when they only have the table
wrong = (a_lo - np.maximum(b_lo * x_mean, b_hi * x_mean).sum(),
         a_hi - np.minimum(b_lo * x_mean, b_hi * x_mean).sum())
right = az.hdi(a_back, prob=0.94).values

print(f"intercept, from summary endpoints (WRONG): [{wrong[0]:8.0f}, {wrong[1]:8.0f}]   width {wrong[1] - wrong[0]:6.0f} g")
print(f"intercept, from transformed draws (RIGHT): [{right[0]:8.0f}, {right[1]:8.0f}]   width {right[1] - right[0]:6.0f} g")
print(f"and the 'intercept' y_mean + y_sd * a = {float(y_mean + y_sd * post_z['a'].mean()):.0f} g"
      " is the mass of an AVERAGE bird, not of one with x = 0")

train = peng[peng.year < 2009]
batch = peng[(peng.year == 2009) & (peng.species == "Gentoo")]
tr_mean, tr_sd = train[PREDICTORS].mean().to_numpy(), train[PREDICTORS].std(ddof=0).to_numpy()
ty_mean, ty_sd = train.body_mass_g.mean(), train.body_mass_g.std(ddof=0)

idata_train = fit_linear(
    (train[PREDICTORS].to_numpy() - tr_mean) / tr_sd, (train.body_mass_g.to_numpy() - ty_mean) / ty_sd, (0, 1), 1, 1
)
draws = az.extract(idata_train, num_samples=1000, random_seed=RANDOM_SEED)


def predict_g(X_new, mean, sd):
    """Posterior-mean prediction in grams, standardising X_new with the given constants."""
    z = (X_new - mean) / sd
    mu_z = draws["a"].values[None, :] + z @ draws["b"].transpose("predictor", "sample").values
    return ty_mean + ty_sd * mu_z.mean(axis=1)


X_batch, y_batch = batch[PREDICTORS].to_numpy(), batch.body_mass_g.to_numpy()
versions = {
    "RIGHT: training mean/sd": predict_g(X_batch, tr_mean, tr_sd),
    "WRONG: mean/sd of the whole table (computed before the split)": predict_g(X_batch, x_mean, x_sd),
    "WRONG: the batch's own mean/sd": predict_g(X_batch, X_batch.mean(axis=0), X_batch.std(axis=0)),
}
print(f"{len(train)} training birds (mean mass {ty_mean:.0f} g),",
      f"{len(batch)} new Gentoos (mean mass {y_batch.mean():.0f} g)\n")
for label, pred_g in versions.items():
    rmse = np.sqrt(np.mean((pred_g - y_batch) ** 2))
    print(f"{label:<62} mean prediction {pred_g.mean():5.0f} g   RMSE {rmse:5.0f} g")

fig, ax = plt.subplots(figsize=(5.5, 4.5))
for (label, pred_g), marker in zip(versions.items(), ["o", "s", "^"]):
    ax.scatter(y_batch, pred_g, s=14, alpha=0.7, marker=marker, label=label.split(" (")[0])
lims = [3000, 6500]
ax.plot(lims, lims, color="k", lw=1)
ax.set(xlabel="observed mass (g)", ylabel="predicted mass (g)", xlim=lims, ylim=lims)
ax.legend(fontsize=7);

log_flipper_c = np.log(peng.flipper_length_mm / 200).to_numpy()   # centred on a meaningful 200 mm
sp_idx = encode(peng.species, SPECIES)

with pm.Model(coords={"species": SPECIES}) as m_loglog:
    a = pm.Normal("a", np.log(4000), 0.5, dims="species")      # log-grams at 200 mm
    elasticity = pm.Normal("elasticity", 3, 2)                  # prior centred on isometric scaling
    sigma = pm.HalfNormal("sigma", 0.5)
    pm.Normal("log_mass", a[sp_idx] + elasticity * log_flipper_c, sigma, observed=np.log(peng.body_mass_g))
    idata_loglog = pm.sample(random_seed=RANDOM_SEED)

ratio = np.exp(idata_loglog.posterior["a"].sel(species="Gentoo") - idata_loglog.posterior["a"].sel(species="Adelie"))
az.summary(xr.Dataset({"elasticity": idata_loglog.posterior["elasticity"], "sigma": idata_loglog.posterior["sigma"],
                       "gentoo_vs_adelie_ratio": ratio}), ci_kind="hdi", ci_prob=0.94, round_to=2)

CONTRACTS = ["Month-to-month", "One year", "Two year"]
billed = telco[telco.tenure > 0].reset_index(drop=True)      # log(0) is the transform telling you something
contract_idx = encode(billed.Contract, CONTRACTS)
assert contract_idx.min() >= 0

with pm.Model(coords={"contract": CONTRACTS}) as m_lognormal:
    mu = pm.Normal("mu", 7, 2, dims="contract")
    sigma = pm.HalfNormal("sigma", 2, dims="contract")
    pm.Normal("log_total", mu[contract_idx], sigma[contract_idx], observed=np.log(billed.TotalCharges))
    idata_ln = pm.sample(random_seed=RANDOM_SEED)
    # posterior predictive on a thinned posterior: 400 replicated datasets of 7032 customers
    ppc = pm.sample_posterior_predictive(idata_ln.isel(draw=slice(None, None, 10)), random_seed=RANDOM_SEED)

log_rep = ppc.posterior_predictive["log_total"].stack(sample=("chain", "draw")).values      # (customer, sample)
post = idata_ln.posterior
retrans = pd.DataFrame(index=CONTRACTS)
retrans["observed mean"] = billed.groupby("Contract").TotalCharges.mean()
retrans["observed median"] = billed.groupby("Contract").TotalCharges.median()
retrans["exp(mean of mu)  [WRONG]"] = np.exp(post["mu"].mean(("chain", "draw"))).values
retrans["exp(mean of log y_rep)  [WRONG]"] = [np.exp(log_rep[contract_idx == k].mean()) for k in range(3)]
retrans["mean of exp(log y_rep)  [RIGHT]"] = [np.exp(log_rep[contract_idx == k]).mean() for k in range(3)]
retrans["sigma"] = post["sigma"].mean(("chain", "draw")).values
retrans.round({"sigma": 2} | {c: 0 for c in retrans.columns[:-1]})

with pm.Model(coords={"contract": CONTRACTS}) as m_gamma:
    log_mean = pm.Normal("log_mean", 7, 2, dims="contract")
    shape = pm.Exponential("shape", 1, dims="contract")
    mean = pm.Deterministic("mean", pm.math.exp(log_mean), dims="contract")     # the estimand, by construction
    pm.Gamma("total", alpha=shape[contract_idx], beta=shape[contract_idx] / mean[contract_idx],
             observed=billed.TotalCharges)
    idata_gamma = pm.sample(random_seed=RANDOM_SEED)

retrans["Gamma, log link: posterior mean of `mean`"] = idata_gamma.posterior["mean"].mean(("chain", "draw")).values
retrans[["observed mean", "mean of exp(log y_rep)  [RIGHT]", "Gamma, log link: posterior mean of `mean`"]].round(0)

fig, axes = plt.subplots(1, 2, figsize=(11, 3.6))
m2m = billed.TotalCharges[billed.Contract == "Month-to-month"]
axes[0].hist(np.log(m2m), bins=50, density=True, alpha=0.5, label="observed")
grid = np.linspace(2.5, 9.5, 200)
mu0 = float(post["mu"].sel(contract="Month-to-month").mean())
s0 = float(post["sigma"].sel(contract="Month-to-month").mean())
normal_pdf = np.exp(-0.5 * ((grid - mu0) / s0) ** 2) / (s0 * np.sqrt(2 * np.pi))
axes[0].plot(grid, normal_pdf, color="C1", label="fitted Normal")
axes[0].set(xlabel="log TotalCharges, month-to-month", title="the log scale is not Normal either")
axes[0].legend(fontsize=8)
axes[1].hist(m2m, bins=60, alpha=0.5)
for col, color, ls in [("exp(mean of mu)  [WRONG]", "C3", "-"), ("mean of exp(log y_rep)  [RIGHT]", "C1", "-"),
                       ("Gamma, log link: posterior mean of `mean`", "C2", "-"), ("observed mean", "k", "--")]:
    axes[1].axvline(retrans.loc["Month-to-month", col], color=color, lw=2, ls=ls, label=col)
axes[1].set(xlabel="TotalCharges ($), month-to-month", title="four answers to 'what is the average?'")
axes[1].legend(fontsize=7);

X_hand = np.column_stack([
    (peng_sexed.species == "Chinstrap").astype(float),
    (peng_sexed.species == "Gentoo").astype(float),
    (peng_sexed.sex == "male").astype(float),
    (peng_sexed.flipper_length_mm - 200) / 10,
])
names_hand = ["species:Chinstrap", "species:Gentoo", "sex:male", "flipper_per_10mm"]
X_hand[:3], names_hand

FORMULA = "0 + C(species, levels=SPECIES) + C(sex, levels=SEXES) + bs(flipper_length_mm, df=4)"
dm = patsy.dmatrix(FORMULA, peng_sexed)


def tidy(name):
    """patsy column names -> something you want to read in az.summary."""
    for old, new in [("C(species, levels=SPECIES)", "species"), ("C(sex, levels=SEXES)", "sex"),
                     ("bs(flipper_length_mm, df=4)", "flipper_spline"), ("[T.", ":"), ("[", ":"), ("]", "")]:
        name = name.replace(old, new)
    return name


X_spline = np.asarray(dm)
spline_names = [tidy(n) for n in dm.design_info.column_names]
print(X_spline.shape, spline_names)

new_rows = pd.DataFrame({"species": ["Gentoo"] * 2, "sex": ["female"] * 2, "flipper_length_mm": [210.0, 225.0]})
print(np.asarray(patsy.build_design_matrices([dm.design_info], new_rows)[0]).round(3))
try:
    patsy.build_design_matrices([dm.design_info], new_rows.assign(species="Emperor"))
except patsy.PatsyError as err:
    print("PatsyError:", str(err).splitlines()[0])

def check_design(X, names):
    """Rank, condition number, and the linear combination of columns that is (nearly) zero."""
    X = np.asarray(X, dtype=float)
    sv = np.linalg.svd(X / np.linalg.norm(X, axis=0), compute_uv=False)    # scale-free: columns to unit length
    rank = np.linalg.matrix_rank(X)
    print(f"columns {X.shape[1]}, rank {rank}, condition number {sv[0] / sv[-1]:.3g}")
    if rank < X.shape[1]:
        null = np.linalg.svd(X)[2][-1]
        null = null / np.abs(null).max()
        terms = " ".join(f"{w:+.2f}*{n}" for w, n in zip(null, names) if abs(w) > 1e-8)
        print("  redundant combination:", terms, "= 0")




n = len(peng_sexed)
species_dummies = np.column_stack([(peng_sexed.species == s).astype(float) for s in SPECIES])
flipper_10 = ((peng_sexed.flipper_length_mm - 200) / 10).to_numpy()[:, None]

designs = {
    "full rank: one column per species": (np.hstack([species_dummies, flipper_10]), SPECIES + ["flipper_per_10mm"]),
    "redundant: intercept + one column per species": (
        np.hstack([np.ones((n, 1)), species_dummies, flipper_10]), ["intercept"] + SPECIES + ["flipper_per_10mm"]),
}
for label, (X, names) in designs.items():
    print(label)
    check_design(X, names)
print("\nnear-collinearity shows up the same way - the uncentred design of section 4:")
check_design(np.column_stack([np.ones(len(X_raw)), X_raw]), ["intercept"] + PREDICTORS)
check_design(np.column_stack([np.ones(len(X_raw)), (X_raw - x_mean) / x_sd]), ["intercept"] + PREDICTORS)

mass_kg = (peng_sexed.body_mass_g / 1000).to_numpy()
fits_design, report_design = {}, {}
for label, (X, names) in designs.items():
    with pm.Model(coords={"column": names}):
        beta = pm.Normal("beta", 0, 5, dims="column")
        sigma = pm.HalfNormal("sigma", 1)
        pm.Normal("mass", pm.math.dot(X, beta), sigma, observed=mass_kg)
        fits_design[label] = pm.sample(random_seed=RANDOM_SEED, progressbar=False)
    report_design[label] = geometry(fits_design[label])

pd.DataFrame(report_design).T.round(2)

bad = fits_design["redundant: intercept + one column per species"]
az.summary(bad, var_names=["beta"], ci_kind="hdi", ci_prob=0.94, round_to=2)

draws = az.extract(bad, num_samples=1500, random_seed=RANDOM_SEED)["beta"]
fig, axes = plt.subplots(1, 2, figsize=(10, 3.8))
axes[0].scatter(draws.sel(column="intercept"), draws.sel(column="Adelie"), s=5, alpha=0.4)
axes[0].set(xlabel="beta[intercept]", ylabel="beta[Adelie]", title="the ridge: only the SUM is identified")
adelie_redundant = (draws.sel(column="intercept") + draws.sel(column="Adelie")).values
axes[1].hist(adelie_redundant, bins=40, alpha=0.6, label="redundant: intercept + Adelie")
good = az.extract(fits_design["full rank: one column per species"], num_samples=1500, random_seed=RANDOM_SEED)["beta"]
axes[1].hist(good.sel(column="Adelie").values, bins=40, alpha=0.6, label="full rank: Adelie")
axes[1].set(xlabel="mass of a 200 mm Adelie (kg)", title="the identified quantity is fine")
axes[1].legend(fontsize=8);

with pm.Model(coords={"column": spline_names}) as m_spline:
    beta = pm.Normal("beta", 0, 5, dims="column")
    sigma = pm.HalfNormal("sigma", 1)
    pm.Normal("mass", pm.math.dot(X_spline, beta), sigma, observed=mass_kg)
    idata_spline = pm.sample(random_seed=RANDOM_SEED)

check_design(X_spline, spline_names)
az.summary(idata_spline, ci_kind="hdi", ci_prob=0.94, round_to=2)

beta_draws = az.extract(idata_spline, num_samples=1000, random_seed=RANDOM_SEED)["beta"]
beta_draws = beta_draws.transpose("column", "sample").values
fig, ax = plt.subplots(figsize=(7, 4))
for k, sp in enumerate(SPECIES):
    sub = peng_sexed[(peng_sexed.species == sp) & (peng_sexed.sex == "female")]
    flipper_grid = np.linspace(sub.flipper_length_mm.min(), sub.flipper_length_mm.max(), 40)
    grid = pd.DataFrame({"species": sp, "sex": "female", "flipper_length_mm": flipper_grid})
    mu_grid = np.asarray(patsy.build_design_matrices([dm.design_info], grid)[0]) @ beta_draws
    ax.fill_between(grid.flipper_length_mm, *np.quantile(mu_grid, [0.03, 0.97], axis=1), color=f"C{k}", alpha=0.3)
    ax.plot(grid.flipper_length_mm, mu_grid.mean(axis=1), color=f"C{k}", label=f"{sp} (female)")
    ax.scatter(sub.flipper_length_mm, sub.body_mass_g / 1000, color=f"C{k}", s=8, alpha=0.5)
ax.set(xlabel="flipper length (mm)", ylabel="body mass (kg)")
ax.legend(fontsize=8);

FACTORS = {
    "Contract": ["Month-to-month", "One year", "Two year"],
    "InternetService": ["DSL", "Fiber optic", "No"],
    "PaymentMethod": ["Bank transfer (automatic)", "Credit card (automatic)", "Electronic check", "Mailed check"],
    "PaperlessBilling": ["No", "Yes"],
    "SeniorCitizen": ["No", "Yes"],
}
for f, levels in FACTORS.items():
    assert set(telco[f]) == set(levels), f

cell_table = telco.groupby(list(FACTORS), as_index=False).agg(
    customers=("churned", "size"), churners=("churned", "sum")
)
n_possible = np.prod([len(v) for v in FACTORS.values()])
print(f"{len(telco)} customers -> {len(cell_table)} non-empty cells (of {n_possible} possible)")
cell_table.sort_values("customers", ascending=False).head()