In [1]:
%run Steppers.ipynb #Import stepper functions (Euler, RK4, etc)
%run ODEs.ipynb     #Import ODEs (Pendulum etc)
%run Tools.ipynb
import matplotlib
matplotlib.use('nbagg')
import matplotlib.pyplot as plt

In [2]:
def InfectionSpread(f,t,a):
    S,I=f
    dSdt=S+I-a*S*I-S**2
    dIdt=a*S*I-2*I
    return(np.array([dSdt,dIdt]))

In [3]:
def SNullCline(S,a):
    I = (S**2-S)/(1-a*S)
    return(I)

def INullCline(I,a):
    S=I*0+2/a
    return(S)

In [14]:
t = np.arange(0,15,0.001)
S0 = 0.2
I0 = 0.5
a=4

iLims=(-0.1,1)
sLims=(-0.1,2/a+0.5)

sol = RK4(InfectionSpread, (S0,I0),t,(a,))
S,I = np.hsplit(sol,2)
plt.xlabel("t")

plt.plot(t,S, label="Susceptable population")
plt.plot(t,I, label="Infected population")
plt.legend()
plt.show()

#plt.plot(S,I)


iDummy = np.linspace(iLims[0],iLims[1], 1000)
sDummy = np.linspace(sLims[0],sLims[1], 1000)

sNullcline = SNullCline(sDummy,a)
iNullcline = INullCline(iDummy,a)

sNullcline[sNullcline>10] = np.inf
sNullcline[sNullcline<-10]=-np.inf


plt.plot(sDummy,sNullcline, label="S Nullcine", color="C0")
plt.plot(iNullcline, iDummy, label="I Nullclines", color="C1")
plt.plot(sDummy, sDummy*0, color="C1")

plt.plot(S, I, label="Phase plot", color="C2")

plt.xlim(sLims)
plt.ylim(iLims)



PlotQuiver(InfectionSpread,sLims,iLims, args=(a,))
plt.xlabel("S (Susceptable population)")
plt.ylabel("I (Infected population)")
plt.legend()
plt.show()