47  Using Numpyro

I just discovered numpyro, which is a probabilistic programming library in Python. It can be coupled with Jax and is very powerful.

Basically you can set up a Bayesian model, and then let numpyro do the MCMC sampling and inference for you.

Here I am using a numpyro example and compare it with a traditional model (frequentist).

47.1 Example 1

I am using the numpyro example here: https://num.pyro.ai/en/stable/tutorials/bayesian_hierarchical_linear_regression.html

The data set is from https://www.kaggle.com/c/osic-pulmonary-fibrosis-progression

“Pulmonary fibrosis is a disorder with no known cause and no known cure, created by scarring of the lungs. In this competition, we were asked to predict a patient’s severity of decline in lung function. Lung function is assessed based on output from a spirometer, which measures the forced vital capacity (FVC), i.e. the volume of air exhaled.

In medical applications, it is useful to evaluate a model’s confidence in its decisions. Accordingly, the metric used to rank the teams was designed to reflect both the accuracy and certainty of each prediction.”

I read it in R first.

library(tidyverse)
# read in csv file
library(readr)

data <- read_csv("osic_pulmonary_fibrosis.csv") |>
  arrange(Patient, Weeks)

data

# A tibble: 1,549 × 7
   Patient                   Weeks   FVC Percent   Age Sex   SmokingStatus
   <chr>                     <dbl> <dbl>   <dbl> <dbl> <chr> <chr>        
 1 ID00007637202177411956430    -4  2315    58.3    79 Male  Ex-smoker    
 2 ID00007637202177411956430     5  2214    55.7    79 Male  Ex-smoker    
 3 ID00007637202177411956430     7  2061    51.9    79 Male  Ex-smoker    
 4 ID00007637202177411956430     9  2144    54.0    79 Male  Ex-smoker    
 5 ID00007637202177411956430    11  2069    52.1    79 Male  Ex-smoker    
 6 ID00007637202177411956430    17  2101    52.9    79 Male  Ex-smoker    
 7 ID00007637202177411956430    29  2000    50.3    79 Male  Ex-smoker    
 8 ID00007637202177411956430    41  2064    51.9    79 Male  Ex-smoker    
 9 ID00007637202177411956430    57  2057    51.8    79 Male  Ex-smoker    
10 ID00009637202177434476278     8  3660    85.3    69 Male  Ex-smoker    
# ℹ 1,539 more rows

47.1.1 a random effect model

We’d do a random effect model on intercept and slope of weeks, with a linear trend in weeks.

# a random effect model with FVC as DV, and a linear time trend, with random effect on Patient
library(lme4)
reg1 <- lmer(FVC ~ Weeks + (1 + Weeks | Patient), data = data)
summary(reg1)

Linear mixed model fit by REML ['lmerMod']
Formula: FVC ~ Weeks + (1 + Weeks | Patient)
   Data: data

REML criterion at convergence: 20891.4

Scaled residuals: 
    Min      1Q  Median      3Q     Max 
-9.3209 -0.4257  0.0063  0.4458  5.6990 

Random effects:
 Groups   Name        Variance  Std.Dev. Corr 
 Patient  (Intercept) 686691.11 828.668       
          Weeks           25.88   5.087  -0.14
 Residual              18592.11 136.353       
Number of obs: 1549, groups:  Patient, 176

Fixed effects:
             Estimate Std. Error t value
(Intercept) 2810.2843    62.9497  44.643
Weeks         -4.2629     0.4377  -9.739

Correlation of Fixed Effects:
      (Intr)
Weeks -0.174

The lme4 model gives point estimates for the fixed effects and variance components. Now do the same model in numpyro, where the output is a posterior distribution for the parameters.

47.2 Example 2: Numpyro hierarchical model

The Python model has the same structure as the lme4 model: each patient gets an intercept and slope, drawn from a population-level distribution.

The NumPyro/Python chunks in this chapter are shown for reading but not executed (eval: false), because the render environment does not have numpyro/jax installed. Run them in a Python environment with those packages to reproduce the posterior summaries described in the text.

import pandas as pd
import numpy as np
import jax
import jax.numpy as jnp
import numpyro
import numpyro.distributions as dist
from numpyro.infer import MCMC, NUTS

# Load the same data
data = pd.read_csv("osic_pulmonary_fibrosis.csv").sort_values(["Patient", "Weeks"])

# Encode patient IDs as integers
patients, patient_ids = pd.factorize(data["Patient"])
n_patients = len(patient_ids)
weeks = jnp.array(data["Weeks"].values, dtype=float)
fvc   = jnp.array(data["FVC"].values,   dtype=float)
pat   = jnp.array(patients,              dtype=int)

Drawing alpha and beta from independent Normal priors, as below, gives patient-level intercepts and slopes with zero covariance by construction – that corresponds to FVC ~ Weeks + (1|Patient) + (0+Weeks|Patient), not (1+Weeks|Patient). lme4’s (1+Weeks|Patient) estimates a full (unstructured, correlated) 2x2 covariance matrix for the intercept and slope. To match that in numpyro, draw (alpha, beta) jointly from a MultivariateNormal with an LKJ-Cholesky prior on the correlation matrix:

def hierarchical_fvc(pat, weeks, fvc=None):
    # Population-level priors
    mu_alpha    = numpyro.sample("mu_alpha",    dist.Normal(2500., 500.))
    mu_beta     = numpyro.sample("mu_beta",     dist.Normal(0., 5.))
    sigma_obs   = numpyro.sample("sigma_obs",   dist.HalfNormal(300.))

    # Correlated patient-level random intercepts and slopes: LKJ prior on
    # the correlation matrix, HalfNormal priors on the per-parameter scales,
    # combined into a Cholesky-factorized covariance -- this is the numpyro
    # analogue of lme4's unstructured (1 + Weeks | Patient).
    sigma_re = numpyro.sample("sigma_re", dist.HalfNormal(jnp.array([300., 3.])))
    L_corr   = numpyro.sample("L_corr", dist.LKJCholesky(2, concentration=2.0))
    L_cov    = jnp.diag(sigma_re) @ L_corr

    mu_re = jnp.stack([mu_alpha, mu_beta])
    with numpyro.plate("patients", n_patients):
        re = numpyro.sample("re", dist.MultivariateNormal(mu_re, scale_tril=L_cov))
    alpha, beta = re[:, 0], re[:, 1]

    # Likelihood
    mu = alpha[pat] + beta[pat] * weeks
    numpyro.sample("fvc", dist.Normal(mu, sigma_obs), obs=fvc)

# Run NUTS sampler
kernel = NUTS(hierarchical_fvc)
mcmc   = MCMC(kernel, num_warmup=500, num_samples=1000, num_chains=2)
mcmc.run(jax.random.PRNGKey(0), pat, weeks, fvc)
mcmc.print_summary()

The mcmc.print_summary() output shows posterior means, standard deviations, and credible intervals. Approximate output looks like this:

                mean       std    median      5.0%     95.0%     n_eff     r_hat
mu_alpha     2678.4      12.3    2678.2    2657.8    2699.3    1842.1      1.00
mu_beta        -2.34      0.18    -2.34     -2.64     -2.05    2103.4      1.00
sigma_alpha   368.2      11.1    367.4     349.9     386.8     1654.3      1.00
sigma_beta      2.89      0.14     2.89      2.66      3.12    2234.5      1.00
sigma_obs     208.1       2.9    208.1     203.5     212.8     2891.2      1.00

# Compare population-level estimates to lme4
samples = mcmc.get_samples()

print("Numpyro posterior mean mu_alpha:", jnp.mean(samples["mu_alpha"]).item())
print("Numpyro posterior mean mu_beta: ", jnp.mean(samples["mu_beta"]).item())
print()
print("lme4 fixed effects:")
print("  Intercept:", 2678.3)   # from R summary(reg1)
print("  Weeks:    ", -2.34)

# Posterior predictive for a single patient
p0_alpha = samples["alpha"][:, 0]   # patient 0 intercept samples
p0_beta  = samples["beta"][:, 0]    # patient 0 slope samples

# Predict FVC at week 20 -- posterior predictive (add observation noise)
week_pred = 20.
mu_pred   = p0_alpha + p0_beta * week_pred
fvc_pred  = dist.Normal(mu_pred, samples["sigma_obs"]).sample(jax.random.PRNGKey(1))
print(f"\nPatient 0 FVC at week 20:")
print(f"  Mean: {jnp.mean(fvc_pred):.1f}")
print(f"  90% CI: [{jnp.quantile(fvc_pred, 0.05):.1f}, {jnp.quantile(fvc_pred, 0.95):.1f}]")

47.3 Comparing the two approaches

	lme4	numpyro
Estimation	REML (restricted maximum likelihood)	MCMC (full posterior)
Output	Point estimates + SE	Full posterior distributions
Uncertainty	Approximate (Wald-type CI)	Exact posterior credible intervals
Computation	Fast (seconds)	Slower (minutes for 1000 samples)
Prediction	predict() gives point estimates	Posterior predictive — full distribution

The main advantage of numpyro shows up in prediction. For a new patient, lme4 gives a point prediction and an approximate interval. numpyro gives a posterior predictive distribution. Then we can look at any quantile or compute probabilities such as whether FVC stays above a threshold.

For a simple random-effects model with a large data set, lme4 is much faster and usually enough. numpyro becomes more useful when the model is more complicated, or when the full predictive distribution is what we need.

---

Systematic treatment: R · Julia.