Kidney: Weibull Regression with Random Effects

This example comes from Volume 1 of the classic BUGS examples and is also documented in the OpenBUGS Kidney example. The data, from McGilchrist and Aisbett (1991), record the times to first and second recurrence of infection in 38 kidney patients on dialysis. Each patient contributes up to two observations, and some of those times are censored: for a censored observation we only know that the patient had gone a certain length of time without a recurrence, not the exact recurrence time.

The question is how the risk of recurrence depends on patient characteristics: age, sex, and the type of underlying disease (coded as "other", GN, AN, or PKD). The model is a Weibull survival regression with a patient-level random effect, so that the two observations from the same patient are allowed to be correlated and unexplained differences between patients are absorbed by a normally distributed "frailty" term. This is the classic Bayesian analogue of a shared-frailty proportional-hazards survival model.

Model

For patient $i$ and recurrence $j$, the recurrence time $t_{ij}$ follows a Weibull distribution whose scale depends on the covariates through a log link, with a per-patient random effect $b_i$:

\[\begin{aligned} t_{ij} &\sim \text{Weibull}(r,\ \mu_{ij}), \quad \text{censored below by } t^{\text{cen}}_{ij} \\ \log \mu_{ij} &= \alpha + \beta_{\text{age}}\, \text{age}_{ij} + \beta_{\text{sex}}\, \text{sex}_i + \beta_{\text{dis}[\text{disease}_i]} + b_i \\ b_i &\sim \text{Normal}(0,\ \tau) \end{aligned}\]

The disease effect uses a corner-point constraint: the first level is fixed to zero and the remaining three levels are estimated. All regression coefficients get flat normal priors, the random-effect precision $\tau$ and the Weibull shape $r$ get gamma priors, and $\sigma = 1/\sqrt{\tau}$ is the standard deviation of the random effects.

using JuliaBUGS
using Distributions

# `censored` comes from Distributions; make it available inside `@bugs`:
JuliaBUGS.@bugs_primitive censored

kidney = @bugs begin
    for i in 1:N
        for j in 1:M
            # Survival times bounded below by censoring times:
            t[i, j] ~ censored(dweib(r, mu[i, j]), var"t.cen"[i, j], nothing)
            mu[i, j] = exp(alpha + var"beta.age" * age[i, j] + var"beta.sex" * sex[i] +
                           var"beta.dis"[disease[i]] + b[i])
        end
        # Random effects:
        b[i] ~ dnorm(0.0, tau)
    end

    # Priors:
    alpha ~ dnorm(0.0, 0.0001)
    var"beta.age" ~ dnorm(0.0, 0.0001)
    var"beta.sex" ~ dnorm(0.0, 0.0001)
    # beta.dis[1] <- 0; # corner-point constraint
    for k in 2:4
        var"beta.dis"[k] ~ dnorm(0.0, 0.0001)
    end
    tau ~ dgamma(1.0E-3, 1.0E-3)
    r ~ dgamma(1.0, 1.0E-3)
    sigma = 1 / sqrt(tau) # s.d. of random effects
end

BUGSModelDef:
begin
    for i = 1:N
        for j = 1:M
            t[i, j] ~ censored(dweib(r, mu[i, j]), var"t.cen"[i, j], nothing)
            mu[i, j] = exp(alpha + var"beta.age" * age[i, j] + var"beta.sex" * sex[i] + var"beta.dis"[disease[i]] + b[i])
        end
        b[i] ~ dnorm(0.0, tau)
    end
    alpha ~ dnorm(0.0, 0.0001)
    var"beta.age" ~ dnorm(0.0, 0.0001)
    var"beta.sex" ~ dnorm(0.0, 0.0001)
    for k = 2:4
        var"beta.dis"[k] ~ dnorm(0.0, 0.0001)
    end
    tau ~ dgamma(0.001, 0.001)
    r ~ dgamma(1.0, 0.001)
    sigma = 1 / sqrt(tau)
end

Names such as var"beta.age" and var"t.cen" are R-style dotted variable names carried over verbatim from the original BUGS program; Julia's var"..." syntax lets us keep the exact original names.

Data

The data set is large, so we load the bundled copy rather than typing it out:

data  = JuliaBUGS.BUGSExamples.VOLUME_1.kidney.data
inits = JuliaBUGS.BUGSExamples.VOLUME_1.kidney.inits
model = kidney(data, inits)

BUGSModel (parameters are in transformed (unconstrained) space, with dimension 43):

  Model parameters:
    beta.sex
    alpha
    b[38], b[37], b[36], b[35], b[34], b[33], b[32], b[31], b[30], b[29], b[28], b[27], b[26], b[25], b[24], b[23], b[22], b[21], b[20], b[19], b[18], b[17], b[16], b[15], b[14], b[13], b[12], b[11], b[10], b[9], b[8], b[7], b[6], b[5], b[4], b[3], b[2], b[1]
    beta.dis[4], beta.dis[3], beta.dis[2]
    beta.age
    tau
    r
    t[38, 2], t[37, 1], t[36, 2], t[36, 1], t[34, 2], t[32, 1], t[26, 1], t[24, 2], t[22, 2], t[20, 2], t[19, 2], t[19, 1], t[16, 2], t[15, 2], t[14, 2], t[14, 1], t[12, 2], t[2, 2]

  Variable sizes and types:
    t.cen: size = (38, 2), type = Matrix{Int64}
    alpha: type = Int64
    b: size = (38,), type = Vector{Float64}
    disease: size = (38,), type = Vector{Int64}
    sex: size = (38,), type = Vector{Int64}
    sigma: type = Float64
    N: type = Int64
    M: type = Int64
    mu: size = (38, 2), type = Matrix{Float64}
    beta.age: type = Int64
    tau: type = Float64
    r: type = Int64
    beta.sex: type = Int64
    beta.dis: size = (4,), type = Vector{Int64}
    age: size = (38, 2), type = Matrix{Int64}
    t: size = (38, 2), type = Matrix{Float64}

We pass the example's published initial values as the second argument to kidney rather than calling kidney(data) alone: with a censored survival likelihood and vague priors, values drawn at random from the priors can land outside the valid range, so a sensible starting point makes construction and sampling reliable.

Here N = 38 is the number of patients and M = 2 is the maximum number of recurrences per patient. t holds the recurrence times (with missing entries where the observation is censored), t.cen holds the corresponding censoring times, age is the patient's age at each observation, sex is a 0/1 indicator, and disease codes the underlying disease type as an integer from 1 to 4. All of the classic examples ship with the package in JuliaBUGS.BUGSExamples, bundling the model, data, initial values, and reference results together.

Sampling

With the model compiled, we attach a gradient backend and draw posterior samples with the No-U-Turn Sampler:

using AbstractMCMC, AdvancedHMC, ADTypes, Mooncake, FlexiChains
using LogDensityProblems

model = kidney(data, inits; adtype=AutoMooncake(; config=nothing))

n_samples, n_adapts = 2000, 1000
D = LogDensityProblems.dimension(model)
chain = AbstractMCMC.sample(
    model, NUTS(0.8), n_samples;
    chain_type=VNChain, n_adapts=n_adapts,
    init_params=rand(D), discard_initial=n_adapts,
)
summarystats(chain)

BUGS-style initial values for this example are available as JuliaBUGS.BUGSExamples.VOLUME_1.kidney.inits and can be applied with initialize!(model, inits).

Results

The source file for this example ships with reference_results = nothing, so there is no bundled table of published posterior summaries to reproduce here. The MultiBUGS and OpenBUGS pages linked above report posterior means and standard deviations for alpha, beta.age, beta.sex, the disease effects beta.dis[2]–beta.dis[4], the Weibull shape r, and the random-effect standard deviation sigma, obtained from a 1000-iteration burn-in followed by 10000 further iterations. A correctly converged chain's summarystats output should match those published values up to Monte Carlo error.

See also: the gallery overview and the getting-started tutorial.