Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- import numpy as np
- import ODESolver
- import matplotlib.pyplot as plt
- class ProblemSIZR():
- def __init__(self, eSS, eSI, alpha, beta, S0, I0, Z0, R0, T, sum, ro):
- self.T = T
- self.sum = sum
- self.S0 = S0
- self.I0 = I0
- self.Z0 = Z0
- self.R0 = R0
- t = np.linspace(0,T)
- if isinstance(eSS, (float,int)):
- self.eSS = lambda t: eSS
- elif callable(eSS):
- self.eSS = eSS
- if isinstance(eSI, (float, int)):
- self.eSI = lambda t: eSI
- elif callable(eSI):
- self.eSI = eSI
- if isinstance(alpha, (float, int)):
- self.alpha = lambda t: alpha
- elif callable(alpha):
- self.alpha = alpha
- if isinstance(ro, (float, int)):
- self.ro = lambda t: ro
- elif callable(ro):
- self.ro = ro
- if isinstance(beta, (float,int)):
- self.beta = lambda t: beta
- elif callable(beta):
- self.beta = beta
- def __call__(self, u, t):
- S, I, Z, R = u
- return [self.sum - self.beta(t)*S*Z - self.eSS(t)*S, self.beta(t)*S*Z -
- self.ro(t)*I - self.eSI(t)*I, self.ro(t)*I - self.alpha(t)*S*Z,
- self.eSS(t)*S + self.eSI(t)*I + self.alpha(t)*S*Z]
- class SolverSIZR(object):
- def __init__(self, problem, dt):
- self.problem, self.dt = problem, dt
- def solveSIZR(self):
- self.solver = ODESolver.RungeKutta4(self.problem)
- ic = [self.problem.S0, self.problem.I0, self.problem.Z0, self.problem.R0]
- self.solver.set_initial_condition(ic)
- n = int(round(self.problem.T/float(self.dt)))
- t = np.linspace(0, self.problem.T, n+1)
- u, self.t = self.solver.solve(t)
- self.S, self.I, self.Z, self.R = u[:,0], u[:,1], u[:,2], u[:,3]
- def plotSIZR(self):
- plt.plot(self.t, self.S)
- plt.plot(self.t, self.I)
- plt.plot(self.t, self.Z)
- plt.plot(self.t, self.R)
- plt.xlabel('Days passed since breakout')
- plt.ylabel('Number of people in a group')
- plt.legend(['Susceptibles', 'Infected', 'Zombies', 'Removed'])
- plt.show()
- problem = ProblemSIZR(eSS=lambda t: 0 if t < 28 else 0.0067,
- eSI=lambda t: 0.0014 if 4 < t < 28 else 0,
- alpha=lambda t: 0 if t < 4 elif 0.0016 if 4 < t < 28 else 0.006,
- beta=lambda t: 0.03 if t < 4 elif 0.0012 if 4 < t < 28 else 0,
- S0=60, I0=0, Z0=1, R0=0, T=33,
- sum=lambda t: 20 if t < 4 elif 2 if 4 < t < 28 else 0, ro=1)
- Solution = SolverSIZR(problem, dt=0.5)
- Solution.solveSIZR()
- Solution.plotSIZR()
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement