unequal-spread/functions.py at master · Sepante/unequal-spread · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
import numpy as np

def init_agents(agents, N, initial_infectious_num, sizes, social_class_num):
    #global agents
    agents['strategy'] = 1 #initially they all choose to go out.
    agents['health'] = 0 #initially they all choose to go out.
    agents['future'] = 0 #initially they all choose to go out.
    #infection_seed = np.random.randint(0, N) #the first infected node

    seed_share = np.round( sizes/sizes.sum() * initial_infectious_num )
    seed_share = seed_share.astype('int')

    if seed_share.sum() != initial_infectious_num:
        print( 'Actually ' + str(seed_share.sum()) + ' infectious seeds have been used.' )

    for social_class in range( social_class_num ):
        people_from_the_class = np.where( agents['social_class'] == social_class )[0]
        infection_seed = np.random.permutation( people_from_the_class )[: seed_share[social_class] ]
        #print(infection_seed)
        agents['health'][infection_seed] = 1
        agents['future'][infection_seed] = 1
        #print(infection_seed)

    #print( np.where(agents['health'] == 1)[0] )


def infect(G, agents, transmit_prob):
    health_list = (agents['health'] > 0)
    #strategy_list = (agents['strategy'] == 1)
    out_and_infected = np.all( [health_list, agents['strategy']], axis = 0 )
    infected_agents = np.where( out_and_infected )[0] #select infectious agents
    #print(len(infected_agents))
    for infectious in infected_agents:
        for neighbor in G.neighbors(infectious):
            if (agents['future'][neighbor] == 0 and agents['strategy'][neighbor] == 1): #if susceptible and out
                #print(neighbor)
                if np.random.random() < transmit_prob:
                    agents['future'][neighbor] = 1 #infect
    return np.sum(health_list)


def recover(agents, recovery_prob):
    infected_agents = np.where( (agents['health'] > 0) )[0]
    #print(len(infected_agents))
    future_status = np.random.choice([-1, 1], len(infected_agents), p=[recovery_prob, 1- recovery_prob] )

    agents['future'][infected_agents] = future_status

    return np.sum(future_status == -1)


def update_infection(agents):
    agents['future'][( agents['health'] == -1 )] = -2
    agents['health'][( agents['health'] == -1 )] = -2

    agents['health'] = agents['future']

def get_newly_recovered_agents(agents, social_class = -1, compartment_to_look_at = 1):
    if social_class == -1:
        return agents['health'] == compartment_to_look_at
    else:
        return np.all( [ agents['social_class'] == social_class , agents['health'] == compartment_to_look_at ], axis = 0 )


def predict_infected_num(N, agents, t_prediction, learning_rate):
    #print('inside = ', t_prediction)

    newly_infected = np.sum( get_newly_recovered_agents(agents, social_class = -1, compartment_to_look_at = 1) ) / N
    #if newly_infected:
        #print("newly: ", newly_infected)

    ## only averages for lowering situation, for increases considers the new value.
#    if newly_infected > t_prediction:
#        new_prediction = newly_infected
#    else:
#        new_prediction = t_prediction * (1-learning_rate ) + newly_infected * learning_rate
#    return new_prediction
    new_prediction = t_prediction * (1-learning_rate ) + newly_infected * learning_rate
    return new_prediction


def update_strategy(agents, exp_stay_home_reward, infection_reward_times_infected_num, beta):
    #survivor_num = np.sum( agents['health'] >= 0 )
    #survivor_num = len(agents)
    survivor_num = 1

    going_out_mean_reward = infection_reward_times_infected_num / survivor_num

    exp_going_out_reward = np.exp(beta * going_out_mean_reward)
    for social_class in range( len(exp_stay_home_reward) ):
        this_class_agents = (agents['social_class'] == social_class)

        staying_home_prob = exp_stay_home_reward[social_class]/ (exp_stay_home_reward[social_class] + exp_going_out_reward)
        #print (social_class, staying_home_prob)

        agents['strategy'][ this_class_agents ] = np.random.choice([0, 1], np.sum(this_class_agents), p=[staying_home_prob, 1- staying_home_prob] )
    #print(infection_reward_times_infected_num , staying_home_prob)
    #print(staying_home_prob)
    return survivor_num

def get_results(agents, social_class_num):
    finally_infected = agents['health'] != 0
    infected_classes = agents[finally_infected]['social_class']
    #print(infected_classes)
    infected_from_each_class = np.zeros(social_class_num, int)
    for social_class in range( social_class_num ):
        infected_from_each_class[social_class] = np.sum( infected_classes == social_class )
    return infected_from_each_class

def get_timed_results(agents, social_class_num):

    currently_infected = agents['health'] > 0
    infected_classes = agents[currently_infected]['social_class']
    infected_from_each_class = np.zeros(social_class_num, int)
    for social_class in range( social_class_num ):
        infected_from_each_class[social_class] = np.sum( infected_classes == social_class )
    return infected_from_each_class