text: import EoN
  import networkx as nx
  import matplotlib.pyplot as plt
  import scipy
  import random
  from scipy import integrate
  
  
  
  '''
  Code to generate figure 2.11.  This is a bit messy because we have to
  define the ODE models.  Since python deals with ODEs by taking 1D arrays
  we have to set up all the variables into a single long vector.
  '''
  
  def star(N):
      G = nx.Graph()
      G.add_node(0)
      for node_id in range(1,N):
          G.add_edge(0,node_id)
      return G
  
  
  
  
  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, tau, gamma, 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
  
  
  def star_graph_dX(X, t, tau, gamma, N):
      '''this system is given in Proposition 2.4, taking Q=S, T=I
      so f_{SI}(k) = f_{QT}(k) = k*tau
      f_{IS}(k) = f_{TQ}(k) = gamma
      X has length 2*(N-1)+2 = 2N'''
  
      #    [[central node infected] + [central node susceptible]]
      #X = [Y_1^1, Y_1^2, ..., Y_1^{N}, Y_2^0, Y_2^1, ..., Y_2^{N-1}]
  
      #Note that in proposition Y^0 is same as Y_2^0
      #and Y^N is same as Y_1^N
  
      #Y1[k]: central node infected, & k-1 peripheral nodes infected
      Y1vec = [0]+list(X[0:N])      #for Y_1^k, use Y1vec[k]
      #pad with 0 to make easier calculations Y_1^0=0
      #the probability of -1 nodes infected is 0
  
      #Y2[k]: central node susceptible & k peripheral nodes infected
      Y2vec = list(X[N:])+[0]   #for Y_2^k use Y2vec[k]
      #padded with 0 to make easier calculations. Y_2^N=0
      #the probability of N (of N-1) peripheral nodes infected is 0
      dY1vec = []
      dY2vec = []
      for k in range(1, N):
          #k-1 peripheral nodes infected, central infected
          dY1vec.append((N-k+1)*tau*Y1vec[k-1] + (k-1)*tau*Y2vec[k-1]
                      +k*gamma*Y1vec[k+1]
                      - ((N-k)*tau + (k-1)*gamma+gamma)*Y1vec[k])
      #now the Y^N equation
      dY1vec.append(tau*Y1vec[N-1] + (N-1)*tau*Y2vec[N-1] - N*gamma*Y1vec[N])
  
      #now the Y^0 equation
      dY2vec.append(gamma*(N-1)*Y1vec[1] + gamma*Y2vec[1]-0)
  
      for k in range(1,N):
          #k peripheral nodes infected, central susceptible
          dY2vec.append(0 + gamma*Y1vec[k+1] + gamma*(k+1)*Y2vec[k+1]
                      - (k*tau + 0 + k*gamma)*Y2vec[k])
  
      return scipy.array(dY1vec + dY2vec)
  
  def star_graph_lumped(N, tau, gamma, I0, tmin, tmax, tcount):
      times = scipy.linspace(tmin, tmax, tcount)
      #    [[central node infected] + [central node susceptible]]
      #X = [Y_1^1, Y_1^2, ..., Y_1^{N}, Y_2^0, Y_2^1, ..., Y_2^{N-1}]
      X0 = scipy.zeros(2*N)  #length 2*N of just 0 entries
      X0[I0]=I0*1./N #central infected, + I0-1 periph infected prob
      X0[N+I0] = 1-I0*1./N #central suscept + I0 periph infected
      X = EoN.my_odeint(star_graph_dX, X0, times, args = (tau, gamma, N))
      #X looks like [[central susceptible,k periph] [ central inf, k-1 periph]] x T
  
      central_inf = X[:,:N]
      central_susc = X[:,N:]
  
      I = scipy.array([ sum(k*central_susc[t][k] for k in range(N))
              + sum((k+1)*central_inf[t][k] for k in range(N))
              for t in range(len(X))])
      S = N-I
      return times, S, I
  
  
  
  
  N=1000
  I0=int(0.1*N)
  iterations = 100 #number of simulations to compare
  gamma = 1
  tmin=0
  tmax=5
  tcount = 1001
  report_times = scipy.linspace(tmin, tmax, 21) #for simulations
  
  
  plt.figure(0)
  tau = 0.005
  G = nx.complete_graph(N)
  t, S, I = complete_graph_lumped(N, tau, gamma, I0, tmin, tmax, tcount)
  plt.plot(t, I/N, label = 'Prediction')
  
  #now check with simulation
  obs_I = 0*report_times
  print("done with complete graph ODE.  Now simulating")
  for counter in range(iterations):
      IC = random.sample(range(N),I0)
      t, S, I = EoN.fast_SIS(G, tau, gamma, initial_infecteds = IC, tmax = tmax)
      obs_I += EoN.subsample(report_times, t, I)
  plt.plot(report_times, obs_I*1./(iterations*N), 'o', label='Simulation')
  plt.axis(ymin=0, ymax=1)
  plt.xlabel('$t$')
  plt.ylabel('$[I]$')
  plt.legend()
  plt.savefig('fig2p11a.png')
  
  print("done with complete graph.  Now star --- warning, this may be slow")
  
  
  
  
  
  
  
  
  plt.clf()
  #for star, if 100 nodes randomly start infected, 1/10 cases have
  #central node infected and 99 peripheral.  9/10 have 100 peripheral.
  
  tau = 4.
  G = star(N)
  
  t, S, I = star_graph_lumped(N, tau, gamma, I0, tmin, tmax, tcount)
  plt.plot(t, I/N, label = 'Prediction')
  print("done with star ODE, now simulating")
  
  obs_I = 0*report_times
  for counter in range(iterations):
      IC = random.sample(range(N),I0)
      t, S, I = EoN.fast_SIS(G, tau, gamma, initial_infecteds = IC, tmax = tmax)
      obs_I += EoN.subsample(report_times, t, I)
  plt.plot(report_times, obs_I*1./(iterations*N), 'o', label='Simulation')
  plt.axis(ymin=0, ymax=1)
  plt.xlabel('$t$')
  plt.ylabel('$[I]$')
  plt.legend()
  plt.savefig('fig2p11b.png')
  