Introduction to Bayesian Power Analysis: Exclude a Null Value

Prepare Notebook

import arviz as az
import matplotlib.pyplot as plt
import matplotlib.ticker as mtick
import numpy as np
import numpy.typing as npt
import pandas as pd
import preliz as pz
import pymc as pm
import pytensor.tensor as pt
import seaborn as sns

az.style.use("arviz-darkgrid")
plt.rcParams["figure.figsize"] = [10, 6]
plt.rcParams["figure.dpi"] = 100
plt.rcParams["figure.facecolor"] = "white"

%load_ext autoreload
%autoreload 2
%config InlineBackend.figure_format = "retina"

seed: int = sum(map(ord, "power_sample_size"))
rng: np.random.Generator = np.random.default_rng(seed=seed)

Description of the Strategy

We strongly recommend reading the original part of the example from the book before continuing (Section 13.2.1). Still, we provide the main ideas below.

Let us assume we have a coin. The goal is to show that the coin is unfair. In other words, we want to show that the \(95\%\) HDI (Highest Density Interval) excludes a ROPE (Region of Practical Equivalence) around \(\theta = 0.50\). To do so, we follow the steps below:

Step 1. Establish the hypothetical distribution of parameter values from which we will generate simulated data. In the example from the book, the hypothetical distribution is a beta distribution with shape parameters of \(0.65 \cdot (2000 - 2) + 1\) and \((1 - 0.65) \cdot (2000 - 2) + 1\). The value of \(2000\) is arbitrary; it’s as if the generating mean of \(0.65\) was based on fictitious previous data containing \(2000\) flips.

fig, ax = plt.subplots()
omega = 0.65
pz.Beta(alpha=omega * (2_000 - 2) + 1, beta=(1 - omega) * (2_000 - 2) + 1).plot_pdf(
    ax=ax
)
ax.set(title=f"omega={omega}");

In practice, we can use many different distributions. For that purpose we store various distribution means in the variable omegas.

omegas = np.linspace(0.5, 0.85, 8)
print(f"omegas: {omegas}")

omegas: [0.5  0.55 0.6  0.65 0.7  0.75 0.8  0.85]

Step 2: We generate random samples of data from the hypothetical distribution and use it to generate synthetic data, i.e. flip coin results. For example, from the hypothetical distribution that is centered on \(\theta = 0.65\), suppose that the selected value is \(0.638\). We can use this value to simulate flipping a coin with that bias \(N\) times. This value of \(N\) is the sample size we want to study.

Let’s generate some sample sizes to consider.

sample_sizes = np.linspace(10, 450, 23, dtype=int)
print(f"sample_sizes: {sample_sizes}")

sample_sizes: [ 10  30  50  70  90 110 130 150 170 190 210 230 250 270 290 310 330 350 370 390 410 430 450]

Step 3: We fit the model we plan to use in practice (with the audience-appropriate priors) to the simulated data above. This will provide us with the posterior distribution of the parameter \(p\in[0,1]\) we are interested in.

Step 4: Tally (assess) whether or not the \(95\%\) HDI excludes a ROPE around the null value of \(\theta = 0.50\). We choose the ROPE to be \([0.48, 0.52]\).

Step 5: Repeat the above steps many times, to approximate the power of the experiment.

We specify the number of simulations to run in the variable n_simulations.

n_simulations = 500

Now that we have clarified the steps, let’s implement the logic in PyMC.

Generate Random Samples of Data

The first step is to generate random samples of data from the hypothetical distribution. We do this by sampling from the beta distribution with the shape parameters we defined earlier.

Remark: Observe that we can vectorize the whole simulation process. Hence, we do not need to loop over any of the parameters.

with pm.Model(
    coords={
        "sample_size": sample_sizes,
        "omega": omegas,
        "simulation": range(n_simulations),
    }
) as data_generating_model:
    # Data container for the sample sizes we want to consider.
    sample_sizes_ = pm.Data(
        "sample_sizes",
        sample_sizes,
        dims=("sample_size",),
    )
    # Expand the sample sizes to the shape of the model.
    sample_sizes_expanded = pt.expand_dims(
        pt.as_tensor_variable(sample_sizes_), axis=[1, 2]
    )
    # Data container for the omegas we want to consider.
    omegas_ = pm.Data("omegas", omegas, dims=("omega",))
    # Expand the omegas to the shape of the model.
    omegas_expanded = pt.expand_dims(pt.as_tensor_variable(omegas_), axis=[0, 2])
    # Sample from the hypothetical distribution.
    p = pm.Beta(
        "p",
        alpha=omegas_expanded * (2_000 - 2) + 1,
        beta=(1 - omegas_expanded) * (2_000 - 2) + 1,
        dims=("sample_size", "omega", "simulation"),
    )
    # Generate the data from the binomial distribution.
    pm.Binomial(
        "data",
        n=sample_sizes_expanded,
        p=p,
        dims=("sample_size", "omega", "simulation"),
    )

pm.model_to_graphviz(data_generating_model)

Now we can simply do a forward pass to generate the data.

with data_generating_model:
    generating_model_idata = pm.sample_prior_predictive(samples=1, random_seed=rng)

generated_data = generating_model_idata["prior"]["data"].sel(chain=0, draw=0).to_numpy()

# Verify the shape of the generated data.
assert generated_data.shape == (len(sample_sizes), len(omegas), n_simulations)

Sampling: [data, p]

Compute Posterior Estimates

Now we fit the model to the generated data. We specify a mild-informative prior distribution for the parameter \(p\) centered on \(0.5\) with a standard deviation of \(0.2\).

fig, ax = plt.subplots()
pz.Beta(mu=0.5, sigma=0.2).plot_pdf(ax=ax)
ax.set(title="Prior distribution of p");

Let’s write the vectorized model.

with pm.Model(
    coords={
        "sample_size": sample_sizes,
        "omega": omegas,
        "simulation": range(n_simulations),
    }
) as model:
    # Data container for the generated data from the previous step.
    generated_data_ = pm.Data(
        "generated_data",
        generated_data,
        dims=("sample_size", "omega", "simulation"),
    )
    # Data container for the sample sizes we want to consider.
    # We keep them as as above as we want to vectorize the process.
    sample_sizes_ = pm.Data(
        "sample_sizes",
        sample_sizes,
        dims=("sample_size",),
    )
    # Expand the sample sizes to the shape of the model.
    sample_sizes_expanded = pt.expand_dims(
        pt.as_tensor_variable(sample_sizes_), axis=[1, 2]
    )
    # Specify the prior distribution for the parameter $p$.
    p = pm.Beta(
        "p",
        mu=0.5,
        sigma=0.2,
        dims=("sample_size", "omega", "simulation"),
    )
    # Condition the likelihood on the generated data.
    pm.Binomial(
        "data",
        n=sample_sizes_expanded,
        p=p,
        observed=generated_data_,
        dims=("sample_size", "omega", "simulation"),
    )

pm.model_to_graphviz(model)

Next, we run our NUTS sampler.

%%time

with model:
    idata = pm.sample(
        draws=1_000,
        chains=4,
        nuts_sampler="numpyro",
        random_seed=rng,
    )

CPU times: user 1h 29s, sys: 13min 19s, total: 1h 13min 49s
Wall time: 11min 32s

We verify we have no divergences.

print(f"Diverging: {idata['sample_stats']['diverging'].sum().item()}")

Diverging: 0

Let’s consider the posterior distribution summary for the specific values of the simulation omega=0.65 and sample_size=150 (which is the particular case from the book).

az.summary(
    idata["posterior"].sel(sample_size=150, omega=0.65, simulation=slice(None, 10)),
    var_names=["p"],
)

The r-hat values are close to \(1.0\), indicating good convergence.

Compute How Many Times the Goal is Attained

Now that we have all the posterior samples at hand, we can easily compute our condition of interest. Namely, we want to know how many times the \(95\%\) HDI excludes the ROPE. Before doing so, let’s visualize the posterior distribution of \(p\) for the particular case of omega=0.65 and sample_size=150 for the first \(100\) simulations.

ax, *_ = az.plot_forest(
    idata["posterior"].sel(sample_size=150, omega=0.65, simulation=slice(None, 100)),
    var_names=["p"],
    combined=True,
    rope=[0.48, 0.52],
    figsize=(7, 20),
)

ax.set(title="100 simulated posterior distribution of p (and rope)");

Computing the \(95\%\) HDI for the parameter \(p\) and evaluating whether it excludes the ROPE is straightforward.

hdi = az.hdi(idata["posterior"], hdi_prob=0.95)["p"]

rope = (0.48, 0.52)

results = ((hdi.sel(hdi="lower") > rope[1]) | (hdi.sel(hdi="higher") < rope[0])).mean(
    dim="simulation"
)

Again, let’s compare the power estimation for the particular case of omega=0.65 and sample_size=150.

results.sel(sample_size=150, omega=0.65).item()

0.884

This is the same result as in the book (well, it is \(.9\), but it is close).

Visualize Results

We can now generate the whole Table 13.1 from the book.

fig, ax = plt.subplots(figsize=(9, 7))

sns.heatmap(
    data=(
        results.to_dataframe()
        .reset_index()
        .pivot_table(index="sample_size", columns="omega", values="p")
    ),
    annot=True,
    cmap="viridis_r",
    linewidths=0.2,
    linecolor="black",
    fmt=".1%",
    cbar_kws={"format": mtick.PercentFormatter(xmax=1)},
    ax=ax,
)
ax.tick_params("y", rotation=0)
ax.set(
    title="Power analysis for different sample sizes",
    xlabel="power",
    ylabel="sample size",
);

The results are very similar to the expected ones.

Remark: In the book, all of these results are computed analytically as the math is not too complex. However, in real applications, we need to use simulation methods.

Finally, let’s visualize the power curve for different sample sizes.

fig, ax = plt.subplots(figsize=(12, 7))

sns.lineplot(
    data=(
        results.to_dataframe()
        .reset_index()
        .assign(omega=lambda x: x["omega"].astype("str"))
    ),
    x="sample_size",
    y="p",
    hue="omega",
    marker="o",
)

ax.axhline(0.8, c="black", ls="--")
ax.legend(loc="center left", title="omega", title_fontsize=14, bbox_to_anchor=(1, 0.5))
ax.set(
    title="Power analysis for different sample sizes",
    xlabel="sample size",
    ylabel="power",
);

Here are some clear remarks about the results from the book:

• As generating mode increases (the omegas), the required sample size decreases. This makes sense intuitively: When the generating mode is large, the sample proportion of heads will tend to be large, and so the HDI will tend to fall toward the high end of the parameter domain. In other words, when the generating mode is large, it doesn’t take a lot of data for the HDI to fall consistently above \(\theta = 0.5\).

• On the other hand, when the generating mode is only slightly above \(\theta = 0.5\), then it takes a large sample for the sample proportion of heads to be consistently above \(0.5\), and for the HDI to be consistently entirely above \(0.5\) (and the ROPE).

Reproducing and coding this vectorized version of the power analysis example was very interesting. This example provides a solid illustration of the power analysis flow, which I plan to extend for more complex applications.