Compartmental models

[pawel-czyz]

@dr-david We discussed using compartmental models a few times, I've decided to write it down here, so I don't forget the equations.

There are several compartmental models.
Two of them are described at the end of this issue.
Using Diffrax we can implement them and model how the fraction of particular variant changes over the time, i.e.,

$$y_k(t) = \frac{i_k(t)}{i_1(t) + \cdots + i_V(t)}$$

and see how well the logistic growth approximates this dynamics over the time.
Note that we can:

1. Solve the model in terms of the usual fractions of the population, and then calculate the proportions $y_k$ just from the data.
2. We can even try to fit the model in the Bayesian manner, backpropagating the gradients through and ODE with Diffrax. Note that, however, this model is not identifiable.

SIR model

In this model, we have functions $s$ and $r$ representing susceptible and recovered (immune) fraction of the population, as well as fraction of population infected with variant $k$, $i_k$, for $k=1, \dotsc, V$.

$$s' = -s( \beta_1 i_1 + \cdots + \beta_V i_V )$$ $$i'_k = \beta_k i_k s - \gamma_k i_k $$ $$r' = \gamma_1 i_1 + \cdots + \gamma_V i_V$$

SIS model

In this model, we do not have the recovered population.

$$s' = -s( \beta_1 i_1 + \cdots + \beta_V i_V ) + (\gamma_1 i_1 + \cdots + \gamma_V i_V)$$ $$i'_k = \beta_k i_k s - \gamma_k i_k $$

Tasks

• Implement the SIR model in Diffrax.
• Implement the SIS model in Diffrax.
• In a notebook (or Snakemake), on simulated data: fit the basic logistic growth model.
• Fit the compartmental models using NumPyro.

[pawel-czyz]

These models are, in principle, not compatible with the current logistic growth due to the used theta, which is a single array. This will require defining more flexible growth models. One possible approach to this is in #20.