def test_last_time_r_threshold(self):
df = pd.DataFrame({
'Time': [1, 2, 3, 5, 6],
'Incidence Number': [0, 0, 0, 0, 0]
})
ser_int = [1, 2]
inference = bp.BranchProPosterior(df, ser_int, 1, 0.2)
inference.run_inference(tau=2)
last_time_r_more_than_1 = inference.last_time_r_threshold('more')
last_time_r_less_than_1 = inference.last_time_r_threshold(
'less', method='Median')
self.assertEqual(len(last_time_r_more_than_1), 3)
self.assertEqual(len(last_time_r_less_than_1), 3)
self.assertEqual(last_time_r_more_than_1[0], 7)
self.assertEqual(last_time_r_more_than_1[1], None)
self.assertEqual(last_time_r_more_than_1[2], 7)
self.assertEqual(last_time_r_less_than_1[0], None)
self.assertEqual(last_time_r_less_than_1[1], 7)
self.assertEqual(last_time_r_less_than_1[2], None)
with self.assertRaises(ValueError):
inference.last_time_r_threshold('=')
inference.last_time_r_threshold('<')
inference.last_time_r_threshold('>')
inference.last_time_r_threshold(2)
inference.last_time_r_threshold('more', method='mean')
def test_get_serial_intervals(self):
df = pd.DataFrame({
'Time': [1, 2, 3, 5, 6],
'Incidence Number': [10, 3, 4, 6, 9]
})
ser_int = [1, 2]
inference = bp.BranchProPosterior(df, ser_int, 1, 0.2)
npt.assert_array_equal(inference.get_serial_intervals(),
np.array([1, 2]))
def test_run_inference(self):
df = pd.DataFrame({
'Time': [1, 2, 3, 5, 6],
'Incidence Number': [10, 3, 4, 6, 9]
})
ser_int1 = [1, 2, 1, 0, 0, 0]
ser_int2 = [1, 2]
inference1 = bp.BranchProPosterior(df, ser_int1, 1, 0.2)
inference1.run_inference(tau=2)
inference2 = bp.BranchProPosterior(df, ser_int2, 1, 0.2)
inference2.run_inference(tau=2)
self.assertEqual(len(inference1.inference_estimates), 3)
self.assertEqual(len(inference1.inference_times), 3)
self.assertEqual(len(inference1.inference_posterior.mean()), 3)
self.assertEqual(len(inference2.inference_estimates), 3)
self.assertEqual(len(inference2.inference_times), 3)
self.assertEqual(len(inference2.inference_posterior.mean()), 3)
def test__init__(self):
df = pd.DataFrame({
'Time': [1, 2, 3, 5, 6],
'Incidence Number': [10, 3, 4, 6, 9]
})
ser_int = [1, 2]
bp.BranchProPosterior(df, ser_int, 1, 0.2)
with self.assertRaises(TypeError) as test_excep:
bp.BranchProPosterior('0', ser_int, 1, 0.2)
self.assertTrue('Incidence data has to' in str(test_excep.exception))
with self.assertRaises(TypeError) as test_excep:
bp.BranchProPosterior(df, 0, 1, 0.2)
self.assertTrue('must be iterable' in str(test_excep.exception))
with self.assertRaises(TypeError) as test_excep:
bp.BranchProPosterior(df, ['zero'], 1, 0.2)
self.assertTrue(
'distribution must contain' in str(test_excep.exception))
with self.assertRaises(ValueError) as test_excep:
bp.BranchProPosterior(df, ser_int, 1, 0.2, time_key='t')
self.assertTrue('No time column' in str(test_excep.exception))
with self.assertRaises(ValueError) as test_excep:
bp.BranchProPosterior(df, ser_int, 1, 0.2, inc_key='i')
self.assertTrue('No incidence column' in str(test_excep.exception))
def test_set_serial_intervals(self):
df = pd.DataFrame({
'Time': [1, 2, 3, 5, 6],
'Incidence Number': [10, 3, 4, 6, 9]
})
ser_int = [1, 2]
new_ser_int = [1, 2, 1]
wrong_ser_int = (1)
inference = bp.BranchProPosterior(df, ser_int, 1, 0.2)
inference.set_serial_intervals(new_ser_int)
npt.assert_array_equal(inference.get_serial_intervals(),
np.array([1, 2, 1]))
with self.assertRaises(ValueError):
inference.set_serial_intervals(wrong_ser_int)
def test_proportion_time_r_more_than_1(self):
df = pd.DataFrame({
'Time': [1, 2, 3, 5, 6],
'Incidence Number': [0, 0, 0, 0, 0]
})
ser_int = [1, 2]
inference = bp.BranchProPosterior(df, ser_int, 1, 0.2)
inference.run_inference(tau=2)
proportions = inference.proportion_time_r_more_than_1(.95)
self.assertEqual(len(proportions), 3)
self.assertEqual(proportions[0], 1)
self.assertEqual(proportions[1], 0)
self.assertEqual(proportions[2], 1)
with self.assertRaises(ValueError):
inference.proportion_time_r_more_than_1(.95, 'mean')
def test_get_intervals(self):
df = pd.DataFrame({
'Time': [1, 2, 3, 5, 6],
'Incidence Number': [0, 0, 0, 0, 0]
})
ser_int = [1, 2]
inference = bp.BranchProPosterior(df, ser_int, 1, 0.2)
inference.run_inference(tau=2)
intervals_df = inference.get_intervals(.95)
self.assertEqual(len(intervals_df['Time Points']), 3)
self.assertEqual(len(intervals_df['Mean']), 3)
self.assertEqual(len(intervals_df['Median']), 3)
self.assertEqual(len(intervals_df['Lower bound CI']), 3)
self.assertEqual(len(intervals_df['Upper bound CI']), 3)
self.assertListEqual(intervals_df['Mean'].to_list(), [5.0] * 3)
self.assertEqual(intervals_df['Central Probability'].to_list(),
[.95] * 3)
def test_add_interval_rt(self):
df = pd.DataFrame({
'Time': [1, 2, 3, 5, 6],
'Incidence Number': [0, 0, 0, 0, 0]
})
ser_int = [1, 2]
inference = bp.BranchProPosterior(df, ser_int, 1, 0.2)
inference.run_inference(tau=2)
intervals_df = inference.get_intervals(.95)
my_plot = bp.ReproductionNumberPlot()
my_plot.add_interval_rt(intervals_df)
npt.assert_array_equal(np.array([my_plot.figure['data'][0]['x']]),
np.array([np.array([4, 5, 6])]))
npt.assert_array_equal(np.array([my_plot.figure['data'][0]['y']]),
np.array([np.array([5.0] * 3)]))
npt.assert_array_equal(np.array([my_plot.figure['data'][1]['x']]),
np.array([np.array([4, 5, 6, 6, 5, 4])]))
npt.assert_array_almost_equal(
np.array([my_plot.figure['data'][1]['y']]),
np.array([np.array([18.444397] * 3 + [0.126589] * 3)]))
with self.assertRaises(TypeError):
bp.ReproductionNumberPlot().add_interval_rt(0)
with self.assertWarns(UserWarning):
df = pd.DataFrame({
'Time Points': [1, 2, 3, 4, 5, 6],
'R_t': [3, 3, 0.5, 0.5, 0.5, 0.5]
})
my_plot = bp.ReproductionNumberPlot()
my_plot.add_ground_truth_rt(df)
dfs1 = pd.DataFrame({
't': [4, 5, 6],
'Mean': [5.0] * 3,
'Lower bound CI': [5.0] * 3,
'Upper bound CI': [5.0] * 3,
'Central Probability': [.95] * 3
})
my_plot.add_interval_rt(dfs1, time_key='t')
with self.assertWarns(UserWarning):
df = pd.DataFrame({
'Time Points': [1, 2, 3, 4, 5, 6],
'R_t': [3, 3, 0.5, 0.5, 0.5, 0.5]
})
my_plot = bp.ReproductionNumberPlot()
my_plot.add_ground_truth_rt(df)
dfs2 = pd.DataFrame({
'Time Points': [4, 5, 6],
'r': [5.0] * 3,
'Lower bound CI': [5.0] * 3,
'Upper bound CI': [5.0] * 3,
'Central Probability': [.95] * 3
})
my_plot.add_interval_rt(dfs2, r_key='r')
def update_posterior(self,
mean,
stdev,
tau,
central_prob,
epsilon=None,
progress_fn=None):
"""Update the posterior distribution based on slider values.
Parameters
----------
mean
(float) updated position on the slider for the mean of
the prior for the Branch Pro model in the posterior.
stdev
(float) updated position on the slider for the standard deviation
of the prior for the Branch Pro model in the posterior.
tau
(int) updated position on the slider for the tau window used in the
running of the inference of the reproduction numbers of the Branch
Pro model in the posterior.
central_prob
(float) updated position on the slider for the level of the
computed credible interval of the estimated R number values.
epsilon
(float) updated position on the slider for the constant of
proportionality between local and imported cases for the Branch Pro
model in the posterior.
progress_fn
Function of integer argument to send to posterior run_inference.
It can be used for dash callbacks set_progress (see
update_posterior_storage in the app script)
Returns
-------
pandas.DataFrame
The posterior distribution, summarized in a dataframe with the
following columns: 'Time Points', 'Mean', 'Lower bound CI' and
'Upper bound CI'
"""
new_alpha = (mean / stdev)**2
new_beta = mean / (stdev**2)
data = self.session_data.get('data_storage')
if data is None:
raise dash.exceptions.PreventUpdate()
time_label, inc_label = data.columns[:2]
num_cols = len(self.session_data.get('interval_storage').columns)
prior_params = (new_alpha, new_beta)
labels = {'time_key': time_label, 'inc_key': inc_label}
if num_cols == 1:
serial_interval = self.session_data.get(
'interval_storage').iloc[:, 0].values
if 'Imported Cases' in data.columns:
# Separate data into local and imported cases
imported_data = pd.DataFrame({
time_label: data[time_label],
inc_label: data['Imported Cases']
})
# Posterior follows the LocImp behaviour
posterior = bp.LocImpBranchProPosterior(
data, imported_data, epsilon, serial_interval,
*prior_params, **labels)
else:
# Posterior follows the simple behaviour
posterior = bp.BranchProPosterior(data, serial_interval,
*prior_params, **labels)
posterior.run_inference(tau)
else:
serial_intervals = self.session_data.get(
'interval_storage').values.T
if 'Imported Cases' in data.columns:
# Separate data into local and imported cases
imported_data = pd.DataFrame({
time_label: data[time_label],
inc_label: data['Imported Cases']
})
# Posterior follows the LocImp behaviour
posterior = bp.LocImpBranchProPosteriorMultSI(
data, imported_data, epsilon, serial_intervals,
*prior_params, **labels)
else:
# Posterior follows the simple behaviour
posterior = bp.BranchProPosteriorMultSI(
data, serial_intervals, *prior_params, **labels)
posterior.run_inference(tau, progress_fn=progress_fn)
return posterior.get_intervals(central_prob)
