A Bayesian Model of NFL Field Goals in Python

%matplotlib inline

import logging

from humanize import intword
from toolz import compose

import arviz as az
import numpy as np
import nutpie
import polars as pl
import preliz as pz
import pymc as pm
import seaborn as sns
from matplotlib import pyplot as plt
from matplotlib import ticker
from pymc_extras.utils.spline import numpy_bspline_basis
from seaborn import objects as so
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import FunctionTransformer, MinMaxScaler, StandardScaler

import nflreadpy as nfl

sns.set(color_codes=True)

logging.getLogger("pytensor.graph.rewriting.basic").setLevel(logging.CRITICAL)

nfl.config.update_config(cache_mode="filesystem")

pbp_df = nfl.load_pbp(seasons=True)

print(f"{intword(pbp_df.shape[0])} x {pbp_df.shape[1]}")

FG_COLS = [
    "play_id",
    "game_id",
    "home_team",
    "away_team",
    "season_type",
    "week",
    "posteam",
    "defteam",
    "half_seconds_remaining",
    "game_seconds_remaining",
    "game_half",
    "qtr",
    "field_goal_result",
    "kicker_player_name",
    "kicker_player_id",
    "kick_distance",
    "score_differential",
    "season",
    "weather",
    "roof",
    "surface",
    "temp",
    "wind",
    "game_stadium",
]

fg_df = pbp_df.filter(pl.col("field_goal_attempt") == 1).select(FG_COLS)

PCT_FORMATTER = lambda: ticker.PercentFormatter(xmax=1, decimals=1)
PCT_LIMIT = (0, 1.05)
PCT_LOCATOR = lambda: ticker.MultipleLocator(0.25)

MADE_LABEL = "Field goals made"

(
    so.Plot(
        fg_df.with_columns(pl.col("field_goal_result").str.to_titlecase()).rename(
            {"field_goal_result": "Result"}
        ),
        x="Result",
    )
    .add(so.Bar(), so.Hist(stat="probability"))
    .scale(
        y=so.Continuous().tick(locator=PCT_LOCATOR()).label(formatter=PCT_FORMATTER())
    )
    .label(y=MADE_LABEL)
    .limit(y=PCT_LIMIT)
)

fg_df = fg_df.filter(pl.col("field_goal_result") != "blocked").with_columns(
    made=pl.col("field_goal_result") == "made"
)

print(intword(fg_df.shape[0]))

fg_df.select(
    pl.col("season").min().alias("first"), pl.col("season").max().alias("last")
)

fg_df = (
    fg_df.with_columns(
        q=pl.col("kick_distance").qcut(30, allow_duplicates=True, include_breaks=True),
    )
    .unnest("q")
    .drop("category")
    .rename({"breakpoint": "bin_max"})
    .with_columns(
        bin_mean=pl.col("kick_distance").mean().over("bin_max").cast(pl.Int64)
    )
)

def ensure_list(maybe_list):
    return maybe_list if isinstance(maybe_list, list) else [maybe_list]


def rolling_mean_by(df, *, name, by, window_size, over=None):
    # centered, row-based window; `by` values are near-evenly spaced
    # (bin_mean, season) so this approximates a centered window over them,
    # avoiding the ~window_size/2 lag a trailing window introduces
    expr = pl.col(name).rolling_mean(window_size, center=True, min_samples=1)
    sort_by = [by]

    if over is not None:
        over = ensure_list(over)
        expr = expr.over(over)
        sort_by += over

    return df.sort(sort_by).with_columns(expr)

DIST_MAJOR_LOCATOR = lambda: ticker.MultipleLocator(5)

DIST_LABEL = "Distance (yards)"

def plot_made_by_dist(df, *, color, color_label, title=None, order=None):
    color_scale = so.Nominal(order=order) if order else so.Nominal()

    return (
        so.Plot(
            df.group_by(color, "bin_mean")
            .agg(pl.col("made").mean())
            .pipe(
                rolling_mean_by, name="made", by="bin_mean", window_size=5, over=color
            ),
            x="bin_mean",
            y="made",
            color=color,
        )
        .add(so.Line())
        .scale(
            x=so.Continuous().tick(locator=DIST_MAJOR_LOCATOR()),
            y=so.Continuous()
            .tick(locator=PCT_LOCATOR())
            .label(formatter=PCT_FORMATTER()),
            color=color_scale,
        )
        .label(x=DIST_LABEL, y=MADE_LABEL, color=color_label, title=title)
    )

(
    so.Plot(
        fg_df.group_by("bin_mean")
        .agg(pl.col("made").mean())
        .pipe(rolling_mean_by, name="made", by="bin_mean", window_size=5),
        x="bin_mean",
        y="made",
    )
    .add(so.Line())
    .scale(
        x=so.Continuous().tick(locator=DIST_MAJOR_LOCATOR()),
        y=so.Continuous().tick(locator=PCT_LOCATOR()).label(formatter=PCT_FORMATTER()),
    )
    .label(x=DIST_LABEL, y=MADE_LABEL)
)

THRESHES = [0.5, 0.75, 0.9]

SEASON_LABEL = "Season"
SEASON_MAJOR_LOCATOR = lambda: ticker.MultipleLocator(5)

(
    so.Plot(
        fg_df.group_by("bin_mean", "season")
        .agg(pl.col("made").mean())
        .join(pl.DataFrame({"thresh": THRESHES}), how="cross")
        .filter(pl.col("made") >= pl.col("thresh"))
        .group_by("season", "thresh")
        .agg(pl.col("bin_mean").max())
        .pipe(
            rolling_mean_by, name="bin_mean", by="season", window_size=5, over="thresh"
        )
        .with_columns(
            pl.col("thresh").map_elements(lambda x: f"{x:.0%}", return_dtype=pl.String),
        ),
        x="season",
        y="bin_mean",
        color="thresh",
    )
    .add(so.Line())
    .scale(
        x=so.Continuous().tick(locator=SEASON_MAJOR_LOCATOR()),
        y=so.Continuous().tick(locator=DIST_MAJOR_LOCATOR()),
        color=so.Nominal(),
    )
    .label(
        x=SEASON_LABEL, y=DIST_LABEL, color="Threshold", title="Field goal success rate"
    )
)

fg_df["wind", "weather"]

fg_df = fg_df.with_columns(is_exposed=pl.col("roof").is_in(["open", "outdoors"]))

plot_made_by_dist(fg_df, color="is_exposed", color_label="Exposed", title="Field goals")

fg_df["surface"].value_counts()

fg_df = (
    fg_df.with_columns(
        surface=pl.col("surface").str.strip_chars().replace("", None),
    )
    .drop_nulls("surface")
    .with_columns(is_turf=~pl.col("surface").is_in(["grass", "dessograss"]))
)

plot_made_by_dist(fg_df, color="is_turf", color_label="Artificial turf")

plot_made_by_dist(
    fg_df.with_columns(home_kick=pl.col("home_team") == pl.col("posteam")),
    color="home_kick",
    color_label="Home team",
)

CLOSE_GAME_MARGIN = 3

fg_df = fg_df.with_columns(
    is_close=pl.col("score_differential").abs() <= CLOSE_GAME_MARGIN,
    score_status=(
        pl.when(pl.col("score_differential") == 0)
        .then(pl.lit("Tied"))
        .when(pl.col("score_differential") > 0)
        .then(pl.lit("Leading"))
        .otherwise(pl.lit("Trailing"))
    ),
)

plot_made_by_dist(
    fg_df,
    color="score_status",
    color_label="Score status",
    title="All games",
    order=["Trailing", "Tied", "Leading"],
)

plot_made_by_dist(
    fg_df.filter(pl.col("is_close")),
    color="score_status",
    color_label="Score status",
    title="Close games",
    order=["Trailing", "Tied", "Leading"],
)

plot_made_by_dist(
    fg_df.filter(pl.col("qtr") <= 4),
    color="qtr",
    color_label="Quarter",
    title="Field goals",
)

plot_made_by_dist(
    fg_df, color="season_type", color_label="Season type", title="Field goals"
)

fg_df.filter(pl.col("game_stadium").str.contains("at Mile High"))[
    "game_stadium"
].value_counts()

STADIUM_DATA_PATH = "https://austinrochford.com/resources/fg_irt/stadiums.csv"

stadium_df = pl.read_csv(
    STADIUM_DATA_PATH, columns=["game_stadium", "building", "elevation_ft"]
).rename({"building": "stadium"})

fg_df = fg_df.join(stadium_df, on="game_stadium", how="left").drop("game_stadium")

ELEV_LABEL = "Elevation (ft)"
ELEV_SCALE = "log"

(
    so.Plot(
        fg_df.with_columns(pl.col("made").cast(pl.Float64)), x="elevation_ft", y="made"
    )
    .add(so.Dot(), so.Agg(), group="stadium")
    .scale(
        x=ELEV_SCALE,
        y=so.Continuous().tick(locator=PCT_LOCATOR()).label(formatter=PCT_FORMATTER()),
    )
    .label(x=ELEV_LABEL, y=MADE_LABEL)
    .limit(x=(1, 10_000))
)

fg_df = fg_df.join(
    pbp_df.sort(["game_id", "qtr", "play_id"])
    .with_columns(
        prev_timeout_team=pl.col("timeout_team").shift(1).over("game_id", "qtr")
    )
    .filter(pl.col("field_goal_attempt") == 1)
    .with_columns(iced=pl.col("prev_timeout_team").eq_missing(pl.col("defteam")))
    .select("play_id", "game_id", "iced"),
    on=["play_id", "game_id"],
)

print(f"{fg_df['iced'].mean():.1%}")

plot_made_by_dist(fg_df, color="iced", color_label="Iced")

ENUM_COL_MAP = {
    "team": ["home_team", "away_team", "posteam", "defteam"],
    "field_goal_result": ["field_goal_result"],
    "kicker_player_name": ["kicker_player_name"],
    "season": ["season"],
    "roof": ["roof"],
    "surface": ["surface"],
    "stadium": ["stadium"],
}

enums = {}

for enum, names in ENUM_COL_MAP.items():
    categories = set()

    for name in names:
        categories |= set(fg_df[name].drop_nulls().cast(pl.String).unique())

    enums[enum] = pl.Enum(sorted(categories))

    for name in names:
        fg_df = fg_df.with_columns(pl.col(name).cast(pl.String).cast(enums[enum]))

make_2d_col = compose(np.transpose, np.atleast_2d, np.asarray)


def make_1d_transform(x, transform):
    return make_pipeline(
        FunctionTransformer(make_2d_col, np.squeeze),
        transform(),
        FunctionTransformer(np.squeeze, make_2d_col),
    ).fit(x)

dist_scaler = make_1d_transform(fg_df["kick_distance"].to_numpy(), MinMaxScaler)
DATA = {"dist_std": dist_scaler.transform(fg_df["kick_distance"].to_numpy())}

DATA["made"] = fg_df["made"].to_numpy()

def unpack_value_dims(value_dims):
    match value_dims:
        case (value, dims):
            pass
        case value:
            dims = "obs"

    return value, dims


def make_data(**data):
    """
    Keyword args should be either:
    * name=value (for dims="obs")
    * name=(value, dims)
    """
    containers = {}

    for name, value_dims in data.items():
        value, dims = unpack_value_dims(value_dims)
        containers[name] = pm.Data(name, value, dims=dims)

    return containers

def make_coefs(*names):
    return {name: pm.Normal(name, 0, 2.5) for name in names}

with pm.Model() as logit_model:
    data = make_data(**DATA)

with logit_model:
    coefs = make_coefs("β0", "β_dist")

with logit_model:
    η = coefs["β0"] + coefs["β_dist"] * data["dist_std"]
    p_made = pm.Deterministic("p_made", pm.math.sigmoid(η), dims="obs")

with logit_model:
    pm.Bernoulli("made_obs", p_made, observed=data["made"], dims="obs")

COMPILE_KWARGS = {"backend": "jax", "gradient_backend": "jax"}
SEED = 123456789

def sample(model, seed=SEED, **sample_kwargs):
    compiled = nutpie.compile_pymc_model(model, **COMPILE_KWARGS)

    trace = nutpie.sample(compiled, seed=seed, **sample_kwargs)
    pm.compute_log_likelihood(trace, model=model, progressbar=False)
    loo = az.loo(trace, pointwise=True)

    return trace, loo

traces = {}
loos = {}

traces["logit"], loos["logit"] = sample(logit_model)

def plot_ax(plotter):
    (ax,) = plotter._figure.axes
    return ax

# match strings that don't end in "_unid" and are not exactly "p_made"
RHAT_EXCLUDE_REGEX = r"^(?!p_made$).*(?<!_unid)$"

RHAT_KWARGS = {"var_names": RHAT_EXCLUDE_REGEX, "filter_vars": "regex"}


def plot_rhat(trace):
    rhat_df = (
        az.rhat(trace, **RHAT_KWARGS)
        .max()
        .to_dataset()
        .to_dataarray(dim="variable")
        .to_dataframe(name="rhat")
        .reset_index()
        .pipe(pl.from_pandas)
    )

    plotter = (
        so.Plot(rhat_df, x="rhat", y="variable")
        .add(so.Bar())
        .scale(x=so.Continuous().tick(every=0.01))
        .limit(x=(0.999, 1.02))
        .label(x=r"$\hat{R}$", y=None)
        .plot()
    )
    plot_ax(plotter).axvline(1.01, c="k", ls="--")

    return plotter

plot_rhat(traces["logit"])

def add_pred_resid(name, df, trace):
    return df.with_columns(
        **{
            f"{name}_pred": trace.posterior["p_made"]
            .mean(dim=("chain", "draw"))
            .to_numpy()
        }
    ).with_columns(pl.col("made").sub(pl.col(f"{name}_pred")).alias(f"{name}_resid"))

fg_df = add_pred_resid("logit", fg_df, traces["logit"])

def bin_resid(name, df, *, by="bin_mean"):
    return (
        df.group_by(by)
        .agg(
            pl.col("made").mean(),
            pl.col(f"{name}_pred").mean().alias("p_hat"),
            pl.len().alias("n"),
        )
        .with_columns(
            (
                (pl.col("made") - pl.col("p_hat"))
                / (pl.col("p_hat") * (1 - pl.col("p_hat")) / pl.col("n")).sqrt()
            ).alias(f"{name}_std_resid")
        )
        .drop("p_hat", "n")
    )

STD_RESID_LABEL = "Success probability\nstandardized residual"
RESID_LABEL = "Success probability\nresidual"

def add_ax_lines(plotter=None, *, val=2, horiz=True, ax=None):
    if ax is None:
        if plotter is None:
            raise ValueError("one of plotter or ax must not be None")
        else:
            ax = plot_ax(plotter)

    axline = ax.axhline if horiz else ax.axvline

    axline(-val, c="k", ls="--")
    axline(val, c="k", ls="--")

    return plotter


def plot_std_resid(
    name,
    df,
    *,
    x="bin_mean",
    x_label=DIST_LABEL,
    y_line=2,
    color=None,
    color_label=None,
    ax=None,
):
    plot_kwargs = {"x": x, "y": f"{name}_std_resid"}

    if color is None:
        by = x
    else:
        by = [x, color]
        plot_kwargs["color"] = color

    plot = (
        so.Plot(bin_resid(name, df, by=by), **plot_kwargs)
        .add(so.Dot(), so.Agg())
        .label(x=x_label, y=STD_RESID_LABEL, color=color_label)
    )

    if ax is None:
        return add_ax_lines(plot.plot(), val=y_line)
    else:
        plot = plot.on(ax).plot()
        add_ax_lines(val=y_line, ax=ax)

        return plot

plot_std_resid("logit", fg_df)

def reindex_array(da, *, obs_dim="obs", **coords):
    da = da.assign_coords({name: (obs_dim, values) for name, values in coords.items()})
    coord_names = list(coords)

    if len(coord_names) == 1:
        return da.swap_dims({obs_dim: coord_names[0]})
    else:
        return da.set_index({obs_dim: coord_names}).unstack(obs_dim)


def sample_posterior_predictive(trace, *, seed=SEED, coords=None):
    pm.sample_posterior_predictive(
        trace,
        var_names="p_made",
        predictions=True,
        random_seed=seed,
        extend_inferencedata=True,
    )

    if coords:
        trace["predictions"]["p_made"] = reindex_array(
            trace["predictions"]["p_made"], **coords
        )

def make_pp_df(*pp_data):
    fst, *rest = pp_data
    df = pl.DataFrame(fst)

    for data in rest:
        df = df.join(pl.DataFrame(data), how="cross")

    return df


def enum_cols(enum, name, values):
    s = pl.Series(values, dtype=enum)
    return {f"{name}_val": s.to_physical(), name: s}

def set_data(df):
    data = {}
    model = pm.modelcontext(None)

    for var in model.data_vars:
        if var.name in df:
            data[var.name] = df[var.name].to_numpy()
        else:
            value, _ = unpack_value_dims(DATA[var.name])
            data[var.name] = np.empty(df.shape[0], dtype=value.dtype)

    pm.set_data(data)

PP_N = 100
PP_DIST = np.linspace(fg_df["kick_distance"].min(), fg_df["kick_distance"].max(), PP_N)
PP_DIST_STD = dist_scaler.transform(PP_DIST)
PP_DIST_DATA = {
    "dist_std": PP_DIST_STD,
    "dist": PP_DIST,
}

pp_df = make_pp_df(PP_DIST_DATA)

with logit_model:
    set_data(pp_df)
    sample_posterior_predictive(
        traces["logit"],
        coords=pp_df[["dist"]].to_dict(as_series=False),
    )

def to_polars(dt):
    ds = dt.to_dataset()
    (name,) = ds.data_vars
    da = ds[name]
    dims = da.dims

    mesh = np.meshgrid(*(ds.coords[dim].to_numpy() for dim in dims), indexing="ij")
    columns = dict(zip(dims, (m.ravel() for m in mesh)))
    columns[name] = da.to_numpy().ravel()

    return pl.DataFrame(columns)

def made_to_float(df):
    return df.with_columns(pl.col("made").cast(pl.Float64))

def plot_axis_ticks(
    plot, *, x_locator=None, x_formatter=None, y_locator=None, y_formatter=None
):
    """Apply real matplotlib locators/formatters post-render.

    Works around a seaborn bug where so.Est()-based layers leave the y-axis
    PseudoAxis incompletely initialized if a matplotlib Locator/Formatter is
    passed through .scale() instead of applied directly to the rendered axes

    See https://github.com/mwaskom/seaborn/issues/3585 for details
    """
    plotter = plot.plot()
    ax = plot_ax(plotter)

    if x_locator is not None:
        ax.xaxis.set_major_locator(x_locator)
    if x_formatter is not None:
        ax.xaxis.set_major_formatter(x_formatter)
    if y_locator is not None:
        ax.yaxis.set_major_locator(y_locator)
    if y_formatter is not None:
        ax.yaxis.set_major_formatter(y_formatter)

    return plotter

def add_posterior_ev(
    plot,
    *,
    data=None,
    ci=True,
    color=None,
    color_by=None,
    linestyle_by=None,
    label=None,
):
    color_kwargs = {} if color is None else {"color": color}
    variables = {}

    if color_by is not None:
        variables["color"] = color_by

    if linestyle_by is not None:
        variables["linestyle"] = linestyle_by

    plot = plot.add(
        so.Line(**color_kwargs), so.Agg(), data=data, label=label, **variables
    )

    if ci:
        plot = plot.add(
            so.Band(**color_kwargs), so.Est(errorbar=("pi", 95)), data=data, **variables
        )

    return plot

POSTERIOR_EV_LABEL = "Posterior expected value"

plot = so.Plot(traces["logit"].predictions.pipe(to_polars), x="dist", y="p_made")
plot = add_posterior_ev(plot, label="Logistic regression")
plot = (
    plot.add(
        so.Line(color="k", linestyle="--"),
        so.Agg(),
        data=fg_df.pipe(made_to_float),
        x="bin_mean",
        y="made",
        label="Actual",
    )
    .scale(x=so.Continuous().tick(locator=DIST_MAJOR_LOCATOR()))
    .label(x=DIST_LABEL, y=MADE_LABEL, legend=POSTERIOR_EV_LABEL)
)

plot_axis_ticks(
    plot,
    y_locator=PCT_LOCATOR(),
    y_formatter=PCT_FORMATTER(),
)

K_BASIS = 12

DATA["dist_basis_val"] = (
    numpy_bspline_basis(DATA["dist_std"], K_BASIS, degree=3),
    ("obs", "dist_basis"),
)

COORDS = {"dist_basis": np.arange(K_BASIS)}

with pm.Model(coords=COORDS) as spline_model:
    data = make_data(**DATA)

    coefs = make_coefs("β0")

def zero_mean(x, *, axis):
    return x - x.mean(axis=axis, keepdims=True)


def dist_spline_coef(name, *, dims="dist_basis"):
    σ = pm.HalfNormal(f"σ_{name}", 2.5)
    coef_unid = pm.GaussianRandomWalk(
        f"{name}_unid", 0, σ, init_dist=pm.Normal.dist(0, 2.5), dims=dims
    )

    return pm.Deterministic(name, zero_mean(coef_unid, axis=0), dims=dims)

with spline_model:
    w = dist_spline_coef("w")
    f_dist = data["dist_basis_val"] @ w

with spline_model:
    η = coefs["β0"] + f_dist
    p_made = pm.Deterministic("p_made", pm.math.sigmoid(η), dims="obs")

    pm.Bernoulli("made_obs", p_made, observed=data["made"], dims="obs")

traces["spline"], loos["spline"] = sample(spline_model)

plot_rhat(traces["spline"])

fg_df = add_pred_resid("spline", fg_df, traces["spline"])

plot_std_resid("spline", fg_df)

fig, (dist_ax, hist_ax) = plt.subplots(ncols=2, figsize=(12, 5))

(
    so.Plot(fg_df, x="spline_resid", color="is_exposed")
    .add(so.Line(), so.KDE(common_norm=False), legend=False)
    .on(hist_ax)
    .scale(y=so.Continuous().tick(at=[]))
    .label(x=RESID_LABEL)
    .plot()
)

plot_std_resid("spline", fg_df, color="is_exposed", color_label="Exposed", ax=dist_ax)
fig.tight_layout();

fig, (dist_ax, hist_ax) = plt.subplots(ncols=2, figsize=(12, 5))

(
    so.Plot(fg_df, x="spline_resid", color="is_turf")
    .add(so.Line(), so.KDE(common_norm=False), legend=False)
    .on(hist_ax)
    .scale(y=so.Continuous().tick(at=[]))
    .label(x=RESID_LABEL)
    .plot()
)

plot_std_resid("spline", fg_df, color="is_turf", color_label="Turf", ax=dist_ax)
fig.tight_layout();

az.plot_compare(az.compare(loos));

PP_DIST_DATA["dist_basis_val"] = numpy_bspline_basis(PP_DIST_STD, K_BASIS, degree=3)

pp_df = make_pp_df(PP_DIST_DATA)

with spline_model:
    set_data(pp_df)

    sample_posterior_predictive(
        traces["spline"],
        coords=pp_df[["dist"]].to_dict(as_series=False),
    )

plot = so.Plot(traces["spline"].predictions.pipe(to_polars), x="dist", y="p_made")
plot = add_posterior_ev(plot, label="Spline")
plot = add_posterior_ev(
    plot,
    data=traces["logit"].predictions.pipe(to_polars),
    color="C1",
    label="Logistic regression",
)

plot_axis_ticks(
    plot.add(
        so.Line(color="k", linestyle="--"),
        so.Agg(),
        data=fg_df.pipe(made_to_float),
        x="bin_mean",
        y="made",
        label="Actual",
    )
    .scale(x=so.Continuous().tick(locator=DIST_MAJOR_LOCATOR()))
    .label(x=DIST_LABEL, y=MADE_LABEL, legend=POSTERIOR_EV_LABEL),
    y_locator=PCT_LOCATOR(),
    y_formatter=PCT_FORMATTER(),
)

(
    so.Plot(fg_df, x="kick_distance")
    .add(so.Bar(), so.Hist())
    .scale(x=so.Continuous().tick(locator=DIST_MAJOR_LOCATOR()))
    .label(x=DIST_LABEL, y="Number of attempts")
)

plotter = (
    so.Plot(fg_df, x="kick_distance")
    .add(so.Bar(), so.Hist())
    .scale(x=so.Continuous().tick(locator=DIST_MAJOR_LOCATOR()))
    .label(x=DIST_LABEL, y="Number of attempts")
    .plot()
)

ax = plot_ax(plotter)

basis_val = numpy_bspline_basis(PP_DIST_STD, K_BASIS, degree=3)
basis_ax = ax.twinx()

for i in range(K_BASIS):
    basis_ax.plot(
        PP_DIST,
        basis_val[:, i],
        c="k",
        ls="--",
        label="Basis functions" if i == 0 else None,
    )

basis_ax.set_ylim(bottom=0)
basis_ax.set_ylabel("Basis function value")
basis_ax.grid(False)

basis_ax.legend(loc="upper left")

plotter

_, ax = plt.subplots()

w_dist = PP_DIST[numpy_bspline_basis(PP_DIST_STD, K_BASIS, degree=3).argmax(axis=0)]

ax.plot(
    w_dist,
    traces["spline"].posterior["w"].std(dim=("chain", "draw")),
    label="$w_i$",
)

ax.set_xlabel(DIST_LABEL)
ax.set_ylabel("Basis function coefficient\nposterior standard deviation")
ax.legend(loc="upper right")

basis_ax = ax.twinx()

for i in range(K_BASIS):
    basis_ax.plot(
        PP_DIST,
        basis_val[:, i],
        c="k",
        ls="--",
        alpha=0.75,
        label="Basis functions" if i == 0 else None,
    )

basis_ax.set_ylim(bottom=0)
basis_ax.set_ylabel("Basis function value")
basis_ax.grid(False)

basis_ax.legend(loc="upper left");

DATA["exposed"] = fg_df["is_exposed"].to_numpy()
DATA["turf"] = fg_df["is_turf"].to_numpy()

with pm.Model(coords=COORDS) as env_model:
    data = make_data(**DATA)

    coefs = make_coefs("β0")

    w = dist_spline_coef("w")
    f_dist = data["dist_basis_val"] @ w

with env_model:
    coefs.update(make_coefs("β_exposed", "β_turf"))

with env_model:
    # fmt: off
    η = sum([
        coefs["β0"],
        coefs["β_exposed"] * data["exposed"],
        coefs["β_turf"] * data["turf"],
        f_dist
    ])
    # fmt: on

with env_model:
    p_made = pm.Deterministic("p_made", pm.math.sigmoid(η), dims="obs")

    pm.Bernoulli("made_obs", p_made, observed=data["made"], dims="obs")

traces["env"], loos["env"] = sample(env_model)

plot_rhat(traces["env"])

fg_df = add_pred_resid("env", fg_df, traces["env"])

plot_std_resid("env", fg_df)

fig, (dist_ax, hist_ax) = plt.subplots(ncols=2, figsize=(12, 5))

(
    so.Plot(fg_df, x="env_resid", color="is_exposed")
    .add(so.Line(), so.KDE(common_norm=False), legend=False)
    .on(hist_ax)
    .scale(y=so.Continuous().tick(at=[]))
    .label(x=RESID_LABEL)
    .plot()
)

plot_std_resid("env", fg_df, color="is_exposed", color_label="Exposed", ax=dist_ax)
fig.tight_layout();

fig, (dist_ax, hist_ax) = plt.subplots(ncols=2, figsize=(12, 5))

(
    so.Plot(fg_df, x="env_resid", color="is_turf")
    .add(so.Line(), so.KDE(common_norm=False), legend=False)
    .on(hist_ax)
    .scale(y=so.Continuous().tick(at=[]))
    .label(x=RESID_LABEL)
    .plot()
)

plot_std_resid("env", fg_df, color="is_turf", color_label="Turf", ax=dist_ax)
fig.tight_layout();

season_to_int = pl.col("season").cast(pl.String).str.to_integer()

add_ax_lines(
    so.Plot(
        bin_resid("env", fg_df, by="season").with_columns(season_to_int).sort("season"),
        x="season",
        y="env_std_resid",
    )
    .add(so.Dot())
    .scale(x=so.Continuous().tick(locator=SEASON_MAJOR_LOCATOR()))
    .label(x=SEASON_LABEL, y=STD_RESID_LABEL)
    .plot()
)

az.plot_compare(az.compare(loos));

az.plot_dist(traces["env"], var_names=["β_exposed", "β_turf"], ci_prob=0.95);

pp_df = make_pp_df(
    PP_DIST_DATA,
    {"exposed": [False, True]},
    {"turf": [False, True]},
)

with env_model:
    set_data(pp_df)

    sample_posterior_predictive(
        traces["env"],
        coords=pp_df[["dist", "exposed", "turf"]].to_dict(as_series=False),
    )

plot = so.Plot(traces["env"].predictions.pipe(to_polars), x="dist", y="p_made")
plot = add_posterior_ev(plot, ci=False, color_by="exposed", linestyle_by="turf")
plot = plot.scale(
    x=so.Continuous().tick(locator=DIST_MAJOR_LOCATOR()), color=so.Nominal()
).label(x=DIST_LABEL, y=MADE_LABEL, color="Exposed", linestyle="Turf")

plot_axis_ticks(plot, y_locator=PCT_LOCATOR(), y_formatter=PCT_FORMATTER())

COORDS["season"] = enums["season"].categories.to_list()

DATA["season_val"] = fg_df["season"].to_physical().to_numpy()

K_SEASON_BASIS = 6

season_basis = numpy_bspline_basis(np.linspace(0, 1, K_BASIS), K_SEASON_BASIS)

COORDS["season_basis"] = np.arange(K_SEASON_BASIS)

φ_prior, _ = pz.maxent(pz.Beta(), np.sqrt(0.25), np.sqrt(0.9), 0.95)

def deviation_coef(name, basis):
    η = pm.HalfNormal(f"η_{name}", 0.5)
    φ = φ_prior.to_pymc(f"φ_{name}")
    σ = pm.Deterministic(f"σ_{name}", η * pm.math.sqrt(1 - φ**2))

    model = pm.modelcontext(None)
    n_season = len(model.coords[name])
    coef_unid = pm.AR(
        f"α_{name}_unid",
        φ,
        σ,
        init_dist=pm.Normal.dist(0, η),
        steps=n_season - 1,
        dims=(f"{name}_basis", name),
    )
    coef = pm.Deterministic(
        f"α_{name}", zero_mean(coef_unid, axis=1), dims=(f"{name}_basis", name)
    )

    return pm.Deterministic(f"w_{name}", (basis @ coef).T, dims=(name, "dist_basis"))

with pm.Model(coords=COORDS) as season_model:
    data = make_data(**DATA)

    coefs = make_coefs("β0", "β_exposed", "β_turf")
    w0 = dist_spline_coef("w0")