## The following will produce a very simple implementation of Euler's
## method for solving the following ODE:
##
##          y'(t) = y(t),   t > 0,
##          y(0) = 0
##
## For now, the code will explicit use the forcing term f(y) = y, but will
## later be generalized.
##

import math
import numpy as np # shortcut name for numpy
from matplotlib import pyplot as plt # shortcut name for pyplot
from matplotlib.animation import FuncAnimation

N = 100 # Number of time steps
t0 = 0 # Initial time
tf = 1 # Final time
h = (tf - t0)/(N-1) # Numerical step size
t = np.linspace(t0,tf,N) # Temporal grid
y = np.zeros([N]) # Empty vector for storing solution values
z = np.zeros([N]) # Empty vector for storing error values
w = np.zeros([N]) # Empty vector for true solution
y[0] = 1 # Initial value of the problem

# The following is the actual implementation of Euler's method
for i in range(1,N):
    y[i] = y[i-1] + h*y[i-1]
    
# The following constructs the true solution
for i in range(0,N):
    w[i] = math.exp(t[i])

# The following computes the error between the approximation
# and the true solution, y = exp(t)
for i in range(0,N):
    z[i] = abs( y[i] - w[i] )

# The following creates strings to generate the titles in the plots
s1 = "Euler Approximation \n for h = "    
s = "%1.2f" % h
title1 = s1 + s
s2 = "Absolute Error \n for h = "
title2 = s2 + s

# The following generates the desired plots

fig, axs = plt.subplots(ncols=2,nrows=2)
gs = axs[1,0].get_gridspec()
# remove the underlying axes
for ax in axs[-1,0:]:
    ax.remove()
axbig = fig.add_subplot(gs[-1,0:])

axs[0,0].plot(t,y,'--',label="Approx.")
axs[0,0].plot(t,w,label="True")
axs[0,0].set_title(title1)
axs[0,0].set_xlabel('t')
axs[0,0].set_ylabel('y')
axs[0,0].legend()

axs[0,1].plot(t,z)
axs[0,1].set_title(title2)
axs[0,1].set_xlabel('t')
axs[0,1].set_ylabel('y')

# The following is a (fairly complicated) bit of code that allows 
# you to make a video of the pointwise error at each of the points 
# in an ever refining domain.

xdata, ydata = [], []
ln, = plt.plot([], [], 'ro')

def out(ii):
	NN = 5*ii
	tt = np.linspace(t0,tf,NN)
	hh = tt[1]-tt[0]
	ww = np.zeros([NN])
	yy = np.zeros([NN])
	zz = np.zeros([NN])
	yy[0] = 1
	for i in range(1,NN):
		yy[i] = yy[i-1] + hh*yy[i-1]
	for i in range(0,NN):
		ww[i] = math.exp(tt[i])
	for i in range(0,NN):
		zz[i] = abs( yy[i] - ww[i] )
	return hh,tt,zz

def init():
	axbig.set_xlim(t0,tf)
	axbig.set_ylim(0,0.15)
	axbig.set_xlabel('t')
	axbig.set_ylabel('Error')
	return ln,
    
def update(ii):
	hh,xdata,ydata = out(ii)
	ln.set_data(xdata, ydata)
	ss = "Error for h = "
	ss1 = "%1.4f" % hh
	title3 = ss + ss1
	axbig.set_title(title3)
	return ln,

ani = FuncAnimation(fig,update,frames=range(1,20),interval=150,init_func=init)

fig.tight_layout()

plt.show()