def propagate_to_stores(env_dic, virus_dic, probability_store_infection_arg):
# Filter on living people because we have a random choice to make in each house
# People who will go to their store (one person per house as imposed by lockdown)
individuals_gotostore = get_random_choice_list([[i for i in env_dic[HA_K][h] if is_alive(i, virus_dic)]
for h in range(len(env_dic[HA_K]))])
# Contagious people who will go to their store
individuals_infected_gotostore = [i for i in individuals_gotostore if is_contagious(i, virus_dic)]
# Stores that will be visited by a contagious person
infected_stores = [env_dic[HS_K][env_dic[IH_K][i]] for i in individuals_infected_gotostore]
# People who live in a house that goes to a contagious store
individuals_attachedto_infected_store = flatten(flatten(
[[env_dic[HA_K][h] for h in env_dic[SH_K][s]] for s in infected_stores]))
# People who did go to that contagious store
individuals_goto_infected_store = list(set(individuals_attachedto_infected_store)
.intersection(set(individuals_gotostore)))
# People who got infected from going to their store
infected_backfromstore = [i for i in individuals_goto_infected_store if get_r() < probability_store_infection_arg]
# INFECTION STATE UPDATE
update_infection_period(infected_backfromstore, virus_dic)
def propagate_to_workplaces(env_dic, virus_dic, probability_work_infection_arg):
# Contagious people who will go to work
infected_gotowork = [i for i in get_infected_people(virus_dic) if i in env_dic[IW_K].keys()
and is_contagious(i, virus_dic)]
# Infected workplaces
infected_workplaces = [env_dic[IW_K][ind] for ind in infected_gotowork]
infected_backfromwork = [i for i in flatten([env_dic[WI_K][k] for k in infected_workplaces])
if get_r() < probability_work_infection_arg]
# INFECTION STATE UPDATE
update_infection_period(infected_backfromwork, virus_dic)
def propagate_to_houses(env_dic, virus_dic, probability_home_infection_arg):
# Houses that contain an infected and contagious person
infected_houses = [env_dic[IH_K][i] for i in get_infected_people(virus_dic)
if is_contagious(i, virus_dic)]
# People infected (not necessarily contagious) from a contagious person living in their house
infected_athome = [i for i in flatten([env_dic[HI_K][hou] for hou in infected_houses])
if get_r() < probability_home_infection_arg]
# INFECTION STATE UPDATE
update_infection_period(infected_athome, virus_dic)
def propagate_to_transportation(env_dic, virus_dic, probability_transport_infection_arg):
# Contagious people who will go to work
infected_who_goto_work = [i for i in get_infected_people(virus_dic) if i in env_dic[IW_K].keys()
and is_contagious(i, virus_dic)]
# Infected public transportation blocks
# No more reduce in python3 ... :(
infected_who_share_blocks = list(set(flatten([env_dic[ITI_K][ind] for ind in infected_who_goto_work])))
infected_bad_luck_transport = [k for k in infected_who_share_blocks
if get_r() < probability_transport_infection_arg]
# INFECTION STATE UPDATE
update_infection_period(infected_bad_luck_transport, virus_dic)
def test_flatten(self):
input_list = [[1, 2], [0], [1, 3]]
result = flatten(input_list)
self.assertEqual(result, [1, 2, 0, 1, 3])