Episimlab is a Python package for rapid development and execution of contagion models. It provides a set of lightweight, reusable, and extensible Python processes which can be combined into models with arbitrary complexity. While Episimlab was originally written with epidemiological use cases in mind, its applications extend to other fields such as computer security. This package specifically supports development of models that:
While Episimlab provides defaults for these behaviors, users can easily customize their models thanks to integration with the xarray-simlab package. For more details, please consult the documentation for Episimlab and xarray-simlab.
See the full example here.
import xsimlab
# Import an existing SIR model
from episimlab.models import ExampleSIR
@xsimlab.process
class CustomRecoveryRate:
"""We modify the model by writing a single custom process. This process
calculates a recovery rate (`rate_I2R`)."""
# Variables output by this process (intent='out')
rate_I2R = xsimlab.variable(
global_name='rate_I2R', groups=['edge_weight'], intent='out',
description="rate of change from compartments I to R")
# Variables ingested by this process (intent='in').
# Setting a `global_name` here allows this process to ingest from an
# upstream CalculateGamma process that would define gamma like:
# gamma = xsimlab.global_ref('gamma', intent='out')
gamma = xsimlab.variable(global_name='gamma', intent='in')
# The state is the N-D array of every individual in the simulation,
# stratified by age, risk group, geospatial identifier, compartment, etc.
state = xsimlab.global_ref('state', intent='in')
def run_step(self):
"""Re-calculate `rate_I2R` at every step of the simulation. We can
instead decide to calculate only once at the beginning by renaming
this method `initialize`. We calculate recovery rate as the product
of `gamma` and the current state of the `I` compartment.
"""
self.rate_I2R = self.gamma * self.state.loc[{'compt': 'I'}]
# Instantiate the included model
model = ExampleSIR()
# Edit the model so it uses our custom process
model.update_processes({'recovery_rate': CustomRecoveryRate})
results = model.run(input_vars={'gamma': 0.6})
model.plot()
# The final state is an N-D array represented in Python as an xarray.DataArray
# object (like a numpy array, but with labeled axes)
final_state = results['compt_model__state']
print(final_state)
# <xarray.DataArray 'compt_model__state' (step: 15, vertex: 4, compt: 4, age: 5, risk: 2)>
# array([[[[[200., 200.],
# [200., 200.],
# ...
# [ 0., 0.],
# [ 0., 0.]]]]])
# Coordinates:
# * age (age) object '0-4' '5-17' '18-49' '50-64' '65+'
# * compt (compt) object 'S' 'I' 'R' 'V'
# * risk (risk) object 'low' 'high'
# * step (step) datetime64[ns] 2020-03-01 2020-03-02 ... 2020-03-15
# * vertex (vertex) object 'Austin' 'Houston' 'San Marcos' 'Dallas'
CustomSetupComptGraph
in the full examplepip install episimlab
To install the most recent (and potentially less stable) development release, install directly from the main
branch of the GitHub repository:
pip install git+https://github.com/UT-Covid/episimlab
Pushing commits to a Pull Request against the main
branch will automatically trigger pytests in GitHub Actions. Tests run in the poetry environment defined by the pyproject.toml
file.
Preferred testing environment runs poetry virtual env within tox.
# Default args
tox
# Pass args to pytest. In this case, we run only the model tests
tox -- tests/test_models.py
If you are interested in contibuting to Episimlab, please see the contribution guide.
If you are having trouble using Episimlab, please first refer to the documentation and examples. If you're still having trouble, please submit a GitHub issue.