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

#*****************************************************************************80
#
## ex52, lab 5 exercise 2.
#
#  Discussion:
#
#    Gradient descent for a function of 2 variables.
#
#  Licensing:
#
#    This code is distributed under the GNU LGPL license.
#
#  Modified:
#
#    18 April 2019
#
#  Author:
#
#    John Burkardt
#
  import matplotlib.pyplot as plt
  import numpy as np
  import platform
  from mpl_toolkits.mplot3d import Axes3D
  from matplotlib import cm

  print ( '' )
  print ( 'ex52:' )
  print ( '  Python version: %s' % ( platform.python_version ( ) ) )
  print ( '  Seek minimizer of a function z(x,y).' )
#
#  Choose a random initial guess.
#
  x = - 3.0 + 6.0 * np.random.rand ( 1 )
  y = - 3.0 + 6.0 * np.random.rand ( 1 )

  x0 = x
  y0 = y
  z0 = f ( x0, y0 )
#
#  Choose a small learning rate.
#
  r = 0.05
#
#  Choose a small tolerance for the derivative.
#
  tol = 0.001
#
#  Loop.
#
  step = 0
  while ( step < 100 ):
    step = step + 1
    z = f ( x, y )
    dzdx, dzdy = fp ( x, y )
    if ( np.sqrt ( dzdx**2 + dzdy**2 ) < tol ):
      break
    x = x - r * dzdx
    y = y - r * dzdy
    print ( '  %d  %g  %g  %g  %g  %g' % ( step, x, y, z, dzdx, dzdy ) )
#
#  Set up table X1, X2, F(X1,X2).
#
  xvec = np.linspace ( -2.0, 2.0, 101 )
  yvec = np.linspace ( -2.0, 2.0, 101 )
  xmat, ymat = np.meshgrid ( xvec, yvec )
  zmat = f ( xmat, ymat )
#
#  Surface plot
#
  fig = plt.figure()
  ax = fig.add_subplot ( 111, projection = '3d' )
  ax.plot_surface ( xmat, ymat, zmat, cmap = cm.coolwarm,
    linewidth = 0, antialiased = False, alpha = 0.5 )
  ax.plot ( x0, y0, z0, 'bo', markersize = 20 )
  ax.plot ( x, y, z, 'ro', markersize = 20 )
  ax.set_xlabel ( '<--- Y --->' )
  ax.set_ylabel ( '<--- X --->' )
  ax.set_zlabel ( '<---Z(X,Y)--->' )
  ax.set_title ( 'Function to be minimized' )
  plt.draw ( )
  filename = 'ex52_surface.png'
  plt.savefig ( filename )
  print ( '' )
  print ( '  Graphics saved in "%s"' % ( filename ) )
  plt.show ( )
  plt.clf ( )
#
#  Terminate.
#
  print ( '' )
  print ( 'ex52:' )
  print ( '  Normal end of execution.' )
  return

def f ( x, y ):

#*****************************************************************************80
#
## F evaluates the function to be minimized.
#
  z = 2.0 * x**2 - 1.05 * x**4 + x**6 / 6.0 + x*y + y**2

  return z

def fp ( x, y ):

#*****************************************************************************80
#
## FP evaluates the derivative of the function to be minimized.
#
  dzdx = 4.0 * x - 4.2 * x**3 + x**5 + y
  dzdy = x + 2.0 * y

  return dzdx, dzdy

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