I am trying to solve a differential equation with Python.
In this two system differential equation if the value of first variable (v) is more than a threshold (30) it should be reset to another value (-65). Below I put my code. The problem is that the value of first variable after reaching 30 remains constant and won't reset to -65. These equations describe the dynamics of a single neuron. The equations are taken from this website and this PDF file.
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.ticker import FormatStrFormatter
from scipy.integrate import odeint
plt.close('all')
a = 0.02
b = 0.2
c = -65
d = 8
i = 0
p = [a,b,c,d,i]
def fun(u,tspan,*p):
du = [0,0]
if u[0] < 30: #Checking if the threshold has been reached
du[0] = (0.04*u[0] + 5)*u[0] + 150 - u[1] - p[4]
du[1] = p[0]*(p[1]*u[0]-u[1])
else:
u[0] = p[2] #reset to -65
u[1] = u[1] + p[3]
return du
p = tuple(p)
y0 = [0,0]
tspan = np.linspace(0,100,1000)
sol = odeint(fun, y0, tspan, args=p)
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
plt.plot(tspan,sol[:,0],'k',linewidth = 5)
plt.plot(tspan,sol[:,1],'r',linewidth = 5)
myleg = plt.legend(['v','u'],\
loc='upper right',prop = {'size':28,'weight':'bold'}, bbox_to_anchor=(1,0.9))
The solution looks like:
Here is the correct solution by Julia, here u1 represent v:
This is the Julia code:
using DifferentialEquations
using Plots
a = 0.02
b = 0.2
c = -65
d = 8
i = 0
p = [a,b,c,d,i]
function fun(du,u,p,t)
if u[1] <30
du[1] = (0.04*u[1] + 5)*u[1] + 150 - u[2] - p[5]
du[2] = p[1]*(p[2]*u[1]-u[2])
else
u[1] = p[3]
u[2] = u[2] + p[4]
end
end
u0 = [0.0;0.0]
tspan = (0.0,100)
prob = ODEProblem(fun,u0,tspan,p)
tic()
sol = solve(prob,reltol = 1e-8)
toc()
plot(sol)
Recommended solution
This uses events and integrates separately after each discontinuity.
import matplotlib.pyplot as plt
import numpy as np
from scipy.integrate import solve_ivp
a = 0.02
b = 0.2
c = -65
d = 8
i = 0
p = [a,b,c,d,i]
# Define event function and make it a terminal event
def event(t, u):
return u[0] - 30
event.terminal = True
# Define differential equation
def fun(t, u):
du = [(0.04*u[0] + 5)*u[0] + 150 - u[1] - p[4],
p[0]*(p[1]*u[0]-u[1])]
return du
u = [0,0]
ts = []
ys = []
t = 0
tend = 100
while True:
sol = solve_ivp(fun, (t, tend), u, events=event)
ts.append(sol.t)
ys.append(sol.y)
if sol.status == 1: # Event was hit
# New start time for integration
t = sol.t[-1]
# Reset initial state
u = sol.y[:, -1].copy()
u[0] = p[2] #reset to -65
u[1] = u[1] + p[3]
else:
break
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
# We have to stitch together the separate simulation results for plotting
ax.plot(np.concatenate(ts), np.concatenate(ys, axis=1).T)
myleg = plt.legend(['v','u'])
Minimum change "solution"
It appears as though your approach works just fine with solve_ivp.
Warning I think in both Julia and solve_ivp, the correct way to handle this kind of thing is to use events. I believe the approach below relies on an implementation detail, which is that the state vector passed to the function is the same object as the internal state vector, which allows us to modify it in place. If it were a copy, this approach wouldn't work. In addition, there is no guarantee in this approach that the solver is taking small enough steps that the correct point where the limit is reached will be stepped on. Using events will make this more correct and generalisable to other differential equations which perhaps have lower gradients before the discontinuity.
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.ticker import FormatStrFormatter
from scipy.integrate import solve_ivp
plt.close('all')
a = 0.02
b = 0.2
c = -65
d = 8
i = 0
p = [a,b,c,d,i]
def fun(t, u):
du = [0,0]
if u[0] < 30: #Checking if the threshold has been reached
du[0] = (0.04*u[0] + 5)*u[0] + 150 - u[1] - p[4]
du[1] = p[0]*(p[1]*u[0]-u[1])
else:
u[0] = p[2] #reset to -65
u[1] = u[1] + p[3]
return du
y0 = [0,0]
tspan = (0,100)
sol = solve_ivp(fun, tspan, y0)
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
plt.plot(sol.t,sol.y[0, :],'k',linewidth = 5)
plt.plot(sol.t,sol.y[1, :],'r',linewidth = 5)
myleg = plt.legend(['v','u'],loc='upper right',prop = {'size':28,'weight':'bold'}, bbox_to_anchor=(1,0.9))
Result
Related
I am going to solve 1D FitzHugh Nagoma with homogeneous Neumann boundary conditions.
How to plot U and V separately. Here a1=2, a0=-0.03 , ep= 0.01 Du= 1, Dv=4 I was confused by plotting the figure
U_t=Du U_xx +U -U^3 - V
V_t=Dv V_xx + ep(U-a1V - a0)
import numpy as np
import matplotlib.pyplot as plt
#matplotlib inline
Du = 0.001
Dv = 0.005
tau = 1
k = -.00
ep = 0.01
a1 = 2
a0 = -0.03
L = 2
N= 10
x = np.linspace(0, L, N+1)
dx = x[1]-x[0]
T = 45 # total time
dt = .01 # time step
size = N
n = int(T / dt) # number of iterations
U = np.random.rand(size)
V = np.random.rand(size)
def laplacian(Z):
Ztop = Z[0:-2]
Zbottom = Z[2:]
Zcenter = Z[1:-1]
return (Ztop + Zbottom -
2 * Zcenter) / dx**2
def show_patterns(U, ax=None):
ax.imshow(U, cmap=plt.cm.copper,
interpolation='bilinear',
extent=[-1, 1])
ax.set_axis_off()
fig, axes = plt.subplots(3, 3, figsize=(16, 16))
step_plot = n // 9
# We simulate the PDE with the finite difference
# method.
for i in range(n):
# We compute the Laplacian of u and v.
deltaU = laplacian(U)
deltaV = laplacian(V)
# We take the values of u and v inside the grid.
Uc = U[1:-1]
Vc = V[1:-1]
# We update the variables.
U[1:-1], V[1:-1] = \
Uc + dt * (Du * deltaU + Uc - Uc**3 - Vc),\
Vc + dt * (Dv * deltaV + ep*(Uc - a1*Vc - a0)) / tau
# Neumann conditions: derivatives at the edges
# are null.
for Z in (U, V):
Z[0] = Z[1]
Z[-1] = Z[-2]
# Z[:, 0] = Z[:, 1]
# Z[:, -1] = Z[:, -2]
# We plot the state of the system at
# 9 different times.
fig, ax = plt.subplots(1, 1, figsize=(8, 8))
show_patterns(U,ax=None)
I got an error 'NoneType' object has no attribute 'imshow'
and could not solve it
This line
show_patterns(U,ax=None)
passed in a None to the ax parameter.
I don't know what ax should be but it needs to be properly initialised.
I've been trying to solve a set of differential equations using solve_ivp. The Jacobian matrix of the system is the A as you can see below. I wanted to enable the option vectorized='True' but unfortunately i do not know how to modify the present code to vectorize the Jacobian matrix A. Does anyone know how this can be done?
# imports
import numpy as np
import scipy.sparse as sp
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt
# grid sizing
R=0.05 #sphere radius
N=1000#number of points
D=0.00002 #diffusion coefficient
k=10 # Arrhenius
Cs=1.0 # Boundary concentration
C0=0.0 # Initial concentration
time_constant=R**2.0/D
dr=R/(N-1)
# Algebra simplification
a=D/dr**2
Init_conc=np.linspace(0,0,N)
B=np.zeros(N)
B[N-1]=Cs*(a+a/(N-1))
#
e1 = np.ones(N)
e2 = np.ones(N)
e3 = np.ones(N)
#
#
#
e1[0]=-k-6*a
e1[1:]=-k-2*a
#
#
e2[1]=6*a
for i in range(2,N) :
e2[i]=a+a/(i-1)
#
#
#
for i in range (0,N-1) :
e3[i]=a-a/(i+1)
A = sp.spdiags([e3,e1,e2],[-1,0,1],N,N,format="csc")
def dc_dt(t,C) :
dc=A.dot(C)+B
return dc
# Solving the system, I want to implement the same thing with vectorized='True'
OutputTimes=np.linspace(0,0.2*time_constant,100)
ans=solve_ivp(dc_dt,(0,0.2*time_constant),Init_conc,method='RK45',t_eval=OutputTimes,vectorized='False')
print (ans)
Please have a look at this answer, the explanation is thorough. For your code in particular, please see below for updated snippet and figure. It is not obvious that vectorize is providing any speed-up. However, providing A for the keyword jac makes a difference. But I guess it is only valid if A is constant?
# imports
import numpy as np
import scipy.sparse as sp
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt # noqa
def dc_dt(t, C):
print(C.shape)
if len(C.shape) == 1:
return np.squeeze(A.dot(C)) + B
else:
return A.dot(C) + np.transpose(np.tile(B, (C.shape[1], 1)))
# return np.squeeze(A.dot(C)) + B
# grid sizing
R = 0.05 # sphere radius
N = 1000 # number of points
D = 0.00002 # diffusion coefficient
k = 10 # Arrhenius
Cs = 1.0 # Boundary concentration
C0 = 0.0 # Initial concentration
time_constant = R**2.0 / D
dr = R / (N - 1)
# Algebra simplification
a = D / dr**2
Init_conc = np.repeat(0, N)
B = np.zeros(N)
B[-1] = Cs * (a + a / (N - 1))
e1 = np.ones(N)
e2 = np.ones(N)
e3 = np.ones(N)
e1[0] = -k - 6 * a
e1[1:] = -k - 2 * a
e2[1] = 6 * a
for i in range(2, N):
e2[i] = a + a / (i - 1)
for i in range(0, N - 1):
e3[i] = a - a / (i + 1)
A = sp.spdiags([e3, e1, e2], [-1, 0, 1], N, N, format="csc")
# Solving the system, I want to implement the same thing with vectorized='True'
OutputTimes = np.linspace(0, 0.2 * time_constant, 10000)
ans = solve_ivp(dc_dt, (0, 0.2 * time_constant), Init_conc,
method='BDF', t_eval=OutputTimes, jac=A, vectorized=True)
plt.plot(np.arange(N), ans.y[:, 0])
plt.plot(np.arange(N), ans.y[:, 1])
plt.plot(np.arange(N), ans.y[:, 10])
plt.plot(np.arange(N), ans.y[:, 20])
plt.plot(np.arange(N), ans.y[:, 50])
plt.plot(np.arange(N), ans.y[:, -1])
plt.show()
I'm converting my code from Matlab to Python and stuck on how to split the state vector such that the result returns the two solution. I have a vector and a single value for the two initial conditions and I expect as the final result a matrix and a vector.
I tried joining the initial conditions (y0 = [c_pt_0, x_0]) in the same manner as the solution (soln = [dfdt,dcdt]) (which is shown below in the code). I also tried a similar approach that is used in matlab, which is concatenating the the initial conditions to one single array and unpacking the results but I think the problem is in the dimensions.
#Basic imports
import numpy as np
import pylab
import matplotlib. pyplot as plt
import scipy
from scipy.integrate import odeint
import matplotlib.pyplot as plt
# define parameters
pi = 3.14159265
V_m = 9.09
m_V__M_Pt = 1e6/195.084
rho = 21.45
R0 = 10**(-8.19)
k_d = 10**(-13)
k_r = 10**(-5)
S = 0.314 #distribution parameter
M = 0.944 #distribution parameter
## geometry
# Finite Volume Method with equidistant elements
r_max = 30.1e-9 #maximum value
n = 301 #number of elements of FVM
dr = r_max/n #length of elements, equidistant
ini_r = np.linspace(5e-10,r_max,n+1) #boundaries of elements
mid_r = ini_r[0:n]+dr/2 #center of elements
## initial conditions
#initial distribution
x0 = 1/(S*np.sqrt(2*pi)*mid_r*1e9)*np.exp((-(np.log(mid_r*1e9)-M)**2)/(2*S**2))
c_pt_0 = 0
y0 = [x0, c_pt_0]
MN_0 = scipy.trapz(np.power(mid_r, 3)*x0,
x=mid_r) # initial mass
M_0 = 4/3*pi*rho*MN_0
def f(y, t):
r = y[0]
c_pt = y[1]
#materials balance
drdt = V_m * k_r * c_pt * np.exp(-R0/ mid_r) - V_m * k_d * np.exp(R0/ mid_r)
dmdt = 4*pi*rho*mid_r**2*drdt
dMdt = np.trapz(r*dmdt, x=mid_r)
dcdt = m_V__M_Pt*(-dMdt)/M_0
dfdt = -(np.gradient(r*drdt, dr))
soln = [dfdt, dcdt]
return soln
#------------------------------------------------------
#define timespace
time = np.linspace(0, 30, 500)
#solve ode system
sln_1 = odeint (f , y0 , time,
rtol = 1e-3, atol = 1e-5)
pylab.plot(mid_r, sln_1[1,:], color = 'r', marker = 'o')
pylab.plot(mid_r, sln_1[-1,:], color = 'b', marker = 'o')
plt.show()
Traceback:
ValueError: setting an array element with a sequence.
Any help is much appreciated.
EDIT: ADDED MATLAB CODE
Here is the MATLAB code that works that I want to convert to python where the state vector is split. I have three files (one main, the f function, and the parameters). Please excuse any face palm coding errors but I do appreciate any suggestions even for this.
modified_model.m:
function modified_model
% import parameters
p = cycling_parameters;
% initial conditions
c_pt_0 = 0;
y0 = [p.x0; c_pt_0];
% call integrator
options_ODE=odeset('Stats','on', 'RelTol',1e-3,'AbsTol',1e-5);
[~, y] = ode15s(#(t,y) f(t, y, p), p.time, y0, options_ODE);
%% Post processing
% split state vector
r = y(:,1:p.n);
c_Pt = y(:,p.n+1);
%% Plot results
figure
hold on;
plot(p.r_m, r(1,:));
plot(p.r_m, r(end,:));
xlabel({'size'},'FontSize',15)
ylabel({'counts'},'FontSize',15)
f.m
function soln = f(~, y, p)
%split state vector
r = y(1:p.n);
c_pt = y(p.n+1);
% materials balance
drdt = p.Vm_Pt.*p.k_rdp.*c_pt.*exp(-p.R0./p.r_m) - p.Vm_Pt.*p.k_dis.*exp(p.R0./p.r_m);
dmdt = 4*pi*p.rho*p.r_m.^2.*drdt;
dMdt = trapz(p.r_m, r.*dmdt);
dcdt = p.I_V*p.m_V__M_Pt*(-dMdt)/p.M_0;
dfdt = - gradient(r.*drdt,p.dr);
soln = [dfdt; dcdt];
and the parameters file: cycling_parameters.m
function p=cycling_parameters
p.M = 195.084;
p.rho = 21.45;
p.time = linspace(0, 30, 500);
p.m_V__M_Pt = 1e6/p.M;
p.Vm_Pt = 9.09;
p.R0_log = -8.1963;
p.k_dis_log = -13;
p.k_rdp_log = -11;
p.R0 = 10^(p.R0_log);
p.k_dis = 10^(p.k_dis_log);
p.k_rdp = 10^(p.k_rdp_log);
%%% geometry %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%% Finite Volume Method with equidistant elements
p.r_max = 10.1e-9; % [m] maximum radius of PRD
p.n = 301; % number of elements of FVM
p.dr = p.r_max/p.n; % [m] length of elements, equidistant
p.r = linspace(5e-10,p.r_max,p.n+1)'; % [m] boundaries of elements
p.r_m = p.r(1:p.n)+p.dr/2; % [m] center of elements
%log normal initial distribution
S = 0.314;
M = 0.944;
p.x0 = 1./(S.*sqrt(2.*pi).*p.r_m*1e9).*exp((-(log(p.r_m*1e9)-M).^2)./(2.*S.^2));
p.r_squared = p.r_m.^2; % [m^2] squares of the radius (center of elements)
p.r_cubed = p.r_m.^3; % [m^3] cubes of the radius (center of elements)
p.MN_0 = trapz(p.r_m, p.r_cubed.*p.x0); % Eq. 2.11 denominator
p.M_0 = 4/3*pi*p.rho*p.MN_0;
p.I_V = 1; %ionomer volume fraction in the catalyst layer
After looking at both codes, the issue is that the odeint solver only takes 1D array inputs and your y0 is [int, array(300,)] and odeint can't work with that. However, you can merge the y0 into a 1D array and then split it up in the function you are integrating over to do the calculation then recombine as the output. Here's a working code of that:
import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt
class P:
def __init__(self, S, M):
self.M = 195.084
self.rho = 21.45
self.m_V__M_Pt = (1*10**6)/self.M
self.Vm_Pt = 9.09
self.R0_log = -8.1963
self.k_dis_log = -13
self.k_rdp_log = -11
self.R0 = 10**(self.R0_log)
self.k_dis = 10**(self.k_dis_log)
self.k_rdp = 10**(self.k_rdp_log)
self.r_max = 10.1*10**(-9)
self.n = 301
self.dr = self.r_max / self.n
self.r = np.linspace(5*10**(-10), self.r_max, self.n)
self.r_m = self.r[0:self.n+1]+self.dr/2
self.x0 = self.compute_x0(S, M)
self.r_squared = np.power(self.r_m, 2)
self.r_cubed = np.power(self.r_m, 3)
self.MN_0 = np.trapz(self.r_m, np.multiply(self.r_cubed, self.x0))
self.M_0 = (4 / 3)* np.pi * self.rho * self.MN_0
self.I_V = 1
def compute_x0(self, S, M):
p1 = np.multiply(2, np.power(S, 2))
p2 = np.multiply(S, np.sqrt(np.multiply(2, np.pi)))
p3 = np.log(self.r_m*1*10**(9)) - M
p4 = np.multiply(p2, self.r_m*10**(9))
p5 = np.power(-p3, 2)
p6 = np.multiply(p4, np.exp(np.divide(p5,p1)))
p7 = np.divide(1, p6)
return p7
def cycling_parameters():
S = 0.314
M = 0.944
p = P(S, M)
return p
def f(y, t):
p = cycling_parameters()
c_pt = y[0]
r = np.delete(y, 0)
p1 = np.multiply(p.Vm_Pt, p.k_rdp)
p2 = np.multiply(p1, c_pt)
p3 = np.multiply(p.Vm_Pt, p.k_dis)
drdt = np.multiply(p2, np.exp(np.divide(-p.R0, p.r_m))) - np.multiply(p3, np.exp(np.divide(p.R0, p.r_m)))
dmdt = np.multiply(4*np.pi*p.rho*np.power(p.r_m, 2), drdt)
p4 = np.multiply(r, dmdt)
dMdt = np.trapz(p.r_m, p4)
dcdt = p.I_V*p.m_V__M_Pt*(-dMdt)/p.M_0
p5 = np.multiply(r, drdt)
dfdt = - np.gradient(p5,p.dr)
ans = np.insert(dfdt, 0, dcdt)
return ans
def modified_model():
p = cycling_parameters()
c_pt_0 = 0
y0 = np.insert(p.x0, 0, c_pt_0)
t = np.linspace(0, 30, 500)
ans = odeint(f, y0, t, rtol = 1e-3, atol = 1e-5)
r = ans[:, 1:p.n+1]
c_Pt = ans[:, 0]
print(r)
print(c_Pt)
plt.plot(p.r_m, r[0, :], color='r', linewidth=0.5)
plt.plot(p.r_m, r[r.shape[0]-1, :], color='b', linewidth=0.5)
plt.show()
if __name__ == '__main__':
modified_model()
Python plot (what this script outputs):
Original Matlab Plot:
I am new to python and in learning stages. I wanted to implement Particle Swarm Optimization(PSO) algorithm which I did by taking help from on-line materials and python tutorials. In PSO, a simple calculus problem is inferred i-e 100 * ((y - (x2))2) + ((1 - (x2))2). This problem is defined in a fitness function.
def fitness(x, y):
return 100 * ((y - (x**2))**2) + ((1 - (x**2))**2)
Now, I want to replace this simple calculus problem by simple first order Ordinary Differential Equation(ODE) by without changing existing function parameters (x,y) and want to return the value of dy_dx,y0 and t for further process.
# Define a function which calculates the derivative
def dy_dx(y, x):
return x - y
t = np.linspace(0,5,100)
y0 = 1.0 # the initial condition
ys = odeint(dy_dx, y0, t)`
In python odeint function is used for ODE which requires three essential parameters i-e func/model, y0( Initial condition on y (can be a vector) and t(A sequence of time points for which to solve for y) Example of odeint parameters.
I don't want to change its parameters because it will be difficult for me to make changes in algorithm.
For simplicity I pasted the full code below and my question is open to anyone if wants to modify the code with further parameters in General Best, Personal Best and r[i].
import numpy as np
from scipy.integrate import odeint
import random as rand
from scipy.integrate import odeint
from numpy import array
import matplotlib.pyplot as plt
def main():
#Variables
n = 40
num_variables = 2
a = np.empty((num_variables, n))
v = np.empty((num_variables, n))
Pbest = np.empty((num_variables, n))
Gbest = np.empty((1, 2))
r = np.empty((n))
for i in range(0, num_variables):
for j in range(0, n):
Pbest[i][j] = rand.randint(-20, 20)
a[i][j] = Pbest[i][j]
v[i][j] = 0
for i in range(0, n):
r[i] = fitness(a[0][i], a[1][i])
#Sort elements of Pbest
Order(Pbest, r, n)
Gbest[0][0] = Pbest[0][0]
Gbest[0][1] = Pbest[1][0]
generation = 0
plt.ion()
fig = plt.figure()
ax = fig.add_subplot(111)
ax.grid(True)
while(generation < 1000):
for i in range(n):
#Get Personal Best
if(fitness(a[0][i], a[1][i]) < fitness(Pbest[0][i], Pbest[1][i])):
Pbest[0][i] = a[0][i]
Pbest[1][i] = a[1][i]
#Get General Best
if(fitness(Pbest[0][i], Pbest[1][i]) < fitness(Gbest[0][0], Gbest[0][1])):
Gbest[0][0] = Pbest[0][i]
Gbest[0][1] = Pbest[1][i]
#Calculate Velocity
Vector_Velocidad(n, a, Pbest, Gbest, v)
generation = generation + 1
print 'Generacion: ' + str(generation) + ' - - - Gbest: ' +str(Gbest)
line1 = ax.plot(a[0], a[1], 'r+')
line2 = ax.plot(Gbest[0][0], Gbest[0][1], 'g*')
ax.set_xlim(-10, 10)
ax.set_ylim(-10, 10)
fig.canvas.draw()
ax.clear()
ax.grid(True)
print 'Gbest: '
print Gbest
def Vector_Velocidad(n, a, Pbest, Gbest, v):
for i in range(n):
#Velocity in X
v[0][i] = 0.7 * v[0][i] + (Pbest[0][i] - a[0][i]) * rand.random() * 1.47 + (Gbest[0][0] - a[0][i]) * rand.random() * 1.47
a[0][i] = a[0][i] + v[0][i]
v[1][i] = 0.7 * v[1][i] + (Pbest[1][i] - a[1][i]) * rand.random() * 1.47 + (Gbest[0][1] - a[1][i]) * rand.random() * 1.47
a[1][i] = a[1][i] + v[1][i]
def fitness(x, y):
return 100 * ((y - (x**2))**2) + ((1 - (x**2))**2)
def Order(Pbest, r, n):
for i in range(1, n):
for j in range(0, n - 1):
if r[j] > r[j + 1]:
#Order the fitness
tempRes = r[j]
r[j] = r[j + 1]
r[j + 1] = tempRes
#Order las X, Y
tempX = Pbest[0][j]
Pbest[0][j] = Pbest[0][j + 1]
Pbest[0][j + 1] = tempX
tempY = Pbest[1][j]
Pbest[1][j] = Pbest[1][j + 1]
Pbest[1][j + 1] = tempY
if '__main__' == main():
main()
I have created some very basic implementations of the mentioned models. However, although graphs seem to look right, the numbers don't add up to a constant. That is for the sum of susceptible/infected/recovered people in each compartment should add up to N (which is total number of people), but it doesn't, for some reason it adds up to some bizarre decimal numbers, and I really don't know how to fix it, after looking at it for 3 days now.
The SI Model:
import matplotlib.pyplot as plt
N = 1000000
S = N - 1
I = 1
beta = 0.6
sus = [] # infected compartment
inf = [] # susceptible compartment
prob = [] # probability of infection at time t
def infection(S, I, N):
t = 0
while (t < 100):
S = S - beta * ((S * I / N))
I = I + beta * ((S * I) / N)
p = beta * (I / N)
sus.append(S)
inf.append(I)
prob.append(p)
t = t + 1
infection(S, I, N)
figure = plt.figure()
figure.canvas.set_window_title('SI model')
figure.add_subplot(211)
inf_line, =plt.plot(inf, label='I(t)')
sus_line, = plt.plot(sus, label='S(t)')
plt.legend(handles=[inf_line, sus_line])
plt.ticklabel_format(style='sci', axis='y', scilimits=(0,0)) # use scientific notation
ax = figure.add_subplot(212)
prob_line = plt.plot(prob, label='p(t)')
plt.legend(handles=prob_line)
type(ax) # matplotlib.axes._subplots.AxesSubplot
# manipulate
vals = ax.get_yticks()
ax.set_yticklabels(['{:3.2f}%'.format(x*100) for x in vals])
plt.xlabel('T')
plt.ylabel('p')
plt.show()
SIS Model:
import matplotlib.pylab as plt
N = 1000000
S = N - 1
I = 1
beta = 0.3
gamma = 0.1
sus = \[\]
inf = \[\]
def infection(S, I, N):
for t in range (0, 1000):
S = S - (beta*S*I/N) + gamma * I
I = I + (beta*S*I/N) - gamma * I
sus.append(S)
inf.append(I)
infection(S, I, N)
figure = plt.figure()
figure.canvas.set_window_title('SIS model')
inf_line, =plt.plot(inf, label='I(t)')
sus_line, = plt.plot(sus, label='S(t)')
plt.legend(handles=\[inf_line, sus_line\])
plt.ticklabel_format(style='sci', axis='y', scilimits=(0,0))
plt.xlabel('T')
plt.ylabel('N')
plt.show()
SIR Model:
import matplotlib.pylab as plt
N = 1000000
S = N - 1
I = 1
R = 0
beta = 0.5
mu = 0.1
sus = []
inf = []
rec = []
def infection(S, I, R, N):
for t in range (1, 100):
S = S -(beta * S * I)/N
I = I + ((beta * S * I)/N) - R
R = mu * I
sus.append(S)
inf.append(I)
rec.append(R)
infection(S, I, R, N)
figure = plt.figure()
figure.canvas.set_window_title('SIR model')
inf_line, =plt.plot(inf, label='I(t)')
sus_line, = plt.plot(sus, label='S(t)')
rec_line, = plt.plot(rec, label='R(t)')
plt.legend(handles=[inf_line, sus_line, rec_line])
plt.ticklabel_format(style='sci', axis='y', scilimits=(0,0))
plt.xlabel('T')
plt.ylabel('N')
plt.show()
I'll look only at the SI model.
Your two key variables are S and I. (You may have reversed the meanings of these two variables, though that does not affect what I write here.) You initialize them so their sum is N which is the constant 1000000.
You update your two key variables in the lines
S = S - beta * ((S * I / N))
I = I + beta * ((S * I) / N)
You apparently intend to add to I and subtract from S the same value, so the sum of S and I is unchanged. However, you actually first change S then use that new value to change I, so the values added and subtracted are not actually the same, and the sum of the variables has not remained constant.
You can fix this by using Python's ability to update multiple variables in one line. Replace those two lines with
S, I = S - beta * ((S * I / N)), I + beta * ((S * I) / N)
This calculates both of the new values before updating the variables, so the same value actually added and subtracted from the two variables. (There are other ways to get the same effect, such as temporary variables for the updated values, or one temporary variable to store the amount to add and subtract, but since you use Python you may as well use its capabilities.)
When I now run the program, I get these graphs:
which I think is what you want.
So the solution above worked for the SIS model as well.
As for the SIR model I had to solve differential equations using odeint, here is a simple solution to the SIR model:
import matplotlib.pylab as plt
from scipy.integrate import odeint
import numpy as np
N = 1000
S = N - 1
I = 1
R = 0
beta = 0.6 # infection rate
gamma = 0.2 # recovery rate
# differential equatinons
def diff(sir, t):
# sir[0] - S, sir[1] - I, sir[2] - R
dsdt = - (beta * sir[0] * sir[1])/N
didt = (beta * sir[0] * sir[1])/N - gamma * sir[1]
drdt = gamma * sir[1]
print (dsdt + didt + drdt)
dsirdt = [dsdt, didt, drdt]
return dsirdt
# initial conditions
sir0 = (S, I, R)
# time points
t = np.linspace(0, 100)
# solve ODE
# the parameters are, the equations, initial conditions,
# and time steps (between 0 and 100)
sir = odeint(diff, sir0, t)
plt.plot(t, sir[:, 0], label='S(t)')
plt.plot(t, sir[:, 1], label='I(t)')
plt.plot(t, sir[:, 2], label='R(t)')
plt.legend()
plt.xlabel('T')
plt.ylabel('N')
# use scientific notation
plt.ticklabel_format(style='sci', axis='y', scilimits=(0,0))
plt.show()