2.24. Custom PDE class: SIR model

This example implements a spatially coupled SIR model with the following dynamics for the density of susceptible, infected, and recovered individuals:

\[\begin{split}\partial_t s &= D \nabla^2 s - \beta is \\ \partial_t i &= D \nabla^2 i + \beta is - \gamma i \\ \partial_t r &= D \nabla^2 r + \gamma i\end{split}\]

Here, \(D\) is the diffusivity, \(\beta\) the infection rate, and \(\gamma\) the recovery rate.

Time: 50.01, Susceptible, Infected, Recovered
  0%|          | 0/50.0 [00:00<?, ?it/s]
Initializing:   0%|          | 0/50.0 [00:00<?, ?it/s]
  0%|          | 0/50.0 [00:00<?, ?it/s]
  0%|          | 0.02/50.0 [00:00<37:48, 45.38s/it]
  0%|          | 0.04/50.0 [00:00<18:59, 22.80s/it]
  1%|          | 0.38/50.0 [00:00<02:07,  2.57s/it]
  3%|3         | 1.7/50.0 [00:01<00:34,  1.38it/s]
  9%|8         | 4.33/50.0 [00:01<00:18,  2.51it/s]
 16%|#6        | 8.06/50.0 [00:02<00:12,  3.31it/s]
 25%|##4       | 12.48/50.0 [00:03<00:10,  3.55it/s]
 33%|###3      | 16.74/50.0 [00:04<00:08,  3.87it/s]
 43%|####2     | 21.46/50.0 [00:05<00:07,  3.93it/s]
 52%|#####1    | 25.89/50.0 [00:06<00:05,  4.10it/s]
 61%|######1   | 30.71/50.0 [00:07<00:04,  4.11it/s]
 70%|#######   | 35.19/50.0 [00:08<00:03,  4.23it/s]
 80%|########  | 40.03/50.0 [00:09<00:02,  4.22it/s]
 89%|########9 | 44.54/50.0 [00:10<00:01,  4.31it/s]
 99%|#########8| 49.4/50.0 [00:11<00:00,  4.39it/s]
 99%|#########8| 49.4/50.0 [00:11<00:00,  4.25it/s]
100%|##########| 50.0/50.0 [00:11<00:00,  4.30it/s]
100%|##########| 50.0/50.0 [00:11<00:00,  4.30it/s]

from pde import FieldCollection, PDEBase, PlotTracker, ScalarField, UnitGrid


class SIRPDE(PDEBase):
    """SIR-model with diffusive mobility"""

    def __init__(
        self, beta=0.3, gamma=0.9, diffusivity=0.1, bc="auto_periodic_neumann"
    ):
        super().__init__()
        self.beta = beta  # transmission rate
        self.gamma = gamma  # recovery rate
        self.diffusivity = diffusivity  # spatial mobility
        self.bc = bc  # boundary condition

    def get_state(self, s, i):
        """generate a suitable initial state"""
        norm = (s + i).data.max()  # maximal density
        if norm > 1:
            s /= norm
            i /= norm
        s.label = "Susceptible"
        i.label = "Infected"

        # create recovered field
        r = ScalarField(s.grid, data=1 - s - i, label="Recovered")
        return FieldCollection([s, i, r])

    def evolution_rate(self, state, t=0):
        s, i, r = state
        diff = self.diffusivity
        ds_dt = diff * s.laplace(self.bc) - self.beta * i * s
        di_dt = diff * i.laplace(self.bc) + self.beta * i * s - self.gamma * i
        dr_dt = diff * r.laplace(self.bc) + self.gamma * i
        return FieldCollection([ds_dt, di_dt, dr_dt])


eq = SIRPDE(beta=2, gamma=0.1)

# initialize state
grid = UnitGrid([32, 32])
s = ScalarField(grid, 1)
i = ScalarField(grid, 0)
i.data[0, 0] = 1
state = eq.get_state(s, i)

# simulate the pde
tracker = PlotTracker(interval=10, plot_args={"vmin": 0, "vmax": 1})
sol = eq.solve(state, t_range=50, dt=1e-2, tracker=["progress", tracker])

Total running time of the script: ( 0 minutes 11.972 seconds)