epiinfer-contact-tracing
Python implementation of epiinfer-contact-tracing ("Forecasting Infection Rates in
Epidemics Given Contact Tracing Information", Kascelan, Yang & Shasha, Algorithms, 2025).
The package, the command and the Docker image are called epitrace. It is packaged so
that anyone can:
Using Docker, the following tasks can be performed:
Simulate an epidemic involving 10,000 individuals;
Generate a 14-day forecast from day 60 and subsequently validate the forecast against observed outcomes;
Produce an out-of-sample forecast extending beyond the available data;
Conduct back-testing by generating forecasts from each day between day 25 and day 79, and measure the associated errors;
Compare epiinfer-contact-tracing (both engines) with the earlier EpiInfer, ARIMA, and Holt-Winters methods;
Convert users’ own case and contact lists for use in forecasting.
New here? Start with the tutorial. It walks through a forecast in
Docker step by step, with the exact output and plots you should get. The documentation is also
published at https://nuzla.github.io/epitrace/.
Install
The project is managed with uv. uv.lock pins every
dependency, so everyone gets the same numerical stack. The package is not on PyPI yet;
from a checkout:
uv sync --all-extras # .venv with the locked versions + plots + baselines
uv run epitrace --help # run any command inside that environment
Other routes:
- Base install.
uv sync gives numpy and pandas only. Add --extra plots or
--extra baselines for the extras.
- Command on your PATH.
uv tool install ".[all]" installs the epitrace command.
- pip.
pip install ".[all]" works too, but it resolves fresh versions instead of the
locked ones.
It is tested on Python 3.11–3.14 (.python-version picks 3.13 by default). The oldest
supported dependencies are also tested: numpy 1.26, pandas 2.0, statsmodels 0.14 and
scipy 1.16. Python 3.10 is not supported: it reaches end of life in October 2026, and no
scipy build for it loads on current macOS.
No Python on the host? Use Docker instead.
60-second tour
uv run epitrace demo # path to a tiny bundled dataset
uv run epitrace describe --data "$(uv run epitrace demo)"
uv run epitrace forecast --data "$(uv run epitrace demo)" --horizon 14 --plot forecast.png
uv run epitrace run --data "$(uv run epitrace demo)" --method unknown_both --max-day 46 \
--horizons 1 5 10 --plots
import epitrace
epitrace.forecast("my_dataset/", horizon=14, seed=0) # mean + 10-90% band
epitrace.run("my_dataset/", "unknown_both").summary() # back-test: RMSE, fRMSE
epitrace.run("my_dataset/", "unknown_both", engine="contact") # the 2025 code instead
Your data is a folder of CSV files; see docs/DATA_FORMAT.md. For the
default engine that is two files: who is symptomatic or recovered each day, and who met
whom. epitrace from-cases builds the folder from a case list and a contact list.
examples/forecast_new_data.py is a runnable template.
Docker
The image bundles Python, the locked dependencies and every epitrace command, so the only
thing you need on the host is Docker. Released versions are published to the GitHub container
registry. CI also builds the image on every push, runs demo, forecast and run --plots in
it as an arbitrary non-root user, and runs the whole tutorial in it.
1. Get the image
Pull a released image (Intel/AMD and ARM, including Apple silicon). Each version tag
publishes one; see the repository's Packages page for the list.
docker pull ghcr.io/nuzla/epitrace:0.3.0
docker tag ghcr.io/nuzla/epitrace:0.3.0 epitrace # the commands below call it "epitrace"
Versions before 0.3.0 were published as ghcr.io/nuzla/epiinfer. Those images stay available
but get no updates; see the changelog for the rename.
For a paper or a script, pin the exact image by digest, which is printed in each release's
workflow summary: ghcr.io/nuzla/epitrace@sha256:…. Version tags can be re-pointed; digests
cannot.
Or build it from this repository:
docker build -f docker/Dockerfile -t epitrace .
docker run --rm epitrace --version
2. How the container is wired
| Host | In the container | Used for |
|---|
./data | /data | datasets you read, and ones simulate / fetch-data write |
./results | /work/results | everything the commands write (the default --out) |
The container runs as a non-root user. Pass --user "$(id -u):$(id -g)" so that the files
it writes into ./results belong to you. Commands below use the shell function
epi() { docker run --rm --user "$(id -u):$(id -g)" \
-v "$PWD/data:/data" -v "$PWD/results:/work/results" epitrace "$@"; }
Paths in the arguments are paths inside the container: datasets are under /data/...
and outputs go under results/....
3. Try it on the bundled demo data (nothing to download)
mkdir -p data results
DEMO="$(epi demo)" # path of the demo dataset inside the image
epi forecast --data "$DEMO" --horizon 14 --out results/demo_forecast.csv
epi run --data "$DEMO" --method unknown_both --start-eval 30 --plots --out results/demo_backtest
epi baseline --data "$DEMO" --which arima --start-eval 30
Open results/demo_backtest/forecast_demo_data_unknown_both.png to see the back-test.
4. Your own data
Put the dataset folder (see docs/DATA_FORMAT.md) in ./data, or
build it from a case list and a contact list with
epi from-cases --cases /data/cases.csv --contacts /data/contacts.csv --out /data/my_dataset:
This is the beginning of the README. Read the rest at the original. Read the full README ↗
