Author

Xiang Ao

Published

October 17, 2025

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.

Code
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.

Code
# 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 data are the OSIC pulmonary-fibrosis panel: 1,549 spirometry readings from 176 patients, with forced vital capacity (FVC) as the outcome and Weeks since baseline as the time variable. The model gives each patient a random intercept and slope.

The fixed effects are an intercept of 2810.3 (SE 62.9) and a slope of \(-4.263\) (SE 0.438): average FVC declines by about 4.3 units per week, and the \(t\) of \(-9.7\) leaves little doubt about the direction. The variance components show where the heterogeneity lives — a patient-intercept standard deviation of 828.7 against a residual of 136.4, so patients differ far more in level than any single measurement deviates from that patient’s own line. The slope SD of 5.09 against a mean slope of \(-4.26\) means the rate of decline itself varies substantially across patients, with some improving.

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.

Code
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:

Code
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
Code
# 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.