Fitting survival data with MortalityTables.jl

begin
    using LsqFit
    using MortalityTables
    using CairoMakie
    using Distributions
    using Optim
    using DataFrames
    using Survival
end

# Apply the site's custom Makie plot theme (cassette futurism).
include(joinpath(@__DIR__, "..", "..", "assets", "themes", "cassette_futurism.jl"))
set_theme!(cassette_futurism_theme())

This tutorial is via PharmCat on Github (this link has similar code with comments in Russian and English).

Fitting a Weibull survival curve

Sample data:

data = let
    survival = [0.99, 0.98, 0.95, 0.9, 0.8, 0.65, 0.5, 0.38, 0.25, 0.2, 0.1, 0.05, 0.02, 0.01]
    times = 1:length(survival)
    DataFrame(; times, survival)
end

14×2 DataFrame

Row	times	survival
	Int64	Float64
1	1	0.99
2	2	0.98
3	3	0.95
4	4	0.9
5	5	0.8
6	6	0.65
7	7	0.5
8	8	0.38
9	9	0.25
10	10	0.2
11	11	0.1
12	12	0.05
13	13	0.02
14	14	0.01

Visualizing the data:

plt = Figure(); ax = Axis(plt[1,1]; xlabel="time", ylabel="survival",
    title="Observed survival proportion")
lines!(ax, data.times, data.survival;
    label="observed", color=CASSETTE_PALETTE[1], linewidth=2)
plt

Observed survival proportions decreasing over time.

Define the two-parameter Weibull model:

x: array of independent variables
p: array of model parameters

model(x, p) will accept the full data set as the first argument x. This means that we need to write our model function so it applies the model to the full dataset. We use @. to apply (“broadcast”) the calculations across all rows.

@. model1(x, p) = survival(MortalityTables.Weibull(; m=p[1], σ=p[2]), x)

model1 (generic function with 1 method)

Fitting the Model

And fit the model with LsqFit.jl:

fit1 = curve_fit(model1, data.times, data.survival, [1.0, 1.0])

lines!(ax, data.times, model1(data.times, fit1.param);
    label="fitted Weibull", color=CASSETTE_PALETTE[2], linewidth=2, linestyle=:dash)
axislegend(ax, position=:rt)
plt

Observed survival proportions and the fitted Weibull survival curve.

Maximum Likelihood estimation

Generate 100 sample datapoints:

# MortalityTables.Weibull(m, σ) has S(x) = exp(-(x/m)^(m/σ)), which in
# Distributions.jl parameterization is Weibull(shape = m/σ, scale = m)
t = rand(Weibull(fit1.param[1] / fit1.param[2], fit1.param[1]), 100)

100-element Vector{Float64}:
  4.399979544008741
  4.693808880233822
 12.318449825867855
  5.394392475814005
  3.930616637905562
  7.324515361499734
  6.0211495938933
  8.212655922356479
  5.398605247730258
  4.224732200821012
  ⋮
  8.339674211714042
 10.838251985737328
  4.364910403782084
  4.42191483320106
 12.03459201049906
  5.062339613254776
  4.201038541988235
  3.036281099171572
  9.709501453759087

Without Censored Data”

#No censored data
fit_mle(Weibull, t)

Distributions.Weibull{Float64}(α=2.9146414287544875, θ=7.4731035169321975)

With Censored Data

Pick some arbitrary observations to censor:

c = collect(trues(100))
c[[1, 3, 7, 9]] .= false

4-element view(::Vector{Bool}, [1, 3, 7, 9]) with eltype Bool:
 0
 0
 0
 0

#ML function
survmle(x) = begin
    ml = 0.0
    for i = 1:length(t)
        if c[i]
            ml += logpdf(Weibull(x[2], x[1]), t[i]) #if not censored log(f(x))
        else
            ml += logccdf(Weibull(x[2], x[1]), t[i]) #if censored log(1-F)
        end
    end
    -ml
end

opt = Optim.optimize(
    survmle,          # function to optimize
    [1.0, 1.0], # lower bound
    [15.0, 15.0],            # upper bound
    [3.0, 3.0]          # initial guess
)

 * Status: success

 * Candidate solution
    Final objective value:     2.270703e+02

 * Found with
    Algorithm:     Fminbox with L-BFGS

 * Convergence measures
    |x - x'|               = 6.27e-08 ≰ 0.0e+00
    |x - x'|/|x'|          = 8.28e-09 ≰ 0.0e+00
    |f(x) - f(x')|         = 1.42e-13 ≰ 0.0e+00
    |f(x) - f(x')|/|f(x')| = 6.26e-16 ≰ 0.0e+00
    |g(x)|                 = 1.62e-09 ≤ 1.0e-08

 * Work counters
    Seconds run:   1  (vs limit Inf)
    Iterations:    4
    f(x) calls:    35
    ∇f(x) calls:   35
    ∇f(x)ᵀv calls: 0

The solution converges to similar values as the function generating the synthetic data:

Optim.minimizer(opt)

2-element Vector{Float64}:
 7.5744579570073265
 2.9027926419820798

Fitting Kaplan Meier

KaplanMeier comes from Survival.jl.

#t- time vector;c - censored events vector
km = fit(Survival.KaplanMeier, t, c)

plt2 = Figure(); ax2 = Axis(plt2[1,1]; xlabel="time", ylabel="survival")
stairs!(ax2, km.events.time, km.survival;
    label="Empirical (Kaplan–Meier)", color=CASSETTE_PALETTE[1], step=:post, linewidth=2)
plt2

Empirical Kaplan-Meier survival estimates over time.

@. model(x, p) = survival(MortalityTables.Weibull(; m=p[1], σ=p[2]), x)

mfit = LsqFit.curve_fit(model, km.events.time, km.survival, [2.0, 2.0])

lines!(ax2, km.events.time, model(km.events.time, mfit.param);
    label="Theoretical Weibull", color=CASSETTE_PALETTE[2], linewidth=2, linestyle=:dash)
axislegend(ax2, position=:rt)
plt2

Kaplan-Meier survival estimates compared with the fitted Weibull curve.