import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pymc as pm
from pymc.math import exp

5. Time-to-event Models: Example 2*#

Consider the model

\[\begin{split} \begin{align*} Y_i & \sim Wei(\nu,\lambda) \\ \nu & = \frac{3}{2} \\ \lambda & \sim Ga(2,3) \\ \end{align*}\end{split}\]

with data

\[ Y_1 = 2, Y_2 = 3, Y_3 = 1^*, Y_4 = \frac{5}{2}, Y_5 = 3^*\]

where *’s represent censored points. The goal is to estimate the parameter \(\lambda\) in Bayesian fashion. The likelihood is the Weibull PDF and Survival function for non-censored and censored points, respectively

\[\begin{split} \begin{align*} L(\lambda | \nu, y_1, ...,y_n) & = \prod_{i=1}^k \nu \lambda y_i ^ {\nu-1} e^{-\lambda y_{i} ^\nu}\times \prod_{i=k+1}^n e^{-\lambda y_i ^nu} \\ & = \nu^k \lambda^k \left ( \prod_{i=1}^k y_i\right )^{\nu-1} e^{-\lambda \sum_{i=1}^n y_i^\nu} \end{align*} \end{split}\]

We can elimate constants, and see that the likelihood is proportional to

\[ L(\lambda | \nu, y_1, ...,y_n) \propto \lambda^k e^{-\lambda \sum_{i=1}^n y_i^\nu} \]

Prior of \(\lambda\) is specified as \(Ga(\alpha,\beta)\), which is also simplified

\[ \pi(\lambda | \alpha, \beta) \propto \lambda^{\alpha-1} e^{-\beta \lambda}\]

With the joint likelihood and prior, we can find the posterior distribution

\[\pi(\lambda | \nu, y_1,...,y_n,\alpha,\beta) \propto \lambda^{k+\alpha-1} e^{-\lambda (\beta + \sum_{i=1}^n y_i^\nu)}\]

And by substituting in the known values, we arrive at

\[\pi(\lambda | \nu, y_1,...,y_n,\alpha,\beta) \propto Ga(3 + 2, 18.1736 +3)\]

The estimator \(\hat{\lambda}_{bayes}\) is the mean of this Gamma distribution

\[ \hat{\lambda}_{bayes} = \frac{5}{21.1736} = 0.2361\]

PyMC implementation#

Now, let’s use PyMC to find the \( \hat{\lambda}_{bayes} \). Note, we are converting between Weibull parameterizations to be compatible with the courses historical solution in WINBUGS.

censored = np.array([np.inf, np.inf, 1, np.inf, 3])
y = np.array([2, 3, 1, 2.5, 3])
with pm.Model() as m:
    λ = pm.Gamma("λ", 2, 3)
    α = 1.5
    β = λ ** (-1 / α)
    obs_latent = pm.Weibull.dist(α, β)
    likelihood = pm.Censored(
        "likelihood", obs_latent, lower=None, upper=censored, observed=y
    )

    trace = pm.sample(1000)

az.summary(trace)
Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [λ]
100.00% [8000/8000 00:00<00:00 Sampling 4 chains, 0 divergences]
Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 1 seconds.
mean sd hdi_3% hdi_97% mcse_mean mcse_sd ess_bulk ess_tail r_hat
λ 0.234 0.104 0.064 0.423 0.002 0.002 1624.0 1838.0 1.0
%load_ext watermark
%watermark -n -u -v -iv -p pytensor
Last updated: Wed Mar 22 2023

Python implementation: CPython
Python version       : 3.11.0
IPython version      : 8.9.0

pytensor: 2.10.1

pandas    : 1.5.3
matplotlib: 3.6.3
pymc      : 5.1.2
numpy     : 1.24.2
arviz     : 0.14.0