Austin Rochford

Decaying Bayesian Updating for Non-stationary Time Series Models

Austin Rochford — Sat, 10 Jan 2026 05:00:00 GMT

In the past, I have both written and spoken about the significant role multi-armed bandits have played in my career. One significant takeaway from my experience so far is the non-stationarity inherent in the realy world. Human behavior, especially online, is constantly shifting and changing in ways that many naive statistical models are not built to handle. There are, of course, many approaches to modeling non-stationary phenomena. In this post I will share a lesser-known approach that I have found occasionaly useful; in fact, I have never been able to find solid references for this approach.

Bayesian Models for Sequential Data¶

I call this approach "decaying Bayesian updating," as it builds on the notion of sequential Bayesian updates as more data arrives over time, while weighting recent data more strongly than older data. Bayesian modeling is theoretically well-suited to modeling sequences of data that arrive over time, as the posterior after time $t$ makes a natural prior for the likelihood of the $(t + 1)$-th observation, whose posterior makes a natural prior for the likelihood of the $(t + 2)$-th observation, and so on.

Throughout most of this post, we will use a simple beta-binomial model for the probability of success in independent trials, though we occasionally will work with more general distributions as appropriate.

First we establish some preliminary results in the stationary case where $X_1, X_2, \ldots, X_T \sim \text{Ber}(p)$ are i.i.d. Bernoulli-distributed random variables with probability of success $p$. Let $\pi_0(p)$ be the prior distribution of $p$ before any data is collected, and $\pi_t(p) = \pi(p\ |\ x_1, x_2, \ldots, x_t)$ be the posterior distribution of $p$ after the first $t$ data points have been collected. There are two plausible methods for calculating $\pi_T(p)$, the posterior distribution of $p$ after all of the data has been collected. Throughout we assume a uniform prior distribution $\pi_0(p) = 1$ on $p$ for simplicity. The following analysis does, however, generalize to other, non-uniform beta-distributed priors.

The first method observes that $\sum_{t = 1}^T X_t \sim \text{Bernoulli}(T, p)$, so by Bayes' theorem

$$ \begin{align} \pi_T(p) & \propto f\left(\sum_{t = 1}^T x_t\ |\ p\right) \cdot \pi_0(p) \\ & \propto p^{\sum_{t = 1}^T x_t} \cdot (1 - p)^{T - \sum_{t = 1}^T x_t}, \end{align} $$

which gives

$$\pi_T(p) \sim \text{Beta}\left(1 + \sum_{t = 1}^T x_t, 1 + T - \sum_{t = 1}^T x_t\right).$$

The second method uses sequential Bayesian updating, noting that the posterior distribution after the first observation is

$$ \begin{align} \pi_1(p) & \propto f(x_1\ |\ p) \cdot \pi_0(p) \\ & = p^{x_1} \cdot (1 - p)^{1 - x_1}. \end{align} $$

Using this posterior as the prior on the second observation, we get

$$ \begin{align} \pi_2(p) & \propto f(x_2\ |\ p) \cdot \pi_1(p) \\ & \propto \left(p^{x_2} \cdot (1 - p)^{1 - x_2}\right) \cdot \left(p^{x_1} \cdot (1 - p)^{1 - x_1}\right) \\ & = p^{x_1 + x_2} \cdot (1 - p)^{2 - (x_1 + x_2)}. \end{align} $$

It is apparent that repeating this process $T - 2$ more times will yield

$$\pi_T(p) \propto p^{\sum_{t = 1}^T x_t} \cdot (1 - p)^{T - \sum_{t = 1}^T x_t},$$

$$\pi_T(p) \sim \text{Beta}\left(1 + \sum_{t = 1}^T x_t, 1 + T - \sum_{t = 1}^T x_t\right),$$

exactly as in the first method.

It is good that these two methods agree in this stationary case; waiting for all of the data to arrive and performing one Bayesian update is equivalent to updating intermediate posterior distributions after each data point arrives. If we were category theorists, we might say the following diagram commutes.

In [1]:

from IPython.display import IFrame

In [2]:

QUIVER_URL = "https://q.uiver.app/#q=WzAsNixbMCwwLCJcXHBpXzAocCkiXSxbMSwxLCJcXHBpXzEocCkiXSxbMiwxLCJcXHBpXzIocCkiXSxbMywxLCJcXGNkb3RzIl0sWzQsMSwiXFxwaV97VCAtMX0ocCkiXSxbNSwwLCJcXHBpX1QocCkiXSxbMCwxLCJmKHhfMVxcIHxcXCBwKSJdLFsxLDIsImYoeF8yXFwgfFxcIHApIl0sWzIsMywiZih4XzNcXCB8XFwgcCkiXSxbMyw0LCJmKHhfe1QgLSAxfVxcIHxcXCBwKSJdLFs0LDUsImYoeF9UXFwgfFxcIHApIl0sWzAsNSwiZlxcbGVmdChcXHN1bV97dCA9IDF9XlR4X3RcXCB8XFwgcFxccmlnaHQpIiwwLHsiY3VydmUiOi0xfV1d&embed"

In [3]:

IFrame(QUIVER_URL, width=889, height=304)

Out[3]:

Throughout this post, we will supplement theoretical analysis with simulations. First we make the necessary Python imports and do some light configuration.

In [4]:

%matplotlib inline

In [5]:

from warnings import filterwarnings

In [6]:

from matplotlib import pyplot as plt, ticker
import numpy as np
import scipy as sp
import seaborn as sns

In [7]:

filterwarnings("ignore", category=RuntimeWarning)

In [8]:

pct_formatter = ticker.StrMethodFormatter("{x:.0%}")

sns.set(color_codes=True)

In [9]:

SEED = 123456789  # for reprodicibility
rng = np.random.default_rng(SEED)

For now we assume that the probability of success is 30% ($p = 0.3$), and that we will collect $T = 5,000$ samples sequentially.

In [10]:

P = 0.3
T = 5_000

We will conduct 1,001 simulations to understand the avarage behavior in this situation.

In [11]:

N_SIM = 1_001

First we generate the $5,000 \times 1,001$ samples.

In [12]:

x = 1 * (rng.uniform(size=(N_SIM, T)) < P)

Next we calculate and visualize the behavior of the maximum likelihood estimator for $p$, $\hat{p}$, in these simulations.

In [13]:

t = 1 + np.arange(T)
p_mle = x.cumsum(axis=1) / t

In [14]:

N_PLOT_TRAJECTORY = 10

In [15]:

def scale_y_axis(ref_val, *, mult=0.1, ax=None):
    if ax is None:
        ax = plt.gca()

    ax.set_ylim((1 - mult) * ref_val, (1 + mult) * ref_val)


def make_pct_axis(*, ref_val=None, mult=0.1, ax=None):
    if ax is None:
        ax = plt.gca()

    ax.yaxis.set_major_formatter(pct_formatter)

    if ref_val is not None:
        scale_y_axis(ref_val, mult=mult, ax=ax)

    ax.set_ylabel("$p$")


def make_time_axis(*, t_max=T, ax=None):
    if ax is None:
        ax = plt.gca()

    ax.set_xlim(1, t_max)
    ax.set_xlabel("$t$")

In [16]:

fig, ax = plt.subplots()

ax.plot(t, p_mle.mean(axis=0), label=r"Average MLE")

ax.plot(t, p_mle[0].T, c="C0", alpha=0.1, label=r"Simulation MLE")
ax.plot(t, p_mle[1:N_PLOT_TRAJECTORY].T, c="C0", alpha=0.2)

ax.axhline(P, c="k", ls="--", label="Actual")

make_time_axis(ax=ax)
make_pct_axis(ref_val=P, ax=ax)

ax.legend();

We see that the average value of $\hat{p}$ across simulations is quite close to the actual value of $p = 0.3$, and the sample simulations plotted converge to this true value as $t$ increases. The former behavior is due to the law of large numbers, and the rate at which simulations' $\hat{p}$s converge is given by the central limit theorem.

Next we calculate the parameters of the posterior distributions at each point in time for each simulation assuming a beta-binomial model with a uniform prior on $p$ at $t = 0$, as above.

In [17]:

α = 1 + x.cumsum(axis=1)
β = 1 + t - x.cumsum(axis=1)

We calculate the posterior expected values of $p$ as well.

In [18]:

def get_bb_post_mean(α, β):
    return α / (α + β)

In [19]:

p_bb = get_bb_post_mean(α, β)

We see that this posterior expected value is quite close to the MLE (on average), especially as $t$ increases.

In [20]:

fig, ax = plt.subplots()

ax.plot(t, p_mle.mean(axis=0), label=r"Average MLE")
ax.plot(t, p_bb.mean(axis=0), label=r"Average posterior EV")

ax.plot(t, p_bb[0].T, c="C1", alpha=0.1, label=r"Simulation posterior EV")
ax.plot(t, p_bb[1:N_PLOT_TRAJECTORY].T, c="C1", alpha=0.2)

ax.axhline(P, c="k", ls="--", label="Actual")

make_time_axis(ax=ax)
make_pct_axis(ref_val=P, ax=ax)

ax.legend();

Decaying Bayesian Updating¶

To motivate decaying Bayesian updating, we consider what happens to the stationary posterior predictions when the probability of success changes over time.

We will use a simple model where

$$ p_t = \begin{cases} 0.3 & \text{if } t < T / 2 \\ 0.7 & \text{if } t \geq T / 2 \end{cases}. $$

First we generate $5,000 \times 1,001$ samples under these assump;tions.

In [21]:

x_ns = 1 * (rng.uniform(size=(N_SIM, T)) < np.where(t < T // 2, P, 1 - P))

Here the suffix ns indicates non-stationarity.

We visualize the MLE and posterior expected value in this situation.

In [22]:

p_mle_ns = x_ns.cumsum(axis=1) / t

In [23]:

α_ns = 1 + x_ns.cumsum(axis=1)
β_ns = 1 + t - x_ns.cumsum(axis=1)

p_bb_ns = get_bb_post_mean(α_ns, β_ns)

In [24]:

fig, ax = plt.subplots()

ax.plot(t, p_mle_ns.mean(axis=0), label="Average MLE")
ax.plot(t, p_bb_ns.mean(axis=0), label="Average posterior EV")

ax.hlines(P, 1, T // 2, color="k", ls="--", label="Actual")
ax.hlines(1 - P, T // 2, T, color="k", ls="--")

make_time_axis(ax=ax)
make_pct_axis(ax=ax)

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

We see that, though these models are not explicitly built to handle nonstationarity, the estimates do start to move towards the later value of $p$ over time, and would continue to do so if we collected more data.

In some situations, however, we might want the estimates to react more quickly to changes in the success probability over time. Decaying Bayesian updates give one method to achieve this, though there are many others. Decaying Bayesian updates can be motivated by writing the standard Bayes' update rule in the somewhat unnatural form

$$ \begin{align} \pi_t(\vartheta) & \propto f(x_t\ |\ \vartheta) \cdot \pi_{t - 1}(\vartheta) \\ & = \left(f(x_t\ |\ \vartheta) \cdot \pi_{t - 1}(\vartheta)\right)^{1 - 0} \cdot \pi_0(\vartheta)^0. \end{align} $$

$\newcommand{\eps}{\varepsilon}$

The following definition and analysis are not specific to the beta-binomial model we have been discussing so far, so we use $\vartheta$ to denote the set set of parameters of the general model.

If we replace $0$ with a small positive number $0 < \eps \ll 1$, we get the decaying Bayesian update rule

$$\pi_t^{\eps}(\vartheta) = \left(f(x_t\ |\ \vartheta) \cdot \pi_{t - 1}^{\eps}(\vartheta)\right)^{1 - \eps} \cdot \pi_0(\vartheta)^{\eps}.$$

With this definition, we see that the decayed posterior at time $t$, $\pi_t^{\eps}(\vartheta),$ is proportional to the weighted geometric mean of the (almost) usual joint distribution $f(x_t\ |\ \vartheta) \cdot \pi_{t - 1}^{\eps}(\vartheta)$ with weight $1 - \eps$ and the no data prior $\pi_0(\vartheta)$ with weight $\eps$.

One advantage of this approach over some other methods of accounting for non-stationarity is that, for many model specifications, there is a closed-form, easily computable form for $\pi_t^{\eps}(\vartheta)$. Throughout the rest of this post, we will frequently transition between probabilities and log-probabilities, depending on which is most convenient for the task at hand. In log-space proprotionality will mean equality up to an additive constant that does not depend on the argument(s) of the function on the left hand side.

Theorem

$$\log \pi_t^{\eps}(\vartheta) \propto \ell_t^{\eps}(\vartheta) + \log \pi_0(\vartheta)$$

where

$$\ell_t^{\eps}(\vartheta) = \sum_{s = 1}^t (1 - \eps)^{t - s + 1} \cdot \log f(x_s\ |\ \vartheta)$$

is a decaying version of the log likelihood function.

Proof

For the first two observations we get

$$ \begin{align} \log \pi_1^{\varepsilon}(\vartheta) & \propto (1 - \eps) \cdot \log f(x_1\ |\ \vartheta) + (1 - \eps) \cdot \log \pi_0(\vartheta) + \eps \cdot \log \pi_0(\vartheta) \\ & = (1 - \eps) \cdot \log f(x_1\ |\ \vartheta) + \log \pi_0(\vartheta) \\ \log \pi_2^{\varepsilon}(p) & \propto (1 - \eps) \cdot \log f(x_2\ |\ \vartheta) + (1 - \eps) \cdot \log \pi_1^{\eps}(\vartheta) + \eps \cdot \log \pi_0(\vartheta) \\ & \propto (1 - \eps)^2 \cdot \log f(x_1\ |\ \vartheta) + (1 - \eps) \cdot \log f(x_2\ |\ \vartheta) + \log \pi_0(\vartheta). \end{align} $$

The emergence of the form in the theorem for arbitrary time $t$ is clear from here.

QED

We see that there is simple, closed-form expression for the contribution of the no data prior to the decayed posterior. For many common distributions, the decayed log likelihood, $\ell_t^{\varepsilon}(\vartheta)$, will also have a relatively simple form. We will return to some specific examples of these distributions throughout this post.

Returning to the beta-binomial case, where

$$\log f(x\ |\ p) = x \log p + (1 - x) \cdot \log(1 - p),$$

we have

$$ \begin{align} \log \pi_t^{\eps}(p) & \propto \ell_t^{\eps}(p) + \log \pi_0(p) \\ & = \sum_{s = 1}^t (1 - \eps)^{t - s + 1} \cdot \log f(x_s\ |\ p) \\ & = \left(\sum_{s = 1}^t (1 - \eps)^{t - s + 1} \cdot x_s\right) \cdot \log p + \left(\sum_{s = 1}^t (1 - \eps)^{t - s + 1} \cdot (1 - x_s)\right) \cdot \log (1 - p), \end{align} $$

$$\pi_t^{\eps}(p) \sim \text{Beta}\left(1 + \sum_{s = 1}^t (1 - \eps)^{t - s + 1} \cdot x_s, 1 + \sum_{s = 1}^t (1 - \eps)^{t - s + 1} \cdot (1 - x_s)\right).$$

In my last post, I showed how to compute sums of this form via matrix multiplication with an upper triangular matrix. We now use this method to calcumate the decayed posterior distributions for a variety of values of $0 < \eps \ll 1$.

In [25]:

def get_decay(T, ε):
    return np.triu((1 - ε) ** sp.linalg.toeplitz(1 + np.arange(T)))

In [26]:

def get_bb_decayed_posterior(x, ε, *, α0=1, β0=1):
    _, T = x.shape
    decay = get_decay(T, ε)

    α = α0 + x @ decay
    β = β0 + (1 - x) @ decay

    return α, β

In [27]:

EPSILONS = [0.1, 0.01, 0.001, 0.0001]

bb_decayed_ns = {ε: get_bb_decayed_posterior(x_ns, ε) for ε in EPSILONS}

We now visualize the posterior expected success rates from these decayed models, along with those of the stationary models calculated previously.

In [28]:

fig, ax = plt.subplots()

ax.plot(t, p_mle_ns.mean(axis=0), label="Average MLE")
ax.plot(t, p_bb_ns.mean(axis=0), label="Average posterior EV")

for ε, (α_decayed_ns, β_decayed_ns) in bb_decayed_ns.items():
    ax.plot(
        t,
        get_bb_post_mean(α_decayed_ns, β_decayed_ns).mean(axis=0),
        label=f"Average posterior EV ($\\varepsilon = {ε}$)",
    )

ax.hlines(P, 1, T // 2, color="k", ls="--", label="Actual")
ax.hlines(1 - P, T // 2, T, color="k", ls="--")

make_time_axis(ax=ax)
make_pct_axis(ax=ax)

ax.legend(loc="upper right", bbox_to_anchor=(1.65, 1));

We see that the decayed models are significantly more responsive to the nonstationarity of the success probability, with exactly how responsive they are depending on the value of $\eps$.

Inuitively, we expect larger values of $\eps$ to lead to faster decay, and therefore faster reactions to changes in the success probability. We do see exactly that in this plot, but there is some nuance. The $\eps = 0.1$ model does indeed react the fastest to the change in success probability, but during the stationary periods, that models does not capture the true sucesss probability well.

Effective sample size¶

This observation leads us to consider the effective sample size of each model. One interpretation of decayed updates is that they reduce the model's effective sample size by a factor controlled by $\eps$ when compared to a stationary model, so that when the true distribution shifts, they have less sample size inertia to overcome when adapting to the change.

We first quantify this phenomenon for the example beta-binomial models we have been working with. In the stationary beta-binomial model, we showed above that $\pi_t(p) \sim \text{Beta}(\alpha_t, \beta_t),$ where

$$ \begin{align} \alpha_t & = 1 + \sum_{s = 1}^t x_s \text{ and} \\ \beta_t & = 1 + t - \sum_{ts= 1}^t x_s. \end{align} $$

We can recover the sample size, $t$, of this model from these parameters by noting that

$$t = \alpha_t + \beta_t - 2.$$

This expression for $t$ motivates the following definition for the effective samples size of a decayed model. If $\pi_t^{\eps}(p) \sim \text{Beta}(\alpha_t^{\eps}, \beta_t^{\eps}),$ we let

$$n_{\text{eff}, t}^{\eps} = \alpha_t^{\eps} + \beta_t^{\eps} - 2$$

be the effective sample size of the model.

We can simplify this expression to

$$ \begin{align} n_{\text{eff}, t}^{\eps} & = \alpha_t^{\eps} + \beta_t^{\eps} - 2 \\ & = \sum_{s = 1}^t (1 - \eps)^{t - s + 1} \\ & = (1 - \eps)^{t + 1} \cdot \left(\frac{1 - (1 - \eps)^{-(t + 1)}}{1 - (1 - \eps)^{-1}} - 1\right) \\ & = \frac{1 - (1 - \eps)^t}{(1 - \eps)^{-1} - 1} \\ & = \frac{1 - \eps}{\eps} \cdot \left(1 - (1 - \eps)^t\right). \end{align} $$

The following plots show both ways of calculating this effective sample size, confirming our simplification is correct.

In [29]:

def get_ess(t, ε):
    return (1 - ε) / ε * (1 - (1 - ε) ** t)

In [30]:

fig, (stat_ax, form_ax) = plt.subplots(
    figsize=(10, 5), ncols=2, sharex=True, sharey=True
)

for ε, (α_decayed, β_decayed) in bb_decayed_ns.items():
    stat_ax.plot(
        t, (α_decayed + β_decayed).mean(axis=0) - 2, label=f"$\\varepsilon = {ε}$"
    )
    form_ax.plot(t, get_ess(t, ε))

stat_ax.plot(t, t, label=r"$\varepsilon = 0$ (stationary)")
form_ax.plot(t, t)

make_time_axis(ax=stat_ax)

stat_ax.set_yscale("log")
stat_ax.set_ylabel("Effective sample size")

stat_ax.set_title(
    r"Average $n_{\text{eff}, t}^{\varepsilon} = \alpha_t^{\varepsilon} + \beta_t^{\varepsilon} - 2$"
)
stat_ax.legend(loc="upper left")

make_time_axis(ax=form_ax)
form_ax.set_title(
    r"$n_{\text{eff}, t}^{\varepsilon} = \frac{1 - \varepsilon}{\varepsilon} \cdot (1 - (1 - \varepsilon)^t)$"
)

fig.tight_layout();

The most striking feature of these plots is that the decayed models have an asymptotic maximum sample size. In fact,

$$n_{\text{eff}, \infty}^\eps = \lim_{t \to \infty} n_{\text{eff}, t}^{\eps} = \frac{1 - \eps}{\eps},$$

so we see that a choice of $0 < \eps \ll 1$ is equivalent to choosing an asymptotic maximum sample size. This relationship between these quantities is visualized below.

In [31]:

fig, ax = plt.subplots()

plot_ε = np.logspace(-5, 0, 10_000)

ax.plot(plot_ε, (1 - plot_ε) / plot_ε)

ax.set_xscale("log")
ax.set_xlabel(r"$\varepsilon$")

ax.set_yscale("log")
ax.set_ylabel(
    r"$n_{\text{eff}, \infty}^{\varepsilon} = \frac{1 - \varepsilon}{\varepsilon}$"
)

ax.set_title("Asymptotic maximum effective sample size");

So far we have not addressed methods for choosing the decay parameter, $\eps$, but this analysis of asymptotic maximum effective sample size provides some intuition. If we have an intuition for the rough time scale (in terms of number of samples) over which we expect the parameters of interest to meaningfully vary, we can choose $\eps$ to produce the corresponding asymptotic maximum effective sample size. We will discuss another perspective on choosing $\eps$ later in this post.

So far our analysis of effective sample size has been theoretical. We now turn to the practical implications of reduced effective sample size. When our (effective) sample size is relatively smaller, we expect more uncertainy in our estimates, corresponding to relatively larger credible intervals. We return to the stationary case to visualize this behavior.

In [32]:

bb_decayed = {ε: get_bb_decayed_posterior(x, ε) for ε in EPSILONS}

In [33]:

CI_WIDTH = 0.95

In [34]:

def get_bb_ci_width(α, β):
    dist = sp.stats.beta(α, β)

    low = dist.ppf((1 - CI_WIDTH) / 2)
    high = dist.isf((1 - CI_WIDTH) / 2)

    return high - low

In [35]:

fig, ax = plt.subplots()

for ε, (α_decayed, β_decayed) in bb_decayed.items():
    ax.plot(
        t,
        get_bb_ci_width(α_decayed.mean(axis=0), β_decayed.mean(axis=0)),
        label=f"$\\varepsilon = {ε}$",
    )

ax.plot(t, get_bb_ci_width(α.mean(axis=0), β.mean(axis=0)), label=r"$\varepsilon = 0$")

make_time_axis(ax=ax)

ax.set_yscale("log")
ax.set_ylabel(f"{CI_WIDTH:.0%}-credible interval width")

ax.legend()
ax.set_title(f"Stationary case ($p = {P:0.1f}$),\naverage posterior parameters");

This plot confirms that as $\eps$ increases and the effective sample size decreases, the width of the credible interval increases. This relationship makes clear a fundamental tradeoff at the heart of model decay; to increase the speed at which the models react to changes in the true parameters, we are introducing uncertainty that makes them less confident in the stationary case.

No data updates, sample time versus clock time¶

Throughout this post so far, we have been implicitly defining the progression of time by the arrival of another piece of data. That is, if we have collected samples $x_1, x_2, \ldots, x_t$, we say that we have reached time $t + 1$ when the next piece of data, $x_{t + 1}$ arrives. We call this definition of time sample time. Depending on the data generating process, the clock time between sample times $t$ and $t + 1$ may be constant or could vary.

While this sample time approach is convenient in many cases, others call for defining the the progression of time independently of data arrival. For simplicity, we will assume data arrives (or not, more about that shortly) at even spaced intervals according to the progression of clock time. If we are guaranteed to have data arrive at even clock time intervals, there is no real difference between samplem time and clock time as we have described them here. It is when data is not guaranteed to arrive at every clock time interval that the difference between these definitions becomes interesting. This sort of situation arrives, for instance, when we are updating models on a regular interval based on visits to certain parts of a website. If the update interval is short (single digit minutes), some portions of a website may not get visited during every time interval.

Assume we are operating by clock time. If no data arrives between time $t$ and $t+1$, we define the decayed posterior at time $t + 1$ as

$$\pi_{t + 1}^{\eps}(\vartheta) \propto \pi_t^{\eps}(\vartheta)^{1 - \eps} \cdot \pi_0(\vartheta)^{\eps}.$$

Note that this is equivalent to assuming the likelihood of no data is one, together with the standard definition of the decayed posterior. Intuitively, when no data arrives, we shrink the posterior back towards the $t = 0$ prior.

Returning to the beta-binomial model for simplicity, assume that at time $t$, $\pi_t^{\eps}(p) \sim \text{Beta}(\alpha_t, \beta_t).$ Assume that no data arrives after time $t$. We will now explore the behavior of the decayed posterior as time goes on.

At time $t + 1$,

$$ \begin{align} \log \pi_{t + 1}^{\eps}(p) & \propto (1 - \eps) \cdot \log \pi_t^{\eps}(p) + \eps \cdot \log \pi_0(p) \\ & \propto (1 - \eps) \cdot \left((\alpha_t - 1) \cdot \log p + (\beta_t - 1) \cdot \log(1 - p)\right) \\ & = (1 - \eps) \cdot (\alpha_t - 1) \cdot \log p + (1 - \eps) \cdot (\beta_t - 1) \cdot \log(1 - p). \end{align} $$

At time $t + 2$,

$$ \begin{align} \log \pi_{t + 2}^{\eps}(p) & \propto (1 - \eps) \cdot \log \pi_{t + 1}^{\eps}(p) + \eps \cdot \log \pi_0(p) \\ & \propto (1 - \eps)^2 \cdot (\alpha_t - 1) \cdot \log p + (1 - \eps)^2 \cdot (\beta_t - 1) \cdot \log(1 - p). \end{align} $$

The pattern is apparent, so at time $t + s$, we have

$$\log \pi_{t + s}^{\eps}(p) \propto (1 - \eps)^s \cdot (\alpha_t - 1) \cdot \log p + (1 - \eps)^s \cdot (\beta_t - 1) \cdot \log(1 - p),$$

$$\pi_{t + s}^{\eps}(p) \sim \text{Beta}\left(1 + (1 - \eps)^s \cdot (\alpha_t - 1), 1 + (1 - \eps)^s \cdot (\beta_t - 1)\right).$$

We see that

$$\lim_{s \to \infty} \pi_{t + s}^{\eps}(p) \sim \text{Beta}(1, 1) = \pi_0(p),$$

so as longer periods of time go without any data arriving, the decayed posterior approaches the $t = 0$ posterior.

Calculating the effective sample size at time $t + s$ we get

$$ \begin{align} n_{\text{eff}, t + s}^{\eps} & = (1 - \eps)^s \cdot (\alpha_t + \beta_t - 2). \end{align} $$

We can get another heuristic for helping to choose $\eps$ by calculating the half-life of information, which is how long it takes the effective sample size to be cut in half after an extended period of no data. We get

$$ \begin{align} n_{\text{eff}, t_{1/2}^{\eps}}^{\eps} & = \frac{1}{2} n_{\text{eff}, t}^{\eps} \\ (1 - \eps)^{t_{1/2}^{\eps}} \cdot (\alpha_t + \beta_t - 2) & = \frac{1}{2} \cdot (\alpha_t + \beta_t - 2) \\ t_{1/2}^{\eps} & = -\frac{\log 2}{\log(1 - \eps)}. \end{align} $$

We visualize how thisn quantity changes with $\eps$ below.

In [36]:

def get_bb_half_life(ε):
    return -np.log(2) / np.log(1 - ε)

In [37]:

fig, ax = plt.subplots()

ax.plot(plot_ε, get_bb_half_life(plot_ε))

ax.set_xscale("log")
ax.set_xlabel(r"$\varepsilon$")

ax.set_yscale("log")
ax.set_ylabel(r"$t_{1/2}^{\varepsilon} = -\frac{\log 2}{\log(1 - \varepsilon)}$")

ax.set_title("Half life of information");

There may be some situations where we have stronger intuition around the half life of information than the time scale over which we expect changes to occur. In this case the above quantity may be used to choose an appropriate value of $\eps$.

Decayed posteriors for exponential family distributions¶

So far we have focused our concrete analysis of decayed Bayesian updated on the simple beta-binomial model. Analogous closed-form expressions exist for the decayed posterior many different distributions. In this section, we will show that this is the case for distributions in the exponential family, a class which includes many common distributions (normal, exponential, $\chi^2$, etc.).

Recall that a distribution is in the exponential family if its probability density function is of the form

$$f(x\ |\ \eta) \propto \exp(\eta \cdot T(x) - A(\eta)),$$

where $T$ is a sufficient statistic of the distribution and $A$ is its log-partition function).

A pleasant property of exponential family distributions for Bayesian statistics is that they have conjugate priors given by

$$\pi(\eta\ |\ \chi, \nu) \propto \exp(\eta \cdot \chi - \nu \cdot A(\eta)).$$

We can quickly verify that the posterior for observation $x$ is

$$ \begin{align} \pi(\eta\ |\ x, \chi, \nu) & \propto f(x\ |\ \eta) \cdot \pi(\eta\ |\ \chi, \nu) \\ & \propto \exp(\eta \cdot T(x) - A(\eta)) \cdot \exp(\eta \cdot \chi - \nu \cdot A(\eta)) \\ & = \exp(\eta \cdot (\chi + T(x)) - (\nu + 1) \cdot A(\eta)), \end{align}$$

so $\pi(\eta\ |\ x, \chi, \nu) = \pi(\eta\ |\ \chi + T(x), \nu + 1),$ and therefore the prior is conjugate.

Theorem

For a distribution in the exponential family with sufficient statistic $T(x)$, log-partition function $A(\eta)$, and conjugate prior $\pi(\eta\ |\ \chi_0, \nu_0)$,

$$\pi_t^{\varepsilon}(\eta\ |\ \chi, \nu) = \pi(\eta\ |\ \chi_t^{\varepsilon}, \nu_t^{\varepsilon}),$$

where

$$\chi_t^{\varepsilon} = \chi_0 + \sum_{s = 1}^t (1 - \varepsilon)^{t - s + 1} \cdot T(x_s),$$

and

$$\nu_t^{\varepsilon} = \nu_0 + \frac{1 - \eps}{\eps} \cdot \left(1 - (1 - \eps)^t\right).$$

Proof

Recall from our first theorem that

$$ \begin{align} \log \pi_t^{\eps}(\eta) & \propto \sum_{s = 1}^t (1 - \eps)^{t - s + 1} \cdot \log f(x_s\ |\ \vartheta) + \log \pi(\eta\ |\ \chi_0, \eta_0) \\ & \propto \sum_{s = 1}^t (1 - \eps)^{t - s + 1} \cdot (\eta \cdot T(x_s) - A(\eta)) + (\eta \cdot \chi_0 - \nu_0 \cdot A(\eta)) \\ & = \eta \cdot\left(\chi_0 + \sum_{s = 1}^t (1 - \eps)^{t - s + 1} \cdot T(x_s)\right) + \left(\nu_0 + \sum_{s = 1}^t (1 - \eps)^{t - s + 1}\right) \cdot A(\eta) \\ & = \eta \cdot \chi_t^{\varepsilon} - \left(\nu_0 + \sum_{s = 1}^t (1 - \eps)^{t - s + 1}\right) \cdot A(\eta). \end{align} $$

We already summed the geometric series present here during our derivation of effective sample size for the beta-binomial model, so we have that

$$ \begin{align} \nu_t^{\varepsilon} & = \nu_0 + \sum_{s = 1}^t (1 - \eps)^{t - s + 1} \\ & = \nu_0 + \frac{1 - \eps}{\eps} \cdot \left(1 - (1 - \eps)^t\right). \end{align} $$

QED

Many of the other calculations we performed for the beta-binomial model above can also be worked out for arbitrary distributions in the exponential family.

This post is available as a Jupyter notebooks here.

In [38]:

%load_ext watermark
%watermark -n -u -v -iv

Last updated: Sat Jan 10 2026

Python implementation: CPython
Python version       : 3.12.5
IPython version      : 8.29.0

matplotlib: 3.9.2
seaborn   : 0.13.2
numpy     : 2.0.2
IPython   : 8.29.0
scipy     : 1.14.1

Cumulative Sums as Matrix Multiplication

Austin Rochford — Wed, 07 Jan 2026 05:00:00 GMT

I am in the process of drafting another post about what I call "decaying Bayes Theorem," which is a method for weighting recent data more heavily than older data during sequential updates of Bayesian statistical models. The implementation of decayed Bayesian updates uses a computational trick that has been useful many times, so I decided break it out into its own short post to avoid it getting lost in the broader points of the forthcoming decaying Bayes post.

The basic insight is that cumulative sums can be written as matrix multiplication by triangular matrices. If we have

$$\mathbf{x} = (x_1, x_2, \ldots, x_n) \in \mathbb{R}^n,$$

the cumulative sum of $x$ is defined as

$$\mathbf{c} = (x_1, x_1 + x_2, \ldots, x_1 + x_2 + \cdots + x_n) \in \mathbb{R}^2.$$ If $U$ is the upper triangular matrix of ones

$$U = \begin{pmatrix} 1 & 1 & \cdots & 1 & 1 \\ 0 & 1 & \cdots & 1 & 1 \\ & & \vdots & & \\ 0 & 0 & \cdots & 1 & 1 \\ 0 & 0 & \cdots & 0 & 1 \end{pmatrix} \in \mathbb{R}^{n \times n},$$

we have that

$$\mathbf{c} = \mathbf{x}^{\top} U.$$ More generally, if we have a matrix of values $X \in \mathbb{R}^{m \times n}$ and let $C \in \mathbb{R}^{m \times n}$ be defined as the cumulative sum of the rows of $X$:

$$C_{i, j} = \sum_{k = 1}^j X_{i, k},$$

we have

$$C = X U.$$

We can validate this calculation numerically as follows.

First we make the necessary Python imports and do some light configuration.

In [1]:

from hypothesis import assume, given
from hypothesis.extra.numpy import arrays, array_shapes
from hypothesis.strategies import floats
from warnings import filterwarnings

In [2]:

import numpy as np
from scipy.linalg import toeplitz

In [3]:

filterwarnings("ignore", category=RuntimeWarning)

We will use the property-based testing library Hypothesis to validate the claims in this post.

Property-based testing works by sampling random plausible inputs to your function and testing invariants on those inputs. The canonical example is that reversing a list twice should result in the original list. In our case, we will calculate the cumulative sum two ways:

with the native NumPy method and
our (possibly generalzed) matrix multplication method,

and ensuring the results are close.

First we define a function that uses matrix multiplication with an upper triangular matrix of ones as described above to calculate the cumulative sum of its argument's rows.

In [4]:

def matmul_cumsum(X):
    _, N = X.shape
    U = np.triu(np.ones((N, N)))

    return X @ U

Next we define Hypothesis strategies for generating the matix X. We ensure that $X$ is two dimensonal and that each of ts rows have finite sum.

In [5]:

matrix_shapes = array_shapes(min_dims=2, max_dims=2)
finite_row_sum_matrices = arrays(np.float64, matrix_shapes).filter(
    lambda x: np.isfinite(x.sum(axis=1)).all()
)

We then define a function that takes a matrix whose rows have finite sum, and checks that the result of calculating the cumulative sum of its columns gives similar results.

In [6]:

@given(finite_row_sum_matrices)
def test_cumsum(X):
    np.testing.assert_allclose(matmul_cumsum(X), X.cumsum(axis=1), rtol=np.inf)

Running this function causes Hypothesis to check the assertion for many random inputs.

In [7]:

test_cumsum()

Since no exceptions were raised, we have confirmed that our matrix multiplication version of the cumulative sum is reasonably accurate.

So far, this reimplementation of cumulative sum is just a mathematical curiosity. This perspective, however, becomes computationally valuable when we want to generalize the idea of cumulative sum. In the forthcoming decaying Bayes post, we will have reason to consider sums of the form

$$C^{\lambda}_{i, j} = \sum_{k = 1}^j \lambda^{j - k + 1} \cdot X_{i, k}$$

for $0 \leq \lambda \leq 1.$ Generalizing the above perspective on cumulative sums as matrix multiplication, if we define

$$U^{\lambda} = \begin{pmatrix} \lambda & \lambda^2 & \cdots & \lambda^{n - 1} & \lambda^n \\ 0 & \lambda & \cdots & \lambda^{n - 2} & \lambda^{n - 1} \\ & & \vdots & & \\ 0 & 0 & \cdots & \lambda & \lambda^2 \\ 0 & 0 & \cdots & 0 & \lambda \end{pmatrix} \in \mathbb{R}^{n \times n},$$

we get that

$$C^{\lambda} = X U^{\lambda}.$$

The following functions implement this operation using matrix multiplication with a suitably-constructed upper diagonal matrix, and test that implementation against naive iteration.

In [8]:

def matmul_decayed_cumsum(X, λ):
    _, N = X.shape
    Uλ = np.triu(λ ** toeplitz(1 + np.arange(N)))

    return X @ Uλ

Note that this test relies on the fact that

$$C_{i, j}^{\lambda} = \lambda \cdot (C_{i, j - 1}^{\lambda} + X_{i, j})$$

for $1 < j \leq n$.

In [9]:

λs = floats(0, 1)

In [10]:

@given(finite_row_sum_matrices, λs)
def test_decayed_cumsum(X, λ):
    _, N = X.shape
    expected = np.empty_like(X)
    expected[:, 0] = λ * X[:, 0]

    for j in range(1, N):
        expected[:, j] = λ * (expected[:, j - 1] + X[:, j])

    np.testing.assert_allclose(matmul_decayed_cumsum(X, λ), expected, rtol=np.inf)

In [11]:

test_decayed_cumsum()

We see that no exceptions were raised, and our matrix multiplication implementation is reasonably accurate.

This post is available as a Jupyter notebook here.

In [12]:

%load_ext watermark
%watermark -n -u -v -iv

Last updated: Wed Jan 07 2026

Python implementation: CPython
Python version       : 3.12.5
IPython version      : 8.29.0

hypothesis: 6.150.0
scipy     : 1.14.1
numpy     : 2.0.2

Implementing the Box-Cox Power Exponential Distribution in Python with PyMC

Austin Rochford — Sat, 13 Dec 2025 05:00:00 GMT

One of my favorite sources of post inspiration is reading about interesting probability/prior distributions and implementing them in PyMC for my own edification (1 2 3 4). I had just such an experience this week reading The Box-Cox power exponential distribution written by Danielle Navarro. This post introduces two distributions motivated by the classic Box-Cox transform, the Box-Cox normal distribution for modeling location, scale, and skewness and the Box-Cox power exponential (BCPE) distribution for modeling location, scale, skewness, and kurtosis. This post will focus on the implementation of and inference for the latter distribution in PyMC. For motivation and further details, please refer to the excellent original post.

First we make the necessary Python imports and do some light configuration.

In [1]:

%matplotlib inline

In [2]:

from pathlib import Path

In [3]:

import arviz as az
from matplotlib import pyplot as plt, rcParams
import numpy as np
import nutpie
from pandas import read_sas
import polars as pl
import pymc as pm
from pytensor import tensor as pt
import seaborn as sns
from seaborn import objects as so

In [4]:

rcParams["figure.figsize"] = (8, 6)

sns.set(color_codes=True)

Load the data¶

We follow the original post by using weight data from the 2021-2023 National Health and Nutrition Examination Survey (NHANES).

In [5]:

DATA_DIR = Path("./data/")

BMX_PATH = DATA_DIR / "BMX_L.xpt"
DEMO_PATH = DATA_DIR / "DEMO_L.xpt"

Demographic data¶

First we load and transform the necessary demographic data for each survey subject.

In [6]:

DEMO_COLS = ["SEQN", "RIAGENDR", "RIDAGEYR"]

In [7]:

GENDER = pl.Enum(["Female", "Male"])

In [8]:

def gender_to_enum(col):
    return (
        pl.when(col == 1)
        .then(pl.lit("Male"))
        .when(col == 2)
        .then(pl.lit("Female"))
        .otherwise(pl.lit("Unknown"))
        .cast(GENDER)
        .alias("gender")
    )

In [9]:

cast_SEQN = pl.col("SEQN").cast(pl.Int64)

In [10]:

demo_df = (
    pl.from_pandas(read_sas(DEMO_PATH))
    .select(DEMO_COLS)
    .with_columns(cast_SEQN, gender_to_enum(pl.col("RIAGENDR")))
    .drop("RIAGENDR")
    .rename({"RIDAGEYR": "age", "SEQN": "seqn"})
)

In [11]:

demo_df

Out[11]:

shape: (11_933, 3)

seqn	age	gender
i64	f64	enum
130378	43.0	"Male"
130379	66.0	"Male"
130380	44.0	"Female"
130381	5.0	"Female"
130382	2.0	"Male"
…	…	…
142306	9.0	"Male"
142307	49.0	"Female"
142308	50.0	"Male"
142309	40.0	"Male"
142310	80.0	"Female"

The columns are

seqn - a unique identifier for the survey subject,
age - the subject's age in years, and
gender - the subject's gender.

Measurement data¶

Next we load and transform the necessary examination data for each survey subject.

In [12]:

BMX_COLS = ["SEQN", "BMXWT"]

In [13]:

bmx_df = (
    pl.from_pandas(read_sas(BMX_PATH))
    .select(BMX_COLS)
    .with_columns(cast_SEQN)
    .rename({"SEQN": "seqn", "BMXWT": "weight"})
)

In [14]:

bmx_df

Out[14]:

shape: (8_860, 2)

seqn	weight
i64	f64
130378	86.9
130379	101.8
130380	69.4
130381	34.3
130382	13.6
…	…
142306	25.3
142307	null
142308	79.3
142309	81.9
142310	72.1

The columns are

seqn - the same unique identifier for the survey subject and
weight - the subject's weight in kilograms.

Joint data¶

We join these dataframes on the unique subject identifier and remove rows with missing data.

In [15]:

full_df = demo_df.join(bmx_df, on="seqn").drop_nulls()

In [16]:

full_df

Out[16]:

shape: (8_754, 4)

seqn	age	gender	weight
i64	f64	enum	f64
130378	43.0	"Male"	86.9
130379	66.0	"Male"	101.8
130380	44.0	"Female"	69.4
130381	5.0	"Female"	34.3
130382	2.0	"Male"	13.6
…	…	…	…
142305	76.0	"Female"	60.4
142306	9.0	"Male"	25.3
142308	50.0	"Male"	79.3
142309	40.0	"Male"	81.9
142310	80.0	"Female"	72.1

Exploratory data analysis¶

Following the original post to validate that we are using the same data set, we visualize the relationship between age, weight and gender below.

In [17]:

(
    so.Plot(full_df, x="age", y="weight", color="gender")
    .add(so.Dot(alpha=0.25))
    .scale(color=so.Nominal(order=GENDER.categories))
    .label(x="Age", y="Weight (kg)", color=str.capitalize)
)

Out[17]:

Finally, we restrict our analysis to subjects 26 to 50 years old and visualize the distribution of weight among these subjects by gender.

In [18]:

df = full_df.filter(pl.col("age").is_between(25, 50, closed="right"))

In [19]:

(
    so.Plot(
        df,
        x="weight",
        color="gender",
    )
    .add(so.Bar(edgewidth=0, width=1), so.Hist())
    .scale(color=so.Nominal(order=GENDER.categories))
    .label(x="Weight (kg)", color=str.capitalize)
)

Out[19]:

As in the original post, we see that these distributions are decidely non-normal, possessing both interesting skewness and kurtosis, requiring a distribution with the flexibility of the BCPE to model accurately.

Modeling¶

Per the original post, the BCPE distribution has four parameters, $\mu \in (-\infty, \infty)$, $\sigma > 0$, $\nu \in (-\infty, \infty)$, and $\tau > 0$. Its probability density is

$$f(y\ |\ \mu, \sigma, \nu, \tau) = \frac{y^{\nu - 1}}{\mu^{\nu} \cdot \sigma} \cdot f(z\ |\ \tau).$$

Here $z$ is related to $y$ by a transformation reminiscient of Box-Cox,

$$z = \begin{cases} \frac{1}{\sigma \cdot \nu} \cdot \left(\left(\frac{y}{\mu}\right)^{\nu} - 1\right) & \textrm{if } \nu \neq 0 \\ \frac{1}{\sigma} \cdot \log\left(\frac{y}{\mu}\right) & \textrm{if } \nu = 0 \end{cases}.$$

The probability density of $z$ is

$$f(z\ |\ \tau) = \frac{\tau}{c \cdot 2^{1 + \tau^{-1}} \cdot \Gamma(\tau^{-1})} \cdot \exp\left(-\frac{1}{2} \cdot \left(\frac{|z|}{c}\right)^{\tau}\right).$$

This distribution has normalizing constant

$$c = \sqrt{\frac{\Gamma(\tau^{-1})}{4^{\tau^{-1}} \cdot \Gamma(3 \tau^{-1})}}.$$

To implement the BCPE distribution in PyMC, we will need to calculate the log probability density $\log f(y\ |\ \mu, \sigma, \nu, \tau)$. To do so, first we note that the log of the normalizing constant $c$ is

$$\log c = \frac{1}{2} \cdot \left(\log \Gamma(\tau^{-1}) - \tau^{-1} \cdot \log 4 - \log \Gamma(3 \tau^{-1})\right),$$

and implement this calculation as a Python function.

In [20]:

def log_c(τ):
    return 0.5 * (
        pt.special.gammaln(τ**-1) - τ**-1 * np.log(4) - pt.special.gammaln(3 * τ**-1)
    )

Next we implement the transformation of $y$ into $z$ as a Python function.

In [21]:

def z(y, μ, σ, ν):
    # ν_ avoids dividing by zero in the third argument to switch when ν is zero
    # it is necessary because both alternatives are evaluated regardless of the
    # value of ν
    ν_ = pt.switch(pt.eq(ν, 0), 1, ν)

    return pt.switch(
        pt.eq(ν, 0), 1 / σ * (pt.log(y) - pt.log(μ)), 1 / (σ * ν_) * ((y / μ) ** ν - 1)
    )

Next we calculate the log density of $z$,

$$\log f(z\ |\ \tau) = \log \tau - \log c - (1 + \tau^{-1}) \cdot \log 2 - \log \Gamma(\tau^{-1}) - \frac{1}{2} \cdot \left(\frac{|z|}{c}\right)^{\tau},$$

and implement it as a Python function.

In [22]:

def log_fz(y, μ, σ, ν, τ):
    log_c_ = log_c(τ)
    z_ = z(y, μ, σ, ν)

    return (
        pt.log(τ)
        - log_c_
        - (1 + τ**-1) * np.log(2)
        - pt.special.gammaln(τ**-1)
        - 0.5 * (pt.abs(z_) / pt.exp(log_c_)) ** τ
    )

Finally we calculate the log probability density of $y$,

$$\log f(y\ |\ \mu, \sigma, \nu, \tau) = (\nu - 1) \cdot \log y - \nu \cdot \log \mu - \log \sigma + \log f(z\ |\ \tau),$$

and implement it as a Python function.

In [23]:

def bcpe_logp(y, μ, σ, ν, τ):
    return (ν - 1) * pt.log(y) - ν * pt.log(μ) - pt.log(σ) + log_fz(y, μ, σ, ν, τ)


def bcpe_density(y, μ, σ, ν, τ):
    return pt.exp(bcpe_logp(y, μ, σ, ν, τ)).eval()

Following the original post, we now set $\mu = 10$ and $\sigma = 10^{-1}$ and visualize the BCPE probability density function for combinations of $\nu \in \{-9, -3, 0, 3, 0\}$ and $\tau \in \{1, 2, 4, 16\}$.

In [24]:

PLOT_NU = np.array([-9, -3, 0, 3, 9], dtype=np.float64)
PLOT_TAU = np.array([1, 2, 4, 16], dtype=np.float64)

In [25]:

y = np.linspace(6, 14, 100)

In [26]:

grid_y, grid_ν, grid_τ = np.meshgrid(y, PLOT_NU, PLOT_TAU, indexing="ij")

grid_fy = bcpe_density(grid_y, 10, 0.1, grid_ν, grid_τ)

In [27]:

grid_df = pl.DataFrame(
    {
        "ν": grid_ν.ravel(),
        "τ": grid_τ.ravel(),
        "y": grid_y.ravel(),
        "density": grid_fy.ravel(),
    }
)

In [28]:

grid = sns.FacetGrid(grid_df, col="ν", row="τ", margin_titles=True)
grid.map(plt.plot, "y", "density")
grid.set_titles(
    row_template="$\\tau = {row_name:.0f}$", col_template="$\\nu = {col_name:.0f}$"
)

for (ax, *_) in grid.axes:
    ax.set_yticks([])

We see that these distributions are quite flexible and result in a wide variety of shapes.

Modeling¶

We now proceed to construct a model of weight as a BCPE distribution in PyMC and perform inference.

First we construct a data container that will let us reuse the same model to infer the distribution of weight for both men and women. As in the original post, we model men's weight first.

In [29]:

def get_gender_weights(df, gender):
    return df.filter(pl.col("gender") == gender).select("weight").to_numpy()

In [30]:

with pm.Model() as model:
    weight = pm.Data("weight", get_gender_weights(df, "Male"))

Next we use the following priors on the paramters of the BCPE distribution of weight.

$$ \begin{align} \mu & \sim N(\bar{w}, 5^2) \\ \sigma & \sim \text{Half-}N(2.5^2) \\ \nu & \sim N(0, 5^2) \\ \tau & \sim \text{Half-}N(2.5^2) \end{align} $$

Here $\bar{w}$ is the mean of the observed weights, making this an empirical Bayesian model.

In [31]:

with model:
    μ = pm.Normal("μ", weight.mean(), 5)
    σ = pm.HalfNormal("σ", 2.5)
    ν = pm.Normal("ν", 0, 5)
    τ = pm.HalfNormal("τ", 2.5)

Finally we use PyMC's CustomDist to add a BCPE likelihood for the observed weights.

In [32]:

with model:
    pm.CustomDist(
        "obs",
        μ,
        σ,
        ν,
        τ,
        logp=bcpe_logp,
        observed=weight,
    )

We now sample from the model's posterior distribution.

In [33]:

SEED = 123456789  # for reproducibility

SAMPLE_KWARGS = {"seed": SEED, "cores": 8}

In [34]:

male_trace = nutpie.sample(nutpie.compile_pymc_model(model), **SAMPLE_KWARGS)

Sampler Progress

Total Chains: 6

Active Chains: 0

Finished Chains: 6

Sampling for now

Estimated Time to Completion: now

Draws	Step Size	Gradients/Draw
1400	0.96	7
1400	0.96	7
1400	0.91	3
1400	0.99	3
1400	0.95	3
1400	0.96	7

The Gelman-Rubin statistics show no cause for concern.

We now compare the maximum likelihood estimates (MLE) of the parameters for men from the original post to their posterior distributions from our model.

In [36]:

REF_VALS = pl.DataFrame(
    {
        "var_name": ["μ", "σ", "ν", "τ"],
        "Female": [76.4064873, 0.2699057, -0.5469284, 2.3326601],
        "Male": [87.6094265, 0.2319320, -0.5234904, 1.6945917],
    }
)

In [37]:

def get_ref_val_dict(gender):
    return {
        var_name: [{"ref_val": val[0][0]}]
        for var_name, val in REF_VALS["var_name", gender]
        .rows_by_key("var_name")
        .items()
    }

In [38]:

PLOT_VARS = ["μ", "σ", "ν", "τ"]

In [39]:

az.plot_posterior(
    male_trace, var_names=PLOT_VARS, grid=(2, 2), ref_val=get_ref_val_dict("Male")
);

We see that the maximum likelihood estimates are quite close to the posterior expected values and medians for each parameter.

Below we plot the BCPE distribution with parameters equal to their posterior expected values.

In [40]:

def get_posterior_ev(trace, vars_=PLOT_VARS):
    return trace.posterior[vars_].mean().to_array().to_numpy()

In [41]:

ax = sns.histplot(get_gender_weights(df, "Male"), stat="density", legend=False)

plot_weight = np.linspace(35, 200)
ax.plot(
    plot_weight,
    bcpe_density(plot_weight, *get_posterior_ev(male_trace)),
    c="k",
    label="BCPE with posterior\nexpected parameters",
)

ax.set_xlabel("Weight (kg)")
ax.set_yticks([])
ax.set_title("Male")
ax.legend();

As in the original post, we see strong agreement between the BCPE density and the histogram estimator.

We now use the data continer to repeat this analysis for women's weights.

In [42]:

with model:
    pm.set_data({"weight": get_gender_weights(df, "Female")})

In [43]:

female_trace = nutpie.sample(nutpie.compile_pymc_model(model), **SAMPLE_KWARGS)

Sampler Progress

Total Chains: 6

Active Chains: 0

Finished Chains: 6

Sampling for now

Estimated Time to Completion: now

Draws	Step Size	Gradients/Draw
1400	0.92	3
1400	0.93	3
1400	0.93	3
1400	0.86	3
1400	0.94	3
1400	0.94	3

In [45]:

az.plot_posterior(
    female_trace, var_names=PLOT_VARS, grid=(2, 2), ref_val=get_ref_val_dict("Female")
);

Again we see that the maximum likelihood estimates are quite close to the posterior expected values and medians for each parameter.

Again we plot the BCPE distribution with parameters equal to their posterior expected values, and see solid agreement with the histogram estimates.

In [46]:

ax = sns.histplot(get_gender_weights(df, "Female"), stat="density", legend=False)

plot_weight = np.linspace(35, 200)
ax.plot(
    plot_weight,
    bcpe_density(plot_weight, *get_posterior_ev(female_trace)),
    c="k",
    label="BCPE with posterior\nexpected parameters",
)

ax.set_xlabel("Weight (kg)")
ax.set_yticks([])
ax.set_title("Female")
ax.legend();

My sincere thanks to Danielle Navarro for the interesting post that inspired this one.

This post is available as a Jupyter notebook here.

In [47]:

%load_ext watermark
%watermark -n -u -v -iv

Last updated: Sat Dec 13 2025

Python implementation: CPython
Python version       : 3.12.5
IPython version      : 8.29.0

nutpie    : 0.15.2
numpy     : 2.0.2
arviz     : 0.20.0
matplotlib: 3.9.2
seaborn   : 0.13.2
polars    : 1.36.1
pytensor  : 2.31.7
pandas    : 2.2.3
pymc      : 5.25.1

Revisiting Multilevel Regression and Poststratification with PyMC

Austin Rochford — Sun, 23 Nov 2025 05:00:00 GMT

A little over eight years ago, I published a post entitled MRPyMC3 - Multilevel Regression and Poststratification with PyMC3, showing how to perform multilevel regression with poststratification (MRP) in Python with PyMC. I periodically enjoy revisiting old posts after both technology and my understanding of the problem advances. This post revisits MRP with a few notable changes:

PyMC is now on major version 5 instead of version 3,
we use nutpie for sampling from the model instead of PyMC's built in sampler, and
we use Polars instead of pandas.

We will not repeat the previous post's full exposition of MRP and will rather focus on the mechanics of its implementation.

First we import the necessary packages and do a bit of light configuration.

In [1]:

%matplotlib inline
%config InlineBackend.figure_format = "retina"

In [2]:

from itertools import zip_longest
import os
from urllib import request
import us
from zipfile import ZipFile

In [3]:

import arviz as az
from matplotlib import cm, pyplot as plt, ticker
from matplotlib.colorbar import ColorbarBase
from matplotlib.colors import Normalize
import numpy as np
import nutpie
import polars as pl
import pymc as pm
from pyreadstat import read_dta
from pytensor import tensor as pt
import seaborn as sns
from seaborn import objects as so
from scipy.special import logit

In [4]:

sns.set_style("darkgrid", {"axes.linewidth": 1, "axes.edgecolor": "black"})

Load and transform the data¶

As in the previous post, we follow Jonathan Kastellec's excellent MRP Primer, which focuses on estimating state-level opinions of gay marriage in 2005/2006 from polling data.

First we download and decompress the data.

In [5]:

DATA_PATH = "./data"
DATA_URI = "https://jkastellec.scholar.princeton.edu/sites/g/files/toruqf3871/files/jkastellec/files/mrp_primer_replication_files.zip"

In [6]:

USER_AGENT = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/141.0.0.0 Safari/537.36"

In [7]:

if not os.path.isdir(DATA_PATH):
    os.mkdir(DATA_PATH)

dest_path = os.path.join(DATA_PATH, os.path.basename(DATA_URI))

if not os.path.exists(dest_path):
    opener = request.build_opener()
    opener.addheaders = [("User-agent", USER_AGENT)]
    request.install_opener(opener)
    request.urlretrieve(DATA_URI, dest_path)

    with ZipFile(dest_path) as src:
        src.extractall(DATA_PATH)

Poll data¶

Next we load and do some light feature engineering on the polling data necessary to build the multilevel model.

In [8]:

UNZIPPED_DIR = "MRP_Primer_Replication_Files"

In [9]:

POLL_PATH = os.path.join(DATA_PATH, UNZIPPED_DIR, "gay_marriage_megapoll.dta")

In [10]:

POLL_COLS = [
    "race_wbh",
    "age_cat",
    "edu_cat",
    "female",
    "region",
    "state",
    "poll",
    "yes_of_all",
]

NN_COLS = ["race_wbh", "age_cat", "edu_cat"]

In [11]:

def to_zero_indexed(name):
    col = pl.col(name)

    return col - col.min()

In [12]:

CAT_COLS = ["age_cat", "edu_cat", "race_wbh"]

cat_col_transforms = [to_zero_indexed(name) for name in CAT_COLS]

In [13]:

GENDER = pl.Enum(["Male", "Female"])
POLL = pl.Enum(
    [
        "ABC 2004Jan15",
        "Gall2004Mar05",
        "Gall2005Aug22",
        "Pew 2004Dec01",
        "Pew 2004Feb11",
    ]
)
RACE = pl.Enum(["White", "Black", "Hispanic"])
REGION = pl.Enum(["dc", "midwest", "northeast", "south", "west"])
STATE = pl.Enum(sorted([state.abbr for state in us.states.STATES] + ["DC"]))

In [14]:

ENUM_CASTS = {
    "female": (GENDER, "gender"),
    "poll": (POLL, None),
    "race_wbh": (RACE, "race"),
    "region": (REGION, None),
    "state": (STATE, None),
}


def cast_enum_cols(df):
    for name, (enum, new_name) in ENUM_CASTS.items():
        if name in df.columns:
            df = df.with_columns(pl.col(name).cast(enum).alias(new_name or name))

            if new_name is not None:
                df = df.drop(name)

    return df

In [15]:

poll_df = (
    pl.from_pandas(read_dta(POLL_PATH)[0])
    .select(POLL_COLS)
    .drop_nulls(NN_COLS)
    .with_columns(*cat_col_transforms)
    .group_by(POLL_COLS[:-1])
    .agg(pl.sum("yes_of_all"), poll_pop=pl.len().cast(pl.Int32))
    .filter(pl.col("state") != "")
    .pipe(cast_enum_cols)
)

In [16]:

poll_df

Out[16]:

shape: (3_693, 9)

age_cat	edu_cat	region	state	poll	yes_of_all	poll_pop	gender	race
i64	i64	enum	enum	enum	i64	i32	enum	enum
1	1	"south"	"DE"	"Gall2005Aug22"	0	1	"Male"	"Black"
3	3	"south"	"FL"	"Pew 2004Dec01"	3	5	"Male"	"White"
0	3	"south"	"VA"	"Pew 2004Feb11"	1	1	"Female"	"White"
0	1	"midwest"	"OH"	"Gall2004Mar05"	0	1	"Male"	"White"
3	1	"south"	"TN"	"ABC 2004Jan15"	0	1	"Male"	"Black"
…	…	…	…	…	…	…	…	…
3	1	"west"	"CA"	"Pew 2004Feb11"	0	1	"Female"	"Black"
1	3	"midwest"	"IA"	"ABC 2004Jan15"	1	1	"Female"	"White"
0	3	"south"	"TN"	"Pew 2004Feb11"	1	1	"Female"	"White"
0	2	"northeast"	"NY"	"Gall2004Mar05"	0	1	"Female"	"White"
3	1	"midwest"	"WI"	"Gall2005Aug22"	0	1	"Female"	"White"

Each row represents a group of people polled:

age_cat represents the age category of that group,
edu_cat represents the education category of that group,
region represents the region of the state in which that group resides,
state represents the state in which that group resides,
poll represents which poll surveyed that group,
yes_of_all represents the number of respondents that indicated support for gay marriage,
poll_pop represents the number of respondents,
gender represents the gender of that group, and
race represents the race of that group.

Census data¶

Next we load and do some light feature engineering on the census data necessary for postratification.

In [17]:

CENSUS_PATH = os.path.join(DATA_PATH, UNZIPPED_DIR, "poststratification 2000.dta")

In [18]:

CENSUS_COLS = [
    "race_wbh",
    "age_cat",
    "edu_cat",
    "female",
    "state",
    "_freq",
    "region",
]

In [19]:

census_df = (
    pl.from_pandas(read_dta(CENSUS_PATH)[0])
    .rename(lambda s: s.lstrip("c").lower())
    .select(CENSUS_COLS)
    .rename({"_freq": "pop"})
    .with_columns(*cat_col_transforms)
    .pipe(cast_enum_cols)
    .sort("state")
)

In [20]:

census_df

Out[20]:

shape: (4_896, 7)

age_cat	edu_cat	state	pop	region	gender	race
i64	i64	enum	i64	enum	enum	enum
0	0	"AK"	467	"west"	"Male"	"White"
1	0	"AK"	377	"west"	"Male"	"White"
2	0	"AK"	419	"west"	"Male"	"White"
3	0	"AK"	343	"west"	"Male"	"White"
0	1	"AK"	958	"west"	"Male"	"White"
…	…	…	…	…	…	…
3	2	"WY"	4	"west"	"Female"	"Hispanic"
0	3	"WY"	8	"west"	"Female"	"Hispanic"
1	3	"WY"	16	"west"	"Female"	"Hispanic"
2	3	"WY"	10	"west"	"Female"	"Hispanic"
3	3	"WY"	1	"west"	"Female"	"Hispanic"

Each row represents a group of people with the given age, education, state, gender, and race combination. All columns have the same meaning as in the poll data. The pop column contains the census population of that group.

State data¶

Finally, we load and do some light feature engineering on data about each state, which we will use, in addition to the poll data, to build our multilevel model.

In [21]:

STATE_PATH = os.path.join(DATA_PATH, UNZIPPED_DIR, "state_level_update.dta")

In [22]:

STATE_COLS = ["sstate", "p_evang", "p_mormon", "kerry_04"]

In [23]:

def to_percentage(name):
    return pl.col(name) / 100

In [24]:

state_df = (
    pl.from_pandas(read_dta(STATE_PATH)[0])
    .select(STATE_COLS)
    .rename(lambda name: name.replace("ss", "s"))
    .with_columns(
        to_percentage("p_evang"), to_percentage("p_mormon"), to_percentage("kerry_04")
    )
    .with_columns(p_relig=pl.col("p_evang") + pl.col("p_mormon"))
    .drop("p_evang", "p_mormon")
    .pipe(cast_enum_cols)
    .join(census_df.group_by("state").agg(pl.first("region")), on="state")
    .sort("state")
)

In [25]:

state_df

Out[25]:

shape: (51, 4)

state	kerry_04	p_relig	region
enum	f64	f64	enum
"AK"	0.355	0.154431	"west"
"AL"	0.368	0.410083	"south"
"AR"	0.446	0.436301	"south"
"AZ"	0.444	0.142887	"west"
"CA"	0.543	0.087176	"west"
…	…	…	…
"VT"	0.589	0.02912	"northeast"
"WA"	0.528	0.128227	"west"
"WI"	0.497	0.129527	"midwest"
"WV"	0.432	0.116178	"south"
"WY"	0.291	0.208844	"west"

There is a row for each state, plus an additional one for the District of Columbia, with the following columns:

state - the state represented by that row,
kerry_04 - the proportion of the state's voters that voted for John Kerry in the 2004 presidential election,
p_relig - the proportion of the state's population that identifies as evangelical Christian or Mormon, and
region - the region of the country the state is in.

Exploratory data analysis¶

Now that we have all the necessary data, we visualize it in order to become familiar with it before building the model.

First we define a few functions that will facilitate intuitive plotting of state-level statistics.

In [26]:

STATE_GRID = [
    ["AK", None, None, None, None, None, None, None, None, None, "ME"],
    [None, None, None, None, None, None, None, None, None, "VT", "NH"],
    ["WA", "ID", "MT", "ND", "MN", None, "MI", None, "NY", "MA", "RI"],
    ["OR", "UT", "WY", "ND", "IA", "WI", "OH", "PA", "NJ", "CT"],
    ["CA", "NV", "CO", "NE", "IL", "IN", "WV", "VA", "MD", "DE"],
    [None, "AZ", "NM", "KS", "MO", "KY", "TN", "SC", "NC"],
    [None, None, None, "OK", "LA", "AR", "MS", "AL", "GA", None, "DC"],
    ["HI", None, None, None, "TX", None, None, None, "FL"],
]

In [27]:

def plot_facecolor(df, col, state, *, ax, norm, cmap, default="gray"):
    is_state = pl.col("state") == state

    if df.select(is_state.any()).item():
        color = cmap(norm(df.filter(is_state)[col]))
    else:
        color = default

    ax.set_facecolor(color)

In [28]:

def plot_state_grid(
    data, *, col, norm, cmap, default="gray", cbar=True, ax_plotter=plot_facecolor
):
    nrows = len(STATE_GRID)
    ncols = max(len(row) for row in STATE_GRID)

    fig, axes = plt.subplots(nrows=nrows, ncols=ncols)

    for state_row, ax_row in zip(STATE_GRID, axes):
        for state, ax in zip_longest(state_row, ax_row):
            if state is None:
                ax.axis("off")
            else:
                ax.set_xticks([])
                ax.set_yticks([])

                ax.annotate(state, xy=(0.1, 0.1), xycoords="axes fraction")

                ax_plotter(
                    data, col, state, ax=ax, norm=norm, cmap=cmap, default=default
                )

    if cbar:
        cbar_ax = fig.add_axes([0.925, 0.25, 0.02, 0.5])
        ColorbarBase(cbar_ax, cmap=cmap, norm=norm)
    else:
        cbar_ax = None

    return fig, axes, cbar_ax

Poll data¶

With the data and these tools in hand, we begin by plotting the number of poll respondents in each state.

In [29]:

fig, _, _ = plot_state_grid(
    poll_df.group_by("state").agg(pl.sum("poll_pop")),
    col="poll_pop",
    norm=Normalize(0, 750),
    cmap=cm.binary,
    default="white",
)

fig.suptitle("Number of responses");

We see immediately that some states have very few respondents, and indeed that Alaska and Hawaii had no respondents.

In [30]:

(
    poll_df.group_by("state")
    .agg(pl.sum("poll_pop"))
    .join(state_df, how="right", on="state")
    .with_columns(pl.col("poll_pop").fill_null(0))
    .sort("poll_pop")
    .select("state", "poll_pop")
    .head()
)

Out[30]:

shape: (5, 2)

state	poll_pop
enum	i32
"AK"	0
"HI"	0
"DC"	6
"WY"	14
"DE"	16

One of the benefits of MRP is the ability to use census data to predict the support for gay marriage in these states with little to no data based on similar respondents in other states.

Next we visualize the empirical support for gay marriage in each state.

In [31]:

PROB_CMAP = sns.diverging_palette(220, 10, as_cmap=True).reversed()
PROB_FORMATTER = ticker.StrMethodFormatter("{x:.1%}")
PROB_LOCATOR = ticker.MultipleLocator(0.2)
PROB_MIN, PROB_MAX = 0.1, 0.70
PROB_NORM = Normalize(0, PROB_MAX)

In [32]:

def make_prob_ax(ax):
    cbar_ax.yaxis.set_major_locator(PROB_LOCATOR)
    cbar_ax.yaxis.set_major_formatter(PROB_FORMATTER)

In [33]:

def rate(num, denom):
    return pl.sum(num) / pl.sum(denom)

In [34]:

disagg_df = poll_df.group_by("state").agg(rate("yes_of_all", "poll_pop"))

In [35]:

fig, _, cbar_ax = plot_state_grid(
    disagg_df,
    col="yes_of_all",
    norm=PROB_NORM,
    cmap=PROB_CMAP,
)

make_prob_ax(cbar_ax)
fig.suptitle("Disaggregation estimate of\nsupport for gay marriage");

Note that the MRP literature, including the Kastellec primer we are following in this post, refer to this quantity as the "disaggregation estimate." We will use that language throughout the rest of this post for consistency. At the end of the post we will see how the MRP estimates of support for gay marriage differ from these disaggregation estimates. The difference between these estimates is largely due to the fact that the polled population is not (and cannot practically be) precisely reflective of the population of each state.

This challenge of representative sampling is further illustrated in the following plot.

In [36]:

(
    so.Plot(
        poll_df.join(
            census_df, on=set(poll_df.columns) & set(census_df.columns), how="right"
        )
        .group_by("state", "gender", "race")
        .agg(pl.col("poll_pop").fill_null(0).sum() == 0)
        .group_by("gender", "race")
        .agg(pl.sum("poll_pop").alias("state_ct")),
        x="race",
        y="state_ct",
        color="gender",
    )
    .add(so.Bar(), so.Dodge())
    .scale(y=so.Continuous().tick(every=5))
    .label(
        x="Race",
        y="Number of states",
        color="Gender",
        title="States with no respondents",
    )
)

Out[36]:

While only two states have zero polled respondents, significantly more states have zero respondents from minority race/gender subpopulations. As in states with zero respondents, a benefit of MRP is the ability to predict support for gay marriage among unsampled subpopulations in these states based on attitudes from similar respondents in other statse.

Modeling¶

With this basic understanding, we now turn to building the multilevel model necessary for MRP.

First we define coordinates that will make our model easier to work with.

In [37]:

n_age = poll_df["age_cat"].n_unique()
n_edu = poll_df["edu_cat"].n_unique()

In [38]:

COORDS = {
    "age": np.arange(n_age),
    "edu": np.arange(n_edu),
    "gender": GENDER.categories.to_numpy(),
    "poll": POLL.categories.to_numpy(),
    "race": RACE.categories.to_numpy(),
    "region": REGION.categories.to_numpy(),
    "state": STATE.categories.to_numpy(),
}

The model includes an intercept with a normal prior,

$$\beta_0 \sim N(0, 2.5^2).$$

In [39]:

with pm.Model(coords=COORDS) as model:
    β0 = pm.Normal("β0", 0, 2.5)

Most of the terms of the model will be hierarchical normal. For example, the prior for the age effect is

$$ \begin{align} \sigma_{\text{age}} & \sim \text{Half-}N(2.5^2), \\ \beta_{\text{age}} & \sim N(0, \sigma_{\text{age}}^2). \end{align} $$

For sampling efficiency, we actually use an equivalent noncentered parameterization of the above priors. The priors for the effects of education, poll, age-education interaction, and gender-race interaction are defined similarly.

In [40]:

HALFNORMAL_SCALE = 1 / np.sqrt(1 - 2 / np.pi)

In [41]:

def noncentered_normal(name, *, dims, μ=0):
    Δ = pm.Normal(f"Δ_{name}", 0, 1, dims=dims)
    σ = pm.HalfNormal(f"σ_{name}", 2.5 * HALFNORMAL_SCALE)

    return pm.Deterministic(name, μ + Δ * σ, dims=dims)

In [42]:

with model:
    β_age = noncentered_normal("β_age", dims="age")
    β_edu = noncentered_normal("β_edu", dims="edu")
    β_poll = noncentered_normal("β_poll", dims="poll")

    β_age_edu = noncentered_normal("β_age_edu", dims=("age", "edu"))
    β_gender_race = noncentered_normal("β_gender_race", dims=("gender", "race"))

The state-level effects are slightly more complex, combining a hierarchical normal region-level effect with the effect of the proportion of its residents that identify as religious and that of the state's support for John Kerry in 2024, . We have

$$ \begin{align} \sigma_{\text{region}} & \sim \text{Half-}N(2.5^2), \\ \alpha_{\text{region}} & \sim N(0, \sigma_{\text{region}}^2) \\ \alpha_{\text{relig}}, \alpha_{\text{kerry}} & \sim N(0, 2.5^2), \\ \beta_{\text{state}} & = \alpha_{\text{region}} + \alpha_{\text{relig}} \cdot x_{\text{relig}} + \alpha_{\text{kerry}} \cdot x_{\text{kerry}}. \end{align} $$

Here $x_{\text{relig}}$ is the proportion of the state's residents that identify as religious and $x_{\text{kerry}}$ is the proportion that voted for John Kerry in 2024.

In [43]:

kerry_04 = state_df["kerry_04"].to_numpy()
p_relig = state_df["p_relig"].to_numpy()
state_region = state_df["region"].cast(pl.Categorical).to_physical().to_numpy()

In [44]:

with model:
    α_region = noncentered_normal("α_region", dims="region")
    α_relig = pm.Normal("α_relig", 0, 2.5)
    α_kerry = pm.Normal("α_kerry", 0, 2.5)

    β_state = (
        α_region[state_region] + α_relig * logit(p_relig) + α_kerry * logit(kerry_04)
    )

Before using these effects to define the likelihood of the observed polling data, we build a number of data containers that will facilitate the posterior predictive sampling that will be necessary to perform poststratification.

In [45]:

with model:
    age = pm.Data("age", poll_df["age_cat"].to_numpy())
    edu = pm.Data("edu", poll_df["edu_cat"].to_numpy())
    gender = pm.Data(
        "gender", poll_df["gender"].to_physical().to_numpy().astype(np.int_)
    )
    poll = pm.Data("poll", poll_df["poll"].to_physical().to_numpy().astype(np.int_))
    pop = pm.Data("pop", poll_df["poll_pop"].to_numpy())
    race = pm.Data("race", poll_df["race"].to_physical().to_numpy().astype(np.int_))
    state = pm.Data("state", poll_df["state"].to_physical().to_numpy())
    yes_of_all = pm.Data("yes_of_all", poll_df["yes_of_all"].to_numpy())

    use_poll = pm.Data("use_poll", True)

Finally we define

$$\eta = \beta_0 + \beta_{\text{age}} + \beta_{\text{edu}} + \beta_{\text{state}} + \beta_{\text{age, edu}} + \beta_{\text{gender, race}} + \beta_{\text{poll}}.$$

Note that to support poststratification, we include the ability to turn off the effect of the poll on the model.

The likelihood is binomial with log-odds $\eta$.

In [46]:

def adv_index(tensor, indices, shape):
    return tensor.ravel()[pt.ravel_multi_index(indices, shape)]

In [47]:

n_gender = GENDER.categories.shape[0]
n_race = RACE.categories.shape[0]

In [48]:

with model:
    η = sum(
        [
            β0,
            β_age[age],
            β_edu[edu],
            β_state[state],
            adv_index(β_age_edu, (age, edu), (n_age, n_edu)),
            adv_index(β_gender_race, (gender, race), (n_gender, n_race)),
            pt.switch(use_poll, β_poll[poll], 0),
        ]
    )

    pm.Binomial("response", pop, pt.sigmoid(η), observed=yes_of_all)

We now sample from the posterior distribution of this model.

In [49]:

SEED = 123456789  # for reproducibility
CHAINS = 8

SAMPLER_KWARGS = {
    "seed": SEED,
    "chains": CHAINS,
    "cores": CHAINS,
    "target_accept": 0.99,
}

In [50]:

trace = nutpie.sample(nutpie.compile_pymc_model(model), **SAMPLER_KWARGS)

Sampler Progress

Total Chains: 8

Active Chains: 0

Finished Chains: 8

Sampling for 3 minutes

Estimated Time to Completion: now

Draws	Step Size	Gradients/Draw
1400	0.05	255
1400	0.04	511
1400	0.04	511
1400	0.04	511
1400	0.05	255
1400	0.05	255
1400	0.04	1023
1400	0.04	255

The Gelman-Rubin statistic, $\hat{R}$, shows no cause for concern.

In [51]:

az.rhat(trace).max().to_array().max()

Out[51]:

<xarray.DataArray ()> Size: 8B
array(1.0036893)

Poststratification¶

We now sample from the model's posterior predictive distribution in order to perform poststratification. To do so, we set the data containers we created above to the values from the census. This step is what allows us to adjust for the non-representativeness of the poll's respondents.

In [52]:

n_census, _ = census_df.shape

In [53]:

with model:
    pm.set_data(
        {
            "age": census_df["age_cat"].to_numpy(),
            "edu": census_df["edu_cat"].to_numpy(),
            "gender": census_df["gender"].to_physical().to_numpy().astype(np.int_),
            "poll": np.zeros(n_census, dtype=np.int_),
            "pop": census_df["pop"].to_numpy(),
            "race": census_df["race"].to_physical().to_numpy().astype(np.int_),
            "state": census_df["state"].to_physical().to_numpy(),
            "yes_of_all": np.zeros(n_census, dtype=np.int_),
            "use_poll": False,
        }
    )

    pm.sample_posterior_predictive(trace, extend_inferencedata=True)

Sampling: [response]

Output()

We use these posterior predictive samples to construct the poststratification estimates of each state's support for gay marriagne.

In [54]:

poststrat_df = (
    census_df.with_columns(
        pp_yes_of_all=trace.posterior_predictive["response"]
        .mean(dim=("chain", "draw"))
        .to_numpy()
    )
    .group_by("state")
    .agg(rate("pp_yes_of_all", "pop").alias("yes_of_all"))
)

Now compare the disaggregation and MRP estimates of support for gay marriage. The disaggregation estimate colors the upper left corner of each state's rectangle, and the MRP estimate colors the lower right corner.

In [55]:

def plot_split_color(dfs, col, state, *, ax, norm, cmap, default="gray"):
    ul_df, lr_df = dfs

    is_state = pl.col("state") == state

    xmin, xmax = ax.get_xlim()
    ymin, ymax = ax.get_ylim()

    # upper left
    if ul_df.select(is_state.any()).item():
        ul_color = cmap(norm(ul_df.filter(is_state)[col]))
    else:
        ul_color = default

    ax.add_patch(plt.Polygon([[0, 0], [xmax, ymax], [0, ymax]], color=ul_color))

    # lower right
    if lr_df.select(is_state.any()).item():
        lr_color = cmap(norm(lr_df.filter(is_state)[col]))
    else:
        lr_color = default

    ax.add_patch(plt.Polygon([[0, 0], [xmax, ymax], [xmax, 0]], color=lr_color))

In [56]:

fig, _, cbar_ax = plot_state_grid(
    (
        poll_df.group_by("state").agg(rate("yes_of_all", "poll_pop")),
        poststrat_df,
    ),
    col="yes_of_all",
    norm=PROB_NORM,
    cmap=PROB_CMAP,
    ax_plotter=plot_split_color,
)

make_prob_ax(cbar_ax)
fig.suptitle("Disaggregation and MRP estimates\nof support for gay marriage");

We see that most states support has moved towards the average with the MRP estimate. Also, MRP produces estimates for Alaska and Hawaii based on their census data, whereas disaggregation cannot.

Finally, we produce another comparison of the two estimates state-by-state.

In [57]:

state_pop = census_df.group_by("state").agg(pl.sum("pop")).sort("state")
state_pop_ = state_pop["pop"].to_numpy().squeeze()

poststrat_state_df = (
    trace.posterior_predictive["response"]
    .rename(response_dim_2="state")
    .assign_coords(state=census_df["state"])
    .groupby("state")
    .sum()
    .pipe(lambda x: x / state_pop_)
)

In [58]:

sorted_states = STATE.categories[
    poststrat_state_df.mean(dim=("chain", "draw")).argsort().to_numpy()
]

In [59]:

(ax,) = az.plot_forest(
    poststrat_state_df,
    coords={"state": sorted_states[::-1]},
    combined=True,
    labeller=az.labels.NoVarLabeller(),
    figsize=(6, 12),
)

disagg_scatter_df = (
    disagg_df.join(
        pl.DataFrame(
            {
                "state": sorted_states.cast(STATE),
                "sort_key": np.arange(sorted_states.shape[0]),
                "y": ax.get_yticks(),
            }
        ),
        on="state",
    )
    .join(state_pop, on="state")
    .sort("sort_key")
)

pop_norm = Normalize(0, disagg_scatter_df["pop"].max())

ax.scatter(
    disagg_scatter_df["yes_of_all"],
    disagg_scatter_df["y"],
    s=200 * pop_norm(disagg_scatter_df["pop"]),
    c="k",
    label="Disaggregation\nestimate",
    zorder=5,
)

ax.axvline(
    (poststrat_state_df.mean(dim=("chain", "draw")) * state_pop_).sum()
    / state_pop_.sum(),
    ls="--",
    c="k",
    label="National average",
)

ax.xaxis.set_major_formatter(PROB_FORMATTER)
ax.set_xlabel("MRP estimate of support for gay marriage")

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

Here the size of the black circles corresponding to the disaggregation estimate correspond to the state's population. We see, by and large, that states with higher population have similar disaggregation and MRP estimates, and states with smaller population have MRP estimates closer to the national (disaggregation) mean than their disaggregation estimates. Notable exceptions to this trend are New York, Florida, and Pennsylvania, among others. These states's deviation from the trend indicates that the respondent population was particularly unrepresentative of the state's demographics.

This post is available as a Jupyter notebook here.

In [60]:

%load_ext watermark
%watermark -n -u -v -iv

Last updated: Sun Nov 23 2025

Python implementation: CPython
Python version       : 3.12.5
IPython version      : 8.29.0

pymc      : 5.25.1
nutpie    : 0.15.2
us        : 3.2.0
seaborn   : 0.13.2
pytensor  : 2.31.7
arviz     : 0.20.0
polars    : 1.14.0
numpy     : 2.0.2
matplotlib: 3.9.2