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.

[milesalanmoore]

@vanAmsterdam Any progress on this? I'm interested in participating if you're still working on this! I've been fitting (thresholding) survival models in Stan for my own work and agree this model class would be great in numpyro.

[vanAmsterdam]

Hi @milesalanmoore thanks for the interest; it took a while for me to get started on this;
I've begun with a minimal implementation here: https://github.com/vanAmsterdam/numpyro/tree/feature/survival
still need to add tests, documentation, examples (started a notebook for this)

in terms of implementation, it would make sense to add:

• RightCensoredDistribution
• LeftCensoredDistribution
• IntervalCensoredDistribution

It took a bit of thinking on how to implement this. A typical implementation (e.g. in R) would be to treat survival analysis as having multivariate observations (i.e. jointly observing times and events)
However, it seems to fit the numpyro format better if implemented when the event is like a mask / truncation point (but then the censoring equivalent, like MaskedDistribution and TruncatedDistribution), so I went with this instead (as I mentioned before)

happy to consider other options, and for you to contribute

[vanAmsterdam]

update: I've now implemented Right / Left / Interval and Doubly censored distributions; you can check the files changed here: vanAmsterdam#1

a couple of tests are still failing and the documentation needs updating

[vanAmsterdam]

Hi @juanitorduz, @milesalanmoore and @fehiepsi I've been making some progress and I think the implementation and api is getting to where I want it to be; API being:

# left/right censored
x = ...
censored = ... (bool)

base_dist = dist.Normal()
cens_dist = RightCensoredDistribution(base_dist, censored)
# or
cens_dist = LeftCensoredDistribution(base_dist, censored)

cens_dist.log_prob(x)

interval censored is somewhat more involved as you need an upper-bound and lower-bound in the observation array, so:

lower_bound = ... (float)
upper_bound = ... (float)
x = jnp.stack([lower_bound, upper_bound], axis=-1)

left_censored = ... (bool)
right_censored = ... (bool)

base_dist = dist.Normal()
cens_dist = IntervalCensoredDistribution(base_dist, left_censored, right_censored)
cens_dist.log_prob(x)

Some things tricky things (also wrt testing):

1. censored data comes in 2-dim (left/right) or 4-dim (interval) data, but we're really modeling a univariate base dist. We can always sample from this base_dist but the result is 1-dim. Due to the api construction this is no problem for left/right censored as the value supplied to log_prob is 1d. For interval censoring there is a mis-match between the dimensions for log_prob evaluations and random sampling; for intervalcensored distributions, obervation shapes are (num_obs, 2). This is a bit of a different view than taken e.g. by PyMC (https://www.pymc.io/projects/docs/en/stable/api/distributions/censored.html). Whereas they say censored distributions are discrete / continuous mixtures, I think I see it more as a way to do correct inference over a 1-d continous base_dist using the data correctly
2. what the PR implements pretty much implements what's done in the current notebooks/source/censoring.ipynb ; though that notebook is instructive, in terms of tutorial we may want to update this. I've made a notebook in the PR that demos all cases of censoring but it's more of a hybrid demo/test than a real instructive notebook (https://github.com/vanAmsterdam/numpyro/blob/feature/survival/notebooks/source/survival_analysis.ipynb); what's also a point is that technically I think the current censoring.ipynb is more of a truncation example as values have a hard cut-off at these limits, not random censoring values as what is typical in e.g. survival analysis
3. i'm also still thinking about some tests like the GOF tests and some censoring specific tests to add to test_distributions.py

My questions are:

1. do you have ideas on how to approach the tests?
2. what to do with the notebooks

progress is here: vanAmsterdam#1

[juanitorduz]

Hey! @vanAmsterdam Thanks for the update! Regarding your points:

1. Do you have ideas on how to approach the tests?

I think (a personal opinion) we can go from simple to complex and work on iterations. We might want to merge the simplest case and then incrementally add more complexity (dimensions?) This is how we developed the HSGP contribution module. It is also easier to review ;)

2. What to do with the notebooks.

I wrote the censoring.ipynb notebook, which I simply ported from my blog (https://juanitorduz.github.io/censoring/). Hence, you can simply replace it with your newer version and add the blog link if people want to go into more details (or not, if it would create more confusion).

We might also think reproducing the fantastic example https://www.pymc.io/projects/examples/en/latest/generalized_linear_models/GLM-truncated-censored-regression.html in NumPyro (and ensuring we do the proper citation of the PyMC original example).

[vanAmsterdam]

ok probably the easiest would be:

• PR 1: left/right censored; with this (and the new truncated cdf) we could reproduce the glm-truncated-censored-regression example
• PR 2: interval censoring;

not sure if it makes sense to have 2 notebooks on censoring / truncation? so maybe we go for reproducing the pymc example and also provide a pointer to your blog in that reproduction?