#! /usr/bin/env python3
#
def ab2 ( f_ode, xRange, yInitial, numSteps ):

#*****************************************************************************80
#
## ab2() uses the Adams-Bashforth second order ODE solver.
#
#  Input:
#    f_ode = evaluates right hand side.
#    xRange = [x1,x2] where the solution is sought on x1<=x<=x2
#    yInitial = column vector of initial values for y at x1
#    numSteps = number of equally-sized steps to take from x1 to x2
#  Output:
#    x = vector of values of x
#    y = matrix whose k-th row is the approximate solution at x(k).
#
  import numpy as np

  x = np.zeros ( numSteps + 1 )
  y = np.zeros ( ( numSteps + 1, np.size ( yInitial ) ) )

  dx = ( xRange[1] - xRange[0] ) / numSteps

  for k in range ( 0, numSteps + 1 ):
    if ( k == 0 ):
      x[k] = xRange[0]
      y[k,:] = yInitial
    elif ( k == 1 ):
      fValue = f_ode ( x[k-1], y[k-1,:] )
      xhalf = x[k-1] + 0.5 * dx
      yhalf = y[k-1,:] + 0.5 * dx * fValue
      fValuehalf = f_ode ( xhalf, yhalf )
      x[k] = x[k-1] + dx
      y[k,:] = y[k-1,:] + dx * fValuehalf
    else:
      fValueold = fValue
      fValue = f_ode ( x[k-1], y[k-1,:] )
      x[k] = x[k-1] + dx
      y[k,:] = y[k-1,:] + dx * ( 3.0 * fValue - fValueold ) / 2.0

  return x, y

def ab2_test ( ):

#*****************************************************************************80
#
## ab2_test() solves the humps ODE using the AB2.
#
#  Licensing:
#
#    This code is distributed under the GNU LGPL license.
#
#  Modified:
#
#    22 April 2020
#
#  Author:
#
#    John Burkardt
#
  import matplotlib.pyplot as plt
  import numpy as np

  print ( '' )
  print ( 'ab2_test():' )
  print ( '  Solve the humps ODE system using ab2().' )

  f = humps_deriv
  tspan = np.array ( [ 0.0, 2.0 ] )
  y0 = 5.1765
  n = 100
  t, y = ab2 ( f, tspan, y0, n )

  plt.plot ( t, y, 'r-', linewidth = 2 )

  a = tspan[0]
  b = tspan[1]
  if ( a <= 0.0 and 0.0 <= b ):
    plt.plot ( [a,b], [0,0], 'k-', linewidth = 2 )

  ymin = min ( y )
  ymax = max ( y )
  if ( ymin <= 0.0 and 0.0 <= ymax ):
    plt.plot ( [0, 0], [ymin,ymax], 'k-', linewidth = 2 )

  plt.grid ( True )
  plt.xlabel ( '<--- T --->' )
  plt.ylabel ( '<--- Y(T) --->' )
  plt.title ( 'ab2_test(): time plot' )

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

  return

def humps_deriv ( x, y ):

#*****************************************************************************80
#
## humps_deriv() evaluates the derivative of the humps function.
#
#  Discussion:
#
#    y = 1.0 / ( ( x - 0.3 )^2 + 0.01 ) \
#      + 1.0 / ( ( x - 0.9 )^2 + 0.04 ) \
#      - 6.0
#
#  Licensing:
#
#    This code is distributed under the GNU LGPL license.
#
#  Modified:
#
#    22 April 2020
#
#  Author:
#
#    John Burkardt
#
#  Input:
#
#    real x[:], y[:]: the arguments.
#
#  Output:
#
#    real yp[:]: the value of the derivative at x.
#
  yp = - 2.0 * ( x - 0.3 ) / ( ( x - 0.3 )**2 + 0.01 )**2 \
       - 2.0 * ( x - 0.9 ) / ( ( x - 0.9 )**2 + 0.04 )**2

  return yp

if ( __name__ == '__main__' ):
  ab2_test ( )