Simulating an outbreak in different social networks.

network_outbreak_sim.py

import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
import seaborn as sns


def sim_outbreak(G, exps, inf_prob, init_infect, lock):
    pop_size = len(G)
    total_infected = []
    for exp in range(exps):
        locked = lock
        A = np.array(nx.adjacency_matrix(G).todense(), dtype="float")
        A[A > 0] = inf_prob * np.random.rand(2 * len(G.edges))

        infected = np.random.randint(pop_size, size=init_infect)
        all_infected = set(infected)
        while len(infected) > 0:
            interact_infect_probs = np.random.rand(len(infected), pop_size)
            inf_contacts = interact_infect_probs < A[infected]
            cand_inf = set(np.where(inf_contacts)[1])
            new_inf = cand_inf - all_infected
            all_infected |= new_inf
            if locked and (len(new_inf) == 0):
                print("Ending lockdown.")
                print(f"Total infected: {len(all_infected)}")
                print(f"Currently infected: {len(infected)}")
                G = nx.complete_graph(len(G))
                A = np.array(nx.adjacency_matrix(G).todense(), dtype="float")
                A[A > 0] = inf_prob * np.random.rand(int(A.sum()))
                locked = False
            else:
                infected = np.array(list(new_inf))

        total_infected.append(len(all_infected))

    print(np.mean(total_infected))
    sns.distplot(total_infected)
    plt.show()


def sim_two_network(friends, exps, inf_prob, init_infect, lock):
    G1 = nx.karate_club_graph()
    A1 = np.array(nx.adjacency_matrix(G1).todense(), dtype="float")
    half_pop_size = len(G1)

    G2 = nx.complete_graph(half_pop_size)
    A2 = np.array(nx.adjacency_matrix(G2).todense(), dtype="float")

    A = np.zeros((2 * half_pop_size, 2 * half_pop_size))
    A[:half_pop_size, :half_pop_size] = A1
    A[half_pop_size:, half_pop_size:] = A2

    g1_inds = np.random.randint(half_pop_size, size=friends)
    g2_inds = np.random.randint(half_pop_size, 2 * half_pop_size, size=friends)

    A[g1_inds, g2_inds] = 1
    A[g2_inds, g1_inds] = 1

    G = nx.from_numpy_array(A)

    pop_size = len(G)
    total_infected = []
    for exp in range(exps):
        locked = lock
        A = np.array(nx.adjacency_matrix(G).todense(), dtype="float")
        A[A > 0] = inf_prob * np.random.rand(2 * len(G.edges))
        if locked:
            A[g1_inds, g2_inds] = 0
            A[g2_inds, g1_inds] = 0

        infected = np.random.randint(pop_size, size=init_infect)
        all_infected = set(infected)
        while len(infected) > 0:
            interact_infect_probs = np.random.rand(len(infected), pop_size)
            inf_contacts = interact_infect_probs < A[infected]
            cand_inf = set(np.where(inf_contacts)[1])
            new_inf = cand_inf - all_infected
            all_infected |= new_inf
            if locked and (len(new_inf) == 0):
                print("Ending lockdown.")
                print(f"Total infected: {len(all_infected)}")
                print(f"Currently infected: {len(infected)}")
                A[g1_inds, g2_inds] = inf_prob * np.random.rand(len(g1_inds))
                A[g2_inds, g1_inds] = inf_prob * np.random.rand(len(g2_inds))
                locked = False
            else:
                infected = np.array(list(new_inf))

        total_infected.append(len(all_infected))

    print(np.mean(total_infected))
    sns.distplot(total_infected)
    plt.show()


def main():
    G = nx.karate_club_graph()
    nx.draw(G)
    plt.show()
    exps = 1000
    inf_prob = 0.25
    init_infect = 3
    sim_outbreak(G, exps, inf_prob, init_infect, False)
    sim_outbreak(G, exps, inf_prob, init_infect, True)

    G = nx.complete_graph(len(G))
    nx.draw(G)
    plt.show()
    sim_outbreak(G, exps, inf_prob, init_infect, False)

    friends = 3
    sim_two_network(friends, exps, inf_prob, 2 * init_infect, False)
    sim_two_network(friends, exps, inf_prob, 2 * init_infect, True)


if __name__ == "__main__":
    main()