The Boston Marathon - So Hot Right Now

Every year I make the short trek one block from Fenway to Kenmore Square to watch exhausted runners grinding through the final mile, occasionally with a cowbell in hand. While I love watching the marathon, I also love running myself, though my longest race to date is just a half marathon. And, unlike most runners, I actually have a small preference for warm-weather running. 70-75°F is my sweet spot. Research, however, does not back me here - most published literature seems to indicate about 50°F is optimum running conditions.

This year marks the fiftieth anniversary of the 1976 Boston Marathon, which peaked at 91.9°F, even too hot for me. It earned the name “Run for the Hoses.” In this race, Jack Fultz won in 2:20:19, which was nine minutes slower than the course record at the time. Bill Rodgers, who’d run 2:09:55 the year before, dropped out, as well as Tom Fleming, who placed third in 1975.

Thermal Fluctuations

Below is the average race time of the top 3 Boston Marathon runners over time.

This confirms that runners have been getting better over time, though non-linearly. Between 1950 and the turn of the century, gains were substantial. Around the turn of the century, the rate of improvement flattened. However, there’s considerable year-over-year variation, which makes a ton of sense. Conditions vary widely. In fact, that 1976 race sticks out like a sore thumb in the middle of the plot.

This made me curious, how much of that variation can be attributed to weather conditions on race day, which I’ll dive into in this post.

Prior work

This is hardly a novel problem. There has been plenty of work done on this topic. Three to highlight:

• Ely, Cheuvront, Roberts & Montain (2007) — pools seven major marathons (including Boston) and divides race-day conditions into WBGT quartiles. For elite men (top-3 finishers), slowdowns go 1.7% → 2.5% → 3.3% → 4.5% as you move from cool to hot. This is the closest thing the literature has to a canonical “how much does heat cost” table.
• Wang et al. (2024) — the world’s top-96 individual marathon athletes (top 16 per continent) followed across events. Models a linear performance degradation above 15°C with slope ~0.39 min/°C for men, 0.71 min/°C for women. This is the first model I attempt to replicate.
• Galloway & Maughan (1997) — not a marathon study, actually. Eight cyclists rode to exhaustion at four temperatures in a lab. Peak endurance came at 10.5°C (~51°F); by 30°C performance had collapsed. This study models the performance degradation as a quadratic function, which I replicate in the second model. See also Maughan, Watson & Shirreffs (2007) for the thermoregulation review that carries the same physiology over to distance running.

My research

These studies give a great groundwork, but I wanted to dig in a bit deeper. Specifically a few questions:

• The literature suggests both linear and quadratic models, which can vary wildly as you get to extreme values. I was curious what this looked like, and what a less constrained fit would look like.
• Both have pretty strongly constrained functional forms. I was interested in a more data-driven approach, namely fitting the weather effect via a spline based model.

So I fit three Bayesian models, each on log-finishing-time. To control for athletic performance, each use a random walk, and each control for precipitation. The models differ in how race-day temperature enters:

1. Wang-replication model: a linear hinge above 59°F, replicating Wang 2024’s functional form.
2. Physiology model: a quadratic hinge above the same knot, motivated by the Galloway study.
3. Spline model: replaces the hinge with a thin-plate spline, effectively letting the data pick the shape.

Each of the three models below is fit in brms with cmdstanr. Response is log_seconds (reader-facing numbers are back-transformed). All three include a state-space level (described below) to absorb improving performance over time, plus precipitation controls. Details on priors, convergence, and specifications live in the collapsible technical details at the bottom.

Controlling for athletic-performance improvement

Runners have gotten faster over the last century. Training methods modernized, footwear evolved, the global elite pool expanded. To measure weather effects, we need to account for that improvement and what performance would have looked like agnostic of weather.

I choose a simple approach to this, a random walk. Each year’s typical finishing time is the previous year’s, plus a small stochastic step. The size of this step is determined as part of the model fitting process. Notably, this does not impose a specific shape, just one wandering step-per-year.

This is a good fit for the following reasons:

• Uncertainty grows when you extrapolate. Predicting ten years past the data accumulates ten years of step variance; predicting one year ahead accumulates one. That matches reality — we know less about the far future than the well-measured middle.
• Year gaps are handled naturally. 1918 (WWI), 2020 (COVID), and 2021 (October race, omitted from the analysis) are missing from the data. The random walk weights its steps by the time elapsed, so the missing years cost us proportional information without breaking anything.

Here’s what the fitted year level looks like:

The red line is the model’s best guess of the typical finishing time for each year absent any weather effect; the bands are 80% and 95% credible intervals. The dots are the actual observed top-3 mean for each year.

If you haven’t worked with random walks often, you might ask: “the line just follows the dots — isn’t this just memorizing the data?” The random walk is regularized: each year’s step is constrained in size by a prior the data has to push against, so the line can track the broad pattern but can’t absorb every wiggle. The years where dots sit visibly above the red band (1976 again, plus 2012 and 2017) are years where something other than the long-run trend is at work, and that residual is what the weather model picks up. If the line went through every dot, there’d be nothing left for the weather term to explain.

In summary, this term captures the bulk year-over-year patterns not explained by the other model terms (like weather), and is temporally regularized so it can’t change too much from one year to the next.

One note is that this plot in particular highlights the 1953 drop from ~156 minutes to ~139 minutes. It seems that marathon was a perfect mix of a few factors: a 25-knot (29 mph) tailwind, 43 degree weather, a fast field, and a course that was found to over 1,000 yards short.

The actual brms implementation does this via a custom Stan block; the code is in the technical details at the bottom.

The Wang-replication model — linear hinge

First I reproduce the Wang model. The “hinge” is a kink at a fixed knot temperature: below the knot, there is no temperature effect — finishing time is flat. Above the knot, finishing time grows linearly with each additional degree. The biological story is that thermoregulation is essentially free until core-temperature rise becomes a problem, after which performance degrades steadily.

Explicitly:

Knot \(k\) at 59°F (15°C). Linear above the knot, exactly flat below by construction; \(\ell_{\text{year}}\) is the random-walk year level.

In brms syntax:

bf(
  log_seconds ~ 0 + I(pmax(tmax_f - KNOT_F, 0)) +
                precip_day + precip_missing,
  sigma ~ s(year, k = 10, bs = "tp")
)

A quick note on how to read this and the next two plots: the y-axis is the predicted slowdown relative to a 50°F reference. Each curve is plotted with the model’s prediction at 50°F subtracted off, so all three pass through zero at 50°F — that’s a reporting anchor, not a model constraint. (The hinge models do have a structural flat region, but it lives below their 59°F knot, not specifically at 50°F.) 50°F is roughly the long-run Patriots’ Day average, so it’s a natural “benign weather” baseline to compare hot or cold years against.

The ribbon is the posterior 95% credible interval. The slope above 15°C comes out to 0.30 min/°C (95% credible interval [0.14, 0.46]), somewhat shallower than Wang’s pooled 0.39.¹ This converts to about 10 seconds per °F above 59. The intervals overlap comfortably and the gap is consistent with the structural differences in the analyses - Boston is a single course, while Wang’s pooled estimate considers variation across many marathons.

The curve is exactly flat below 59°F because the model says it has to be: there’s no below-knot coefficient. By design, this model only makes claims above the hinge. The other two models allow for varying behavior for cold weather too.

The Physiology model — quadratic hinge

Next we do a Physiology-based model, which assumes temperature response is quadratic. Similar to Wang, we pin a knot at 59 degrees, however we allow for variation both above and below that knot:

This approach is motivated by Galloway & Maughan (1997), a cycling study which found an inverted-U relationship with performance peaking at ~10.5°C. We shift the knot to 15°C to align with Wang’s runner-based study, but keep the quadratic functional form on both sides of it.

bf(
  log_seconds ~ 0 + I(pmax(KNOT_F - tmax_f, 0)) +
                I(pmax(tmax_f - KNOT_F, 0)^2) +
                precip_day + precip_missing,
  sigma ~ s(year, k = 10, bs = "tp")
)

Despite allowing variation below the knot, it remains approximately flat below it. Above, we see strong quadratic behavior, however it is a bit alarming how close it gets to the 92°F day, which smells a bit overfit to me.

The Spline model

Last, I fit a model more of my own crafting, with a bit less structure than the literature, which more freely allows the data to drive the curve rather than strict functional form.

Here, we just have the performance and precipitation controls, and a thin-plate spline on temperature. A thin-plate spline is a flexible curve fit through the data with a built-in penalty for getting too wiggly. One small change from the previous two models: the spline takes temperature as an anomaly (race-day max minus the day-of-year climatological average) rather than the raw °F. The hinge models needed a raw-scale variable to anchor the 59°F knot; the spline doesn’t have a fixed knot, so the climate-normalized version is the natural input (Boston’s Patriots’-Day climatology only varies by a few degrees across April, so the practical difference is small).

This model also adds a second equation letting the residual variance depend smoothly on year: the within-year top-3 SD shrinks across the century as the elite field densifies and pacing professionalizes.

bf(
  log_seconds ~ 0 + s(tmax_anomaly_f, k = 6, bs = "tp") +
                precip_day + precip_missing,
  sigma ~ s(year, k = 10, bs = "tp")
)

The Spline model agrees with the other two below ~80°F where the data is dense. Above 85°F the curve bends less aggressively than the Physiology model’s quadratic — with one data point above 90°F, the thin-plate prior pulls toward a simpler shape.

So What’s the Answer?

So we’ve got three different formulations to answer our question. However, we want one final answer, not three. Often, modelers will just select a “best” model. If we go by that heuristic, the Spline model wins through Leave One Out (LOO) cross-validation, a measure of how well the model predicts held-out observations. The “LOO Δ vs best (in SE)” shown below is the gap to the best model expressed in standard errors of that gap. Spline wins individual leave-one-out: Physiology is modestly worse (~2 SE), Wang-replication is clearly worse (~6 SE).

However, I wanted to be a bit careful here. Notably, we’re in the area of outliers and extrapolation, dangerous waters for statistical models. In this regime, I err on conservatism where possible. There is plenty of literature that shows a many-models approach often beats a single, highly predictive model. That is what I chose to do here, known as Model Stacking: build a weighted mixture where each model’s weight is chosen to maximize the predictive performance of the combined model, following Yao, Vehtari, Simpson & Gelman (2018).

Stacking can favor a model that isn’t the best individual predictor, and we see that here — the largest contributor to the blend is the quadratic Physiology model, not the Spline. Stacking asks “if I’m allowed to combine the three models into one weighted prediction, what weights minimize the combined prediction error?” A model can be a slightly worse solo predictor and still earn high mixture weight if it makes its mistakes in different places than the others. Uncorrelated errors cancel out when blended. The Physiology model is more aggressive at the tails, predicting a higher heat penalty than the other two. The Spline and linear hinge models make similar predictions to each other, so they contribute redundant information. The Physiology model’s distinct behavior is what earns it a higher weight in the blend.

That being said, the data-rich region appropriately is the most impactful in fitting the models, but we assume it cleanly extrapolates to a less saturated region. The mixture’s weights reflect predictive accuracy in mild conditions, and we’re assuming that accuracy carries to the heat tail, which may not be a perfect assumption. For the rest of this post I report the stacked mixture as the headline. The Spline is the natural single-model starting point (it’s the best individual predictor), but for an extrapolation question with a single anchoring data point at 92°F, leaning on a blend that incorporates the mechanism-committed Physiology and Wang-replication models is the more defensible read.

Effect of Weather on Final Race Times

For every race-year in the panel, the stacked model estimates the weather-only contribution to that year’s top-3 average finish time: the difference between what the model predicts happened and what it would have predicted on a 50°F day. Note that by construction, this contribution is bounded below at zero: each model’s marginal curve has its minimum at the 50°F reference so any departure from the reference is by construction non-negative.

Most years cluster within a minute of zero. Two stand clearly above, 1976 and 2017. 2018 in particular is interesting: the model predicts a slowdown as a result of cold weather, rather than heat.

The 1976 counterfactual

In an alternate universe where it was a 50°F Patriots’ Day in 1976, the stacked mixture estimates the 1976 top-3 mean would have been faster by 6:27 (95% credible interval 2:47–10:52).

For context, by model: the Physiology model alone says 7:21 (the largest of the three), the Wang-replication model says 5:44, and the Spline model says 5:20.

The figure shows each model’s full posterior density side-by-side. The model identifies a cohort-level heat effect — every finisher in the field gets shifted by the same amount, so the counterfactual time is the same for each place by construction:

Observed top-3 mean: 2:21:14; the model’s 50°F counterfactual top-3 mean is 2:14:47.

2026 prediction

I still have 40 minutes until I’m technically allowed to burn the midnight oil, but I did do my last data pull as I wrapped up this piece at 11:20 pm in Boston, the night before the race. The NWS forecast for Hopkinton is 48°F, near perfect conditions. Under that weather, the stacked mixture predicts basically no temperature-related penalty and a top-3 mean of 2:05:10 with a 95% posterior predictive interval of 1:57:54–2:12:29. Obviously this is an incredibly wide interval. It would be remarkable to beat the world record of 2:00:35 on Boston’s famously difficult course, and 11% of our posterior probability claims the average of the top-3 finishers would be below that. However, that is to be expected given we formulated this as a weather model, not a performance forecasting model.

Of course, these curves all look similar, which is a result of the forecast being near-perfect.

Good luck to all the runners out there!

Caveats

Temperature is a proxy. The best marathon-performance studies use WBGT (wet-bulb globe temperature), which incorporates humidity, radiation, and wind. Blue Hill doesn’t have continuous dewpoint data until 2006 or wind data until 1998 at Logan. I kept maximum temperature + a precipitation indicator for simplicity.

DNF selection. The 1976 top-3 is the 3 best finishers, not the 3 best in the field. Rodgers and Fleming dropped out — athletes who ran 2:09 and 2:12 on 50°F in 1975 and 1977. This selection bias affects the study, but I kept the top-3 framing following the precedent set by Ely 2007.

Technical details

parameters {
  real ll_init;
  vector[Y] tilde_l;
  real<lower=0> sigma_l;
}
transformed parameters {
  vector[Y] ll;
  ll[1] = ll_init;
  for (y in 2:Y) {
    ll[y] = ll[y-1] + sigma_l * sqrt(dt[y]) * tilde_l[y];
  }
}
model {
  ll_init  ~ normal(9, 0.3);
  tilde_l  ~ std_normal();
  sigma_l  ~ normal(0, 0.05);
}

bf(
  log_seconds ~ 0 + s(tmax_anomaly_f, k = 6, bs = "tp") +
                precip_day +
                precip_missing,
  sigma ~ s(year, k = 10, bs = "tp")
)

bf(
  log_seconds ~ 0 + I(pmax(knot_f - tmax_f, 0)) +
                I(pmax(tmax_f - knot_f, 0)^2) +
                precip_day +
                precip_missing,
  sigma ~ s(year, k = 10, bs = "tp")
)

bf(
  log_seconds ~ 0 + I(pmax(tmax_f - knot_f, 0)) +
                precip_day +
                precip_missing,
  sigma ~ s(year, k = 10, bs = "tp")
)

Acknowledgments and sources

• NOAA GHCN-Daily for Blue Hill weather records and 1991–2020 Climate Normals for the climatology baseline.
• adrian3/Boston-Marathon-Data-Project for pre-2019 race results.
• The Boston Athletic Association for historical results 2019–2025 and race-date records.
• Paul Bürkner’s brms and the Stan developers.
• Matt Porter for getting my wheels spinning on the problem.

For live updates throughout the tournament, see the dashboard.

Background

Every March, 68 college basketball teams from 31 conferences get thrown into a single-elimination bracket, and \(2^{63}\) (roughly 9.2 quintillion) possible outcomes follow. I find bracketology to be a really fun statistical problem: individual games are often fairly predictable, but with that many possible outcomes, low-probability upsets are virtually guaranteed to sneak through somewhere. Everyone’s bracket breaks. That’s why even an octopus can put together competitive predictions for single-elimination tournaments.

For this year’s tournament, I fit a model to forecast the NCAA March Madness tournaments. The approach builds on the same philosophy as my NHL team strength model, a hierarchical Bayesian model for handling head-to-head matchups where we have historical scores. This model decomposes team quality into attack and defense parameters, and the same decomposition is a natural fit for basketball. By decomposing offense and defense, we can see how Houston’s defense-first approach is fundamentally different from Alabama’s strong offensive profile. Simulating bracket knowing how a team wins, not just that it wins provides a richer, more digestible understanding of the predictions.

The Model

Step 1: Bradley-Terry (Win/Loss)

The classic Bradley-Terry model assigns each team \(i\) a latent strength \(\theta_i\) and models the probability that team \(i\) beats team \(j\) as a function of the difference in their strengths. In its typical form, the only outcome is binary: who won. This is a workhorse model for pairwise comparisons (Elo ratings are a special case) and it works fine for generating bracket probabilities.

But basketball gives us more than just wins and losses. Beating a team by 30 tells you something very different than winning by 1, but the basic Bradley-Terry model sees those as equivalent.

Step 2: Score Differential

A natural extension is to model the score margin directly:

Now \(\theta_i\) captures how many points team \(i\) is above or below average, and the likelihood uses the full margin rather than collapsing it to a binary outcome. Home court advantage \(\alpha\) enters linearly. I stood this up initially, and it achieved a Brier score of 0.189 across 10 seasons of held-out tournament predictions.

But this model still assigns each team a single “strength” number. Michigan and Houston might have the same \(\theta\), but they get there in completely different ways. The single strength parameter can’t see that.

Step 3: Offense-Defense Decomposition

In this approach, each game produces two observations, the score produced by each team. By observing each team’s score separately, you can naturally split team strength into offense and defense:

Each team gets two parameters: an offensive strength \(\text{off}_i\) (how many points they generate above average) and a defensive strength \(\text{def}_i\) (how many points they prevent above average). The global intercept \(\mu\) anchors the average score per team per game, estimating at about 70 points. Home court advantage \(\alpha\) applies in the usual way.

Note that if you subtract the second equation from the first, you recover the margin model. Any information the margin model could learn, this model can as well. What it gains is the ability to distinguish Houston’s defense-first 77-55 wins from Michigan’s offense-first 101-83 wins, even when both are comfortable victories.

This is the same framework Baio & Blangiardo used for modeling the Premier League, adapted for college basketball. The key insight from that work carries over: by modeling scores directly instead of margins, you can identify teams that win by outscoring opponents versus teams that win by shutting them down.

Hierarchical Structure and the LKJ Correlation

Not all teams are created equal in terms of schedule strength. A mid-major team that went 28-3 against weak competition: how good are they really? I handle this with a hierarchical prior. Each team’s offense and defense are drawn from conference-level distributions:

where \(L\) is the Cholesky factor of a \(2 \times 2\) covariance matrix with an LKJ(\(\eta=2\)) prior on the correlation, and \(z_i \sim \text{Normal}(0, 1)\). Conference-level means for offense (\(\mu^{\text{off}}_c\)) and defense (\(\mu^{\text{def}}_c\)) are estimated separately. The SEC might produce strong defenses while the Big Ten generates potent offenses, or vice versa.

The LKJ prior lets the model learn whether offense and defense trade off or co-occur within a conference. The non-centered parameterization via \(z_i\) is the standard trick for sampling efficiency in hierarchical models.

The ultimate result is that we get conference-level effects: a dominant team in a weak conference gets pulled down slightly, and an underperforming team in a strong conference gets a boost.

On the likelihood

One decision to make in this process is the likelihood for number of points scored in a game. Here we balance pragmatism and expressiveness.

In basketball, points are discrete units so a natural question is why not use a Poisson or negative binomial likelihood? A couple of reasons:

• Basketball scores are sums of 1/2/3-point possessions, not a genuine counting process. A team’s score of 78 isn’t “78 events occurred” the way 3 goals in a hockey game is.
• At a mean of ~70 points, the Gaussian puts essentially zero mass below zero, and by the central limit theorem it’s a good approximation to a sum of many small discrete contributions anyway.
• The continuous approximation (Gaussian) is simpler and easier to fit, which was helpful given I had one night to put this together.

With individual scores instead of margins, outliers become a concern. Blowout games could pull the model around. To handle this empirically, I fit both a Gaussian and a Student-t likelihood (which is more robust to outliers) and compared them via LOO cross-validation:

The ELPD difference was negligible in favor of Gaussian. The Student-t model estimated \(\nu \approx 53\), which is basically indistinguishable from a Gaussian, so I kept the Gaussian for simplicity.

Full Model Specification

with pm.Model(coords=coords) as model:
    mu_intercept = pm.Normal("mu_intercept", mu=70, sigma=10)

    sigma_off_conf = pm.HalfNormal("sigma_off_conf", sigma=5)
    sigma_def_conf = pm.HalfNormal("sigma_def_conf", sigma=5)
    mu_off_conf = pm.Normal("mu_off_conf", mu=0, sigma=sigma_off_conf, dims="conference")
    mu_def_conf = pm.Normal("mu_def_conf", mu=0, sigma=sigma_def_conf, dims="conference")

    sd_dist = pm.HalfNormal.dist(sigma=5)
    chol, corr, stds = pm.LKJCholeskyCov("lkj", n=2, eta=2, sd_dist=sd_dist,
                                          compute_corr=True)
    z = pm.Normal("z", mu=0, sigma=1, shape=(n_teams, 2))

    team_effects = pt.stack([mu_off_conf[conf_of_team],
                             mu_def_conf[conf_of_team]], axis=1) + pt.dot(z, chol.T)
    off = pm.Deterministic("off", team_effects[:, 0], dims="team")
    deff = pm.Deterministic("def", team_effects[:, 1], dims="team")

    alpha = pm.Normal("alpha", mu=3.5, sigma=2.0)
    sigma = pm.HalfNormal("sigma", sigma=15)

    mu_score_i = mu_intercept + off[team_i] - deff[team_j] + alpha * home
    mu_score_j = mu_intercept + off[team_j] - deff[team_i] - alpha * home

    pm.Normal("score_i", mu=mu_score_i, sigma=sigma, observed=score_i, dims="game")
    pm.Normal("score_j", mu=mu_score_j, sigma=sigma, observed=score_j, dims="game")

Sampled with nutpie: 4 chains, 2,000 draws each after 2,000 tuning steps. Zero divergences, \(\hat{R} \leq 1.01\) for all parameters.

Data

The model is fit on the 2025-26 regular season: all 5,647 Division I games across 365 teams and 31 conferences. Data comes from the Kaggle March ML Mania 2026 competition dataset. Each game provides: which teams played, both scores, and whether it was a home game, away game, or neutral site.

Posterior Predictive Check

The posterior predictive distribution matches the observed score distribution well.

Results

Overall Team Strength

The overall strength of a team is the sum of its offensive and defensive contributions: \(\text{off}_i + \text{def}_i\).

Michigan and Duke sit at the top, followed by Arizona and Florida. These are the 4 number one seeds. Some may say this is chalky, I say it’s a great sanity check.

The Offense-Defense Decomposition

Every dot is a team, grey ones are those that didn’t make the tournament. The x-axis is offensive strength (higher = generates more points above average), the y-axis is defensive strength (higher = allows fewer points than average). Teams in the upper-right are good at both; teams in the lower-left are bad at both.

A few teams to highlight:

• Houston (overall 23.6): off = 7.0, def = 16.6. This is a defense-first team by a wide margin. Elite defensive rating, merely decent offense.
• Alabama (overall 21.0): off = 20.7, def = 0.3. The highest offensive rating of any tournament team, however paired with defense essentially at the league average.
• Duke (overall 28.0): off = 12.0, def = 16.0. Also defense-first, but with a more balanced profile than Houston.

The side-by-side rankings make the decomposition concrete. The list of top-20 offenses and top-20 defenses are quite different. A team can rank in the top 5 offensively while barely cracking the top 20 defensively, or vice versa.

Physically, \(\text{off}_i\) is how many points above average team \(i\) produces per game, and \(\text{def}_i\) is how many points below average they hold opponents to. These are per-game quantities, not per-possession, so they bake in pace alongside efficiency.

Conference Effects: Offense vs Defense

Conferences generally don’t lean dramatically in one direction or the other - good conferences are good, bad conferences are bad. The correlation in offensive and defensive effects is 0.83, very strong.

Shooters Shoot - Simpson’s Paradox in the Wild

One thing worth noting in the plot above is the marginal correlation between offense and defense across all 365 teams is +0.38. That’s positive: teams that are good at offense tend to also be good at defense. Intuitively, this makes sense: programs with good coaching, recruiting, and resources tend to be good at everything.

But the LKJ model parameter, the within-conference correlation that the model learned, is -0.20. That’s negative. Within a conference, offense and defense trade off. This is Simpson’s paradox. The relationship reverses when you condition on the confounding variable (conference membership). Here’s why:

Between conferences, the correlation is +0.83. Strong conferences (SEC, Big 12, Big Ten) produce teams that are above average on both offense and defense. Weak conferences produce teams below average on both. This between-group correlation is strong and positive, and when you pool all teams together ignoring conference, it dominates the marginal relationship.

Within a conference, though, there’s also a structural component: games are zero-sum at the score level. When Team A runs up the score on Team B, that same game hurts Team B’s defensive numbers. One team’s strong offensive showing is simultaneously a poor defensive showing for the opponent, and conference opponents play each other repeatedly, effectively inducing a negative correlation. The data suggests a modest negative correlation of about -0.20 within conferences, largely shaking out from this.

The LKJ model parameter of -0.20 captures precisely this within-conference structure. That’s what it’s designed to do: after the conference means (\(\mu^{\text{off}}_c\), \(\mu^{\text{def}}_c\)) absorb the between-conference variation, the LKJ correlation models the residual relationship between offense and defense at the team level.

Home Court Advantage

The model estimates \(\alpha\) at 1.44 ± 0.11 points. At first glance this looks low, but remember that \(\alpha\) appears in both score equations with opposite signs: the home team’s expected score goes up by \(\alpha\), and the away team’s goes down by \(\alpha\). The net effect on the margin is \(2\alpha \approx 2.9\) points, consistent with the published literature and KenPom’s estimate of ~3.5 raw points (our estimate is slightly lower because the model accounts for team quality differences simultaneously).

This parameter matters for training but zeroes out for tournament predictions, since all men’s tournament games are on neutral courts. I’ll come back to why this distinction matters for the women’s tournament later.

Simulating the Tournament

Win probabilities come from the Gaussian CDF applied to the strength difference. For two teams \(i\) and \(j\) on a neutral court:

The \(\sqrt{2}\) comes from the fact that we’re comparing the difference of two independent score random variables, each with variance \(\sigma^2\).

For each of 10,000 simulations, I draw a complete set of team strengths from the posterior and play through the 68-team bracket. If only I could spam all of them to my bracket pool.

Championship Odds

Again, the four number one seeds rise to the top, monopolizing nearly 50% of the championship probability. However, that does mean a 50% chance none of them win the tournament.

Advancement Probabilities

The advancement heatmap tells the full story. Notice how probability drops off dramatically at each round. Even #1 seeds only have ~50% probability of making the Final Four. This is because as the tournament progresses, you’re playing better competition, and you have more and more hurdles to clear. If the first round you have a 10% chance of dropping, the second you have 20%, that’s a total of a 28% chance of dropping out in the first two rounds, worse than one-in-four.

Below is the breakdown by region:

Upset Watch

Worth noting that the model is fairly chalky here. It will generally agree with the committee’s seeding, since it’s estimating overall team strength from the same regular season data. However, it does see Texas A&M as the one 10 seed to beat a 7 in the first round (though, it’s just over 50%, so effectively a coin flip).

Tail Outcomes

The most fun part of running 10,000 brackets: the tail outcomes. These are things that could happen, the model assigns them nonzero probability, even if they’re unlikely.

Deepest Cinderella runs across 10,000 simulations:

• Northern Iowa (12) won the championship in 2 simulations
• NC State (11) and SMU (11) each won the championship
• Cal Baptist (13) reached the Championship game
• Tennessee St (15) reached the Championship game
• Siena (16) reached the Elite Eight

Historical Validation

I validated by fitting the offense-defense model on each regular season from 2015–2025 (skipping 2020’s COVID cancellation) and predicting that year’s tournament.

Results across 669 tournament games:

• Brier score: 0.189
• Accuracy: 70.4%
• Log loss: 0.555

It is worth calling out that this Brier score is equivalent to a single strength parameter model. Sort of a no-free-lunch thing going on here, we’re a bit constrained by using just the score margin as our outcome. That being said, the decomposition paints a much clearer picture of team structure, making the model interpretable and useful in a way not achievable by a single parameter model.

The calibration plot shows the model is reasonably well-calibrated. There is a slight bit of the characteristic “s-curve” shape which comes from overfitting that could be improved upon in the future.

The Women’s Tournament

I fit an independent model on the women’s regular season using the identical offense-defense specification. The results tell a very different story from the men’s tournament.

UConn and the Chalkiness Gap

The men’s tournament has two co-favorites separated by 1.3 percentage points. The women’s tournament has UConn and everyone else:

UConn’s 34-0 record translates to a posterior strength distribution that clearly leads the pack, they are elite on both offense and defense.

The structural difference between the men’s and women’s fields is stark. The men’s top 30 is a smooth gradient where each team’s HDI overlaps with the teams around it. The women’s field has a clear break: five teams in one tier, Texas sort of in-between, then below that, the rest of the tournament field packs into a narrow band.

Talent in women’s college basketball is less equitably distributed. The top programs separate themselves from the field by a wider margin than their men’s counterparts, and the model sees it directly in the regular season results. For one, the correlation between offense and defense is much stronger in the women’s tournament.

We can quantify the concentration with a Gini coefficient over championship probabilities (0 means every team has equal odds, 1 means a single team wins every simulation). The women’s field is substantially more top-heavy: a Gini of 0.93 compared to 0.79 for the men’s. Put differently, only 29 women’s teams won the title in any of 10,000 simulations, compared to 46 on the men’s side.

Women’s Advancement Probabilities

The advancement heatmap makes the concentration visible. The women’s top seeds are a sea of deep red through the Sweet 16 and beyond, while probability drops off far more sharply for everyone else.

Home Court Advantage: This Time it Actually Matters

The women’s model estimates \(\alpha\) at 1.37 ± 0.10, almost identical to the men’s 1.44. The posteriors overlap almost entirely, suggesting that home court is roughly the same phenomenon across genders (a margin effect of about 2.7-2.9 points).

It is worth noting that this affects predictions differently than in the men’s tournament. In the men’s tournament, every game is neutral-site, so home court zeros out entirely. In the women’s tournament, the top 16 seeds host rounds 1 and 2 on their home courts. That ~2.8-point margin advantage compounds on top of the strength differential that already favors higher seeds, making the early rounds even more lopsided.

The Brackets

Based on the most likely outcome at each matchup:

Men’s

Final Four: Duke, Florida, Michigan, Arizona

Championship Game: Michigan vs. Duke

Champion: Michigan (16.2%, meaning we’re 83.8% sure this is wrong)

Women’s

Final Four: UConn, South Carolina, UCLA, Texas

Championship Game: UConn vs. South Carolina

Champion: UConn (33.8%, meaning we’re 66.2% sure this is wrong)

The contrast captures the structural difference between the two fields. The men’s champion is nearly a coin-flip among three or four teams. The women’s champion is the clearest favorite the model produces, but 33.8% still means there’s a 66.2% chance that UConn loses. Single-elimination tournaments are designed to produce uncertainty, and even a dominant team can only be so dominant across six consecutive games.

Following Along

I’ve put together an interactive dashboard that will update daily as the tournament progresses, so you can watch the model’s predictions shift as games are played and teams are eliminated. It is worth noting that these predictions run the statistical model live every morning, meaning there may be relatively small differences between predictions here and the dashboard, since I do not fix the random seed.

I also submitted these predictions to the Kaggle March ML Mania 2026 competition to see how they stack up against other approaches. By no means do I expect this to win, since it’s more a pedagogy-first exercise and more feature-dense approaches will likely perform better. That being said, it will be fun for a performance comparison.

Code

All code for this analysis is available on GitHub.

This is the second of two notebooks I wrote and contributed to Bambi’s example documentation back in 2024. The first post, covering polynomial regression basics, is here. This one goes deeper into what happens when you use the poly keyword in a Bambi formula. Specifically, looking at the orthogonalization that happens under the hood.

The original notebook lives in the Bambi docs. What follows is the content, lightly adapted for this blog.

Orthogonal Polynomial Regression

While the content here can stand alone, it is a companion to the polynomial regression post, which contains additional useful examples.

import warnings

import arviz as az
import bambi as bmb
import formulae
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import scipy
import seaborn as sns
from typing import Optional

SEED = 1234
az.style.use("arviz-darkgrid")
warnings.filterwarnings("ignore")

Revisiting Polynomial Regression

To start, we’ll recreate the projectile motion data defined in the polynomial regression post with \(x_0 = 1.5\) \(m\) and \(v_0 = 7\) \(m\)/\(s\). This will follow:

Where \(g\) will be the acceleration of gravity on Earth, \(-9.81\) \(m\)/\(s^2\). First we’ll generate the data.

g = -9.81
v0 = 7
x0 = 1.5
t = np.linspace(0, 2, 100)
x_projectile = (1/2) * g * t**2 + v0 * t + x0
noise = np.random.normal(0, 0.2, x_projectile.shape)
x_obs_projectile = x_projectile + noise
df_projectile = pd.DataFrame({"t": t, "x": x_obs_projectile, "x_true": x_projectile})
df_projectile = df_projectile[df_projectile["x"] >= 0]

plt.scatter(df_projectile.t, df_projectile.x, label='Observed Displacement', color="C0")
plt.plot(df_projectile.t, df_projectile.x_true, label='True Function', color="C1")
plt.xlabel('Time (s)')
plt.ylabel('Displacement (m)')
plt.ylim(bottom=0)
plt.legend()
plt.show()

Putting this into Bambi, we set \(\beta_2 = \frac{g}{2}\), \(\beta_1 = v_0\), and \(\beta_0 = x_0\), then perform the following regression:

We expect to recover \(\beta_2 = -4.905\), \(\beta_1 = 7\), \(\beta_0 = 1.5\) from our fit. We start with the approach from the other notebook where we explicitly tell formulae to calculate coefficients on \(t^2\) and \(t\).

model_projectile_all_terms = bmb.Model("x ~ I(t**2) + t + 1", df_projectile)
fit_projectile_all_terms = model_projectile_all_terms.fit(
    idata_kwargs={"log_likelihood": True}, random_seed=SEED
)
az.summary(fit_projectile_all_terms)

            mean     sd  hdi_3%  hdi_97%  mcse_mean  mcse_sd  ess_bulk  ess_tail  r_hat
sigma      0.173  0.014   0.148    0.199      0.000    0.000    2978.0    2514.0    1.0
Intercept  1.455  0.057   1.346    1.558      0.001    0.001    2553.0    2663.0    1.0
I(t ** 2) -4.999  0.097  -5.181   -4.817      0.002    0.002    2067.0    1996.0    1.0
t          7.186  0.161   6.872    7.483      0.003    0.003    2120.0    1967.0    1.0

The parameters are recovered as anticipated.

If you want to include all terms of a variable up to a given degree, you can also use the keyword poly. So if we want the linear and quadratic effects, as in this case, we would designate poly(t, 2).

model_projectile_poly = bmb.Model("x ~ poly(t, 2) + 1", df_projectile)
fit_projectile_poly = model_projectile_poly.fit(
    idata_kwargs={"log_likelihood": True}, random_seed=SEED
)
az.summary(fit_projectile_poly)

                mean     sd  hdi_3%  hdi_97%  mcse_mean  mcse_sd  ess_bulk  ess_tail  r_hat
sigma          0.173  0.014   0.147    0.201      0.000    0.000    4781.0    3156.0    1.0
Intercept      2.883  0.019   2.849    2.919      0.000    0.000    6952.0    3211.0    1.0
poly(t, 2)[0] -3.792  0.173  -4.109   -3.458      0.002    0.003    6298.0    2977.0    1.0
poly(t, 2)[1] -8.987  0.171  -9.307   -8.669      0.002    0.003    6499.0    3087.0    1.0

Now there are fitted coefficients for \(t\) and \(t^2\), but wait, those aren’t the parameters we used! What’s going on here?

The poly Keyword

To fully understand what’s going on under the hood, we must wade into some linear algebra. When the poly keyword is used, instead of directly using the values of \(x, x^2, x^3, \dots, x^n\), it converts them into orthogonal polynomials. When including the effect from multiple polynomial terms, there will generally be correlation between them. Including all of these into a model can be a problem from the fitting perspective due to multicollinearity. By orthogonalizing, the correlation is removed by design.

As it turns out, it’s difficult to get any information on how the orthogonalization is performed. Here is the implementation for poly in formulae, but to fully understand, I went into the source code for the R Stats library where poly is defined as a function for use on any vector, and took a look at its code.

Here’s a step-by-step summary, along with a toy example for \(x^4\).

• The data is first centered around the mean for stability

X = np.array([1, 2, 3, 4, 5])

mean = np.mean(X)
X_centered = X - mean
print(f"Array: {X}, mean: {mean}.\nCentered: {X_centered}")

Array: [1 2 3 4 5], mean: 3.0.
Centered: [-2. -1.  0.  1.  2.]

• A Vandermonde matrix is created. This just takes the input data and generates a matrix where columns represent increasing polynomial degrees. In this example, the first column is \(x^0\), a constant term. The second is \(x^1\), or the centered data. The third column is \(x^2\), the fourth is \(x^3\), the last is \(x^4\).

degree = 4
simple_vander = np.vander(X_centered, N=degree+1, increasing=True)
simple_vander

array([[ 1., -2.,  4., -8., 16.],
       [ 1., -1.,  1., -1.,  1.],
       [ 1.,  0.,  0.,  0.,  0.],
       [ 1.,  1.,  1.,  1.,  1.],
       [ 1.,  2.,  4.,  8., 16.]])

• QR decomposition is performed. There are several methods to doing this in practice, the most common being the Gram-Schmidt process. Here I just take advantage of the Numpy implementation. We take the above matrix and convert it into two components, an orthogonal matrix \(Q\), and an upper triangular matrix \(R\).

q, r = np.linalg.qr(simple_vander)
print("Orthogonal matrix Q:\n", q.round(4))
print("\nUpper triangular matrix R:\n", r.round(4))

Orthogonal matrix Q:
 [[-0.4472 -0.6325  0.5345 -0.3162 -0.1195]
 [-0.4472 -0.3162 -0.2673  0.6325  0.4781]
 [-0.4472 -0.     -0.5345  0.     -0.7171]
 [-0.4472  0.3162 -0.2673 -0.6325  0.4781]
 [-0.4472  0.6325  0.5345  0.3162 -0.1195]]

Upper triangular matrix R:
 [[ -2.2361  -0.      -4.4721  -0.     -15.2053]
 [  0.       3.1623   0.      10.7517   0.    ]
 [  0.       0.       3.7417   0.      16.5702]
 [  0.       0.       0.       3.7947   0.    ]
 [  0.       0.       0.       0.      -2.8685]]

• Last take the dot product of \(Q\) with the diagonal elements of \(R\). \(Q\) is then scaled to the magnitude of the polynomial degrees in \(R\). This serves as our transformation matrix which transforms input data into the space defined by the orthogonal polynomials.

diagonal = np.diag(np.diag(r))  # First call gets elements, second creates diag matrix
transformation_matrix = np.dot(q, diagonal)
print(transformation_matrix.round(4))

[[ 1.     -2.      2.     -1.2     0.3429]
 [ 1.     -1.     -1.      2.4    -1.3714]
 [ 1.     -0.     -2.      0.      2.0571]
 [ 1.      1.     -1.     -2.4    -1.3714]
 [ 1.      2.      2.      1.2     0.3429]]

• From the transformation matrix, we get squared norms (norm2), which give us the scale of each polynomial. We also get the value by which we need to shift each polynomial to match the centered data (alpha).

norm2 = np.sum(transformation_matrix**2, axis=0)

weighted_sums = np.sum(
    (transformation_matrix**2) * np.reshape(X_centered, (-1, 1)),
    axis=0
)
normalized_sums = weighted_sums / norm2
adjusted_sums = normalized_sums + mean
alpha = adjusted_sums[:degree]

print(f"Norm2: {norm2}\nalpha: {alpha}")

Norm2: [ 5.         10.         14.         14.4         8.22857143]
alpha: [3. 3. 3. 3.]

• Finally, we iteratively apply this to all desired polynomial degrees, shifting the data and scaling by the squared norms appropriately to maintain orthogonality with the prior term.

transformed_X = np.full((len(X), degree+1), np.nan)
transformed_X[:,0] = 1
transformed_X[:, 1] = X - alpha[0]
for i in range(1, degree):
    transformed_X[:, i + 1] = (
        (X - alpha[i]) * transformed_X[:, i] -
        (norm2[i] / norm2[i - 1]) * transformed_X[:, i - 1]
    )

transformed_X /= np.sqrt(norm2)
transformed_X

array([[ 4.47213595e-01, -6.32455532e-01,  5.34522484e-01,
        -3.16227766e-01,  1.19522861e-01],
       [ 4.47213595e-01, -3.16227766e-01, -2.67261242e-01,
         6.32455532e-01, -4.78091444e-01],
       [ 4.47213595e-01,  0.00000000e+00, -5.34522484e-01,
         2.34055565e-16,  7.17137166e-01],
       [ 4.47213595e-01,  3.16227766e-01, -2.67261242e-01,
        -6.32455532e-01, -4.78091444e-01],
       [ 4.47213595e-01,  6.32455532e-01,  5.34522484e-01,
         3.16227766e-01,  1.19522861e-01]])

This is now a matrix of orthogonalized polynomials of X. The first column is just a constant. The second column corresponds to the input \(x\), the next is \(x^2\) and so on. In most implementations, the constant term is eliminated, giving us the following final matrix.

transformed_X[:,1:]

array([[-6.32455532e-01,  5.34522484e-01, -3.16227766e-01,
         1.19522861e-01],
       [-3.16227766e-01, -2.67261242e-01,  6.32455532e-01,
        -4.78091444e-01],
       [ 0.00000000e+00, -5.34522484e-01,  2.34055565e-16,
         7.17137166e-01],
       [ 3.16227766e-01, -2.67261242e-01, -6.32455532e-01,
        -4.78091444e-01],
       [ 6.32455532e-01,  5.34522484e-01,  3.16227766e-01,
         1.19522861e-01]])

The approach shown in this derivation has been reproduced below as a Scikit-Learn style class, where the fit method calculates the coefficients and the transform method returns orthogonalized data. It is also at this gist, including the typical BaseEstimator, TransformerMixin inheritances.

class OrthogonalPolynomialTransformer:
    """Transforms input data using orthogonal polynomials."""

    def __init__(self, degree: int = 1) -> None:
        self.degree = degree + 1  # Account for constant term
        self.norm2 = None
        self.alpha = None

    def fit(self, X, y=None):
        """Calculate transformation matrix, extract norm2 and alpha."""
        self.norm2 = None
        self.alpha = None

        X = np.asarray(X).flatten()
        if self.degree >= len(np.unique(X)):
            raise ValueError(
                "'degree' must be less than the number of unique data points."
            )

        mean = np.mean(X)
        X_centered = X - mean

        vandermonde = np.vander(X_centered, N=self.degree + 1, increasing=True)
        Q, R = np.linalg.qr(vandermonde)

        diagonal = np.diag(np.diag(R))
        transformation_matrix = np.dot(Q, diagonal)
        self.norm2 = np.sum(transformation_matrix**2, axis=0)

        weighted_sums = np.sum(
            (transformation_matrix**2) * np.reshape(X_centered, (-1, 1)),
            axis=0
        )
        normalized_sums = weighted_sums / self.norm2
        adjusted_sums = normalized_sums + mean
        self.alpha = adjusted_sums[:self.degree]
        return self

    def transform(self, X):
        """Iteratively apply up to 'degree'."""
        X = np.asarray(X).flatten()
        transformed_X = np.empty((len(X), self.degree + 1))

        transformed_X[:, 0] = 1
        if self.degree > 0:
            transformed_X[:, 1] = X - self.alpha[0]

        if self.degree > 1:
            for i in range(1, self.degree):
                transformed_X[:, i + 1] = (
                    (X - self.alpha[i]) * transformed_X[:, i] -
                    (self.norm2[i] / self.norm2[i - 1]) * transformed_X[:, i - 1]
                )

        transformed_X /= np.sqrt(self.norm2)
        return transformed_X[:, 1:self.degree]

    def fit_transform(self, X, y=None):
        self.fit(X, y)
        return self.transform(X)

An example call is shown below. It’s worth noting that in this implementation, the constant term is not returned, the first column corresponds to \(x\), the second to \(x^2\), and the third to \(x^3\).

X = np.array([1, 2, 3, 4, 5])
poly3 = OrthogonalPolynomialTransformer(degree=3).fit(X)
poly3.transform(X)

array([[-6.32455532e-01,  5.34522484e-01, -3.16227766e-01],
       [-3.16227766e-01, -2.67261242e-01,  6.32455532e-01],
       [ 0.00000000e+00, -5.34522484e-01,  2.34055565e-16],
       [ 3.16227766e-01, -2.67261242e-01, -6.32455532e-01],
       [ 6.32455532e-01,  5.34522484e-01,  3.16227766e-01]])

This matches what you may get when calling the same function in R:

> poly(X, 4)
                 1          2             3          4
[1,] -6.324555e-01  0.5345225 -3.162278e-01  0.1195229
[2,] -3.162278e-01 -0.2672612  6.324555e-01 -0.4780914
[3,] -3.288380e-17 -0.5345225  9.637305e-17  0.7171372
[4,]  3.162278e-01 -0.2672612 -6.324555e-01 -0.4780914
[5,]  6.324555e-01  0.5345225  3.162278e-01  0.1195229

or, most relevant, from formulae,

formulae_poly = formulae.transforms.Polynomial()
formulae_poly(X, 4)

array([[-0.63245553,  0.53452248, -0.31622777,  0.11952286],
       [-0.31622777, -0.26726124,  0.63245553, -0.47809144],
       [ 0.        , -0.53452248, -0.        ,  0.71713717],
       [ 0.31622777, -0.26726124, -0.63245553, -0.47809144],
       [ 0.63245553,  0.53452248,  0.31622777,  0.11952286]])

For an example, applying this function to x over a domain from 0-10,

x = np.linspace(0, 10, 100)
x2 = x**2

transformer = OrthogonalPolynomialTransformer(degree=2)
x_orthogonalized = transformer.fit_transform(x)
x_orth = x_orthogonalized[:, 0]
x2_orth = x_orthogonalized[:, 1]

data = np.vstack([x, x2, x_orth, x2_orth]).T
df = pd.DataFrame(data, columns=['x', '$x^2$', '$x$ Orth', '$x^2$ Orth'])
correlation_matrix = df.corr()
sns.heatmap(correlation_matrix, annot=True, cmap='Reds')
plt.xticks(rotation=45)
plt.show()

We now see that the orthogonalized version of \(x\) and \(x^2\) are no longer correlated to each other. Next, we construct a response variable and plot against it.

y = 3 * x2 + x

fig, axs = plt.subplots(2, 2, figsize=(8, 6), sharey='row')

plots = [
    (x, 'x', False),
    (x2, '$x^2$', False),
    (x_orth, 'Orthogonalized $x$', True),
    (x2_orth, 'Orthogonalized $x^2$', False)
]

for ax, plot_data in zip(axs.flat, plots):
    x_val, xlabel = plot_data[:2]
    if len(plot_data) == 3 and plot_data[2]:
        sns.regplot(x=x_val, y=y, ax=ax, line_kws={"color": "C1"})
    else:
        sns.scatterplot(x=x_val, y=y, ax=ax)
    ax.set(xlabel=xlabel, ylabel='y')

    if plot_data[1] == 'Orthogonalized $x^2$':
        ax.axvline(0, color='k', linestyle='--')

plt.tight_layout()
plt.show()

The top half shows the response variable against \(x\) and \(x^2\), this should look familiar.

The bottom half shows the new orthogonalized polynomial terms. First, you’ll notice the domain is centered at 0 and more compressed than the original scale, which is done within the orthogonalization process. Otherwise, the \(x\) term is the same. Remember in the construction, the first order is untouched, then subsequent terms are built orthogonal to the first degree polynomial.

I’ve shown a linear fit on top of the first order term. What you’ll notice is that the orthogonalized \(x^2\) correspond to the residuals of this line. At the lowest values of \(y\), the fit is poor, and this is where the orthogonalized \(x^2\) is highest. As the first order term crosses the linear fit, you see the orthogonalized \(x^2\) cross zero, then go to negative values as it dips under the linear fit. It crosses 0 one more time and then is once again poor at the highest values shown. Since the \(x^2\) is proportional to the residuals of the first order term, if we plot the orthogonalized \(x^2\) term against the residuals, we should see a linear trend.

slope, intercept, r_value, p_value, std_err = scipy.stats.linregress(x_orth, y)

y_pred = intercept + slope * x_orth
residuals = y - y_pred

plt.figure(figsize=(10, 5))
plt.subplot(1, 2, 1)
plt.scatter(x_orth, y, label='Original data')
plt.plot(x_orth, y_pred, color='C1', label='Fitted line')
plt.xlabel('$x$ Orth')
plt.ylabel('y')
plt.title('$x$ Orth vs y with Linear Fit')
plt.legend()

plt.subplot(1, 2, 2)
plt.scatter(x2_orth, residuals)
plt.xlabel('$x^2$ Orth')
plt.ylabel('Residuals')
plt.title('$x^2$ Orth vs Residuals')
plt.axhline(0, color='black', linestyle='--')

plt.tight_layout()
plt.show()

And, in fact, the linear trend bears out when plotting the orthogonal \(x^2\) vs the residuals.

We can take this a degree higher and look at a cubic term.

x3 = x**3
x2 = x**2
y_cubic = 2.5 * x3 - 15 * x2 + 55 * x

transformer = OrthogonalPolynomialTransformer(degree=3)
x_orthogonalized = transformer.fit_transform(x)
x_orth = x_orthogonalized[:, 0]
x2_orth = x_orthogonalized[:, 1]
x3_orth = x_orthogonalized[:, 2]

fig, axs = plt.subplots(2, 3, figsize=(12, 8), sharey='row')

plots = [
    (x, 'x', 'x vs y', False),
    (x2, '$x^2$', '$x^2$ vs y', False),
    (x3, '$x^3$', '$x^3$ vs y', False),
    (x_orth, '$x$ Orth', '$x$ Orth vs y', True),
    (x2_orth, '$x^2$ Orth', '$x^2$ Orth vs y', True),
    (x3_orth, '$x^3$ Orth', '$x^3$ Orth vs y', False)
]

for ax, plot_data in zip(axs.flat, plots):
    x_val, xlabel, title = plot_data[:3]
    if len(plot_data) == 4 and plot_data[3]:
        sns.regplot(x=x_val, y=y_cubic, ax=ax, line_kws={"color": "C1"})
    else:
        sns.scatterplot(x=x_val, y=y_cubic, ax=ax)
    ax.set(xlabel=xlabel, ylabel='y', title=title)

    if title in ('$x^2$ Orth vs y', '$x^3$ Orth vs y'):
        ax.axvline(0, color='k', linestyle='--')

plt.tight_layout()
plt.show()

At a cubic level, it’s a bit more difficult to see the trends, however, the procedure is still the same. We can model each subsequent term against the residuals of the prior, and we can see that since this data was constructed from a cubic function, the \(x^3\) plot against the residuals of the \(x^2\) term is linear.

slope, intercept, r_value, p_value, std_err = scipy.stats.linregress(x_orth, y_cubic)
y_pred = intercept + slope * x_orth
residuals = y_cubic - y_pred

slope_res, intercept_res, r_value_res, p_value_res, std_err_res = scipy.stats.linregress(
    x2_orth, residuals
)
residuals_pred = intercept_res + slope_res * x2_orth
second_order_residuals = residuals - residuals_pred

plt.figure(figsize=(15, 5))

plt.subplot(1, 3, 1)
sns.scatterplot(x=x_orth, y=y_cubic, hue=np.arange(len(x_orth)),
                palette="viridis", legend=False)
plt.plot(x_orth, y_pred, color='black', label='Linear Model')
plt.xlabel('$x$ Orth')
plt.ylabel('y')
plt.title('$x$ Orth vs y with Linear Fit')
plt.legend()

plt.subplot(1, 3, 2)
sns.scatterplot(x=x2_orth, y=residuals, hue=np.arange(len(x2_orth)),
                palette="viridis", legend=False)
plt.plot(x2_orth, residuals_pred, color='black')
plt.xlabel('$x^2$ Orth')
plt.ylabel('Residuals')
plt.title('$x^2$ Orth vs Residuals')
plt.axhline(0, color='grey', linestyle='--', zorder=-1)

plt.subplot(1, 3, 3)
sns.scatterplot(x=x3_orth, y=second_order_residuals, hue=np.arange(len(x3_orth)),
                palette="viridis", legend=False)
plt.xlabel('$x^3$ Orth')
plt.ylabel('Second Order Residuals')
plt.title('$x^3$ Orth vs Second Order Residuals')
plt.axhline(0, color='grey', linestyle='--', zorder=-1)
plt.annotate('Point hue denotes index',
             xy=(0.99, 0.05), ha='right', xycoords='axes fraction',
             fontsize=14, color='black')

plt.tight_layout()
plt.show()

The main takeaway of this deep dive is the following: The poly keyword when used in a formula creates orthogonal polynomials. This is well-suited for fitting statistical models, since it eliminates the risk of multicollinearity between terms.

This wasn’t used in the other notebook since we were trying to recover parameters associated with each term. However, if you’re building a statistical model, especially one in which prediction is the focus, they may be the appropriate approach.

As one final note, the formulae version of poly does include a raw argument, which allows you to get the non-orthogonalized versions of each polynomial term. You can call that in Bambi like bmb.Model("y ~ poly(x, 4, raw=True)", df).

Orthogonal Polynomials in Practice

In order to see the poly keyword in action, we’ll take a look at the cars dataset. This dataset, preloaded into Seaborn, includes information on cars manufactured between 1970-1982. First we’ll load it in and take a look at the included variables.

df_mpg = sns.load_dataset("mpg")
df_mpg.head()

    mpg  cylinders  displacement  horsepower  weight  acceleration  model_year origin                       name
0  18.0          8         307.0       130.0    3504          12.0          70    usa  chevrolet chevelle malibu
1  15.0          8         350.0       165.0    3693          11.5          70    usa          buick skylark 320
2  18.0          8         318.0       150.0    3436          11.0          70    usa         plymouth satellite
3  16.0          8         304.0       150.0    3433          12.0          70    usa              amc rebel sst
4  17.0          8         302.0       140.0    3449          10.5          70    usa                ford torino

In this example, we’ll take a look at how a car’s fuel efficiency (mpg) relates to its horsepower (hp).

To start, we’ll just plot the joint distribution, as well as the distribution of the response variable.

df_mpg = df_mpg.dropna(subset=["horsepower", "mpg"])

plt.figure(figsize=(14, 6))
plt.subplot(1, 2, 1)
sns.regplot(data=df_mpg, x="horsepower", y="mpg", line_kws={"color": "firebrick"})

plt.subplot(1, 2, 2)
sns.histplot(df_mpg["mpg"], edgecolor="black", kde=True)
plt.xlabel('MPG')
plt.ylabel('Count')
plt.title('Histogram of MPG')

plt.tight_layout()
plt.show()

Immediately, we see that the linear fit doesn’t seem to model this data perfectly, it exhibits some nonlinearity. We’ll use a polynomial regression in order to see if we can improve that fit and capture the curvature. We will first fit a linear model as a benchmark.

mpg_hp_linear_mod = bmb.Model("mpg ~ horsepower", df_mpg)
mpg_hp_linear_fit = mpg_hp_linear_mod.fit(
    idata_kwargs={"log_likelihood": True}, random_seed=SEED
)
mpg_hp_linear_mod.predict(mpg_hp_linear_fit, kind="response")

fig = plt.figure()
for p in [.68, .95]:
    bmb.interpret.plot_predictions(
        mpg_hp_linear_mod,
        mpg_hp_linear_fit,
        "horsepower",
        pps=True,
        legend=True,
        prob=p,
        ax=plt.gca()
    )
sns.scatterplot(data=df_mpg, x="horsepower", y="mpg", color='blue', label='True Data');

Looking at this plot with the 68% and 95% CIs shown, the fit looks okay. Most notably, at about 160 hp, the data diverge from the fit pretty drastically. The fit at low hp values isn’t particularly good either, there’s quite a bit that falls outside of our 95% CI. This can be accented pretty heavily by looking at the residuals from the mean of the model.

predicted_mpg = mpg_hp_linear_fit.posterior["mu"].mean(("chain", "draw"))
residuals = df_mpg["mpg"] - predicted_mpg
sns.scatterplot(data=df_mpg, x="horsepower", y=residuals)
plt.axhline(0, color='black', lw=2)
plt.ylabel("Residuals")
plt.title('Residuals for linear model')

This is definitely not flat like we would ideally like it.

Next we fit a polynomial regression, including a square term.

mpg_hp_sq_mod = bmb.Model("mpg ~ poly(horsepower, 2)", df_mpg)
mpg_hp_sq_fit = mpg_hp_sq_mod.fit(
    idata_kwargs={"log_likelihood": True}, random_seed=SEED
)
mpg_hp_sq_mod.predict(mpg_hp_sq_fit, kind="response")

fig = plt.figure()
for p in [.68, .95]:
    bmb.interpret.plot_predictions(
        mpg_hp_sq_mod,
        mpg_hp_sq_fit,
        "horsepower",
        pps=True,
        legend=True,
        prob=p,
        ax=plt.gca()
    )
sns.scatterplot(data=df_mpg, x="horsepower", y="mpg", color='blue', label='True Data')
plt.title("Quadratic Fit")

Visually, this seems to look better. Particularly at high values, the model follows the pattern in the data much, much better, since we allow for curvature by including the polynomial term. Generating the same residual plot gives the following,

predicted_mpg = mpg_hp_sq_fit.posterior["mu"].mean(("chain", "draw"))
residuals = df_mpg["mpg"] - predicted_mpg
sns.scatterplot(data=df_mpg, x="horsepower", y=residuals)
plt.axhline(0, color='black', lw=2)
plt.ylabel("Residuals")
plt.title('Residuals for quadratic model')

This is far closer to flat than before.

For a true comparison, we can look at the elpd difference between the models.

az.compare({"Linear": mpg_hp_linear_fit, "Quadratic": mpg_hp_sq_fit})

           rank     elpd_loo     p_loo  elpd_diff    weight         se       dse  warning scale
Quadratic     0 -1137.318163  4.249899    0.00000  0.915146  18.118029   0.00000    False   log
Linear        1 -1181.836583  3.362876   44.51842  0.084854  15.124216  10.37422    False   log

The quadratic model performs better by LOO-CV.

Cautionary Tales

Last, we’re going to investigate a couple of pitfalls with polynomial regression.

Fitting too many polynomial degrees

Typically, when fitting a statistical model, you want to come to your data with a hypothesis and motivate your polynomial degree based on domain knowledge and expertise with the data. Instead of being principled, we’re going to throw caution to the wind and iteratively fit models from degree 1-10 and then see which performs best.

poly_fits, poly_models = {}, {}
for degree in range(1, 10):
    model = bmb.Model(f"mpg ~ poly(horsepower, {degree})", df_mpg)
    fit = model.fit(
        idata_kwargs={"log_likelihood": True}, random_seed=SEED, progressbar=False
    )
    poly_models[f"Poly{degree}"] = model
    poly_fits[f"Poly{degree}"] = fit

cmp = az.compare(poly_fits)
cmp

       rank     elpd_loo      p_loo  elpd_diff    weight         se       dse  warning scale
Poly7     0 -1133.197215   9.458617   0.000000  0.649372  18.880097   0.000000    False   log
Poly6     1 -1134.193525   8.909358   0.996310  0.000000  18.614197   1.827800     True   log
Poly8     2 -1134.296370  10.653330   1.099154  0.000000  18.876017   0.609013    False   log
Poly5     3 -1134.866504   7.696208   1.669289  0.000000  18.554120   3.508807    False   log
Poly9     4 -1135.238197  12.004611   2.040982  0.000000  18.951758   1.579663    False   log
Poly2     5 -1137.318129   4.249865   4.120914  0.000000  18.118010   6.509055    False   log
Poly3     6 -1137.990983   5.376918   4.793768  0.285760  18.402214   7.003322    False   log
Poly4     7 -1138.858924   7.061455   5.661709  0.000000  18.308157   6.200712    False   log
Poly1     8 -1181.870350   3.400270  48.673135  0.064868  15.124141  11.007970    False   log

A 7th-degree polynomial seems to do better than the quadratic one we fit before. But we must also notice that most ELPD values are very similar. Let’s do a plot, so we can more easily grasp how different models are according to the ELPD. We are going to use az.plot_compare, and we are going to add a blue band to indicate models that have an ELPD difference of less than 4 with respect to the first-ranked model. Essentially models that are that close cannot be distinguished when using ELPD.

ax = az.plot_compare(cmp, figsize=(12, 4), plot_ic_diff=False, legend=False);
best_loo = cmp["elpd_loo"].iloc[0]
ax.axvspan(best_loo-4, best_loo, color="C0", alpha=0.2);

We can see that Poly6, Poly8, Poly5 and Poly9 are all very similar (within a difference of 4 units). Even more, all models except Poly1 have overlapping standard errors.

Overall, this is telling us that there is no clear gain in predictive performance once we move beyond a quadratic model. If we want to pick a single model, then we need another criteria to decide. If we have no reason to prefer a more complex model, choosing the simpler one (Poly2 in this example) is a good heuristic.

Before deciding let’s do a couple more plots. First, see what those residuals look like!

best_model = poly_models["Poly7"]
best_fit = poly_fits["Poly7"]
best_model.predict(best_fit, kind="response")

predicted_mpg = best_fit.posterior["mu"].mean(("chain", "draw"))
residuals = df_mpg["mpg"] - predicted_mpg
sns.scatterplot(data=df_mpg, x="horsepower", y=residuals)
plt.axhline(0, color='black', lw=2)
plt.ylabel("Residuals")
plt.title('Residuals for degree 7 model');

Hey, that looks pretty good, the residuals appear nice and flat. Before we go full steam ahead with this model, let’s take a look at the posterior predictive distribution.

fig = plt.figure()
for p in [.68, .95]:
    bmb.interpret.plot_predictions(
        best_model,
        best_fit,
        "horsepower",
        pps=True,
        legend=True,
        prob=p,
        ax=plt.gca()
    )
sns.scatterplot(data=df_mpg, x="horsepower", y="mpg", color='blue', label='True Data')
plt.title("Best Fit Model: 7th Degree Polynomial");

Uh-oh. You can see that while this gave the best elpd, and had a nice residual plot, it’s obviously overfit, as expected given that we already showed that the difference with the quadratic model is small. Given our knowledge about how cars operate, we expect a decreasing trend of fuel efficiency at higher horsepower. The 7th degree polynomial absolutely is not consistent with that. First, looking at the low values, it increases before starting the decreasing trend. Second, it starts to go back up at the high end of the data, strongly latching onto a couple of points that are likely driven by noise.

This behavior evokes the classic quote,

“With four parameters I can fit an elephant, and with five I can make him wiggle his trunk.” - John von Neumann

The takeaway here is that as you fit higher polynomial degrees, you increase the risk of overfitting.

Extrapolation of polynomial models

With any model, we should be careful when extrapolating and ensure our assumptions hold, but this particularly applies when considering polynomial regression. Since we consider the higher order polynomials, terms can quickly blow up outside of our domain.

For example, with the quadratic fit, we see that the drop in mpg flattens out at higher horsepower. However, if you look closely at the posterior predictive of the quadratic model, you can start to see the fit rise again at the end. But, if we extend this beyond the bounds, due to the curvature of a second degree polynomial, we see a reversal of the negative effect on horsepower, where our quadratic model implies a higher horsepower leads to better mpg.

extrapolate_x_hp = np.linspace(0, 500, 250)
mpg_hp_sq_mod.predict(mpg_hp_sq_fit, data=pd.DataFrame({"horsepower": extrapolate_x_hp}))

sns.scatterplot(data=df_mpg, x="horsepower", y="mpg", color='blue', label='True Data')
plt.plot(
    extrapolate_x_hp,
    mpg_hp_sq_fit.posterior["mu"].mean(("chain", "draw")),
    color="red",
    label="Extrapolated Fit",
)
plt.xlim(left=0, right=extrapolate_x_hp.max())
plt.legend(frameon=False)

This is strictly untrue based on what we know about cars and what we’ve seen in the data, so you would not want to use the model outside of the intended domain. If that is the goal, you would want to find a more appropriate specification. Something like an exponential or inverse fit may be appropriate, in order to make sure the fit approaches 0, while still forbidding predictions below 0.

Extrapolation issues are not unique to polynomial regression, for example we run into forbidden values with linear regression when extrapolating too.

mpg_hp_linear_mod.predict(
    mpg_hp_linear_fit, data=pd.DataFrame({"horsepower": extrapolate_x_hp})
)
sns.scatterplot(data=df_mpg, x="horsepower", y="mpg", color='blue', label='True Data')

plt.plot(
    extrapolate_x_hp,
    mpg_hp_linear_fit.posterior["mu"].mean(("chain", "draw")),
    color="red",
    label="Predicted"
)
plt.fill_between(
    extrapolate_x_hp, plt.ylim()[0], 0,
    color='grey', alpha=0.5, label="MPG Forbidden region"
)
plt.xlim(left=0, right=extrapolate_x_hp.max())
plt.ylim(bottom=mpg_hp_linear_fit.posterior["mu"].mean(("chain", "draw")).min())
plt.legend(frameon=False);

However, it is highlighted in this notebook because, due to the nature of polynomial regression, it can be very sensitive outside the fitting domain. Just for fun to wrap this notebook up, we will take a look at what the 7th order “best model” does outside of where we fit the model.

extrapolate_x_hp = np.linspace(0, 300, 250)
best_model.predict(best_fit, data=pd.DataFrame({"horsepower": extrapolate_x_hp}))

sns.scatterplot(data=df_mpg, x="horsepower", y="mpg", color='blue', label='True Data')
plt.plot(
    extrapolate_x_hp,
    best_fit.posterior["mu"].mean(("chain", "draw")),
    color="red",
    label="Extrapolated Fit",
)
plt.fill_between(
    extrapolate_x_hp, plt.ylim()[0], 0,
    color='grey', alpha=0.5, label="MPG Forbidden region"
)
plt.xlim(left=0, right=extrapolate_x_hp.max())
plt.ylim(bottom=best_fit.posterior["mu"].mean(("chain", "draw")).min())
plt.legend(frameon=False);

Yikes.

Back in 2024, I wrote a couple of example notebooks that got merged into the Bambi documentation. For those unfamiliar, Bambi is a library for fitting Bayesian regression models using a formulaic interface on top of PyMC (the closest thing in python to brms, in my opinion). I realized I never migrated the content here, so I thought it was time to do so.

This post covers polynomial regression. The original notebook lives in the Bambi docs.

What follows is the content from the notebook, lightly adapted for this blog format.

Polynomial Regression

Unlike many other examples shown in Bambi, there aren’t specific polynomial methods or families implemented – most of the interesting behavior for polynomial regression occurs within the formula definition. Regardless, there are some nuances that are useful to be aware of.

This example uses the kinematic equations from classical mechanics as a backdrop. Specifically, an object in motion experiencing constant acceleration can be described by the following:

where \(x_0\) and \(x_f\) are the initial and final locations, \(v_0\) is the initial velocity, and \(a\) is acceleration.

A falling ball

First, we’ll consider a simple falling ball, released from 50 meters. In this situation, \(v_0 = 0\) \(m\)/\(s\), \(x_0 = 50\) \(m\) and \(a = g\), the acceleration due to gravity, \(-9.81\) \(m\)/\(s^2\). So dropping out the \(v_0 t\) component, the equation takes the form:

We’ll start by simulating data for the first 2 seconds of motion. We will also assume some measurement error with a gaussian distribution of \(\sigma = 0.3\).

import warnings

import arviz as az
import bambi as bmb
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

SEED = 1234
az.style.use("arviz-darkgrid")
warnings.filterwarnings("ignore")

g = -9.81  # acceleration due to gravity (m/s^2)
t = np.linspace(0, 2, 100)  # time in seconds
inital_height = 50
x_falling = 0.5 * g * t**2 + inital_height

rng = np.random.default_rng(SEED)
noise = rng.normal(0, 0.3, x_falling.shape)
x_obs_falling = x_falling + noise
df_falling = pd.DataFrame({"t": t, "x": x_obs_falling})

fig, ax = plt.subplots(figsize=(10, 6))
ax.scatter(t, x_obs_falling, label="Observed Displacement", color="C0")
ax.plot(t, x_falling, label="True Function", color="C1")
ax.set(xlabel="Time (s)", ylabel="Displacement (m)")
ax.legend();

Casting the equation \(x_f = \frac{1}{2} g t^2 + x_0\) into a regression context, fitting:

We let time (\(t\)) be the independent variable, and final location (\(x_f\)) be the response/dependent variable. This allows our coefficients to be proportional to \(g\) and \(x_0\). The intercept, \(\beta_0\) corresponds exactly to \(x_0\), the initial height. Letting \(\beta_1 = \frac{1}{2} g\) gives \(g = 2\beta_1\) when \(x_1 = t^2\), meaning we’re doing polynomial regression. We can put this into Bambi via the following, optionally including the + 1 to emphasize that we choose to include the coefficient.

model_falling = bmb.Model("x ~ I(t**2) + 1", df_falling)
results_falling = model_falling.fit(idata_kwargs={"log_likelihood": True}, random_seed=SEED)

The term I(t**2) indicates to evaluate inside the I. For including just the \(t^2\) term, you can express it in any of the following ways:

• I(t**2)
• {t**2}
• Square the data directly, and pass it as a new column

To verify, we’ll fit the other two versions as well.

model_falling_variation1 = bmb.Model(
    "x ~ {t**2} + 1",  # Using {t**2} syntax
    df_falling
)
results_variation1 = model_falling_variation1.fit(random_seed=SEED)

model_falling_variation2 = bmb.Model(
    "x ~ tsquared + 1",  # Using data with the t variable squared
    df_falling.assign(tsquared=t**2)
)
results_variation2 = model_falling_variation2.fit(random_seed=SEED)

print("I{t**2} coefficient: ", round(results_falling.posterior["I(t ** 2)"].values.mean(), 4))
print("{t**2} coefficient: ", round(results_variation1.posterior["I(t ** 2)"].values.mean(), 4))
print("tsquared coefficient: ", round(results_variation2.posterior["tsquared"].values.mean(), 4))

I{t**2} coefficient:  -4.8476
{t**2} coefficient:  -4.8476
tsquared coefficient:  -4.8476

Each of these provides identical results, giving -4.9, which is \(g/2\). This makes the acceleration exactly the \(-9.81\) \(m\)/\(s^2\) acceleration that generated the data. Looking at our model summary,

az.summary(results_falling)

             mean     sd  hdi_3%  hdi_97%  mcse_mean  mcse_sd  ess_bulk  ess_tail  r_hat
sigma       0.336  0.025   0.289    0.381      0.000    0.000    5977.0    2861.0    1.0
Intercept  49.961  0.051  49.870   50.058      0.001    0.001    5997.0    3145.0    1.0
I(t ** 2)  -4.848  0.028  -4.899   -4.799      0.000    0.000    5704.0    2844.0    1.0

We see that both \(g/2 = -4.9\) (so \(g=-9.81\)) and the original height of \(x_0 = 50\) \(m\) are recovered, along with the injected noise.

We can then use the model to answer some questions, for example, when would the ball land? This would correspond to \(x_f = 0\).

calculated_x0 = results_falling.posterior["Intercept"].values.mean()
calculated_g = -2 * results_falling.posterior["I(t ** 2)"].values.mean()
calculated_land = np.sqrt(2 * calculated_x0 / calculated_g)
print(f"The ball will land at {round(calculated_land, 2)} seconds")

The ball will land at 3.21 seconds

Or if we want to account for our measurement error and use the full posterior,

calculated_x0_posterior = results_falling.posterior["Intercept"].values
calculated_g_posterior = -2 * results_falling.posterior["I(t ** 2)"].values
calculated_land_posterior = np.sqrt(2 * calculated_x0_posterior / calculated_g_posterior)
lower_est = round(np.quantile(calculated_land_posterior, 0.025), 2)
upper_est = round(np.quantile(calculated_land_posterior, 0.975), 2)
print(f"The ball landing will be measured between {lower_est} and {upper_est} seconds")

The ball landing will be measured between 3.2 and 3.23 seconds

Projectile Motion

Next, instead of a ball strictly falling, instead imagine one thrown straight upward. In this case, we add the initial velocity back into the equation.

We will envision the ball tossed upward, starting at 1.5 meters above ground level. It will be tossed at 7 m/s upward. It will also stop when hitting the ground.

v0 = 7
x0 = 1.5
x_projectile = (1/2) * g * t**2 + v0 * t + x0
noise = rng.normal(0, 0.2, x_projectile.shape)
x_obs_projectile = x_projectile + noise
df_projectile = pd.DataFrame({"t": t, "tsq": t**2, "x": x_obs_projectile, "x_true": x_projectile})
df_projectile = df_projectile[df_projectile["x"] >= 0]

fig, ax = plt.subplots(figsize=(10, 6))
ax.scatter(df_projectile.t, df_projectile.x, label="Observed Displacement", color="C0")
ax.plot(df_projectile.t, df_projectile.x_true, label='True Function', color="C1")
ax.set(xlabel="Time (s)", ylabel="Displacement (m)", ylim=(0, None))
ax.legend();

Modeling this using Bambi, we must include the linear term on time to capture the initial velocity. We’ll do the following regression,

which then maps the solved coefficients to the following: \(\beta_0 = x_0\), \(\beta_1 = v_0\), and \(\beta_2 = \frac{g}{2}\).

model_projectile_all_terms = bmb.Model("x ~ I(t**2) + t + 1", df_projectile)
fit_projectile_all_terms = model_projectile_all_terms.fit(
    idata_kwargs={"log_likelihood": True}, target_accept=0.9, random_seed=SEED
)

az.summary(fit_projectile_all_terms)

            mean     sd  hdi_3%  hdi_97%  mcse_mean  mcse_sd  ess_bulk  ess_tail  r_hat
sigma      0.202  0.017   0.171    0.234      0.000    0.000    2723.0    2328.0    1.0
Intercept  1.561  0.066   1.441    1.687      0.001    0.001    2058.0    2550.0    1.0
I(t ** 2) -4.867  0.114  -5.079   -4.649      0.003    0.002    1667.0    1966.0    1.0
t          6.909  0.189   6.553    7.262      0.005    0.003    1694.0    2039.0    1.0

hdi = az.hdi(fit_projectile_all_terms.posterior, hdi_prob=0.95)
print(f"Initial height: {hdi['Intercept'].sel(hdi='lower'):.2f} to "
      f"{hdi['Intercept'].sel(hdi='higher'):.2f} meters (True: {x0} m)")
print(f"Initial velocity: {hdi['t'].sel(hdi='lower'):.2f} to "
      f"{hdi['t'].sel(hdi='higher'):.2f} meters per second (True: {v0} m/s)")
print(f"Acceleration: {2*hdi['I(t ** 2)'].sel(hdi='lower'):.2f} to "
      f"{2*hdi['I(t ** 2)'].sel(hdi='higher'):.2f} meters per second squared (True: {g} m/s^2)")

Initial height: 1.43 to 1.69 meters (True: 1.5 m)
Initial velocity: 6.54 to 7.28 meters per second (True: 7 m/s)
Acceleration: -10.16 to -9.27 meters per second squared (True: -9.81 m/s^2)

We once again are able to recover all our input parameters.

In addition to directly calculating all terms, to include all polynomial terms up to a given degree you can use the poly keyword. We don’t do that in this notebook for two reasons. First, by default it orthogonalizes the terms making it ill-suited to this example since the coefficients have physical meaning and we want to directly interpret them. The orthogonalization process can be disabled by the raw argument of poly, but we still elect not to use poly here because in later examples we decide to use different effects on the \(t\) term vs the \(t^2\) term, and doing so is not easy when using poly. However, just to show that the results match when using the raw = True argument, we’ll fit the same model as above.

model_poly_raw = bmb.Model("x ~ poly(t, 2, raw=True)", df_projectile)
fit_poly_raw = model_poly_raw.fit(idata_kwargs={"log_likelihood": True}, random_seed=SEED)
az.summary(fit_poly_raw)

                          mean     sd  hdi_3%  hdi_97%  mcse_mean  mcse_sd  ess_bulk  ess_tail  r_hat
sigma                    0.201  0.017   0.172    0.234      0.000    0.000    3066.0    2205.0    1.0
Intercept                1.561  0.067   1.437    1.682      0.001    0.001    2535.0    2154.0    1.0
poly(t, 2, raw=True)[0]  6.911  0.196   6.556    7.280      0.004    0.004    2092.0    2075.0    1.0
poly(t, 2, raw=True)[1] -4.870  0.118  -5.095   -4.653      0.003    0.002    2059.0    2166.0    1.0

We see the same results, where poly(t, 2, raw=True)[0] corresponds to the coefficient on \(t\) (\(v_0\) in our example), and poly(t, 2, raw=True)[1] is the coefficient on \(t^2\) (\(\frac{g}{2}\)).

Measuring gravity on a new planet

In the next example, you’ve been recruited to join the space program as a research scientist, looking to directly measure the gravity on a new planet, PlanetX. You don’t know anything about this planet or its safety, so you have time for one, and only one, throw of a ball. However, you’ve perfected your throwing mechanics, and can achieve the same initial velocity wherever you are. To baseline, you make a toss on planet Earth, warm up your spacecraft and stop at Mars to make a toss, then travel far away, and make a toss on PlanetX.

First we simulate data for this experiment.

def simulate_throw(v0, g, noise_std, time_step=0.25, max_time=10, seed=1234):
    rng = np.random.default_rng(seed)
    times = np.arange(0, max_time, time_step)
    heights = v0 * times - 0.5 * g * times**2
    heights_with_noise = heights + rng.normal(0, noise_std, len(times))
    valid_indices = heights_with_noise >= 0
    return times[valid_indices], heights_with_noise[valid_indices], heights[valid_indices]

# Define the parameters
v0 = 20  # Initial velocity (m/s)
g_planets = {"Earth": 9.81, "Mars": 3.72, "PlanetX": 6.0}
noise_std = 1.5

# Generate data
records = []
for planet, g in g_planets.items():
    times, heights, heights_true = simulate_throw(v0, g, noise_std)
    for time, height, height_true in zip(times, heights, heights_true):
        records.append([planet, time, height, height_true])

df = pd.DataFrame(records, columns=["Planet", "Time", "Height", "Height_true"])
df["Planet"] = df["Planet"].astype("category")

And drawing those trajectories,

fig, ax = plt.subplots(figsize=(10, 6))

for i, planet in enumerate(df["Planet"].cat.categories):
    subset = df[df["Planet"] == planet]
    ax.plot(subset["Time"], subset["Height_true"], alpha=0.7, color=f"C{i}")
    ax.scatter(subset["Time"], subset["Height"], alpha=0.7, label=planet, color=f"C{i}")

ax.set(
    xlabel="Time (seconds)", ylabel="Height (meters)",
    title="Trajectory Comparison", ylim=(0, None)
)
ax.legend(title="Planet");

We now aim to model this data. We again use the following equation (calling displacement \(h\) for height):

where \(g_p\) now has a subscript to indicate the planet that we’re throwing from.

In Bambi, we’ll do the following:

Height ~ I(Time**2):Planet + Time + 0

which corresponds one-to-one with the above formula. The intercept is eliminated since we start from \(x=0\).

planet_model = bmb.Model("Height ~ I(Time**2):Planet + Time + 0", df)
planet_model.build()
planet_fit = planet_model.fit(chains=4, idata_kwargs={"log_likelihood": True}, random_seed=SEED)

The model has fit. Let’s look at how we did recovering the simulated parameters.

az.summary(planet_fit)

                                mean     sd  hdi_3%  hdi_97%  mcse_mean  mcse_sd  ess_bulk  ess_tail  r_hat
sigma                          1.759  0.147   1.498    2.044      0.003    0.003    2054.0    1938.0    1.0
I(Time ** 2):Planet[Earth]    -4.998  0.075  -5.145   -4.865      0.002    0.001    1833.0    2431.0    1.0
I(Time ** 2):Planet[Mars]     -1.884  0.022  -1.925   -1.844      0.001    0.000    1428.0    1763.0    1.0
I(Time ** 2):Planet[PlanetX]  -3.017  0.036  -3.087   -2.953      0.001    0.001    1519.0    1729.0    1.0
Time                          20.128  0.166  19.827   20.449      0.004    0.003    1393.0    1714.0    1.0

Getting the gravities back to the physical value,

hdi = az.hdi(planet_fit.posterior, hdi_prob=0.95)
print(f"g for Earth: {2*hdi['I(Time ** 2):Planet'].sel({'I(Time ** 2):Planet_dim':'Earth', 'hdi':'lower'}):.2f} "
      f"to {2*hdi['I(Time ** 2):Planet'].sel({'I(Time ** 2):Planet_dim':'Earth', 'hdi':'higher'}):.2f} "
      f"meters (True: -9.81 m)")
print(f"g for Mars: {2*hdi['I(Time ** 2):Planet'].sel({'I(Time ** 2):Planet_dim':'Mars', 'hdi':'lower'}):.2f} "
      f"to {2*hdi['I(Time ** 2):Planet'].sel({'I(Time ** 2):Planet_dim':'Mars', 'hdi':'higher'}):.2f} "
      f"meters (True: -3.72 m)")
print(f"g for PlanetX: {2*hdi['I(Time ** 2):Planet'].sel({'I(Time ** 2):Planet_dim':'PlanetX', 'hdi':'lower'}):.2f} "
      f"to {2*hdi['I(Time ** 2):Planet'].sel({'I(Time ** 2):Planet_dim':'PlanetX', 'hdi':'higher'}):.2f} "
      f"meters (True: -6.0 m)")
print(f"Initial velocity: {hdi['Time'].sel(hdi='lower'):.2f} to {hdi['Time'].sel(hdi='higher'):.2f} "
      f"meters per second (True: 20 m/s)")

g for Earth: -10.29 to -9.71 meters (True: -9.81 m)
g for Mars: -3.85 to -3.68 meters (True: -3.72 m)
g for PlanetX: -6.18 to -5.90 meters (True: -6.0 m)
Initial velocity: 19.80 to 20.45 meters per second (True: 20 m/s)

We can see that we’re pretty close to recovering most the parameters, but the fit isn’t great. Plotting the posteriors for \(g\) against the true values,

earth_posterior = -2 * planet_fit.posterior["I(Time ** 2):Planet"].sel(
    {"I(Time ** 2):Planet_dim": "Earth"})
planetx_posterior = -2 * planet_fit.posterior["I(Time ** 2):Planet"].sel(
    {"I(Time ** 2):Planet_dim": "PlanetX"})
mars_posterior = -2 * planet_fit.posterior["I(Time ** 2):Planet"].sel(
    {"I(Time ** 2):Planet_dim": "Mars"})

fig, axs = plt.subplots(1, 3, figsize=(12, 6))
az.plot_posterior(earth_posterior, ref_val=9.81, ax=axs[0])
axs[0].set_title("Posterior $g$ on Earth")
az.plot_posterior(mars_posterior, ref_val=3.72, ax=axs[1])
axs[1].set_title("Posterior $g$ on Mars")
az.plot_posterior(planetx_posterior, ref_val=6.0, ax=axs[2])
axs[2].set_title("Posterior $g$ on PlanetX");

The fit seems to work, more or less, but certainly could be improved.

Adding a prior

But, we can do better! We have a very good idea of the acceleration due to gravity on Earth and Mars, so why not use that information? From an experimental standpoint, we can consider these throws from a calibration mindset, allowing us to get some information on the resolution of our detector, and our throwing apparatus. With informative priors constraining the Earth and Mars gravity parameters, the model can more precisely estimate the unknown PlanetX gravity, as there will be less uncertainty propagating from the calibration planets.

For Earth, at the extremes, \(g\) takes values as low as 9.78 \(m\)/\(s^2\) (at the Equator) up to 9.83 (at the Poles). So we can add a very strong prior,

For Mars, we know the mean value is about 3.72 \(m\)/\(s^2\). There’s less information on local variation readily available by a cursory search, however we know that the radius of Mars is about half that of Earth, so \(\sigma = \frac{0.025}{2} = 0.0125\) might make sense, but to be conservative we’ll round that up to \(\sigma = 0.02\).

For PlanetX, we must use a very loose prior. We might say that we know the ball took longer to fall than Earth, but not as long as on Mars, so we can split the difference. Then set a very wide \(\sigma\) value.

Since these correspond to \(g/2\), we’ll divide all values by 2 when putting them into Bambi. Additionally, we know the balls landed eventually, so \(g\) must be negative. We’ll truncate the upper limit of the distribution at 0.

Now, for defining this in Bambi, the term of interest is I(Time ** 2):Planet. Often, you set one prior that applies to all groups, however, if you want to set each group individually, you can pass a list to the bmb.Prior definition. The broadcasting rules from PyMC apply here, so it could equivalently take a numpy array. You’ll notice that the priors are passed alphabetically by group name.

priors = {
    "I(Time ** 2):Planet": bmb.Prior(
        "TruncatedNormal",
        mu=[
            -9.81/2,  # Earth
            -3.72/2,  # Mars
            -6.77/2   # PlanetX
        ],
        sigma=[
            0.025/2,  # Earth
            0.02/2,   # Mars
            3/2       # PlanetX
        ],
        upper=[0, 0, 0]
    )}

planet_model_with_prior = bmb.Model(
    'Height ~ I(Time**2):Planet + Time + 0',
    df,
    priors=priors
)

planet_model_with_prior.build()
idata = planet_model_with_prior.prior_predictive()
az.summary(idata.prior, kind="stats")

                                 mean       sd   hdi_3%  hdi_97%
sigma                          14.466   13.809    0.025   36.595
I(Time ** 2):Planet[Earth]     -4.905    0.012   -4.928   -4.883
I(Time ** 2):Planet[Mars]      -1.860    0.010   -1.880   -1.841
I(Time ** 2):Planet[PlanetX]   -3.622    1.509   -6.360   -0.915
Time                            0.520   14.788  -26.565   27.992

Here we’ve sampled the prior predictive and can see that our priors are correctly specified to the associated planets.

Next we fit the model.

planet_fit_with_prior = planet_model_with_prior.fit(
    chains=4, idata_kwargs={"log_likelihood": True}, random_seed=SEED
)
planet_model_with_prior.predict(planet_fit_with_prior, kind="pps");

az.summary(planet_fit_with_prior)[0:5]

                                mean     sd  hdi_3%  hdi_97%  mcse_mean  mcse_sd  ess_bulk  ess_tail  r_hat
sigma                          1.759  0.142   1.495    2.024      0.002    0.002    3333.0    2373.0    1.0
I(Time ** 2):Planet[Earth]    -4.907  0.012  -4.929   -4.884      0.000    0.000    4360.0    2943.0    1.0
I(Time ** 2):Planet[Mars]     -1.862  0.009  -1.879   -1.847      0.000    0.000    2054.0    2614.0    1.0
I(Time ** 2):Planet[PlanetX]  -2.985  0.023  -3.025   -2.940      0.000    0.000    2282.0    2772.0    1.0
Time                          19.960  0.075  19.827   20.103      0.002    0.001    2025.0    2249.0    1.0

We see some improvements here! Off the cuff, these look better, you’ll notice the \(v_0\) coefficient on Time covers the true value of 20 m/s.

Now taking a look at the effects before and after adding the prior on the gravities,

earth_posterior_2 = -2 * planet_fit_with_prior.posterior["I(Time ** 2):Planet"].sel(
    {"I(Time ** 2):Planet_dim": "Earth"})
mars_posterior_2 = -2 * planet_fit_with_prior.posterior["I(Time ** 2):Planet"].sel(
    {"I(Time ** 2):Planet_dim": "Mars"})
planetx_posterior_2 = -2 * planet_fit_with_prior.posterior["I(Time ** 2):Planet"].sel(
    {"I(Time ** 2):Planet_dim": "PlanetX"})

fig, axs = plt.subplots(2, 3, figsize=(12, 6), sharex='col')
az.plot_posterior(earth_posterior, ref_val=9.81, ax=axs[0,0])
axs[0,0].set_title("Earth $g$ - No Prior")
az.plot_posterior(mars_posterior, ref_val=3.72, ax=axs[0,1])
axs[0,1].set_title("Mars $g$ - No Prior")
az.plot_posterior(planetx_posterior, ref_val=6.0, ax=axs[0,2])
axs[0,2].set_title("PlanetX $g$ - No Prior")

az.plot_posterior(earth_posterior_2, ref_val=9.81, ax=axs[1,0])
axs[1,0].set_title("Earth $g$ - Priors Used")
az.plot_posterior(mars_posterior_2, ref_val=3.72, ax=axs[1,1])
axs[1,1].set_title("Mars $g$ - Priors Used")
az.plot_posterior(planetx_posterior_2, ref_val=6.0, ax=axs[1,2])
axs[1,2].set_title("PlanetX $g$ - Priors Used");

Adding the prior gives smaller uncertainties for Earth and Mars by design, however, we can see the estimate for PlanetX has also improved by injecting our knowledge into the model.

A little over a decade ago, toward the start of my Ph.D., I started programming for real work for the first time. I was doing data analysis in C++ using ROOT (and yes, data analysis in C++ is as awful as it sounds). At the time, My advisor was the first person to introduce me to a terminal and to emacs. I’m pretty sure when I wasn’t looking, he aliased emacs = emacs -nw just so I wouldn’t even have to use the emacs GUI.

In 2018, I switched to VSCode. The integrated terminal made it feel like a one-stop-shop - everything in one program, no context switching, plus a rich editing experience. But I’ve always missed parts of a terminal-only workflow, and lately I’ve been drifting back to it.

This has made me a sucker for terminal-based tooling. One friction point I noticed recently: there isn’t a good terminal user interface (TUI) for Linear project management, which forces me to context-switch into the native app to log progress, comment on issues, etc. So this past weekend I hacked together my own.

linear-term

linear-term is the TUI I put together (repo here). It was built using textual in python. The original design was intended to look similar to the native app, but within the terminal.

It’s a 3-panel layout: the center shows your issues, the right panel shows issue details once selected, and the left panel has filtering options. You can toggle through each via TAB or F1, F2, F3.

I also created a kanban board view, accessible via b, where you can look at issues by their status.

I also added some CLI tools. There are existing Linear CLIs out there, but I wanted this to be enough of a one-stop-shop that you didn’t have to install a bunch of other tools.

For example:

$ linear-term list --mine
TJB-1 [Backlog] --- Get familiar with Linear @Tyler Burch
TJB-4 [Done] --- Import your data @Tyler Burch
TJB-3 [In Progress] --- Connect your tools @Tyler Burch
TJB-2 [Todo] --- Set up your teams @Tyler Burch

I mainly put this together for my own use, but please feel free to use it if you’re interested. Happy to hear feedback, or take contributions too.

Several times this year, I’ve had an idea for a post that never came to fruition. I figured a good way to wrap up the year would be to collect some of those loose threads into a single lightning-round collection of thoughts.

So, without further ado, here’s what I did in 2025:

I wrote a lot of R

A little over a year ago, I switched groups within our analytics team, and now have a strong requirement to use R for building projects. As a result, 2025 was the first year where I wrote more lines of R than any other language.

A few reflections after a year of working in R full-time:

1. I used to advise people who wanted to break into baseball analytics that it didn’t matter what language they learned, just pick R or Python and learn as much as you can. If a job requires you to switch, you can learn the other on the job, but prioritize knowledge depth while learning. This year has reinforced that opinion; good design patterns are ubiquitous, and learning on the fly has not been arduous. Living in a world with LLMs helps too (more on that later).

2. That being said, I’m now of the opinion that it’s near impossible to be a proper statistician and not engage with some R. There are several non-negotiable statistics libraries that plainly don’t have reasonable non-R alternatives (cough mgcv cough). Certainly, there are ways around it if you go the Python route, but often it’s fitting a square peg in a round hole.

3. I still prefer working in Python, especially for production code. R has so many idiosyncrasies that still frustrate me. The pain point that continues to plague me more than any is silent failures and tragically uninformative error messages. (Seriously, what does object of type 'closure' is not subsettable even mean?)

I made some small contributions to tidymodels

Speaking of R, after some initial skepticism, I’ve grown pretty fond of the tidymodels framework.

Following my personal policy of helping out the libraries that help me, I made a couple of small PRs to tidymodels, probably the most interesting being adding fold weighting to the tune package.

For context, tune handles hyperparameter tuning within tidymodels. For each candidate hyperparameter set, you evaluate performance across resampling folds and take the set with the best average performance. Prior to this change, each fold contributed equally to that average, regardless of how much data it represented. That assumption is fine in a classic K-Fold setup, but if you switch to something like expanding window CV where folds have variable sizes, it seemed to miss the mark.

With this change, folds can now be weighted (naturally by training set size in the expanding window example) so later, more informative folds carry more influence when selecting hyperparameters. It’s a very small tweak, but it fixed a real issue I ran into on a project.

I’ve thought a lot about forecasting problems

Related to the above, more than before, I’ve found myself approaching problems through a forecasting lens. The largest project I worked on this year wasn’t explicitly a forecasting task, but required conditioning on time due to distributional shifts over time.

One of the clearest implications was cross-validation. I have increasingly used rolling or expanding window setups for cases when time could even plausibly matter, even if the problem isn’t explicitly framed as forecasting.

Lately, I’ve found myself defaulting to the assumption that time is relevant unless proven otherwise. Instead of asking whether I should use time-aware cross-validation, I enter with the posture that I need to be convinced it’s safe not to. Stationarity assumptions are helpful when true, but frequently break down in real-world problems, and ignoring that can lead to misleading performance estimates. Perhaps it can be a bit overcautious, but it’s one of those cases where a bit more care here has made me more comfortable with the results I’m delivering.

I used a lot of tokens

Agentic AI was impossible to ignore in 2025. I’ve long been bought in on codebase-aware LLMs. I find copy/pasting from ChatGPT both painfully slow and prone to stripping useful information. Because of this, I was a pretty early adopter of Cursor; I started using it in August of 2023. The tab complete was enough for me to buy in, and when agent-based edits came along, it became an important part of my workflow.

These days I’m using some combination of Cursor, Codex, and Claude Code for scaffolding boilerplate, generating prototypes (especially front-ends), quickly testing hypotheses, and making publication-quality plots faster than I could myself. The domain of tasks where it’s faster to prompt than to tweak, dissect, debug, etc. has grown way faster than I expected.

I don’t have any novel insights in this domain that haven’t been said elsewhere. The advice that I think about most day-to-day in my workflow is:

• Be meticulous about context. Keep context window usage under 50% if possible. The smaller the haystack, the easier the needle is to find.
• Spend time on prompting well. An extra 5 minutes on a good prompt can save an hour of debugging.
• Be judicious about correctness. One “wrong” line in the context window can yield hundreds of bad lines of code, or subtle unexpected bugs. Clear, correct prompts are key.
• Let the tool actually run code and see output. They’re better about this in other languages, but I find particularly in R they have to be pushed to do this. Add logging statements so it can iterate with itself and find bugs.

One thing I have noticed is that most of the public conversation around LLM tooling is focused more in the domain of traditional software engineering, where specs are far clearer than data analysis workflows. The projects I work on are fuzzier, and the road from question to answer (or from idea to predictions) is not typically well-defined from the outset. I haven’t seen nearly as much written about what works well in this environment and am still figuring it out day-by-day for myself.

I loosened my grip on statistical dogma

My first experience reading an honest stats text was McElreath’s Statistical Rethinking back in 2020 (10/10 recommend). Much of my early experience with doing statistics was strongly through a Bayesian lens. I used to feel pretty strongly that this was the best way to do things when possible.

In 2025, I let go of a lot of those biases. At the end of the day, I’m a practitioner; I need answers to problems. I’m not debating or writing about the philosophy of statistics, and in many of the problems I work on trying to wedge Bayesian inference into the solution can be a hindrance more than a value add.

If a frequentist framework can get me to effectively the same answer more quickly, I’ve become much more comfortable using it. In a lot of applied settings, the difference between a Bayesian posterior and a well-validated frequentist estimate is functionally negligible, while the difference in iteration speed isn’t. The Bayesian paradigm still shapes how I think about uncertainty, but I don’t reach for that machinery unless it provides something that will be used and I’m willing to wait for MCMC chains to finish sampling before I can provide an answer.

One of my repeated lines throughout this year has been “just use the right tool for the job,” which could even be a hidden thread behind this post. For me, the “right” tool is the one that answers the real question under real constraints and produces results that stakeholders can understand and act on. In practice, that’s often the approach that generates the most business value (or in my case, wins the most games) as quickly and correctly as possible.

I changed a lot of diapers

Above all else, I welcomed a second child into the world in May. At the time of writing this, I’m parenting a teething 7-month-old who gives the best smiles, as well as a curious and fiercely independent daughter who turned 3 yesterday.

I have a tendency to value myself entirely based on my work and the things I produce. Fatherhood is a constant forcing function to get out of that mindset and to enjoy life outside of sheer production. While it has exhausting moments, it has made me appreciate the day-to-day moments so much more.

Right now I’m most appreciating the sense of wonder from my toddler. I dread air travel, airports are awful, logistics are a nightmare, I could go on for hours. But hearing my daughter say “I’m so excited” while getting on a plane, and watching her stare with awe out the window during takeoff, has made me stop for a few moments and appreciate how cool life really can be. There are countless places where getting a chance to look through her lens has made me far more appreciative of the little things in life.

See you soon, hopefully before another 2 years go by, but no promises.

• How Not to Be Wrong: The Power of Mathematical Thinking - Jordan Ellenberg
• Fluent in 3 Months: How Anyone at Any Age Can Learn to Speak Any Language from Anywhere in the World - Benny Lewis
• Zak George’s Dog Training Revolution - Zak George
• Winning Fixes Everything: How Baseball’s Brightest Minds Created Sports’ Biggest Mess - Evan Drellich
• Heaven and Hell: A History of the Afterlife - Bart D. Ehrman
• My Life as a Quant: Reflections on Physics and Finance - Emanuel Derman
• The Checklist Manifesto: How to Get Things Right - Atul Gawande
• The Big Short: Inside the Doomsday Machine - Michael Lewis

Background

Earlier this NHL season, I posted a Bayesian hierarchical model for NHL scoring in an aim to understand the skill of the Bruins based on the first 21 games (in which they went 18-3). This model has been expanded to better model NHL games (specifically the overtime structure), fit on all of 2022-2023 data to get the goal creation and suppression parameters for each team, and then used to project the remainder of the playoffs, which can be found here.

Methodology

Original Model

The base of the model has remained unchanged, based on the Baio and Blangiardo model. For regulation scoring, I fit the following model for goal scoring, \(y = (y_{g0}, y_{g1})\), a Poisson process:

for game \(g\), with \(j = {0, 1}\) an indicator variable for home ice. In their paper the rate parameter \(\theta_{gj}\) is given by the following:

Where \(h\) is the home ice advantage, \(a\) is the “attack strength,” \(d\) is the “defense strength,” and \(h\) and \(v\) denote “home” or “visitor” respectively. Last, \(\alpha\) is a flat intercept. In words: the home scoring rate is proportional to the attack skill of the home team, minus the defense of the away team, plus home ice advantage. The away scoring rate is the opposite, with no advantage.

Overtime

However, some crucial updates have been made, specifically accounting for the overtime mechanisms in hockey. For games that reach overtime, the $\theta$ parameters are scaled down to the relative time frame

\(K\) represents the scaling factor for overtime expectations. For the regular season, \(K = 12\), since overtime is 5 minutes. In playoff games, \(K\) is set to 3, as each overtime period lasts one-third of the time of a regulation period. This assumption implies that the goal creation and suppression parameters remain the same during overtime, a necessary compromise given the relatively small dataset of OT games.

I also introduce a custom likelihood function, which compares the observed home and away overtime goals to the expected overtime goal rates (ot_home_theta and ot_away_theta). This function is only applied to games that went into overtime. It allows only 3 possible outcomes:

1. No goals scored by either team: (0, 0)
2. Home team scores 1 goal and away team scores 0 goals: (1, 0)
3. Home team scores 0 goals and away team scores 1 goal: (0, 1)

For each allowed outcome \((h\_goals, a\_goals)\), we calculate the log-likelihood of the observed home and away overtime goals as follows:

Here, \(P(y_{h,ot})\) and \(P(y_{a,ot})\) represent the probabilities of observing the home and away overtime goals, given their respective expected goal rates and the allowed outcome. And recall, these are Poisson distributed.

Then for custom likelihood function, the log-sum-exp of the log-likelihoods is as follows:

Here, \(\exp(\text{loglikelihood}_{h_g, a_g})\) simply represents the likelihood of observing the home and away overtime goals, given their respective expected goal rates and the allowed outcome.

Shootouts

For regular season games, if the score is still the same after the overtime consideration, a shootout model is then introduced, modeling the probability of the home team winning the using a familiar logistic regression. I introduce team-specific coefficients for shootout success and failure, denoted by \(so_o\) (success) and \(so_d\) (failure), as well as an intercept term \(so_i\) and a home advantage term \(so_{adv}\). Then, we calculate the probability of the home team winning the shootout using the logistic function:

Finally, we model the shootout_winner variable as a Bernoulli random variable with probability \(so_{P_\text{home}}\):

This shootout model is conditioned only on games that went to a shootout.

Frankly, I believe all this is far more transparent with code. So without further ado,

def overtime_goals_likelihood(observed_ot_h_goals, observed_ot_a_goals, ot_h_theta, ot_a_theta):
    allowed_outcomes = [(0, 0), (1, 0), (0, 1)]
    likelihoods = []

    for h_goals, a_goals in allowed_outcomes:
        h_likelihood = pm.logp(pm.Poisson.dist(mu=ot_h_theta), observed_ot_h_goals * h_goals)
        a_likelihood = pm.logp(pm.Poisson.dist(mu=ot_a_theta), observed_ot_a_goals * a_goals)
        likelihoods.append(h_likelihood + a_likelihood)

    return pm.math.logsumexp(pm.math.stack(likelihoods), axis=0)

    
home_idx, teams = pd.factorize(data["home_team"], sort=True)
away_idx, _ = pd.factorize(data["away_team"], sort=True)

coords = {
    "team": teams,
    "match": np.arange(len(data)),
}

with pm.Model(coords=coords) as model:
    # Global model parameters
    intercept = pm.Normal("intercept", mu=0, sigma=2)
    home = pm.Normal("home", mu=0, sigma=0.2)

    # Hyperpriors for attacks and defs
    sd_att = pm.HalfCauchy("sd_att", 0.2)
    sd_def = pm.HalfCauchy("sd_def", 0.2)

    # Team-specific model parameters
    atts_star = pm.Normal("atts_star", mu=0, sigma=sd_att, dims="team")
    defs_star = pm.Normal("defs_star", mu=0, sigma=sd_def, dims="team")

    # Demeaned team-specific parameters
    atts = pm.Deterministic("atts", atts_star - at.mean(atts_star), dims="team")
    defs = pm.Deterministic("defs", defs_star - at.mean(defs_star), dims="team")

    # Expected goals for home and away teams during regulation
    home_theta = at.exp(intercept + home + atts[home_idx] - defs[away_idx])
    away_theta = at.exp(intercept + atts[away_idx] - defs[home_idx])

    # Likelihood (Poisson distribution) for regulation goals
    home_points = pm.Poisson("home_points", mu=home_theta, observed=data['home_goals'], dims="match")
    away_points = pm.Poisson("away_points", mu=away_theta, observed=data['away_goals'], dims="match")

    # Overtime and shootout deterministics
    overtime = data['home_goals'] == data['away_goals']
    shootout = (data['home_goals_ot'] == data['away_goals_ot']) & overtime

    # Expected goals for home and away teams during overtime (scaled down by 1/12)
    ot_home_theta = home_theta * (1 / 12)
    ot_away_theta = away_theta * (1 / 12)

    # Likelihood (custom likelihood function) for overtime goals
    if overtime.sum() > 0:
        pm.Potential("ot_goals_constraint",
                    overtime_goals_likelihood(data.home_goals_ot, data.away_goals_ot, ot_home_theta, ot_away_theta))

    # Shootout model (conditioned on games that went to shootout)
    so_coeff_o = pm.Normal("so_coeff_o", mu=0, sigma=1, dims="team")  # Offensive shootout coefficient
    so_coeff_d = pm.Normal("so_coeff_d", mu=0, sigma=1, dims="team")  # Defensive shootout coefficient
    so_coeff_h = pm.Normal("so_coeff_h", mu=0, sigma=1)  # Home advantage coefficient
    so_intercept = pm.Normal("so_intercept", mu=0, sigma=1)  # Intercept term

    so_logit = (so_intercept +
                so_coeff_o[home_idx[shootout]] - so_coeff_o[away_idx[shootout]] +
                so_coeff_d[home_idx[shootout]] - so_coeff_d[away_idx[shootout]] +
                so_coeff_h * home)

    if shootout.sum() > 0:
        so_prob = pm.math.invlogit(so_logit)
        shootout_winner = pm.Bernoulli("shootout_winner", p=so_prob, observed=data['shootout_winner'][shootout])

    trace = pm.sample(4000, tune=3000)
return model, trace