#!/usr/bin/env python
# coding: utf-8

# ## Program 2.1 (SIR model) - original Fortran code

# In[1]:


get_ipython().run_cell_magic('writefile', 'Program_2_1.f', 'C\nC This is the FORTRAN version of program 2.1 from page 19 of\nC "Modeling Infectious Disease in humans and animals"\nC by Keeling & Rohani.\nC \nC It is the simple SIR epidemic without births or deaths.\nC\nC This code is written to be simple, transparent and readily compiled.\nC Far more elegant and efficient code can be written.\nC\nC This code can be compiled using the intel fortran compiler:\nC ifort -Vaxlib  -o Program_2_1  Program_2_1.f  \nC \n\n\nC Main program starts here.\n\n      program main\n\n      REAL beta\n      REAL gamma\n\n      REAL S,I,R\n\n      REAL S0\n      REAL I0\n      REAL MaxTime\n      REAL EVERY, step, t\n\n      INTEGER GivesName\n\n      CHARACTER*2000 str, FileName\n\n      COMMON /parameters/ beta, gamma\n      COMMON /variables/ S, I, R\n\n      GivesName=iargc()\n\n      if (GivesName.eq.0) then\n         beta=1.4247\n         gamma=0.14286\n         S0=1 - 1.0d-6\n         I0=1.0d-6\n         MaxTime=70\n      else\n \nc     \nc     READ IN ALL THE VARIABLES\nc\n\n         call getarg(1,FileName)\n\n         open(1,file=FileName,STATUS=\'OLD\',ACCESS=\'SEQUENTIAL\')\n  \n         \n         read(1,*) str\n         read(1,*) beta\n         \n         read(1,*) str\n         read(1,*) gamma\n         \n         read(1,*) str\n         read(1,*) S0\n         \n         read(1,*) str\n         read(1,*) I0\n         \n         read(1,*) str\n         read(1,*) MaxTime \n         \n         close(1)\n      endif\nC\nC     Check all variables are OK & set up intitial conditions */\nC     \n      \n      if ( S0.le.0) then\n         write(*,*) "ERROR: Initial level of susceptibles (",S0,") is \n     .        less than or equal to zero."\n         STOP\n      endif\n\n      if ( I0.le.0) then\n         write(*,*) "ERROR: Initial level of infecteds (",I0,") is \n     .        less than or equal to zero."\n         STOP\n      endif\n\n      if ( beta.le.0) then\n         write(*,*) "ERROR: Transmission rate beta (",beta,") is \n     .        less than or equal to zero."\n         STOP\n      endif\n\n      if ( gamma.le.0) then\n         write(*,*) "ERROR: Recovery rate gamma (",gamma,") is \n     .        less than or equal to zero."\n         STOP\n      endif\n\n      if ( MaxTime.le.0) then\n         write(*,*) "ERROR: Maximum run time (",MaxTime,") is \n     .        less than or equal to zero."\n         STOP\n      endif\n\n      if (S0+I0.ge.1) then\n        write(*,*) "WARNING: Initial level of susceptibles+infecteds\n     .        (",S0,"+",I0,"=",S0+I0,") is greater than one."\n      endif\n\n      if (beta.lt.gamma) then\n        write(*,*) "WARNING: Basic reproductive ratio (R_0=",\n     .        beta/gamma,") is less than one."\n      endif\n\n\n      S=S0\n      I=I0\n      R=1-S0-I0\n     \nC     \nC     Find a suitable time-scale for outputs \nC     \n      step=0.01/((beta+gamma)*S0)\n\n      Every=1.0/((beta+gamma)*S0)\n      if (Every.gt.1) then\n         Every=10.0**INT(log10(Every))\n      else\n         Every=10.0**INT(log10(Every)-1)\n      endif\n      DO WHILE (MaxTime/Every.gt.10000)\n         Every=Every*10.0  \n      ENDDO\n\n      open(1,recl=3000,file=\'Program_2_1_f.out\',ACCESS=\'SEQUENTIAL\')\nC     for F77 use\nC      open(1,file=\'Output_Risk\',ACCESS=\'SEQUENTIAL\')\nC\n\nC\nC     The main iteration routine\nC     \n\n      t=0\n      write(1,*) t,S,I,R\n      DO WHILE (t.lt.MaxTime)\n         \n\n         CALL Runge_Kutta(step)\n         \n         t=t+step\n         \nC     If time has moved on sufficiently, output the current data\n         \n         if( INT(t/Every).gt.INT((t-step)/Every) ) then\n            write(1,*) t,S,I,R\n         endif\n      \n      ENDDO\n      \n      END\n  \n\n\n      SUBROUTINE Runge_Kutta(step)\n\n      REAL InitialPop(3), tmpPop(3)\n      REAL dPop1(3), dPop2(3), dPop3(3), dPop4(3)\n\n      REAL S,I,R, step\n      COMMON /variables/ S, I, R\n\nC\nC     Integrates the equations one step, using Runge-Kutta 4 \nC     Note: we work with arrays rather than variables to make the\nC     coding easier\nC\n      InitialPop(1)=S\n      InitialPop(2)=I\n      InitialPop(3)=R\n\n      CALL Diff(InitialPop,dPop1)\n      do k=1,3\n         tmpPop(k)=InitialPop(k)+step*dPop1(k)/2.0\n      ENDDO \n\n      CALL Diff(tmpPop,dPop2)\n      do k=1,3\n         tmpPop(k)=InitialPop(k)+step*dPop2(k)/2.0\n      ENDDO\n\n      CALL Diff(tmpPop,dPop3)\n      do k=1,3\n         tmpPop(k)=InitialPop(k)+step*dPop3(k)\n      ENDDO\n\n      CALL Diff(tmpPop,dPop4)\n\n      do k=1,3\n         tmpPop(k)=InitialPop(k)+step*(dPop1(k)/6 + dPop2(k)/3 +\n     .        dPop3(k)/3 + dPop4(k)/6)\n      ENDDO\n\n      S=tmpPop(1)\n      I=tmpPop(2)\n      R=tmpPop(3)\n\n      RETURN\n      END\n\n\n \nC The Main Differential Equations.\n\n      SUBROUTINE Diff(Pop, dPop)\n\n      REAL Pop(3), dPop(3)\n\n      REAL beta\n      REAL gamma\n\n      COMMON /parameters/ beta, gamma\n\n\nC     Set up temporary variables to make the equations look neater\n\n      REAL tmpS, tmpI, tmpR\n\n      tmpS=Pop(1)\n      tmpI=Pop(2)\n      tmpR=Pop(3)\n\nC\nC     The differential equations\nC\n\nC     dS/dt =\n      dPop(1) = - beta*tmpS*tmpI\n\nC     dI/dt =\n      dPop(2) = beta*tmpS*tmpI - gamma*tmpI\n\nC     dR/dt =\n      dPop(3) = gamma*tmpI\n      \n      RETURN\n      END\n')


# In[2]:


get_ipython().system('gfortran -O3 -o Program_2_1_f Program_2_1.f')


# In[3]:


get_ipython().system('./Program_2_1_f parameters.txt')


# In[4]:


import pandas as pd
import matplotlib.pyplot as plt
plt.style.use("ggplot")


# In[5]:


sir_out = pd.read_csv("Program_2_1_f.out",sep="\s+",engine="python",header=None,names=["t","S","I","R"],index_col=False)


# In[6]:


sline = plt.plot("t","S","",data=sir_out,color="red",linewidth=2)
iline = plt.plot("t","I","",data=sir_out,color="green",linewidth=2)
rline = plt.plot("t","R","",data=sir_out,color="blue",linewidth=2)
plt.xlabel("Time",fontweight="bold")
plt.ylabel("Number",fontweight="bold")
legend = plt.legend(title="Population",loc=5,bbox_to_anchor=(1.25,0.5))
frame = legend.get_frame()
frame.set_facecolor("white")
frame.set_linewidth(0)