#! /usr/bin/env python3
#
def stiff_direction ( LAMBDA ):

#*****************************************************************************80
#
## stiff_direction() displays the direction field for stiff_ode().
#
#  Modified:
#
#    23 January 2023
#
  from stiff_ode import stiff_ode, stiff_solution, stiff_lambda
  import matplotlib.pyplot as plt
  import numpy as np
#
#  Set the value of LAMBDA.
#
  stiff_lambda ( LAMBDA )
#
#  Evaluate the right hand side over a grid, and display the vector field.
#
  x = np.linspace ( 0.0, 2.0 * np.pi, 16 )
  y = np.linspace ( -1.5, +1.5, 9 )
  X, Y = np.meshgrid ( x, y )
  PX = np.ones ( X.shape )
  PY = stiff_ode ( X, Y )
  scale = 5.0
  plt.quiver ( X, Y, PX, PY, scale ) 
#
#  Solve the ODE, and display it as a red line.
#
  x2 = np.linspace ( 0.0, 2.0 * np.pi, 101 )
  y2 = stiff_solution ( x2 )
  plt.plot ( x2, y2, 'r' )
#
#  Finish the graph.
#
  plt.grid ( True )
  plt.axis ( 'equal' )
  plt.title ( 'Direction field for stiff ODE, LAMBDA = ' + str ( LAMBDA ) )
  plt.savefig ( "stiff_direction.png" )
  plt.show ( )
  plt.close ( )

if ( __name__ == '__main__' ):
  LAMBDA = 4.0
  stiff_direction ( LAMBDA )