import EoN
import networkx as nx
import matplotlib.pyplot as plt
import scipy
from scipy import integrate
def complete_graph_dX(X, t, tau, gamma, N):
r'''This system is given in Proposition 2.3, taking Q=S, T=I
f_{SI}(k) = f_{QT}= k*\tau
f_{IS}(k) = f_{TQ} = \gamma
\dot{Y}^0 = \gamma Y^1 - 0\\
\dot{Y}^1 = 2\gamma Y^2 + 0Y^0 - (\gamma + (N-1)\tau)Y^1
\dot{Y}^2 = 3\gamma Y^3 + (N-1)\tau Y^1 - (2\gamma+2(N-2))Y^2
...
\dot{Y}^N = (N-1)\tau Y^{N-1} - N\gamma Y^N
Note that X has length N+1
'''
#X[k] is probability of k infections.
dX = []
dX.append(gamma*X[1])
for k in range(1,N):
dX.append((k+1)*gamma*X[k+1]+ (N-k+1)*(k-1)*tau*X[k-1]
- ((N-k)*k*tau + k*gamma)*X[k])
dX.append((N-1)*tau*X[N-1] - N*gamma*X[N])
return scipy.array(dX)
def complete_graph_lumped(N, I0, tmin, tmax, tcount):
times = scipy.linspace(tmin, tmax, tcount)
X0 = scipy.zeros(N+1) #length N+1 of just 0 entries
X0[I0]=1. #start with 100 infected.
X = integrate.odeint(complete_graph_dX, X0, times, args = (tau, gamma, N))
#X[t] is array whose kth entry is p(k infected| time=t).
I = scipy.array([sum(k*Pkt[k] for k in range(len(Pkt))) for Pkt in X])
S = N-I
return times, S, I
N=200
gamma = 1
k = N-1.
tau_c = gamma/k
rho = 0.1
for tau, label in zip([0.9*tau_c, tau_c, 1.2*tau_c, 1.5*tau_c],['a', 'b', 'c', 'd']):
plt.clf()
t, S, I = complete_graph_lumped(N, int(N*rho), 0 , 20, 1001)
plt.plot(t, I)
S0 = (1-rho)*N
I0 = rho*N
t, S, I = EoN.SIS_homogeneous_meanfield(S0, I0, k, tau, gamma, tmin=0, tmax=20,
tcount=1001)
plt.plot(t, I, '--')
S0 = (1-rho)*N
I0 = rho*N
SI0 = (1-rho)*N*k*rho
SS0 = (1-rho)*N*k*(1-rho)
t, S, I = EoN.SIS_homogeneous_pairwise(S0, I0, SI0, SS0, k, tau, gamma, tmin = 0,
tmax=20, tcount=1001)
plt.plot(t, I, ':')
plt.xlabel('$t$')
plt.ylabel('Prevalence')
plt.savefig('fig4p5{}.png'.format(label))
