add branching flow example to epidemiology models

[ibulu]

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]

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]

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 ([email protected]). Main reasons why I'd like to implement branching are:

• to account for asymptomatic transmission (very relevant for covid)
• add distributed-delay to the deaths compartment (I believe this is called Erlang distributed models)
• and finally, model R0 by a neural network using mobility data...hence the need for separate D compartment and distributed-delay. Mobility effects I and its effect is observed later on deaths.

[fritzo]

@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]

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

[fritzo]

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]

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