# -*- coding: utf-8 -*-
"""
Goal: check the order of the Explicit Euler scheme 
for the example of the SIR system as in the exercice.
@author: turinici
"""

import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt
from scipy.stats import linregress

T,N=150,150
h=T/N
S0,I0,Rinit=1.e+6,10.,0.
Ntotal=S0+I0+Rinit

beta=0.5
gamma=1./3.
U0=np.array([S0,I0,Rinit])

def sir(t,X):
    '''the function that defines the ODE    '''
    S,I,R=X
    return np.array([-beta*S*I/Ntotal,beta*S*I/Ntotal-gamma*I,
	gamma*I])

t0,tf=52,60	
sol=odeint(sir,U0,[0,t0,57,tf],tfirst=True,atol=1.e-8)
#sol.shape (4,3)
S52,I52,R52=sol[1,:]#this will be the initial state, considered exact
S60,I60,R60=sol[3,:]#this will be the final state, considered exact

error_list=[]
hlist=[0.01,0.05,0.1,0.5,1.,2.,4.]#list of h values
for h in hlist:
	N=int((tf-t0)/h)
	U=sol[1,:].copy()#initial data
	for jj in range(N):#this is the Explicit Euler scheme
		U=U+h*sir(jj*h+t0,U)
	S_EE_T_h=U[0]	
	error_list.append(np.abs(S_EE_T_h-S60))

plt.figure('order for EE')
plt.loglog(hlist,error_list,'r-o')
linear_regression_fit=linregress(np.log(hlist), np.log(error_list))
print('empirical order is=', linear_regression_fit.slope)