import numpy as np
from scipy.integrate import odeint
from .benchmark_simulator import BenchmarkSimulator
[docs]
class SIR(BenchmarkSimulator):
def __init__(
self,
N: float = 1e6,
T: int = 160,
I0: float = 1.0,
R0: float = 0.0,
subsample: int | str = "original",
total_count: int = 1000,
scale_by_total: bool = True,
rng: np.random.Generator = None,
):
"""SIR simulated benchmark
See: https://arxiv.org/pdf/2101.04653.pdf, Task T.9
NOTE: the simulator scales outputs between 0 and 1.
Parameters
----------
N: float, optional, default: 1e6
The size of the simulated population.
T: int, optional, default: 160
The duration (time horizon) of the simulation.
The last time-point is not included.
I0: float, optional, default: 1.0
The number of initially infected individuals.
R0: float, optional, default: 0.0
The number of initially recovered individuals.
subsample: int, str or None, optional, default: 'original'
The number of evenly spaced time points to return. If `None`,
no subsampling will be performed, all `T` timepoints will be returned
and a trailing dimension will be added. If an integer is provided,
subsampling is performed and no trailing dimension will be added.
'original' reproduces the original benchmark task subsampling of 10 points.
total_count: int, optional, default: 1000
The `N` parameter of the binomial noise distribution. Used just
for scaling the data and magnifying the effect of noise, such that
`max infected == total_count`.
scale_by_total: bool, optional, default: True
Scales the outputs by `total_count` if set to True.
rng: np.random.Generator or None, optional, default: None
An optional random number generator to use.
"""
self.N = N
self.T = T
self.I0 = I0
self.R0 = R0
self.subsample = subsample
self.total_count = total_count
self.scale_by_total = scale_by_total
self.rng = rng
if self.rng is None:
self.rng = np.random.default_rng()
def _deriv(self, x, t, N, beta, gamma):
"""Helper function for scipy.integrate.odeint."""
s, i, r = x
dS = -beta * s * i / N
dI = beta * s * i / N - gamma * i
dR = gamma * i
return dS, dI, dR
[docs]
def prior(self):
"""Generates a random draw from a 2-dimensional (independent) lognormal prior
which represents the contact and recovery rate parameters of a basic SIR model.
Returns
-------
params : np.ndarray of shape (2, )
A single draw from the 2-dimensional prior.
"""
return self.rng.lognormal(mean=[np.log(0.4), np.log(1 / 8)], sigma=[0.5, 0.2])
[docs]
def observation_model(self, params: np.ndarray):
"""Runs a basic SIR model simulation for T time steps and returns `subsample` evenly spaced
points from the simulated trajectory, given disease parameters (contact and recovery rate) `params`.
Parameters
----------
params : np.ndarray of shape (2,)
The 2-dimensional vector of disease parameters.
Returns
-------
x : np.ndarray of shape (subsample, ) or (T, 1) if subsample=None
The time series of simulated infected individuals. A trailing dimension of 1 should
be added by a BayesFlow configurator if the data is (properly) to be treated as time series.
"""
# Create vector (list) of initial conditions
x0 = self.N - self.I0 - self.R0, self.I0, self.R0
# Unpack parameter vector into scalars
beta, gamma = params
# Prepare time vector between 0 and T of length T
t_vec = np.arange(0, self.T)
# Integrate using scipy and retain only infected (2-nd dimension)
irt = odeint(self._deriv, x0, t_vec, args=(self.N, beta, gamma))[:, 1]
# Subsample evenly the specified number of points, if specified
if self.subsample == "original":
irt = irt[::17]
elif self.subsample is not None:
irt = irt[:: (self.T // self.subsample)]
else:
irt = irt[:, None]
# Truncate irt, so that small underflow below zero becomes zero
irt = np.maximum(irt, 0.0)
# Add noise and scale, if indicated
x = self.rng.binomial(n=self.total_count, p=irt / self.N)
if self.scale_by_total:
x = x / self.total_count
return x
