def test_transition_matrix_shape_checking_correct(self):
graph = nx.Graph()
graph.add_nodes_from(range(50))
state_names = ['susceptible', 'infected']
# Same shape shouldn't raise.
transition_matrix = np.array([[0, 0], [0, 1]])
treatment_transition_matrix = np.array([[1, 0], [1, 0]])
_ = infectious_disease.Params(
population_graph=graph,
transition_matrix=transition_matrix,
treatment_transition_matrix=treatment_transition_matrix,
state_names=state_names,
healthy_index=0,
healthy_exit_index=1,
infectious_index=1,
infection_probability=0.1,
num_treatments=0,
max_treatments=0)
# Different shape should raise.
treatment_transition_matrix = np.array([[1, 0, 0], [1, 0, 0], [1, 0, 0]])
with self.assertRaises(ValueError):
_ = infectious_disease.Params(
population_graph=graph,
transition_matrix=transition_matrix,
treatment_transition_matrix=treatment_transition_matrix,
state_names=state_names,
healthy_index=0,
healthy_exit_index=1,
infectious_index=1,
infection_probability=0.1,
num_treatments=0,
max_treatments=0)
def test_param_equality_correct(self):
graph = nx.Graph()
graph.add_nodes_from(range(50))
state_names = ['susceptible', 'infected']
transition_matrix = np.ones((2, 2)) * 0.5
treatment_transition_matrix = np.ones((2, 2)) * 0.5
num_treatments = max_treatments = 0
p1 = infectious_disease.Params(
population_graph=graph,
transition_matrix=transition_matrix,
treatment_transition_matrix=treatment_transition_matrix,
state_names=state_names,
healthy_index=0,
healthy_exit_index=1,
infectious_index=1,
infection_probability=0.1,
num_treatments=num_treatments,
max_treatments=max_treatments)
p2 = infectious_disease.Params(
population_graph=graph,
transition_matrix=transition_matrix,
treatment_transition_matrix=treatment_transition_matrix,
state_names=state_names,
healthy_index=0,
healthy_exit_index=1,
infectious_index=1,
infection_probability=0.1,
num_treatments=num_treatments,
max_treatments=max_treatments)
self.assertEqual(p1, p2)
def test_unnormalized_transition_in_params_raises(self):
graph = nx.Graph()
graph.add_nodes_from(range(50))
state_names = ['susceptible', 'infected']
transition_matrix = np.ones((2, 2))
treatment_transition_matrix = np.ones((2, 2))
with self.assertRaises(ValueError):
_ = infectious_disease.Params(
population_graph=graph,
transition_matrix=transition_matrix,
treatment_transition_matrix=treatment_transition_matrix,
state_names=state_names,
healthy_index=0,
healthy_exit_index=1,
infectious_index=1,
infection_probability=0.1,
num_treatments=0,
max_treatments=0)
