NetSIR.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat May 11 08:56:41 2019

@author: nolte

D. D. Nolte, Introduction to Modern Dynamics: Chaos, Networks, Space and Time, 2nd ed. (Oxford,2019)
"""

# https://www.python-course.eu/networkx.php
# https://networkx.github.io/documentation/stable/tutorial.html
# https://networkx.github.io/documentation/stable/reference/functions.html

import numpy as np
from scipy import integrate
from matplotlib import pyplot as plt
import networkx as nx
import time
from random import random

tstart = time.time()

plt.close('all')

r = 0.0001    # 0.00015
k = 10        # connections  50
dill = 14     # days ill
dpq = 14      # days shelter in place
fnq = 0.5     # fraction NOT sheltering in place
mr0 = 0.01    # mortality rate
mr1 = 0.03     # extra mortality rate if exceeding hospital capacity
P = 330       # population of US in Millions
HC = 0.003    # hospital capacity

dinf = fnq*dill + (1-fnq)*np.exp(-dpq/dill)*dill;

#betap = r*k*dinf;
#mu = 1/dill;
betap = 0.014;
mu = 0.13;

print('beta = ',betap)
print('dinf = ',dinf)
print('betap/mu = ',betap/mu)


N = 128      # 50

# model_case 1 = Complete Graph
# model_case 2 = Barabasi
# model_case 3 = Power Law
# model_case 4 = D-dimensional Hypercube
# model_case 5 = Erdos Renyi
# model_case 6 = Random Regular
# model_case 7 = Strogatz
# model_case 8 = Hexagonal lattice
# model_case 9 = Tree

model_case = int(input('Input Model Case (1-9)'))
if model_case == 1: # Complete Graph
    facoef = 0.5
    nodecouple = nx.complete_graph(N)
elif model_case == 2: # Barabasi
    facoef = 2
    k = 1
    nodecouple = nx.barabasi_albert_graph(N, k, seed=None)
elif model_case == 3: # Power law
    facoef = 3
    k = 2
    triangle_prob = 0.05
    nodecouple = nx.powerlaw_cluster_graph(N, k, triangle_prob)        
elif model_case == 4:
    Dim = 6
    facoef = 3
    nodecouple = nx.hypercube_graph(Dim)
    N = nodecouple.number_of_nodes()
elif model_case == 5: # Erdos-Renyi
    facoef = 5
    prob = 0.03
    nodecouple = nx.erdos_renyi_graph(N, prob, seed=None, directed=False)
elif model_case == 6: # Random
    facoef = 10
    nodecouple = nx.random_regular_graph(3, N, seed=None)
elif model_case == 7: # Watts
    facoef = 7
    k = 4;              # nearest neighbors
    rewire_prob = 0.2   # rewiring probability
    nodecouple = nx.watts_strogatz_graph(N, k, rewire_prob, seed=None)
elif model_case == 8:   # Hexagonal lattice
    facoef = 8
    rows = 4
    colm = 8
    nodecouple = nx.hexagonal_lattice_graph(rows, colm, periodic=True, with_positions=False)
    N = nodecouple.number_of_nodes()
elif model_case == 9: # k-fold tree
    facoef = 16
    k = 3  # degree
    h = 4  # height
    sm = 0
    for lp in range(h+1):
        sm = sm + k**lp
    N = sm
    nodecouple = nx.balanced_tree(k, h)

indhi = 0
deg = 0
for omloop in nodecouple.node:
    degtmp = nodecouple.degree(omloop)
    if degtmp > deg:
        deg = degtmp
        indhi = omloop
print('highest degree node = ',indhi)
print('highest degree = ',deg)

plt.figure(1)
colors = [(random(), random(), random()) for _i in range(10)]
#nx.draw(nodecouple)
nx.draw_circular(nodecouple,node_size=75, node_color=colors)
#nx.draw_spring(nodecouple,node_size=75, node_color=colors)
#nx.draw_spectral(nodecouple)

#print('Number of nodes = ',nodecouple.number_of_nodes())
#print('Number of edges = ',nodecouple.number_of_edges())
#print('Average degree = ',nx.k_nearest_neighbors(nodecouple))
print(nx.info(nodecouple))
        
# function: omegout, yout = coupleN(G)
def coupleN(G,tlim):

    # function: yd = flow_deriv(x_y)
    def flow_deriv(x_y,t0):
        
        N = np.int(x_y.size/2)
        yd = np.zeros(shape=(2*N,))
        ind = -1
        for omloop in G.node:
            ind = ind + 1
            temp1 = -mu*x_y[ind] + betap*x_y[ind]*x_y[N+ind]
            temp2 =  -betap*x_y[ind]*x_y[N+ind]
            linksz = G.node[omloop]['numlink']
            for cloop in range(linksz):
                cindex = G.node[omloop]['link'][cloop]
                indx = G.node[cindex]['index']
                g = G.node[omloop]['coupling'][cloop]
                
                temp1 = temp1 + g*betap*x_y[indx]*x_y[N+ind]
                temp2 = temp2 - g*betap*x_y[indx]*x_y[N+ind]
            
            yd[ind] = temp1
            yd[N+ind] = temp2
                
        return yd
    # end of function flow_deriv(x_y)
    x0 = x_y
    t = np.linspace(0,tlim,tlim)      # 600  300
    y = integrate.odeint(flow_deriv, x0, t)        
    
    return t,y
    # end of function: omegout, yout = coupleN(G)

lnk = np.zeros(shape = (N,), dtype=int)
ind = -1
for loop in nodecouple.node:
    ind = ind + 1
    nodecouple.node[loop]['index'] = ind
    nodecouple.node[loop]['link'] = list(nx.neighbors(nodecouple,loop))
    nodecouple.node[loop]['numlink'] = len(list(nx.neighbors(nodecouple,loop)))
    lnk[ind] = len(list(nx.neighbors(nodecouple,loop)))

gfac = 0.1

ind = -1
for nodeloop in nodecouple.node:
    ind = ind + 1
    nodecouple.node[nodeloop]['coupling'] = np.zeros(shape=(lnk[ind],))
    for linkloop in range (lnk[ind]):
        nodecouple.node[nodeloop]['coupling'][linkloop] = gfac*facoef
            
x_y = np.zeros(shape=(2*N,))   
for loop in nodecouple.node:
    x_y[loop]=0
    x_y[N+loop]=nodecouple.degree(loop)
    #x_y[N+loop]=1
x_y[N-1]=1
x_y[2*N-1] = x_y[2*N-1] - 0.1
N0 = np.sum(x_y[N:2*N]) - x_y[indhi] - x_y[N+indhi]
     
tlim0 = 600
t0,yout0 = coupleN(nodecouple,tlim0)                           # Here is the subfunction call for the flow


plt.figure(2)
plt.yscale('log')
plt.gca().set_ylim(1e-3, 1)
for loop in range(N):
    lines1 = plt.plot(t0,yout0[:,loop])
    lines2 = plt.plot(t0,yout0[:,N+loop])
    lines3 = plt.plot(t0,N0-yout0[:,loop]-yout0[:,N+loop])

    plt.setp(lines1, linewidth=0.5)
    plt.setp(lines2, linewidth=0.5)
    plt.setp(lines3, linewidth=0.5)
    

Itot = np.sum(yout0[:,0:127],axis = 1) - yout0[:,indhi]
Stot = np.sum(yout0[:,128:255],axis = 1) - yout0[:,N+indhi]
Rtot = N0 - Itot - Stot
plt.figure(3)
plt.plot(t0,Itot,t0,Stot,t0,Rtot)
plt.legend(('Infected','Susceptible','Removed'))
plt.hold




# Repeat but innoculate 
x_y = np.zeros(shape=(2*N,))   
for loop in nodecouple.node:
    x_y[loop]=0
    x_y[N+loop]=nodecouple.degree(loop)
    #x_y[N+loop]=1
x_y[N-1]=1
x_y[2*N-1] = x_y[2*N-1] - 0.1
N0 = np.sum(x_y[N:2*N]) - x_y[indhi] - x_y[N+indhi]
     
tlim0 = 50
t0,yout0 = coupleN(nodecouple,tlim0)                           # Here is the subfunction call for the flow


# remove all edges from highest-degree node
ee = list(nodecouple.edges(indhi))
nodecouple.remove_edges_from(ee)
print(nx.info(nodecouple))

#nodecouple.remove_node(indhi)        
lnk = np.zeros(shape = (N,), dtype=int)
ind = -1
for loop in nodecouple.node:
    ind = ind + 1
    nodecouple.node[loop]['index'] = ind
    nodecouple.node[loop]['link'] = list(nx.neighbors(nodecouple,loop))
    nodecouple.node[loop]['numlink'] = len(list(nx.neighbors(nodecouple,loop)))
    lnk[ind] = len(list(nx.neighbors(nodecouple,loop)))

ind = -1
x_y = np.zeros(shape=(2*N,)) 
for nodeloop in nodecouple.node:
    ind = ind + 1
    nodecouple.node[nodeloop]['coupling'] = np.zeros(shape=(lnk[ind],))
    x_y[ind] = yout0[tlim0-1,nodeloop]
    x_y[N+ind] = yout0[tlim0-1,N+nodeloop]
    for linkloop in range (lnk[ind]):
        nodecouple.node[nodeloop]['coupling'][linkloop] = gfac*facoef

#N1 = np.sum(yout0[tlim0-1,:])
#x_y = yout0[tlim0-1,:]
    
tlim1 = 500
t1,yout1 = coupleN(nodecouple,tlim1)

t = np.zeros(shape=(tlim0+tlim1,))
yout = np.zeros(shape=(tlim0+tlim1,2*N))
t[0:tlim0] = t0
t[tlim0:tlim1+tlim0] = tlim0+t1
yout[0:tlim0,:] = yout0
yout[tlim0:tlim1+tlim0,:] = yout1


plt.figure(4)
plt.yscale('log')
plt.gca().set_ylim(1e-3, 1)
for loop in range(N):
    lines1 = plt.plot(t,yout[:,loop])
    lines2 = plt.plot(t,yout[:,N+loop])
    lines3 = plt.plot(t,N0-yout[:,loop]-yout[:,N+loop])

    plt.setp(lines1, linewidth=0.5)
    plt.setp(lines2, linewidth=0.5)
    plt.setp(lines3, linewidth=0.5)
    

Itot = np.sum(yout[:,0:127],axis = 1) - yout[:,indhi]
Stot = np.sum(yout[:,128:255],axis = 1) - yout[:,N+indhi]
Rtot = N0 - Itot - Stot
plt.figure(3)
plt.plot(t,Itot,t,Stot,t,Rtot,linestyle='dashed')
plt.legend(('Infected','Susceptible','Removed'))
plt.hold()


elapsed_time = time.time() - tstart
print('elapsed time = ',format(elapsed_time,'.2f'),'secs')