daily pastebin goal
66%
SHARE
TWEET

Untitled

a guest Nov 17th, 2018 68 Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. import numpy as np
  2. import ODESolver
  3. import matplotlib.pyplot as plt
  4. import piecewise
  5.  
  6. class ProblemSIZR():
  7.     def __init__(self, eSS, eSI, alpha, beta, S0, I0, Z0, R0, T, sum, ro):
  8.  
  9.         self.T = T
  10.         self.sum = sum
  11.         self.S0 = S0
  12.         self.I0 = I0
  13.         self.Z0 = Z0
  14.         self.R0 = R0
  15.         t = np.linspace(0,T)
  16.  
  17.         if isinstance(eSS, (float,int)):
  18.             self.eSS = lambda t: eSS
  19.         elif callable(eSS):
  20.             self.eSS = eSS
  21.  
  22.         if isinstance(eSI, (float, int)):
  23.             self.eSI = lambda t: eSI
  24.         elif callable(eSI):
  25.             self.eSI = eSI
  26.  
  27.         if isinstance(alpha, (float, int)):
  28.             self.alpha = lambda t: alpha
  29.         elif callable(alpha):
  30.             self.alpha = alpha
  31.  
  32.         if isinstance(ro, (float, int)):
  33.             self.ro = lambda t: ro
  34.         elif callable(ro):
  35.             self.ro = ro
  36.  
  37.         if isinstance(beta, (float,int)):
  38.             self.beta = lambda t: beta
  39.         elif callable(beta):
  40.             self.beta = beta
  41.  
  42.  
  43.  
  44.     def __call__(self, u, t):
  45.         S, I, Z, R = u
  46.         return [self.sum - self.beta(t)*S*Z - self.eSS(t)*S, self.beta(t)*S*Z -
  47.         self.ro(t)*I  - self.eSI(t)*I, self.ro(t)*I - self.alpha(t)*S*Z,
  48.         self.eSS(t)*S + self.eSI(t)*I + self.alpha(t)*S*Z]
  49.  
  50.  
  51. class SolverSIZR(object):
  52.     def __init__(self, problem, dt):
  53.         self.problem, self.dt = problem, dt
  54.  
  55.     def solveSIZR(self):
  56.         self.solver = ODESolver.RungeKutta4(self.problem)
  57.         ic = [self.problem.S0, self.problem.I0, self.problem.Z0, self.problem.R0]
  58.         self.solver.set_initial_condition(ic)
  59.         n = int(round(self.problem.T/float(self.dt)))
  60.         t = np.linspace(0, self.problem.T, n+1)
  61.         u, self.t = self.solver.solve(t)
  62.         self.S, self.I, self.Z, self.R = u[:,0], u[:,1], u[:,2], u[:,3]
  63.  
  64.  
  65.  
  66.     def plotSIZR(self):
  67.         plt.plot(self.t, self.S)
  68.         plt.plot(self.t, self.I)
  69.         plt.plot(self.t, self.Z)
  70.         plt.plot(self.t, self.R)
  71.         plt.xlabel('Days passed since breakout')
  72.         plt.ylabel('Number of people in a group')
  73.         plt.legend(['Susceptibles', 'Infected', 'Zombies', 'Removed'])
  74.         plt.show()
  75.  
  76.  
  77. problem = ProblemSIZR(eSS=lambda t: 0 if t < 28  else 0.0067,
  78.                       eSI=lambda t: 0.0014 if 4 < t < 28 else 0,
  79.                       alpha=lambda t: 0 if t < 4 elif 0.0016 if 4 < t < 28 else 0.006,
  80.                       beta=lambda t: 0.03 if t < 4 elif 0.0012 if 4 < t < 28 else 0,
  81.                       S0=60, I0=0, Z0=1, R0=0, T=33,
  82.                       sum=lambda t: 20 if t < 4 elif 2 if 4 < t < 28 else 0, ro=1)
  83.  
  84.  
  85. Solution = SolverSIZR(problem, dt=0.5)
  86. Solution.solveSIZR()
  87. Solution.plotSIZR()
RAW Paste Data
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy. OK, I Understand
 
Top