#! /usr/bin/env python3
#
from back_euler_lam import back_euler_lam
from stiff_ode import stiff_lambda, stiff_solution
import matplotlib.pyplot as plt
import numpy as np

print ( '' )
print ( 'exercise4:' )
print ( 'Solve the stiff ODE system using back_euler_lam().' )

LAMBDA = 10000.0
stiff_lambda ( LAMBDA )

xRange = np.array ( [ 0.0, 2.0 * np.pi ] )
yInit = 0.0
numSteps = 40
x1, y1 = back_euler_lam ( LAMBDA, xRange, yInit, numSteps )

print ( 'First 4 values of Y are ', y1[0:4] )

x2 = np.linspace ( 0.0, 2.0 * np.pi, 41 )
y2 = stiff_solution ( x2 )

plt.plot ( x1, y1, 'ro' )
plt.plot ( x2, y2, 'b-', linewidth = 2 )
plt.grid ( True )
plt.xlabel ( '<--- X --->' )
plt.ylabel ( '<--- Y(X) --->' )
plt.legend ( [ 'Backward Euler', 'Exact' ] )
plt.title ( 'Stiff ODE' )

filename = 'exercise4.jpg'
plt.savefig ( filename )
print ( 'Graphics saved as "%s"' % ( filename ) )
plt.show ( )
plt.close ( )