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)} &= \frac{h_0(t) exp(b_1 az)}{h_0(t) exp(b_1 z)} \\ 125&= \frac{exp(b_1 az)}{exp(b_1 z)} \\ 126&= 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) &= h_j(t) \times exp(b_1 (a -1 )z)\\ 133h_i(t) &= 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">>>></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.