pyro-ppl / pyro

Deep universal probabilistic programming with Python and PyTorch
http://pyro.ai
Apache License 2.0
8.5k stars 982 forks source link

add branching flow example to epidemiology models #2552

Open ibulu opened 4 years ago

ibulu commented 4 years ago

It'd be quite helpful to add some examples for more complicated models with complicated flows such as SEIRD or SEIRDS, where, for example, E branches into ID (infectious individuals who'll die) and IR (infectious individuals who'll recover). It's very straightforward to write down the transition() function. However, it is not very obvious to write compute_flows()

fritzo commented 4 years ago

HI @ibulu I agree it would be helpful, and I'd like to have an example. However I haven't been able to figure out how to implement a .compute_flows() method for models with loops. The crux is that, in directed graph of state transitions, flows around (undirected) loops cannot be identified from mere time series of counts of each compartment.

As a simple example to aid our discussion, consider compartments S = susceptible, I = infected and symptomatic, A as infected but asymptomatic, and R as recovered, with transitions:

S ---> A
|      |
V      V
I ---> R

Consider four time series over three time steps

S: 9 8 7
I: 1 1 1
A: 0 1 1
R: 0 0 1

At the first time step, the flux is unambiguously S2A=1. However at the second time step it is unclear whether the flux was S2I=1,I2R=1 or S2A=1,A2R=1. This ambiguity corresponds to flow around the cycle S->I->R->A->S.

A simpler example is an SIS model.

One idea to implement .compute_flows() might be to treat those cyclic flows as latent variables, and to sample them (via pyro.sample) within the .compute_flows() method. I'm open to other ideas if you have any. I'd also be happy to zoom and brainstorm if you'd like.

Addresses #2426

ibulu commented 4 years ago

yeah. I'd be happy to brainstorm together on this. I am available anytime tomorrow from 10:30AM to 4:30PM central time for Zoom call if that works on your end (ibulu77@gmail.com). Main reasons why I'd like to implement branching are:

fritzo commented 4 years ago

@ibulu I sent an invite for 1pm Central tomorrow. I'd also like to hear more about your models, or read more if you have any preprints. (anyone else, feel free to ping me to attend the meeting)

ibulu commented 4 years ago

Thank you. RSVP'ed and invited few more people on our end.

fritzo commented 4 years ago

Hi @ibulu, I think your SEIRD model is actually easier than the looping models I described (and easier than your SEIRDS model). I believe you can simply define:

class SEIRDModel(CompartmentalModel):
    def __init__(self, ...):
        compartments = ("S", "E", "IR", "ID", "D")
        super().__init__(compartments, duration, population)

    def compute_flows(self, prev, curr, t):
        S2E = prev["S"] - curr["S"]
        ID2D = curr["D"] - prev["D"]
        E2ID = curr["ID"] - prev["ID"] + ID2D
        E2IR = prev["E"] - curr["E"] + S2E - E2ID
        IR2R = prev["IR"] - curr["IR"] + E2IR
        return {
            "S2E_{}".format(t): S2E,
            "ID2D_{}".format(t): ID2D,
            "E2ID_{}".format(t): E2ID,
            "E2IR_{}".format(t): E2IR,
            "IR2R_{}".format(t): IR2R,
        }

Let me know if you need any help getting that working, and I'll try to add an example when I get a chance.

ibulu commented 4 years ago

This is so cool! will combine this with boxcar model!