Flepi 3.0: Generalizing the FOI of n-dimensional compartmental models.

[twallema]

After bouncing of an interesting chat with @pearsonca about flepimop's future, I'm proud to say I finally understand the force-of-infection in compartmental models with $n$ strata, where each strata has a 2D within-dimension interaction matrix and across-dimension broadcasts. I hope it may serve as a mathematical and computational blueprint for FlepiMoP 3.0 and spark discussion, and I hope you enjoy this read.

Force of infection in $n$ dimensions

For a model with $n$ strata labeled $k: 1 \to n$, where every strata's interactions are captured by a 2D interaction matrix $N^k_{i_k j_k}$, which can be broadcast across any combination of the remaining $n-1$ dimensions by means of a broadcasting tensor $B^k$ (dimensionality can range from 0 to $n-1$), to obtain an $n+1$ dimensional interaction tensor $N^k_{i_1 j_1 i_2 \cdots i_n}$, a unified contact tensor $C_{i_1 j_1, \cdots, i_n, j_n}$ can be defined as follows,

$$ C_{i_1 j_1, \cdots, i_n, j_n} = N^1_{i_1 j_1 i_2 \cdots i_n} \cdot N^2_{i_2 j_2 i_1 i_3 \cdots i_n} \cdots N^n_{i_n j_n i_1 \cdots i_{n-1}}, $$

and,

$$ N^k_{i_1 j_1 i_2 \cdots i_n} = N^k_{ij} \cdot B^k_{{(i_l)_{l=1, l\neq k}^n}} $$

Because the position of the broadcasting tensor in the expanded product does not matter, the unified contact tensor can be written as,

$$ C_{i_1 j_1, \cdots, i_n, j_n} = N^1_{i_1 j_1} \cdot N^2_{i_2 j_2} \cdots N^n_{i_n j_n} $$

$$ \cdot B^1_{{(i_k)_{k=2}^n}} $$

$$ \cdot B^2_{{(i_k)_{k=1,k\neq2}^n}} $$

$$ \cdot \cdots \cdot B^n_{{(i_k)_{k=0}^{n-1}}} $$

For some reason these three terms above don't want to fit on one line but use your imagination. In this equation the terms $N$ represent the within-dimension interactions, and the terms $B$ represent the across-dimension broadcasts of $N$. Using the unified contact tensor the force of infection can be neatly written as,

$$ \lambda_{i_1 i_2 \cdots i_n} = \beta \sum\limits_{j_1 j_2 \cdots j_n} C_{i_1 j_1, \cdots, i_n, j_n} \dfrac{I_{j_1 j_2 \cdots j_n}}{T_{j_1 j_2 \cdots j_n}}, $$

the coolest thing is how easily this can be implemented in the computer using Einstein summation (for the sake of brevity omitting any broadcasting B),

C = np.einsum('i1j1, i2j2, ..., injn->i1j1i2j2...injn', N1, N2, ..., Nn)
l = beta * np.einsum('i1j1i2j2...injn, j1j2...jn -> i1i2...in', C, I/T)

Using $\mathbf{s}$ and $\mathbf{r}$ to group the multi-dimensional indices representing the different stratifications,

$$ \mathbf{\lambda_s = \beta \sum_{r} C_{s,r} \dfrac{I_r}{N_r}}, $$

which is consistent with the one-dimensional specification commonly found in textbooks (see f.i. Keeling and Rohani p.238).

Two-dimensional example: age and space stratification

Assuming there are no broadcasts of the social contact matrix $N_{ab}$ and the mobility matrix $M_{cd}$. The unified contact tensor can be constructed to accommodate individuals having a fraction $f_v$ of their contacts on their visited patch, while having $(1-f_v)$ contacts on their home patch (a commonly used way to incorporate temporary human movement as a driver of infectious disease; see f.i. Keeling & Rohani section 7.2.1.4),

$$ C_{ab, cd} = (1-f_v) N_{ab} \cdot \delta_{cd} + f_v N_{ab} \cdot M_{cd}, $$

by exploiting Kronecker delta (equivalent to identity matrix for this case) as a "mobility matrix of people staying". If $f_v = 1$, individuals make all their contacts on visited patches. The force of infection for an SIR model using the unified contact tensor is,

$$ \lambda_{ac} = \beta \sum\limits_{b=0}^N \sum\limits_{d=0}^G C_{ab,cd} \dfrac{I_{bd}}{T_{bd}}. $$

Which is coded up neatly using np.einsum,

C =  ((1 - f_v) * np.einsum('ab,cd->abcd', N, np.eye(M.shape[0])) + f_v * np.einsum('ab,cd->abcd', N, M))
l = beta * np.einsum('abcd, bd->ac', C, I/T)

and yields the same results as the reference model used in the other discussion, for an epidemic seeded in municipality Aarlen,

Adding broadcasting

To incorporate the scenario where the mobility matrix is different for individuals of different ages in both the home patch and visited patch, we introduce a 1D vector $\alpha_a$ that scales the mobility matrix for each age group $a$. The unified contact tensor for this case becomes,

$$ C_{ab, cd} = (1-f_v) N_{ab} \cdot \delta_{cd} + f_v N_{ab} \cdot (\alpha_a \cdot M_{cd}), $$

and is again coded up rather neatly using np.einsum,

C =  ((1 - f_v) * np.einsum('ab,cd->abcd', N, np.eye(M.shape[0])) + f_v * np.einsum('ab,a, cd->abcd', N, np.ones(N.shape[0]), M))

whose implementation as $\alpha_a=1$ does not alter the simulation, while for $\alpha_a=0$ there is no mobility. The position of $\alpha_a$ in the Einstein summation does not alter the results, further highlighting the syntax' flexibility,

C =  ((1 - f_v) * np.einsum('ab,cd->abcd', N, np.eye(M.shape[0])) + f_v * np.einsum('ab,cd,a->abcd', N, M, np.ones(N.shape[0])))

is equivalent to,

C =  ((1 - f_v) * np.einsum('ab,cd->abcd', N, np.eye(M.shape[0])) + f_v * np.einsum('a,ab,cd->abcd', np.ones(N.shape[0]), N, M))

Three dimensions

Extending this age- and space stratified model (without above broadcast) with an additional stratification representing race, with racial segregation modeled using a two-dimensional racial segregation matrix $\mathbf{R}$ is straightforward (have to computationally test). Mathematically, the unified contact tensor becomes,

$$ C_{ab, cd, ef} = (1-f_v) N_{ab} \cdot \delta_{cd} \cdot R_{ef} + f_v N_{ab} \cdot M_{cd} \cdot R_{ef}, $$

and the force of infection is,

$$ \lambda_{ace} = \beta \sum\limits_{bdf} C_{ab, cd, ef} \dfrac{I_{bdf}}{T_{bdf}}. $$

Homogeneous mixing among races can implemented by making $\mathbf{R}$ a matrix full of ones, while perfect racial segregation is implemented using the identity matrix. Another example highlighted by Carl is implementing vaccines in this way, where the interaction matrix $\mathbf{V}$ could be used to represent clustering of people who do and don't get vaccinated.

Computational complexity

The above simulation has 3 disease states, 581 spatial patches, 4 age groups resulting in 6972 compartments, simulated over 120 days using scipy's RK45 using a 2024 MacBook Pro M3 Max.

• Using np.einsum the simulation takes 3.6 s and 4.0 s without and with broadcasting respectively.
• As we are now using tensors, tensorflow's einsum function is more efficient, taking 1.2 s on the CPU.
• The use of tensors may open the way to simulations of very large systems on the GPU.
• Simulating a similar number of compartments using FlepiMoP took 28.0 s, resulting in a 25x speed up on a CPU.

Implications

• The combination of unified contact tensors and Einstein summation provide a mathematical and computational blueprint for building compartmental epidemiological models with $n$ strata.
• The framework can be implemented generically, users should be able to provide high-level instructions used by a software to construct a unified contact tensor. This is then translated in the appropriate string representation of the np.einsum function.
• I've been browsing literature and the 1D case is in most textbooks, there's plenty of papers on 2D models (including one I wrote myself), however, I so far have not found a good generalization to $n$ dimensions that is not "in too deep" for a practical audience (that I belong to myself).
• I thus think this would make for a good "Applied Mathematics and Computation" paper if pitched appropriately at an applied readership.

[pearsonca]

Quick note re:

I've been browsing literature and the 1D case is in most textbooks, there's plenty of papers on 2D models (including one I wrote myself), however, I so far have not found a good generalization to $n$ dimensions.

I think there are likely some opportunities here to show that systems like SIR+V can be approached instead in terms of stratification on intervention status.

[twallema]

I might have done this for COVID-19 by modeling the vaccine status as a stratification: https://www.sciencedirect.com/science/article/pii/S0307904X23002810.

[twallema]

I like the idea of using vaccines as a test case. We could demonstrate how an SIR might work for age and space, and then how extending the model with vaccines can be done using a stratification? Perhaps benchmark computational speeds of adding V as a disease state vs. adding it as a dimension?

[twallema]

Meeting 29/08/2024 - Carl & Tijs

@pearsonca

Latest version generalization FOI: FlepiMoP_3_0.pdf

Powerpoint with meeting summary: n_dimensional_compartmental_models.pptx

[twallema]

@alsnhll: expressive languages for chemical reaction networks are very similar to epidemics: https://kappalanguage.org