def test_beta_search(self):
expected_beta = 0.5
model_init = sir.Model(alpha=0.3, beta=expected_beta)
experimantal = model_init.run()
actual_under_beta = sir.model_approx(model_init,
experimantal,
var_name='beta',
var_init=0.1,
learning_rate=0.01,
max_iter=100)
actual_over_beta = sir.model_approx(model_init,
experimantal,
var_name='beta',
var_init=1,
learning_rate=0.01,
max_iter=100)
self.assertTrue(close(actual_under_beta, expected_beta))
self.assertTrue(close(actual_over_beta, expected_beta))
def test_susceptible_search(self):
exp_susceptible = 750
model_init = sir.Model(alpha=0.5, beta=0.3, s_init=exp_susceptible)
experimantal = model_init.run()
actual_under = sir.model_approx(model_init,
experimantal,
var_name='s_init',
var_init=300,
learning_rate=2,
max_iter=300)
actual_over = sir.model_approx(model_init,
experimantal,
var_name='s_init',
var_init=1200,
learning_rate=2,
max_iter=300)
self.assertTrue(close(actual_under, exp_susceptible))
self.assertTrue(close(actual_over, exp_susceptible))
def test_small_beta(self):
expected_beta = 0.0005
model_init = sir.Model(alpha=0.03, beta=expected_beta)
experimantal = model_init.run()
actual_under_beta = sir.model_approx(model_init,
experimantal,
var_name='beta',
var_init=0,
learning_rate=0.0005,
max_iter=500,
threshold=0)
actual_over_beta = sir.model_approx(model_init,
experimantal,
var_name='beta',
var_init=0.001,
learning_rate=0.00001,
max_iter=500)
self.assertTrue(close(actual_under_beta, expected_beta, eps=0.0001))
self.assertTrue(close(actual_over_beta, expected_beta, eps=0.0001))
def update(val=0):
alpha = s_alpha.val / 100000
beta = s_beta.val / 500000
model = sir.Model(alpha, beta, s_sInit.val, s_iInit.val, 0)
model.time_unit = 10
data = model.run()
basic_reproduction = 2 # sir['Ro']
if basic_reproduction >= 0:
print('epidemic')
else:
print('none epidemic')
print(basic_reproduction)
print('alpha: ' + str(alpha) + ' beta: ' + str(beta))
s_line.set_data(range(len(data[0])), data[0])
i_line.set_data(range(len(data[1])), data[1])
r_line.set_data(range(len(data[2])), data[2])
fig.canvas.draw_idle()
from mpl_toolkits.mplot3d import axes3d import matplotlib.pyplot as plt import sir fig = plt.figure() ax = fig.add_subplot(111, projection='3d') model = sir.Model(0.03, 0.0008) X, Y, Z = model.run() ax.plot_wireframe(X, Y, Z) plt.show()