.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "examples_gallery/advanced_pdes/pde_sir.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_examples_gallery_advanced_pdes_pde_sir.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_examples_gallery_advanced_pdes_pde_sir.py:


Custom PDE class: SIR model
===========================

This example implements a `spatially coupled SIR model
<https://en.wikipedia.org/wiki/Compartmental_models_in_epidemiology>`_ with
the following dynamics for the density of susceptible, infected, and recovered
individuals:

.. math::

    \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

Here, :math:`D` is the diffusivity, :math:`\beta` the infection rate, and
:math:`\gamma` the recovery rate.

.. GENERATED FROM PYTHON SOURCE LINES 19-69



.. image-sg:: /examples_gallery/advanced_pdes/images/sphx_glr_pde_sir_001.png
   :alt: Time: 50, Susceptible, Infected, Recovered
   :srcset: /examples_gallery/advanced_pdes/images/sphx_glr_pde_sir_001.png
   :class: sphx-glr-single-img


.. rst-class:: sphx-glr-script-out

 .. code-block:: none


      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<23:07, 27.76s/it]
      0%|          | 0.06/50.0 [00:00<07:45,  9.32s/it]
      1%|          | 0.62/50.0 [00:00<00:48,  1.01it/s]
      6%|▌         | 3.03/50.0 [00:00<00:13,  3.61it/s]
     16%|█▌        | 8.11/50.0 [00:01<00:06,  6.14it/s]
     31%|███       | 15.42/50.0 [00:02<00:04,  6.96it/s]
     46%|████▋     | 23.14/50.0 [00:03<00:03,  7.37it/s]
     62%|██████▏   | 31.19/50.0 [00:04<00:02,  7.62it/s]
     79%|███████▉  | 39.42/50.0 [00:04<00:01,  8.07it/s]
     97%|█████████▋| 48.69/50.0 [00:05<00:00,  8.14it/s]
     97%|█████████▋| 48.69/50.0 [00:06<00:00,  7.74it/s]
    100%|██████████| 50.0/50.0 [00:06<00:00,  7.95it/s] 
    100%|██████████| 50.0/50.0 [00:06<00:00,  7.95it/s]






|

.. code-block:: Python


    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(interrupts=10, plot_args={"vmin": 0, "vmax": 1})
    sol = eq.solve(state, t_range=50, dt=1e-2, tracker=["progress", tracker])


.. rst-class:: sphx-glr-timing

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


.. _sphx_glr_download_examples_gallery_advanced_pdes_pde_sir.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: pde_sir.ipynb <pde_sir.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: pde_sir.py <pde_sir.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: pde_sir.zip <pde_sir.zip>`