PageSourceSearch

https://sidravi1.github.io/blog/2021/10/11/cox-proportional-hazard-model

html sidravi1.github.io collected 2026-10-03 09:57:02 UTC 19,996 bytes, 248 lines download raw bytes

1<!DOCTYPE html>
2<html lang="en">
3<head>
4<!-- Global site tag (gtag.js) - Google Analytics -->
vendor: 64 bytes, lines 4-5
4
5<script async src="https://www.googletagmanager.com/gtag/js?id=
5UA-157344899-1
vendor: 12 bytes, line 5
5"></script>
6<script>
7  
vendor: 134 bytes, lines 7-11
7window.dataLayer = window.dataLayer || [];
8  function gtag(){dataLayer.push(arguments);}
9  gtag('js', new Date());
10
11  gtag('config', '
11UA-157344899-1
vendor: 4 bytes, line 11
11');
12</script>
12
13
14
15<meta charset="UTF-8">
16<meta name="viewport" content="width=device-width, initial-scale=1.0">
17<meta http-equiv="X-UA-Compatible" content="ie=edge">
18<!-- For altair plots -->
19<script src="https://cdn.jsdelivr.net/npm/vega@3"></script>
vendor: 1 bytes, line 19
19
20<script src="https://cdn.jsdelivr.net/npm/vega-lite@2"></script>
vendor: 1 bytes, line 20
20
21<script src="https://cdn.jsdelivr.net/npm/vega-embed@3"></script>
21
22<link href="https://fonts.googleapis.com/css?family=Nunito+Sans:400,400i,700&display=swap" rel="stylesheet">
23<link rel="stylesheet" href="/assets/css/style.css">
24<title>[Survival models] Cox proportional hazard model</title>
25<!-- Begin Jekyll SEO tag v2.8.0 -->
26<title>[Survival models] Cox proportional hazard model | Sid Ravinutala</title>
27<meta name="generator" content="Jekyll v3.10.0" />
28<meta property="og:title" content="[Survival models] Cox proportional hazard model" />
29<meta name="author" content="Sid Ravinutala" />
30<meta property="og:locale" content="en_US" />
31<meta name="description" content="In the previous blogpost, we talked about the hazard function - the probability that the event happens at time $t$ given that you it hasn’t till time $t$. In the Cox proportional hazard model, we assume the hazard function is:" />
32<meta property="og:description" content="In the previous blogpost, we talked about the hazard function - the probability that the event happens at time $t$ given that you it hasn’t till time $t$. In the Cox proportional hazard model, we assume the hazard function is:" />
33<link rel="canonical" href="https://sidravi1.github.io/blog/2021/10/11/cox-proportional-hazard-model" />
34<meta property="og:url" content="https://sidravi1.github.io/blog/2021/10/11/cox-proportional-hazard-model" />
35<meta property="og:site_name" content="Sid Ravinutala" />
36<meta property="og:type" content="article" />
37<meta property="article:published_time" content="2021-10-11T10:42:00+00:00" />
38<meta name="twitter:card" content="summary" />
39<meta property="twitter:title" content="[Survival models] Cox proportional hazard model" />
40<script type="application/ld+json">
41{"@context":"https://schema.org","@type":"BlogPosting","author":{"@type":"Person","name":"Sid Ravinutala"},"dateModified":"2021-10-11T10:42:00+00:00","datePublished":"2021-10-11T10:42:00+00:00","description":"In the previous blogpost, we talked about the hazard function - the probability that the event happens at time $t$ given that you it hasn’t till time $t$. In the Cox proportional hazard model, we assume the hazard function is:","headline":"[Survival models] Cox proportional hazard model","mainEntityOfPage":{"@type":"WebPage","@id":"https://sidravi1.github.io/blog/2021/10/11/cox-proportional-hazard-model"},"url":"https://sidravi1.github.io/blog/2021/10/11/cox-proportional-hazard-model"}</script>
41
42<!-- End Jekyll SEO tag -->
43
44<script>
45MathJax = {
46  tex: {
47    inlineMath: [['$', '$'], ['\\(', '\\)']],
48    displayMath: [['$$', '$$'], ['\\[', '\\]']]
49  },
50  svg: {
51    fontCache: 'global'
52  },
53  tags: 'ams',
54  processEscapes: true,
55  options: {
56    renderActions: {
57      /* add a new named action not to override the original 'find' action */
58      find_script_mathtex: [10, function (doc) {
59        for (const node of document.querySelectorAll('script[type^="math/tex"]')) {
60          const display = !!node.type.match(/; *mode=display/);
61          const math = new doc.options.MathItem(node.textContent, doc.inputJax[0], display);
62          const text = document.createTextNode('');
63          node.parentNode.replaceChild(text, node);
64          math.start = {node: text, delim: '', n: 0};
65          math.end = {node: text, delim: '', n: 0};
66          doc.math.push(math);
67        }
68      }, '']
69    }
70  }
71};
72</script>
vendor: 1 bytes, line 72
72
73<script type="text/javascript" id="MathJax-script" async
74  src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js">
75</script>
75
76</head>
77<body>
78  <header class="texture-red">
79    <div class="container"><div class="navbar">
80	<ul>
81        <a href="/about.html"><li>About</li></a>
82        <a href="/resume.html"><li>Resume</li></a>
83        <a href="/blog"><li>Blog</li></a>
84	</ul>
85</div></div><div class="container">
86	<h1>[Survival models] Cox proportional hazard model</h1>
87	<h4 class="post-description"></h4>
88	<div class="post-date" style="margin-top:20px">
89		Published on Oct 11, 2021
90	</div>
91	<ul class="post-tags"></ul>
92</div>
93</header>
94  <main>
95    <div class="container">
96      <div class="post-container">
97          <p>In the <a href="/blog/2021/10/01/censored-data-and-intro-to-survival-models">previous blogpost</a>, we talked about the <em>hazard</em> function - the probability that the event happens at time $t$ given that you it hasn’t till time $t$. In the Cox proportional hazard model, we assume the hazard function is:</p>
98
99\[h(t) = h_0(t) \times exp(b_1x_1 + b_2x_2 + ... + b_px_p)\]
100
101<p>Some interesting things to point out here:</p>
102<ul>
103  <li>time $t$ only occurs in $h_0(t)$</li>
104  <li>covariates, $x_i$ only occur in the $exp$ bit</li>
105</ul>
106
107<p>So there is some base hazard rate $h_0(t)$ and the only way it varies for units with different characteristics is through the $exp$ bit. So why is it called “proportional” hazard model?</p>
108
109<p>Imagine a simple model where:</p>
110
111\[h(t) = h_0(t) \times exp(b_1x_1)\]
112
113<p>And we have a unit $i$ with $x_{i1} = z$ and another unit $j$ with $x_{ji} = a \cdot z$ so let’s see the hazard for each:</p>
114
115\[h_i(t) = h_0(t) \times exp(b_1z)\]
116
117<p>And</p>
118
119\[h_j(t) = h_0(t) \times exp(b_1 az)\]
120
121<p>So the hazard ratio of the two is:</p>
122
123\[\begin{aligned}
124\frac{h_i(t)}{h_j(t)} &amp;= \frac{h_0(t) exp(b_1 az)}{h_0(t) exp(b_1 z)} \\
125&amp;= \frac{exp(b_1 az)}{exp(b_1 z)} \\
126&amp;= exp(b_1 (a - 1)z)
127\end{aligned}\]
128
129<p>So no matter what the $t$ is, the hazard ration between these two is always some constant. Another way is to note that:</p>
130
131\[\begin{aligned}
132h_i(t) &amp;= h_j(t) \times exp(b_1 (a -1 )z)\\
133h_i(t) &amp;= c h_j(t)
134\end{aligned}\]
135
136<p>where $c = exp(b_1 (a -1 )z)$. So $h_i(t)$ is proportional to $h_j(t)$ - times does not make an appearance in the coefficient. This assumption may not be true and there are a bunch of tests you can do to check it.</p>
137
138<p>Alright. Enough with the theory.</p>
139
140<h2 id="partial-likelihood">Partial likelihood</h2>
141
142<p>I’m not going to derive it. The internet is your friend.</p>
143
144\[PL(\beta | D) = \prod_{j=1}^{d} \frac{\exp(x_j'\beta)}{\sum_{l \in R_j} \exp(x'_l\beta)}\]
145
146<p>where $R_j$ is the <em>risk set</em> - people at risk of death/event when the event for $j$ occurs. This should look familiar - it’s softmax. The intuition follows from there. So we want to pick the betas that maximise the likelihood of the event happening out of all the units at risk.</p>
147
148<p>Log-likelihood just needs some $algebra$:</p>
149
150\[pll(\beta | D) = \sum_{j=1}^{d} \left( x_j'\beta - \log\left(\sum_{l \in R_j} \exp(x'_l\beta)\right)\right)\]
151
152<p>Now all we need to code it up with <code class="language-plaintext highlighter-rouge">jax</code> for some sweet sweet autograd.</p>
153
154<h2 id="minimize-negative-log-likelihood">Minimize negative log likelihood</h2>
155
156<figure class="highlight"><pre><code class="language-python" data-lang="python"><span class="kn">import</span> <span class="nn">jax.scipy</span> <span class="k">as</span> <span class="n">jsp</span>
157<span class="kn">import</span> <span class="nn">jax</span>
158
159<span class="o">@</span><span class="n">jax</span><span class="p">.</span><span class="n">jit</span>
160<span class="k">def</span> <span class="nf">neglogp</span><span class="p">(</span><span class="n">betas</span><span class="p">,</span> <span class="n">riskset</span><span class="o">=</span><span class="n">riskset</span><span class="p">,</span> <span class="n">observed</span><span class="o">=</span><span class="n">observed</span><span class="p">):</span>
161    <span class="n">betas_x</span> <span class="o">=</span> <span class="n">betas</span> <span class="o">@</span> <span class="n">x</span><span class="p">.</span><span class="n">T</span>
162    <span class="n">riskset_beta</span> <span class="o">=</span> <span class="n">betas_x</span> <span class="o">*</span> <span class="n">riskset</span>
163    <span class="n">ll_matrix</span> <span class="o">=</span> <span class="p">(</span><span class="n">betas_x</span> <span class="o">-</span> <span class="n">jsp</span><span class="p">.</span><span class="n">special</span><span class="p">.</span><span class="n">logsumexp</span><span class="p">(</span><span class="n">riskset_beta</span><span class="p">,</span> <span class="n">b</span><span class="o">=</span><span class="n">riskset</span><span class="p">,</span> <span class="n">axis</span><span class="o">=</span><span class="mi">1</span><span class="p">))</span>
164    <span class="k">return</span> <span class="o">-</span><span class="p">(</span><span class="n">observed</span> <span class="o">*</span> <span class="n">ll_matrix</span><span class="p">).</span><span class="nb">sum</span><span class="p">()</span></code></pre></figure>
165
166<p>If you know <code class="language-plaintext highlighter-rouge">numpy</code> and <code class="language-plaintext highlighter-rouge">scipy</code>, you know a lot of <code class="language-plaintext highlighter-rouge">jax</code>.</p>
167
168<p>P.S: Took me forever to realise that <code class="language-plaintext highlighter-rouge">logsumexp</code> take a weight argument as <code class="language-plaintext highlighter-rouge">b</code>. No idea why <code class="language-plaintext highlighter-rouge">statsmodels</code> doesn’t use this.</p>
169
170<p>Ok. Now all we need is to minimize it. Here comes the magic of autograd</p>
171
172<figure class="highlight"><pre><code class="language-python" data-lang="python"><span class="kn">from</span> <span class="nn">jax</span> <span class="kn">
172import</span> <span class="n">grad</span><span class="p">,</span> <span class="n">hessian</span>
173<span class="kn">from</span> <span class="nn">scipy.optimize</span> <span class="kn">import</span> <span class="n">minimize</span>
174
175<span class="n">dlike</span> <span class="o">=</span> <span class="n">grad</span><span class="p">(</span><span class="n">neglogp</span><span class="p">)</span>
176<span class="n">dlike2</span> <span class="o">=</span> <span class="n">hessian</span><span class="p">(</span><span class="n">neglogp</span><span class="p">)</span>
177<span class="n">res</span> <span class="o">=</span> <span class="n">minimize</span><span class="p">(</span><span class="n">neglogp</span><span class="p">,</span> <span class="n">np</span><span class="p">.</span><span class="n">ones</span><span class="p">(</span><span class="mi">2</span><span class="p">),</span> <span class="n">method</span><span class="o">=</span><span class="s">'Newton-CG'</span><span class="p">,</span> <span class="n">jac</span><span class="o">=</span><span class="n">dlike</span><span class="p">,</span> <span class="n">hess</span><span class="o">=</span><span class="n">dlike2</span><span class="p">)</span>
178<span class="n">res</span><span class="p">.</span><span class="n">x</span>
179
180<span class="o">&gt;&gt;&gt;</span> <span class="n">DeviceArray</span><span class="p">([</span><span class="o">-</span><span class="mf">0.14952305</span><span class="p">,</span> <span class="o">-</span><span class="mf">0.16527627</span><span class="p">],</span> <span class="n">dtype</span><span class="o">=</span><span class="n">
180float32</span><span class="p">)</span></code></pre></figure>
181
182<p>You can compare this to <code class="language-plaintext highlighter-rouge">statsmodels</code> or <code class="language-plaintext highlighter-rouge">lifelines</code> package. It’s in the notebook - you can check it out.</p>
183
184<h2 id="lets-get-bayesian">Let’s get Bayesian</h2>
185
186<p>Blackjax is a pretty sweet package with a lot of sampling methods. Let’s use the NUTS sampler (yes - i know, that’s calling it a “sampler sampler”). Let’s use it to get some samples. This is basically the example in the blackjax documentation.</p>
187
188<p>First we need the posterior probability which is simply the prior $\times$ the likelihood. In log terms:</p>
189
190<figure class="highlight"><pre><code class="language-python" data-lang="python"><span class="k">def</span> <span class="nf">logprior</span><span class="p">(</span><span class="n">betas</span><span class="p">):</span>
191    <span class="k">return</span> <span class="n">jsp</span><span class="p">.</span><span class="n">stats</span><span class="p">.</span><span class="n">norm</span><span class="p">.</span><span class="n">logpdf</span><span class="p">(</span><span class="n">betas</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mf">1.0</span><span class="p">).</span><span class="nb">sum</span><span class="p">()</span>
192
193<span class="n">logprob</span> <span class="o">=</span> <span class="k">lambda</span> <span class="n">x</span><span class="p">:</span> <span class="n">logprior</span><span class="p">(</span><span class="o">**</span><span class="n">x</span><span class="p">)</span> <span class="o">+</span> <span class="n">neglogp3</span><span class="p">(</span><span class="o">**</span><span class="n">x</span><span class="p">)</span></code></pre></figure>
194
195<p>Ok. Let’s setup the initial values and the kernel.</p>
196
197<figure class="highlight"><pre><code class="language-python" data-lang="python"><span class="n">initial_position</span> <span class="o">=</span> <span class="p">{</span><span class="s">"betas"</span><span class="p">:</span> <span class="n">jnp</span><span class="p">.</span><span class="n">ones</span><span class="p">(</span><span class="mi">2</span><span class="p">)}</span>
198<span class="n">initial_state</span> <span class="o">=</span> <span class="n">nuts</span><span class="p">.</span><span class="n">new_state</span><span class="p">(</span><span class="n">initial_position</span><span class="p">,</span> <span class="n">logprob</span><span class="p">)</span>
199
200<span class="n">inv_mass_matrix</span> <span class="o">=</span> <span class="n">np</span><span class="p">.</span><span class="n">array</span><span class="p">([</span><span class="mf">0.5</span><span class="p">,</span> <span class="mf">0.5</span><span class="p">])</span>
201<span class="n">step_size</span> <span class="o">=</span> <span class="mf">1e-3</span>
202
203<span class="n">nuts_kernel</span> <span class="o">=</span> <span class="n">jax</span><span class="p">.</span><span class="n">jit</span><span class="p">(</span><span class="n">nuts</span><span class="p">.</span><span class="n">kernel</span><span class="p">(</span><span class="n">logprob</span><span class="p">,</span> <span class="n">step_size</span><span class="p">,</span> <span class="n">inv_mass_matrix</span><span class="p">))</span></code></pre></figure>
204
205<p>and the inference loop for samples:</p>
206
207<figure class="highlight"><pre><code class="language-python" data-lang="python"><span class="k">def</span> <span class="nf">inference_loop</span><span class="p">(</span><span class="n">rng_key</span><span class="p">,</span> <span class="n">kernel</span><span class="p">,</span> <span class="n">initial_state</span><span class="p">,</span> <span class="n">num_samples</span><span class="p">):</span>
208    <span class="k">def</span> <span class="nf">one_step</span><span class="p">(</span><span class="n">state</span><span class="p">,</span> <span class="n">rng_key</span><span class="p">):</span>
209        <span class="n">state</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">kernel</span><span class="p">(</span><span class="n">rng_key</span><span class="p">,</span> <span class="n">state</span><span class="p">)</span>
210        <span class="k">return</span> <span class="n">state</span><span class="p">,</span> <span class="n">state</span>
211
212    <span class="n">keys</span> <span class="o">=</span> <span class="n">jax</span><span class="p">.</span><span class="n">random</span><span class="p">.</span><span class="n">split</span><span class="p">(</span><span class="n">rng_key</span><span class="p">,</span> <span class="n">num_samples</span><span class="p">)</span>
213    <span class="n">_</span><span class="p">,</span> <span class="n">states</span> <span class="o">=</span> <span class="n">jax</span><span class="p">.</span><span class="n">lax</span><span class="p">.</span><span class="n">scan</span><span class="p">(</span><span class="n">one_step</span><span class="p">,</span> <span class="n">initial_state</span><span class="p">,</span> <span class="n">keys</span><span class="p">)</span>
214
215    <span class="k">return</span> <span class="n">states</span>
216
217    <span class="n">rng_key</span> <span class="o">=</span> <span class="n">jax</span><span class="p">.</span><span class="n">random</span><span class="p">.</span><span class="n">PRNGKey</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
218    <span class="n">states</span> <span class="o">=</span> <span class="n">inference_loop</span><span class="p">(</span><span class="n">rng_key</span><span class="p">,</span> <span class="n">nuts_kernel</span><span class="p">,</span> <span class="n">initial_state</span><span class="p">,</span> <span class="mi">1000</span><span class="p">)</span>
219
220    <span class="n">beta_samples</span> <span class="o">=</span> <span class="n">states</span><span class="p">.</span><span class="n">position</span><span class="p">[</span><span class="s">"betas"</span><span class="p">].</span><span class="n">block_until_ready</span><span class="p">()</span></code></pre></figure>
221
222<p>Now let’s get some samples:</p>
223
224<figure class="highlight"><pre><code class="language-python" data-lang="python"><span class="n">rng_key</span> <span class="o">=</span> <span class="n">jax</span><span class="p">.</span><span class="n">random</span><span class="p">.</span><span class="n">PRNGKey</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
225<span class="n">states</span> <span class="o">=</span> <span class="n">inference_loop</span><span class="p">(</span><span class="n">rng_key</span><span class="p">,</span> <span class="n">nuts_kernel</span><span class="p">,</span> <span class="n">initial_state</span><span class="p">,</span> <span class="mi">1000</span><span class="p">)</span>
226
227<span class="n">beta_samples</span> <span class="o">=</span> <span class="n">states</span><span class="p">.</span><span class="n">position</span><span class="p">[</span><span class="s">"betas"</span><span class="p">].</span><span class="n">block_until_ready</span><span class="p">()</span></code></pre></figure>
228
229<p>Here’s the trace from it:
230<img src="https://sidravi1.github.io/assets/20211012_trace.png" alt="trace" /></p>
231
232<p>And the posterior:</p>
233
234<p><img src="https://sidravi1.github.io/assets/20211012_posterior.png" alt="posterior" /></p>
235
236<p>And there you have it folks - some nice confidence intervals.</p>
237
238<h2 id="final-words">Final words</h2>
239
240<p>Ok - I cheated a little bit. I removed ties in the database. If you have ties you need a slight change to the likelihood - check out the <a href="https://github.com/statsmodels/statsmodels/blob/main/statsmodels/duration/hazard_regression.py">statsmodels</a> code for <code class="language-plaintext highlighter-rouge">efron_loglike</code> or <code class="language-plaintext highlighter-rouge">breslow_loglike</code>.</p>
241
242<p>The point of this was also to play around with <code class="language-plaintext highlighter-rouge">jax</code> and <code class="language-plaintext highlighter-rouge">blackjax</code>. Mission accomplished. Let’s look at some parametric models next time around.</p>
243
244      </div>
245
246        <!-- Configure Disqus --></div>
247  </main></body>
248</html>

Line numbers count LF bytes from the start of the resource, as the search results do. Vendor segments are library code the classifier recognised; they are stored but not indexed. Bytes are shown as Latin1 characters, one per byte.