#! /usr/bin/env python3
#
def predator_deriv ( t, y ):

#*****************************************************************************80
#
## predator_deriv() returns the right hand side of the predator test problem.
#
#  Licensing:
#
#    This code is distributed under the MIT license.
#
#  Modified:
#
#    04 January 2023
#
#  Author:
#
#    John Burkardt
#
  import numpy as np

  r = y[0]
  f = y[1]

  drdt =   2.0 * r - 0.001 * r * f
  dfdt = -10.0 * f + 0.002 * r * f

  value = np.array ( [ drdt, dfdt ] )

  return value

def predator_rk4 ( ):

#*****************************************************************************80
#
## predator_rk4() calls rk4() with specific values of n, a, b, and y0.
#
#  Licensing:
#
#    This code is distributed under the MIT license.
#
#  Modified:
#
#    04 January 2023
#
#  Author:
#
#    John Burkardt
#
  import matplotlib.pyplot as plt
  import numpy as np

  f = predator_deriv
  tspan = np.array ( [ 0.0, 5.0 ] )
  y0 = np.array ( [ 5000, 100 ] )
  n = 1000
  
  t, y = rk4 ( f, tspan, y0, n )

  dt = ( tspan[1] - tspan[0] ) / n

  plt.clf ( )
  plt.plot ( t, y, linewidth = 3 )
  plt.grid ( True )
  plt.xlabel ( '<-- T -->' )
  plt.ylabel ( '<-- Y(T) -->' )
  label = 'Predator rk4 solution with dt = ' + str ( dt )
  plt.title ( label )
  filename = 'predator_rk4_plot.png'
  plt.savefig ( filename )
  print ( '  Graphics saved as "' + filename + '"' )

  plt.clf ( )
  plt.plot ( y[:,0], y[:,1], linewidth = 3 )
  plt.grid ( True )
  plt.xlabel ( '<-- Prey(T) -->' )
  plt.ylabel ( '<-- Predators(T) -->' )
  label = 'Predator rk4 solution with dt = ' + str ( dt )
  plt.title ( label )
  filename = 'predator_rk4_phase.png'
  plt.savefig ( filename )
  print ( '  Graphics saved as "' + filename + '"' )

  return

def predator_rk5 ( ):

#*****************************************************************************80
#
## predator_rk5() calls rk5() with specific values of n, a, b, and y0.
#
#  Licensing:
#
#    This code is distributed under the MIT license.
#
#  Modified:
#
#    04 January 2023
#
#  Author:
#
#    John Burkardt
#
  import matplotlib.pyplot as plt
  import numpy as np

  f = predator_deriv
  tspan = np.array ( [ 0.0, 5.0 ] )
  y0 = np.array ( [ 5000, 100 ] )
  n = 1000
  
  t, y = rk5 ( f, tspan, y0, n )

  dt = ( tspan[1] - tspan[0] ) / n

  plt.clf ( )
  plt.plot ( t, y, linewidth = 3 )
  plt.grid ( True )
  plt.xlabel ( '<-- T -->' )
  plt.ylabel ( '<-- Y(T) -->' )
  label = 'Predator rk5 solution with dt = ' + str ( dt )
  plt.title ( label )
  filename = 'predator_rk5_plot.png'
  plt.savefig ( filename )
  print ( '  Graphics saved as "' + filename + '"' )

  plt.clf ( )
  plt.plot ( y[:,0], y[:,1], linewidth = 3 )
  plt.grid ( True )
  plt.xlabel ( '<-- Prey(T) -->' )
  plt.ylabel ( '<-- Predators(T) -->' )
  label = 'Predator rk5 solution with dt = ' + str ( dt )
  plt.title ( label )
  filename = 'predator_rk5_phase.png'
  plt.savefig ( filename )
  print ( '  Graphics saved as "' + filename + '"' )

  return

def rk45 ( yprime, tspan, y0, n ):

#*****************************************************************************80
#
## rk45() uses explicit Runge-Kutta schemes of order 4 and 5.
#
#  Discussion:
#
#    This function approximates the solution of a differential
#    equation of the form:
#      y' = yprime(t,y)
#      y(tspan(1)) = y0
#
#  Licensing:
#
#    This code is distributed under the MIT license.
#
#  Modified:
#
#    04 January 2023
#
#  Author:
#
#    John Burkardt
#
#  Input:
#
#    function handle yprime: defines the derivative function of the form
#      value = yprime ( t, y )
#
#    real tspan(2): contains the initial and final times.
#
#    real y0(m): contains the initial condition.
#
#    integer n: the number of steps to take.
#
#  Output:
#
#    real t(n+1): the times.
#
#    real y(n+1,m): the estimated solution values.
#
#    real e(n+1,m): an estimate for the additional error
#    made at each step.
# 
  import numpy as np

  m = len ( y0 )
  t = np.zeros ( n + 1 )
  y = np.zeros ( [ n + 1, m ] )
  e = np.zeros ( [ n + 1, m ] )

  dt = ( tspan[1] - tspan[0] ) / n

  t[0] = tspan[0]
  y[0,:] = y0.copy ( )
  e[0,:] = 0.0
 
  a = np.array ( [ \
       [           0.0,            0.0,            0.0,           0.0, 0.0 ], \
       [          0.25,            0.0,            0.0,           0.0, 0.0 ], \
       [      3.0/32.0,       9.0/32.0,            0.0,           0.0, 0.0 ], \
       [ 1932.0/2197.0, -7200.0/2197.0,  7296.0/2197.0,           0.0, 0.0 ], \
       [   439.0/216.0,           -8.0,   3680.0/513.0, -845.0/4104.0, 0.0 ], \
       [     -8.0/27.0,            2.0, -3544.0/2565.0, 1859.0/4104.0, -11.0/40.0 ] ] )
  b = np.array ( [ 16.0/135.0, 0.0, 6656.0/12825.0, 28561.0/56430.0, -9.0/50.0, 2.0/55.0 ] )
  c = np.array ( [ 0.0, 0.25, 3.0/8.0, 12.0/13.0, 1.0, 0.5 ] )
  d = np.array ( [ 25.0/216.0, 0.0, 1408.0/2565.0, 2197.0/4104.0, -1.0/5.0, 0.0 ] )

  for i in range ( 0, n ):

    k1 = dt * ( yprime ( t[i] + c[0] * dt, y[i,:] ) )
    k2 = dt * ( yprime ( t[i] + c[1] * dt, y[i,:] + a[1,0] * k1 ) )
    k3 = dt * ( yprime ( t[i] + c[2] * dt, y[i,:] + a[2,0] * k1 \
      + a[2,1] * k2 ) )
    k4 = dt * ( yprime ( t[i] + c[3] * dt, y[i,:] + a[3,0] * k1 \
      + a[3,1] * k2 + a[3,2] * k3 ) )
    k5 = dt * ( yprime ( t[i] + c[4] * dt, y[i,:] + a[4,0] * k1 \
      + a[4,1] * k2 + a[4,2] * k3 + a[4,3] * k4 ) )
    k6 = dt * ( yprime ( t[i] + c[5] * dt, y[i,:] + a[5,0] * k1 \
      + a[5,1] * k2 + a[5,2] * k3 + a[5,3] * k4 + a[5,4] * k5 ) )

    y4 = y[i,:] + d[0] * k1 + d[1] * k2 + d[2] * k3 \
                + d[3] * k4 + d[4] * k5 + d[5] * k6
    y5 = y[i,:] + b[0] * k1 + b[1] * k2 + b[2] * k3 \
                + b[3] * k4 + b[4] * k5 + b[5] * k6

    t[i+1] = t[i] + dt
    y[i+1,:] = y5.copy ( )
    e[i+1,:] = y5 - y4

  return t, y, e

def rk45_test ( ):

#*****************************************************************************80
#
## rk45_test() tests rk45().
#
#  Licensing:
#
#    This code is distributed under the MIT license.
#
#  Modified:
#
#    04 January 2023
#
#  Author:
#
#    John Burkardt
#
  import platform

  print ( '' )
  print ( 'rk45_test()' )
  print ( '  Python version: ' + platform.python_version ( ) )
  print ( '  Test rk45().' )

  predator_rk4 ( )
  predator_rk5 ( )
#
#  Terminate.
#
  print ( '' )
  print ( 'rk45_test():' )
  print ( '  Normal end of execution.' )

  return

def rk4 ( yprime, tspan, y0, n ):

#*****************************************************************************80
#
## rk4() uses an explicit Runge-Kutta scheme of order 4.
#
#  Discussion:
#
#    This function approximates the solution of a differential
#    equation of the form:
#      y' = yprime(t,y)
#      y(tspan(1)) = y0
#
#  Licensing:
#
#    This code is distributed under the MIT license.
#
#  Modified:
#
#    04 January 2023
#
#  Author:
#
#    John Burkardt
#
#  Input:
#
#    function handle yprime: defines the derivative function of the form
#      value = yprime ( t, y )
#
#    real tspan(2): contains the initial and final times.
#
#    real y0(m): contains the initial condition.
#
#    integer n: the number of steps to take.
#
#  Output:
#
#    real t(n+1): the times.
#
#    real y(n+1,m): the estimated solution values.
#
  import numpy as np

  m = len ( y0 )
  t = np.zeros ( n + 1 )
  y = np.zeros ( [ n + 1, m ] )

  dt = ( tspan[1] - tspan[0] ) / n

  t[0] = tspan[0]
  y[0,:] = y0.copy ( )

  a = np.array ( [ \
      [       0.0,          0.0,          0.0,        0.0 ], \
      [      0.25,          0.0,          0.0,        0.0 ], \
      [  4.0/81.0,    32.0/81.0,          0.0,        0.0 ], \
      [ 57.0/98.0, -432.0/343.0, 1053.0/686.0,        0.0 ], \
      [   1.0/6.0,          0.0,    27.0/52.0, 49.0/156.0 ] ] )

  b = np.array ( [ 43.0/288.0, 0.0, 243.0/416.0, 343.0/1872.0, 1.0/12.0 ] )
  c = np.array ( [ 0.0, 0.25, 4.0/9.0, 6.0/7.0, 1.0 ] )
 
  for i in range ( 0, n ):

    k1 = dt * yprime ( t[i] + c[0] * dt, y[i,:] )
    k2 = dt * yprime ( t[i] + c[1] * dt, y[i,:] + a[1,0] * k1 )
    k3 = dt * yprime ( t[i] + c[2] * dt, y[i,:] + a[2,0] * k1 \
      + a[2,1] * k2 )
    k4 = dt * yprime ( t[i] + c[3] * dt, y[i,:] + a[3,0] * k1 \
      + a[3,1] * k2 + a[3,2] * k3 )
    k5 = dt * yprime ( t[i] + c[4] * dt, y[i,:] + a[4,0] * k1 \
      + a[4,1] * k2 + a[4,2] * k3 + a[4,3] * k4 )

    y4 = y[i,:] + b[0] * k1 + b[1] * k2 + b[2] * k3 + b[3] * k4 + b[4] * k5

    t[i+1] = t[i] + dt
    y[i+1,:] = y4.copy ( )

  return t, y

def rk5 ( yprime, tspan, y0, n ):

#*****************************************************************************80
#
## rk5() uses an explicit Runge-Kutta scheme of order 5.
#
#  Discussion:
#
#    This function approximates the solution of a differential
#    equation of the form:
#      y' = yprime(t,y)
#      y(tspan(1)) = y0
#
#  Licensing:
#
#    This code is distributed under the MIT license.
#
#  Modified:
#
#    04 January 2023
#
#  Author:
#
#    John Burkardt
#
#  Input:
#
#    function handle yprime: defines the derivative function of the form
#      value = yprime ( t, y )
#
#    real tspan(2): contains the initial and final times.
#
#    real y0(m): contains the initial condition.
#
#    integer n: the number of steps to take.
#
#  Output:
#
#    real t(n+1): the times.
#
#    real y(n+1,m): the estimated solution values.
#
  import numpy as np

  m = len ( y0 )
  t = np.zeros ( n + 1 )
  y = np.zeros ( [ n + 1, m ] )

  dt = ( tspan[1] - tspan[0] ) / n

  t[0] = tspan[0]
  y[0,:] = y0.copy ( )
 
  a = np.array ( [ \
       [           0.0,            0.0,            0.0,           0.0, 0.0 ], \
       [          0.25,            0.0,            0.0,           0.0, 0.0 ], \
       [      3.0/32.0,       9.0/32.0,            0.0,           0.0, 0.0 ], \
       [ 1932.0/2197.0, -7200.0/2197.0,  7296.0/2197.0,           0.0, 0.0 ], \
       [   439.0/216.0,           -8.0,   3680.0/513.0, -845.0/4104.0, 0.0 ], \
       [     -8.0/27.0,            2.0, -3544.0/2565.0, 1859.0/4104.0, -11.0/40.0 ] ] )

  b = np.array ( [ 16.0/135.0, 0.0, 6656.0/12825.0, 28561.0/56430.0, -9.0/50.0, 2.0/55.0 ] )
  c = np.array ( [ 0.0, 0.25, 3.0/8.0, 12.0/13.0, 1.0, 0.5 ] )
  d = np.array ( [ 25.0/216.0, 0.0, 1408.0/2565.0, 2197.0/4104.0, -1.0/5.0, 0.0 ] )

  for i in range ( 0, n ):

    k1 = dt * yprime ( t[i] + c[0] * dt, y[i,:] )
    k2 = dt * yprime ( t[i] + c[1] * dt, y[i,:] + a[1,0] * k1 )
    k3 = dt * yprime ( t[i] + c[2] * dt, y[i,:] + a[2,0] * k1 \
      + a[2,1] * k2 )
    k4 = dt * yprime ( t[i] + c[3] * dt, y[i,:] + a[3,0] * k1 \
      + a[3,1] * k2 + a[3,2] * k3 )
    k5 = dt * yprime ( t[i] + c[4] * dt, y[i,:] + a[4,0] * k1 \
      + a[4,1] * k2 + a[4,2] * k3 + a[4,3] * k4 )
    k6 = dt * yprime ( t[i] + c[5] * dt, y[i,:] + a[4,0] * k1 \
      + a[5,1] * k2 + a[5,2] * k3 + a[5,3] * k4 + a[5,4] * k5 )

    y5 = y[i,:] + b[0] * k1 + b[1] * k2 + b[2] * k3 \
                + b[3] * k4 + b[4] * k5 + b[5] * k6

    t[i+1] = t[i] + dt
    y[i+1,:] = y5.copy ( )

  return t, y

def timestamp ( ):

#*****************************************************************************80
#
## timestamp() prints the date as a timestamp.
#
#  Licensing:
#
#    This code is distributed under the MIT license. 
#
#  Modified:
#
#    21 August 2019
#
#  Author:
#
#    John Burkardt
#
  import time

  t = time.time ( )
  print ( time.ctime ( t ) )

  return

if ( __name__ == '__main__' ):
  timestamp ( )
  rk45_test ( )
  timestamp ( )