SEIRD: A Complex Contagion Model#

This is a simple model combining concepts and ideas from the well known SEIR (Susceptible-Exposed-Infected-Recovered) and SIRD (Susceptible-Infected-Recovered-Deceased) models, and adapting them to a spatial 2D grid model.

Fundamentals#

We model a population of agents on a two-dimensional grid of cells that can move randomly or away from infected agents. Each cell can be in one of the following states:

  • empty: there is no agent on the cell.

  • susceptible: the agent on the cell is healthy but susceptible to the disease and can be exposed if in contact with the disease.

  • exposed: the agent on the cell is exposed to the disease, meaning that it can already infect other agents in contact with the disease; however, there are no symptoms yet, thus, the agent gets noticed as infected only after the incubation period.

  • infected: the agent on the cell is infected and can infect neighboring agents.

  • recovered: the agent on the cell is recovered, thus, it is immune to the disease. However, it can lose its immunity after a while and become susceptible again.

  • deceased: the agent on the cell is deceased. The cell will again be empty in the next time step.

Special Cell States#

Additionally, cells can also have the following “special” states:

  • source: the cells are infection sources, thus, they can transition neighboring agents from the susceptible to the exposed state.

  • inert: inert cells are cells that do not partake in any of the model dynamics. They can be used to model spatial heterogeneities like compartmentalization, much like stones in the Forest Fire model.

Cells that are initialized in one of these states should not be regarded as representing agents: there is no movement for these cells, nor can these cells change their state.

Implementation#

The implementation allows for a range of different storylines by changing the parameters. Keep in mind that individual processes can often be disabled by setting probabilities to zero.

Update Rules#

In each time step the cells update their respective states asynchronously, but randomly shuffled to reduce artefacts, according to the following rules:

  1. A living (susceptible, exposed, infected, or recovered) cell becomes empty with probability p_empty.

  2. An empty cell turns into a susceptible one with probability p_susceptible. With probability p_immune the cell is immune to being infected.

  3. A susceptible cell either becomes randomly exposed with probability p_exposed or becomes exposed with probability p_transmit * (1 - p_random_immunity) if a neighboring cell is exposed or infected, with p_transmit being the neighboring cell’s probability of transmitting the disease. Disease transmission happens only if the cell is _not_ immune.

  4. An exposed cell becomes infected with probability p_infected.

  5. An infected cell recovers with probability p_recovered and becomes immune, becomes deceased with probability p_deceased, or else stays infected.

  6. A recovered cell can lose its immunity with p_lose_immunity and becomes susceptible again.

  7. A deceased cell turns into an empty cell.

Movement#

In each time step, the agents on the cells can move to empty neighboring cells according to the following rules:

  1. A living (susceptible, exposed, infected, or recovered) cell moves with probability p_move_randomly to a randomly chosen empty neighboring cell, if there is any.

  2. A living cell moves away from an infected neighboring cell to a randomly selected neighboring empty cell if there is any.

Note

Movement is meant to represent an agent moving from one cell to another. Thus, the agents also take with them their agent-specific properties immune and p_transmit. This is implemented using a swap operation of the corresponding State objects.

Heterogeneities#

As in the Forest Fire model, there is the possibility to introduce heterogeneities into the grid that are implemented as two additional possible cell states:

  • source: these are constant exposure sources. They spread the infection like normal infected or exposed cells, but don not revert to the empty state. If activated, they are per default at the lower boundary of the grid, though this can be changed in the configuration.

  • inert: inert cells are cells that do not partake in the dynamics of the model, and hence can be used to represent barriers. If enabled, the default mode is clustered_simple, which leads to randomly distributed inert cells whose neighbors have a certain probability to also be inert.

Both make use of the entity selection interface.

Immunity Control#

Via the immunity_control parameter in the model configuration, additional immunities can be introduced at desired times, thereby manipulating a cell’s immune state. This feature can be used e.g. to investigate the effect of vaccination. The immunities are introduced before the update rule above is carried out.

Exposure Control#

Via the exposure_control parameter in the model configuration, additional exposures can be introduced at desired times. The exposures are introduced before the update rule above is carried out.

Transmission Control#

Via the transmission_control parameter in the model configuration, the cell-specific state p_transmit can be manipulated. The cell state manipulation happens before the update rule above is carried out.

Data Output#

The following data is stored alongside the simulation:

  • kind: the state of each cell:

    • 0: empty

    • 1: susceptible

    • 2: exposed

    • 3: infected

    • 4: recovered

    • 5: deceased

    • 6: source, is constantly infectious

    • 7: inert, does not take part in any interaction

  • age: the age of each cell, reset after a cell turns empty.

  • cluster_id: a number identifying to which cluster a cell belongs; is 0 for non-living cells. Recovered cells do not count towards it.

  • exposed_time: the time steps a living cell has already been exposed to the disease, for each cell.

  • immunity: whether or not a cell is immune, for each cell.

  • densities: the densities of each of the kind of cells over time; this is a labeled 2D array with the dimensions time and kind.

  • counts: cumulative counters for a number of events, e.g. state transitions. This is a 2D array with the dimensions time and label, where the latter describes the name of the counted event:

    • empty_to_susceptible, i.e. “birth”

    • living_to_empty, i.e. “random death”

    • susceptible_to_exposed_contact, via local contact with a neighbor

    • susceptible_to_exposed_random, via a random point exposure, controlled by p_exposed

    • susceptible_to_exposed_controlled, via Exposure Control

    • exposed_to_infected, as controlled by p_infected

    • infected_to_recovered, as controlled by p_recovered

    • infected_to_deceased, as controlled by p_deceased

    • recovered_to_susceptible, which happens when losing immunity, as controlled by p_lose_immunity

    • move_randomly, as controlled by p_move_randomly

    • move_away_from_infected, as enabled by move_away_from_infected

Hint

When setting the write_ca_data parameter to False, only the densities and counts are stored. If spatially resolved information is not required, it is recommended to use this option, as it improves the run time and reduces the amount of disk space needed (which can be very large for large grids).

If spatial information is required, the write_start and write_every options can help cut down on disk space usage.

Default configuration parameters#

Below are the default configuration parameters for the SEIRD model:

# --- Space parameters --------------------------------------------------------
# The physical space this model is embedded in
space:
  periodic: true


# --- CellManager and cell initialization -------------------------------------
cell_manager:
  grid:
    structure: square
    resolution: 64      # in cells per unit length of physical space

  neighborhood:
    mode: Moore

  # Cell initialization parameters
  cell_params:
    # Initial susceptible density, value in [0, 1]
    # Probability that a cell is initialized as susceptible (instead of empty)
    p_susceptible: 0.3

    # Initial immune density, value in [0, 1]
    # With this probability, a cell is initialized as immune; however, only if
    # it was initialized as susceptible already (determined by p_susceptible)
    p_immune: &p_immune 0.001

    # The cell-specific probability p_transmit to transmit the disease to other
    # cells. This parameter sets the initial value for each cell.
    # This probability can be manipulated using the transmission_control to
    # either manipulate it for all cells or a subset of cells at specified
    # times. This allows, for example, to set some cells to be superspreader.
    p_transmit: &p_transmit
      mode: value
      value:
        default: 1.
      uniform:
        range: [0., 1.]


# --- Model Dynamics ----------------------------------------------------------
# NOTE that the default parameters below are chosen according to the
#      Covid-19 storyline roughly approximating the current knowledge, which
#      of course still changes and updates weekly (updated June 16th, 2020).
#      Parameters are chosen such that one time-step corresponds to one day.

# Probability per site and time step to transition from state empty to
# susceptible.
# As default, there are no randomly appearing susceptible cells coming into
# the system.
p_susceptible: 0.

# Probability to be immune per transition from an empty to a susceptible cell
# via p_susceptible.
p_immune: *p_immune

# Cell specific probability of transmitting the disease set at transition
# from an empty to a susceptible cell via p_susceptible. It sets the cell
# state `p_transmit` used in the model dynamics.
p_transmit: *p_transmit

# Probability per site and time step for a susceptible cell to _not_ become
# infected if an infected cell is in the neighborhood. This probability
# applies per event, so it does _not_ mean that an immune cell is also immune
# in the next iteration step.
p_random_immunity: 0.

# Probability per susceptible cell and time step for a random point exposure
# NOTE This is affected by the exposure control; see below.
p_exposed: 0.001

# Probability per exposed cell and time step to transition to infected state
# The default corresponds to a mean incubation period of 5 time steps
p_infected: 0.2

# Probability to recover if infected
# Note that p_recovered + p_deceased need to be smaller than 1
# The default is set to 1/14=0.0714, modeling that, on average, infections
# require approximately 14 days to go away.
p_recovered: 0.0714

# Probability to decease if the cell is infected
# Note that p_recovered + p_deceased need to be smaller than 1
# The default is set to 2%, approximating the infection fatality rate
# (see e.g. https://ourworldindata.org/mortality-risk-covid for a rough
# explanation).
p_deceased: 0.02

# Probability per site and time step to transition to empty
# As default, no living cells are vanishing by chance i.e. the system
# is approximated to be closed for the given time scales.
p_empty: 0.

# The probability to lose immunity if a cell is recovered
# This value is quite uncertain because it is not known whether and with what
# probability immunity is lost in the case of Covid-19. Currently, the effect
# of losing immunity through virus mutation seems to be smaller than for the
# common influenza virus with roughly one year. However, there are studies
# suggesting that especially for light Covid-19 infections, new infections
# can happen with a probably non-negligible rate.
# The given default tries to incorporate and estimate both effects roughly.
p_lose_immunity: 0.01

# ... Movement ...............................................................
# Whether to allow cells to move away from infected neighboring cells
# If a neighbor is infected a cell searches for a random empty neighboring
# place and moves towards it. If there is no space, do nothing.
move_away_from_infected: false

# Probability to move in a random direction
p_move_randomly: 0.2


# --- Exposure Control -------------------------------------------------------
# Exposure control to investigate the time-dependent influence of the
# disease driving force. Note that exposure control is applied at the
# beginning of an iteration step. Its effect is seen in the following
# time step
exposure_control:
  enabled: false

  # The number of additional exposures to be placed on the grid
  num_additional_exposures: 10

  # Add the additional exposures at the given times
  # Note that the !listgen tag creates a list from the parameters
  # (start, stop, step_size)
  # To disable, pass an empty sequence.
  at_times: !listgen [0, 100, 20]

  # Change the probability of random exposure.
  # The expected value is a list of [iteration_step, new_value] pairs, e.g.
  #   - [10, .5]
  #   - [42, 0.]
  # ... will set p_expose from the default value to .5 at time 10 and set it
  # back to 0. at time 42.
  # To disable, pass an empty sequence.
  change_p_exposed: []


# --- Immunity Control --------------------------------------------------------
# Immunity control to investigate the time-dependent influence of actively
# provided immunities. Note that immunity control is applied at the
# beginning of an iteration step but after the exposure control. Its effect
# is seen in the following time step
immunity_control:
  enabled: false

  # The number of additional immunities to be placed on the grid
  num_additional_immunities: 10

  # Add the additional immunities at the given times
  # Note that the !listgen tag creates a list from the parameters
  # (start, stop, step_size)
  # To disable, pass an empty sequence.
  at_times: !listgen [0, 100, 20]

  # Change the probability of a random immunity when new susceptible cells
  # appear through p_susceptible.
  # The expected value is a list of [iteration_step, new_value] pairs, e.g.
  #   - [10, .5]
  #   - [42, 0.]
  # ... will set p_immune from the default value to .5 at time 10 and set it
  # back to 0. at time 42.
  # To disable, pass an empty sequence.
  change_p_immune: []


# --- Transmission Control ----------------------------------------------------
# Transmission control to investigate the time-dependent influence of actively
# changing decease transmission probabilities.
# Note that transmission control is applied at the beginning of an iteration
# step after the exposure and immunity control. Its effect is seen only in the
# following time step
transmission_control:
  enabled: false
  # Change the probability to transmit a decease of some randomly
  # selected exposed or infected cells.
  # The expected value is a list of mappings
  # change_p_transmit:
  #   - time: 10
  #     num_cells: 6
  #     cell_kind: susceptible
  #     p_transmit: 0.5
  #   - time: 42
  #     num_cells: 2
  #     cell_kind: exposed
  #     p_transmit: 0.
  # ... will set p_transmit from the default value to .5 for 6 randomly
  # selected susceptible cells at time 10 and set it to 0. for 2 randomly
  # selected exposed cells at time 42.
  # If num_cells exceeds the current number of present cells with the specified
  # kind all of them are chosen to reset their p_transmit value.
  # To disable, pass an empty sequence.
  change_p_transmit: []


# --- Heterogeneities ---------------------------------------------------------
# Some cells can be permanently infected or turned into inert cells.
# Both these features are using the `select_entities` interface; consult the
# documentation regarding the information on available selection modes.

# Make some cells inert: these do not take part in any of the processes
inert_cells:
  enabled: false
  mode: clustered_simple

  # Clustering parameters
  p_seed: .02                # Probability with which a cell is a cluster seed
  p_attach: .1               # Attachment probability (per neighbor)
  num_passes: 5              # How many attachment procedures to perform

# Set some cells to be permanently infected (invoked after inert cells are set)
infection_source:
  enabled: false
  mode: boundary

  # Boundary selection parameters (requires space to be set to NON-periodic!)
  boundary: bottom



# --- Output Configuration ----------------------------------------------------
# Whether to write out spatially resolved data from the CA
# If false, will write only the non-spatial `densities` and `counts` datasets.
# Setting this to false can be useful if no spatial analysis is required or if
# using huge grids.
write_ca_data: true

# HDF5 Compression level for all datasets
# A value of 1-3 is a good default. Choose a lower value if speed is limited by
# the CPU or a higher value if speed is limited by data writing.
compression: 3

Available plots#

The following plot configurations are available for the SEIRD model:

Default Plot Configuration#

# -- Plot of all densities over time ------------------------------------------
densities:
  based_on: densities.uni

densities_facet_grid:
  based_on: densities.mv


# -- Counter-based ------------------------------------------------------------

# Transition plots ............................................................
transitions/raw:
  based_on: transitions/raw

transitions/combined:
  based_on: transitions/combined

transitions/smoothed:
  based_on: transitions/smoothed


# Movement events .............................................................
movement:
  based_on: movement


# -- Distribution plots -------------------------------------------------------
age_distribution/final:
  based_on: age_distribution/final

age_distribution/time_series:
  based_on: age_distribution/time_series
  enabled: false

age_distribution/deceased:
  based_on: age_distribution/deceased



# -- Phase diagrams from two densities ----------------------------------------
phase_diagram/SI:
  based_on: phase_diagram/SI

phase_diagram/SE:
  based_on: phase_diagram/SE

phase_diagram/EI:
  based_on: phase_diagram/EI


# ... and their facet grid equivalent
phase_diagram_facet_grid/SI:
  based_on: phase_diagram_facet_grid/SI

phase_diagram_facet_grid/SE:
  based_on: phase_diagram_facet_grid/SE

phase_diagram_facet_grid/EI:
  based_on: phase_diagram_facet_grid/EI




# -- CA plots -----------------------------------------------------------------
ca/state:
  based_on: ca/state
  enabled: false

ca/state_final:
  based_on:
    - ca/state
    - .plot.ca.snapshot
  enabled: false


ca/age:
  based_on: ca/age
  enabled: false

ca/age_final:
  based_on:
    - ca/age
    - .plot.ca.snapshot
  enabled: false


ca/clusters:
  based_on: ca/clusters
  enabled: false

ca/clusters_final:
  based_on:
    - ca/clusters
    - .plot.ca.snapshot
  enabled: false


ca/combined:
  based_on:
    - ca/state
    - ca/age
    - ca/clusters

ca/combined_final:
  based_on:
    - ca/combined
    - .plot.ca.snapshot

Base Plot Configuration#

_variables:
  base_path: &base_path  data/SEIRD

  cmap: &cmap
    # NOTE Order is important here, mapping to `kind` values 0 ... 7
    empty:        &color_empty                darkkhaki           # 0
    susceptible:  &color_susceptible          forestgreen         # 1
    exposed:      &color_exposed              darkorange          # 2
    infected:     &color_infected             firebrick           # 3
    recovered:    &color_recovered            slategray           # 4
    deceased:     &color_deceased             black               # 5
    source:       &color_source               maroon              # 6
    inert:        &color_inert                moccasin            # 7

  labels:
    label_susceptible:    &label_susceptible  Susceptible Density [1/A]
    label_infected:       &label_infected     Infected Density [1/A]
    label_exposed:        &label_exposed      Exposed Density [1/A]
    label_recovered:      &label_recovered    Recovered Density [1/A]



# =============================================================================
#  ╔╦╗╔═╗╔╦╗╔═╗╦  ╔═╗╔╦╗╔═╗╔═╗
#   ║ ║╣ ║║║╠═╝║  ╠═╣ ║ ║╣ ╚═╗
#   ╩ ╚═╝╩ ╩╩  ╩═╝╩ ╩ ╩ ╚═╝╚═╝
# =============================================================================
# -- Overloads ----------------------------------------------------------------
# Overload some configs to insert model-specific settings

# Model-specific defaults
.defaults:
  based_on: .defaults

  # Optionally define something here ...


# -- Creators -----------------------------------------------------------------
.creator.universe:
  based_on:
    - .creator.universe
    - .defaults

  dag_options:
    select_path_prefix: *base_path

.creator.multiverse:
  based_on:
    - .creator.multiverse
    - .defaults

  select_and_combine:
    base_path: *base_path



# -- Styles -------------------------------------------------------------------
# .. Property cycles ..........................................................
# Using these as base configurations allows to have consistent kind colors
# NOTE The order in the evaluated format string needs to match the order in
#      which the cycling happens! This is an implicit mapping.

.style.prop_cycle.ESEIRDSI:
  style:
    axes.prop_cycle: !format
      fstr: "cycler('color', ['{cmap[empty]:}', '{cmap[susceptible]:}', '{cmap[exposed]:}', '{cmap[infected]:}', '{cmap[recovered]:}', '{cmap[deceased]:}', '{cmap[source]:}', '{cmap[inert]:}'])"
      cmap: *cmap

.style.prop_cycle.seird:
  style:
    axes.prop_cycle: !format
      fstr: "cycler('color', ['{cmap[susceptible]:}', '{cmap[exposed]:}', '{cmap[infected]:}', '{cmap[recovered]:}', '{cmap[deceased]:}'])"
      cmap: *cmap

.style.prop_cycle.deceased:
  style:
    axes.prop_cycle: !format
      fstr: "cycler('color', ['{cmap[deceased]:}'])"
      cmap: *cmap

.style.prop_cycle.default:
  style:
    axes.prop_cycle: &default_prop_cycle "cycler('color', ['1f77b4', 'ff7f0e', '2ca02c', 'd62728', '9467bd', '8c564b', 'e377c2', '7f7f7f', 'bcbd22', '17becf'])"


# =============================================================================
#  ╔╦╗╔═╗╔═╗  ╔╦╗╔═╗╔╦╗╔═╗  ╔═╗╔═╗╔═╗╦═╗╔═╗╔╦╗╦╔═╗╔╗╔╔═╗
#   ║║╠═╣║ ╦  ║║║║╣  ║ ╠═╣  ║ ║╠═╝║╣ ╠╦╝╠═╣ ║ ║║ ║║║║╚═╗
#  ═╩╝╩ ╩╚═╝  ╩ ╩╚═╝ ╩ ╩ ╩  ╚═╝╩  ╚═╝╩╚═╩ ╩ ╩ ╩╚═╝╝╚╝╚═╝
# =============================================================================
# The following entries can be included into a plot configuration to make
# certain meta-operations available for the data transformation framework.

.dag.meta_ops.build_densities_dataset:
  dag_options:
    meta_operations:
      #
      # Selects the relevant densities for the phase diagram plot
      #
      # Args:
      #   0:             The `densities` data container
      #
      # Returns:
      #   xr.Dataset with variables:  susceptible, exposed, infected, recovered
      #
      build_densities_dataset:
        - .sel: [!arg 0, {kind: susceptible}]
          kwargs: {drop: true}
          tag: susceptible
        - .sel: [!arg 0, {kind: exposed}]
          kwargs: {drop: true}
          tag: exposed
        - .sel: [!arg 0, {kind: infected}]
          kwargs: {drop: true}
          tag: infected
        - .sel: [!arg 0, {kind: recovered}]
          kwargs: {drop: true}
          tag: recovered

        - xr.Dataset:
            data_vars:
              susceptible: !dag_tag susceptible
              exposed: !dag_tag exposed
              infected: !dag_tag infected
              recovered: !dag_tag recovered




# =============================================================================
#  ╔═╗╦  ╔═╗╔╦╗╔═╗
#  ╠═╝║  ║ ║ ║ ╚═╗
#  ╩  ╩═╝╚═╝ ╩ ╚═╝
# =============================================================================
# -- Densities time series ----------------------------------------------------

densities.uni:
  based_on:
    - .creator.universe
    - .plot.facet_grid.line
    - .hlpr.kind.time_series
    - .style.prop_cycle.ESEIRDSI

  select:
    data: densities

  x: time
  hue: kind

  helpers: &helpers_densities
    set_limits:
      y: [0., 1.]
    set_labels:
      y: Density [1/A]
      only_label_outer: true
    set_title:
      title: Densities
    set_legend:
      loc: best

densities.mv:
  based_on:
    - .creator.multiverse
    - .plot.facet_grid.with_auto_encoding
    - .plot.facet_grid.errorbands
    - .hlpr.kind.time_series
    - .hlpr.legend.hide
    - .animation.disabled
    - .style.prop_cycle.ESEIRDSI

  expected_multiverse_ndim: [1, 2, 3]

  select_and_combine:
    fields:
      _data: densities

  transform:
    - xr.full_like: [!dag_tag _data, !expr nan]
      tag: nans

    # Try to compute mean and std over ceratain dimensions, falling back to
    # a NaN array if that dimension is not available
    - .mean: [!dag_tag _data, seed]
      allow_failure: silent
      fallback: !dag_tag _data
      tag: mean
    - .std: [!dag_tag _data, seed]
      allow_failure: silent
      fallback: !dag_tag nans  # use NaN instead
      tag: std

    # Combine into a dataset
    - xr.Dataset:
        data_vars:
          mean: !dag_tag mean
          std: !dag_tag std
      tag: data

  y: mean
  yerr: std
  x: time
  hue: kind

  helpers:
    <<: *helpers_densities
    set_title:
      enabled: false
    set_suptitle:
      title: Densities


# -- Any kind of phase plot ---------------------------------------------------

# The base plot, defining styles etc.
.plot.phase_diagram:
  based_on:
    - .defaults
    - .plot.facet_grid.scatter
    - .dag.meta_ops.build_densities_dataset

  cmap: viridis_r
  hue: time
  helpers:
    set_limits:
      x: [0, ~]
      y: [0, ~]
    set_labels:
      only_label_outer: true

  # Parameters that are passed on to plt.scatter
  s: 10


# Universe phase diagram
phase_diagram:
  based_on:
    - .creator.universe
    - .plot.phase_diagram

  select:
    data:
      path: densities
      transform:
        - build_densities_dataset

  helpers:
    set_title:
      title: Phase Diagram


# Multiverse phase diagram, with
phase_diagram_facet_grid:
  based_on:
    - .creator.multiverse
    - .plot.facet_grid.with_auto_encoding
    - .plot.phase_diagram

  expected_multiverse_ndim: [1, 2, 3]

  select_and_combine:
    fields:
      _data:
        path: densities
        transform:
          - build_densities_dataset


  transform:
    # Try to compute mean over certain dimensions, falling back if it fails
    - .mean: [!dag_tag _data, seed]
      allow_failure: silent
      fallback: !dag_tag _data

    - pass: !dag_prev
      tag: data


# .. Specializations ..........................................................
# .. Universe . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

phase_diagram/SI:
  based_on: phase_diagram

  x: susceptible
  y: infected
  helpers:
    set_labels:
      x: *label_susceptible
      y: *label_infected

phase_diagram/SE:
  based_on: phase_diagram
  x: susceptible
  y: exposed
  helpers:
    set_labels:
      x: *label_susceptible
      y: *label_exposed

phase_diagram/EI:
  based_on: phase_diagram
  x: exposed
  y: infected
  helpers:
    set_labels:
      x: *label_exposed
      y: *label_infected

# .. Universe . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

phase_diagram_facet_grid/SI:
  based_on: phase_diagram_facet_grid
  x: susceptible
  y: infected
  helpers:
    set_labels:
      x: *label_susceptible
      y: *label_infected

phase_diagram_facet_grid/SE:
  based_on: phase_diagram_facet_grid
  x: susceptible
  y: exposed
  helpers:
    set_labels:
      x: *label_susceptible
      y: *label_exposed

phase_diagram_facet_grid/EI:
  based_on: phase_diagram_facet_grid
  x: exposed
  y: infected
  helpers:
    set_labels:
      x: *label_exposed
      y: *label_infected




# --- Counter-based -----------------------------------------------------------

counters:
  based_on:
    - .creator.universe
    - .plot.facet_grid

  select:
    counts: counts
    d_counts:
      path: counts
      transform:
        - .diff: [!dag_prev , time]
  dag_options:
    select_path_prefix: data/SEIRD

  style:
    # Set the default prop cycle; somehow lost for categorical label dimension
    axes.prop_cycle: *default_prop_cycle

  helpers:
    set_limits:
      x: [0, max]
      y: [0, ~]

    # Only set the x labels on the bottom row, as x-axes are shared.
    axis_specific:
      bottom_row:
        axis: [~, -1]
        set_labels:
          x: Time [Iteration Steps]


# .. Movement stats ...........................................................
movement:
  based_on: counters

  transform:
    - .sel: [!dag_tag d_counts]
      kwargs:
        label:
          - move_randomly
          - move_away_from_infected
      tag: data

  kind: line
  x: time
  hue: label

  helpers:
    set_labels:
      y: Movement Events


# .. State transitions ........................................................
transitions/raw:
  based_on: counters

  transform:
    - &op_sel_transitions
      .sel: [!dag_tag d_counts]
      kwargs:
        label:
          - empty_to_susceptible
          - living_to_empty
          - susceptible_to_exposed_contact
          - susceptible_to_exposed_random
          - susceptible_to_exposed_controlled
          - exposed_to_infected
          - infected_to_recovered
          - infected_to_deceased
          - recovered_to_susceptible

    # Define as data to be plotted
    - define: !dag_prev
      tag: data

  kind: line
  x: time
  col: label
  col_wrap: 3
  sharey: true

transitions/combined:
  based_on: counters

  transform:
    - *op_sel_transitions
    - define: !dag_prev
      tag: data

  kind: line
  x: time
  hue: label

  helpers:
    set_labels:
      y: Counts

transitions/smoothed:
  based_on:
    - transitions/combined
    - .dag.meta_ops.rolling

  dag_options:
    define:
      smoothed_by: &smoothed_by 7

  transform:
    - *op_sel_transitions

    # Apply rolling mean over a number of time coordinates
    - rolling: [!dag_prev , mean, {time: *smoothed_by}]
      kwargs:
        # Configure rolling window such that early values are included
        min_periods: 2
        center: false
      tag: data

  helpers:
    set_labels:
      y: !format ["Counts (smoothed over {} data points)", *smoothed_by]



# --- Distributions -----------------------------------------------------------

age_distribution:
  based_on:
    - .defaults
    - .creator.universe
    - .plot.histogram

  select:
    age: age
    kind: kind

  x: age
  hue: kind

  helpers:
    set_labels:
      x: Age
      y: Counts
    set_title:
      title: Age Distribution

age_distribution/final:
  based_on:
    - age_distribution
    - .style.prop_cycle.seird
  transform:
    - SEIRD.compute_age_distribution:
        age: !dag_tag age
        kind: !dag_tag kind
        coarsen_by: 10
    # Select the last time step
    - .isel: [!dag_prev , {time: -1}]
      tag: counts

  helpers:
    set_title:
      title: Final Age Distribution


age_distribution/time_series:
  based_on:
    - age_distribution
    - .style.prop_cycle.seird
    - .animation.defaults

  transform:
    - SEIRD.compute_age_distribution:
        age: !dag_tag age
        kind: !dag_tag kind
        coarsen_by: 10
        normalize: false  # optional. Default: false
      tag: counts

  # Represent the time dimension as frames of the animation
  frames: time


age_distribution/deceased:
  based_on:
    - age_distribution
    - .style.prop_cycle.deceased
  transform:
    - SEIRD.compute_age_distribution:
        age: !dag_tag age
        kind: !dag_tag kind
        coarsen_by: 10
        compute_for_kinds: [deceased]
    - .sum: [!dag_prev , time]
    - print: !dag_prev
      tag: counts

  helpers:
    set_title:
      title: Deceased Age Distribution





# --- CA Plots ----------------------------------------------------------------
ca/state:
  based_on:
    - .creator.universe
    - .plot.ca

  select:
    kind: kind

  to_plot:
    kind:
      title: State
      cmap: *cmap


ca/age:
  based_on:
    - .creator.universe
    - .plot.ca

  select:
    age: age

  to_plot:
    age:
      title: Age
      cmap: YlGn


ca/clusters:
  based_on:
    - .creator.universe
    - .plot.ca

  select:
    cluster:
      path: cluster_id
      with_previous_result: true
      transform:
        - where: ['!=', 0]       # 0 is masked: not part of a cluster
        - np.fmod: 20

  to_plot:
    cluster:
      title: Clusters
      vmin: 0
      vmax: 20
      cmap: tab20
      no_cbar_markings: true

For available base plots, see Base Plot Configuration Pool.