Implement LeftCensoredDistribution and RightCensoredDistribution

[vanAmsterdam]

Implements LeftCensoredDistribution and RightCensoredDistribution

These distributions take any base_dist with a cdf method implemented and allow for handling left-or-right censored data. See Issue #1909.

Current status

• implementation
• standard tests
• censored distribution specific tests

[vanAmsterdam]

I've removed the tutorial / notebook from this PR to chop up the work

[vanAmsterdam]

anyone up for a review? @juanitorduz @fehiepsi ?

[juanitorduz]

Great! I will take a look! In the meantime I think we need to add these distribution into the docs: https://github.com/pyro-ppl/numpyro/blob/master/docs/source/distributions.rst (you can use make docs to generate the docs locally)

Also, I think @fehiepsi 's review is key ;)

[juanitorduz]

Thank you @vanAmsterdam ! This looks great! Here are some initial comments :)

Question:
This PR handles left and right censoring. Are there plans for interval censoring (where we know the event occurred in an interval [a, b])? Or can we add a uni-test so see if we can do this already from this implementation?

[vanAmsterdam]

Question: This PR handles left and right censoring. Are there plans for interval censoring (where we know the event occurred in an interval [a, b])? Or can we add a uni-test so see if we can do this already from this implementation?

I've code for this on a branch that's now out of sync (see https://github.com/vanAmsterdam/numpyro/blob/95fd78dee8fe794d6e59517ad4e3c09804fcdc5a/numpyro/distributions/censored.py#L220)
I wanted to split the work up in 2 PRs for easier review, but happy to go either way
the interval censoring needs some additional work, especially on tests

[juanitorduz]

Question: This PR handles left and right censoring. Are there plans for interval censoring (where we know the event occurred in an interval [a, b])? Or can we add a uni-test so see if we can do this already from this implementation?

I've code for this on a branch that's now out of sync (see https://github.com/vanAmsterdam/numpyro/blob/95fd78dee8fe794d6e59517ad4e3c09804fcdc5a/numpyro/distributions/censored.py#L220) I wanted to split the work up in 2 PRs for easier review, but happy to go either way the interval censoring needs some additional work, especially on tests

perfect! makes total sense 🙌

[juanitorduz]

@vanAmsterdam I created a PR to your branch to try to fix the docstrings (it now works fine for me) vanAmsterdam#2

[juanitorduz]

You can try ruining make doctest locally to make the debugging easier :)

[vanAmsterdam]

You can try ruining make doctest locally to make the debugging easier :)

yes figured this one out a little too late :) but it's passing now so I think we're set for a review. thanks for the help!

[juanitorduz]

@fehiepsi I think this one is ready for your review :)

[fehiepsi]

LGTM, thanks @vanAmsterdam and @juanitorduz!

[fehiepsi]

Actually, I have some minor comments.

[fehiepsi]

I think you need self.base_dict.expand(self.batch_shape).sample(...).

[fehiepsi]

You'll need self.base_dict.expand(self.batch_shape).sample(...). promote_shapes does not broadcast the parameters, it only adds some missing singleton dimensions.

[vanAmsterdam]

good catch; updated in newest commit and also added a test