def test_chose_draw_plot(self, mock_plt):
chose_draw_plot(["pop"], self.get_stats(is_empty=False), True)
assert mock_plt.called
chose_draw_plot(["R0"], self.get_stats(is_empty=False), True)
assert mock_plt.called
chose_draw_plot(["summ"], self.get_stats(is_empty=False), True)
assert mock_plt.called
chose_draw_plot(["lock"], self.get_stats(is_empty=False), True)
assert mock_plt.called
chose_draw_plot(["new"], self.get_stats(is_empty=False), True)
assert mock_plt.called
chose_draw_plot(["hos"], self.get_stats(is_empty=False), True)
assert mock_plt.called
joblib.dump(env_dic, env_file)
scenario_day_dic = {
-1: scx_base_just_a_flu.do_parallel_run,
0: sc0_base_lockdown.do_parallel_run,
1: sc1_simple_lockdown_removal.do_parallel_run,
2: sc2_yoyo_lockdown_removal.do_parallel_run,
3: sc3_loose_lockdown.do_parallel_run,
4: sc4_rogue_citizen.do_parallel_run,
5: sc5_rogue_neighborhood.do_parallel_run,
6: sc6_travelers.do_parallel_run,
7: sc7_nominal_lockdown_removal.do_parallel_run,
8: sc8_innoculation.do_parallel_run,
9: sc9_vaccination.do_parallel_run
}
scenario_variant_dic = {
10: sc10_variant.do_parallel_run
}
if params[scenario_id_key] in scenario_day_dic:
stats_result = launch_parallel_byday(params, env_dic, scenario_day_dic[params[scenario_id_key]],
params[ncpu_key])
elif params[scenario_id_key] in scenario_variant_dic:
stats_result = launch_parallel_byvariant(params, env_dic, scenario_variant_dic[params[scenario_id_key]],
params[ncpu_key])
else:
raise KeyError("Unknown scenario")
print("It took : %.2f seconds" % (time.time() - t_start))
chose_draw_plot(params[draw_graph_key], stats_result, params[show_plot_key])
import random
import time
from scenario.example import sc0_base_lockdown
from simulator.constants.keys import random_seed_key, draw_graph_key
from simulator.helper.parser import get_parser
from simulator.helper.plot import chose_draw_plot
from simulator.helper.simulation import get_default_params
from simulator.helper.environment import get_environment_simulation
if __name__ == '__main__':
params = get_default_params()
args = get_parser().parse_args()
for arg in vars(args):
v = getattr(args, arg)
if arg in params and v is not None:
params[arg] = v
random.seed(params[random_seed_key])
env_dic = get_environment_simulation(params)
t_start = time.time()
stats_result = sc0_base_lockdown.launch_run(params, env_dic)
print("It took : %.2f seconds" % (time.time() - t_start))
chose_draw_plot(params[draw_graph_key], stats_result)