survival analysis / time-to-event analysis in numpyro

[vanAmsterdam]

in my research domain (healthcare) many outcomes are right-censored (aka survival analysis / time-to-event); the data for an individual patient are basically a tuple of (observed_time, event / censoring indicator).

What would be a good way to incorporate this into numpyro? There are several distributions already implemented in the library that are often used for (bayesian) survival analysis (e.g. exponential, weibull), but calculating the likelihood correctly for censored data would require some extra work.

are others interested in adding support for survival analysis to numpyro?

an example implemented in problow is here: https://probflow.readthedocs.io/en/stable/examples/time_to_event.html

[juanitorduz]

@vanAmsterdam There is an example where this is done by hand https://num.pyro.ai/en/latest/tutorials/censoring.html (motivated by the code in https://kylejcaron.github.io/posts/censored_demand/2024-02-06-censored-demand.html) . Would that be enough? Or are you looking into something like the PyMC wrapper https://www.pymc.io/projects/docs/en/latest/api/distributions/censored.html ?

[vanAmsterdam]

thanks for these pointers! I'd be thinking more along the lines of the pymc implementation indeed. This would make it easier for people to use numpyro for survival analysis with different choices of distributions.

I'd be happy to try to port this pymc approach to numpyro if we think this is a worthwhile addition and agree on how it should be implemented

• based on the implementation, add a tutorial

[vanAmsterdam]

e.g. we might also think about it as an effect handler just like mask, but then a censored variant, and even potentially a truncated variant

so the interface could be something like

t = jnp.array([.72, 1.23, .45])
censored = jnp.array([True, False, True])

with handlers.censored(censored):
   sample('obs_t', dist.Normal(), obs=t)

btw stan docs also have some nice background information on this topic: https://mc-stan.org/docs/stan-users-guide/truncation-censoring.html#censored-data

[juanitorduz]

pymc approach to numpyro if we think this is a worthwhile

IMO, that would be great! This can enable users to use this more directly (also creating awareness about it).

But I guess @fehiepsi and his team are the ones to decide whether this is something they want :)

[fehiepsi]

Yeah,we welcome the contribution.