#! /usr/bin/env python3
#
def exercise4 ( ):

  from expm_ode import expm_ode
  from rk2 import rk2
  from rk3 import rk3
  import numpy as np

  print ( "exercise4:" )
#
#  Demonstrate the expm_ode returns a vector for vector input.
#
  x = 1.0
  y = np.array ( [ 5.0, 6.0 ] )
  fValue = expm_ode ( x, y )
  print ( "expm_ode(x,y) returns ", fValue )
#
#  Solve a system of two ODE's.
#
  f_ode = expm_ode
  xRange = np.array ( [ 0.0, 2.0 ] )
  yInit = np.array ( [ 5.0, 6.0 ] )
  numSteps = 40
  xs, ys = rk3 ( f_ode, xRange, yInit, numSteps )
  print ( '  Solution at', xs[-1], 'is', ys[-1,:] )
#
#  Solve two scalar ODE's.
#
  f_ode = expm_ode
  xRange = np.array ( [ 0.0, 2.0 ] )
  yInit = 5.0
  numSteps = 40
  x1, y1 = rk3 ( f_ode, xRange, yInit, numSteps )

  f_ode = expm_ode
  xRange = np.array ( [ 0.0, 2.0 ] )
  yInit = 6.0
  numSteps = 40
  x2, y2 = rk3 ( f_ode, xRange, yInit, numSteps )
#
#  Compare results
#
  print ( 'y1[-1,0] = ', y1[-1,0], 'ys[-1,0]=', ys[-1,0] )
  print ( 'y2[-1,0] = ', y2[-1,0], 'ys[-1,1]=', ys[-1,1] )
#
#  Compare all values.
#
  print ( '||y1[:] - ys[:,0]|| = ', np.linalg.norm ( y1[:,0] - ys[:,0] ) )
  print ( '||y2[:] - ys[:,1]|| = ', np.linalg.norm ( y2[:,0] - ys[:,1] ) )

if ( __name__ == "__main__" ):
  exercise4 ( )