Gradient descent implementation

gradient_descent.py

from matplotlib import cm #For the colors
import numpy as np
import matplotlib.pyplot as plt

def f(x,y):
    return x**2 + y**2

def derivate(point_with_increment, point, h):
    return (f(point_with_increment[0], point_with_increment[1]) - f(point[0], point[1])) / h

def gradient(point, h):
    grad = np.zeros(2)
    for i, v  in enumerate(point):
        cp = np.copy(point)
        cp[i] = cp[i] + h

        dp = derivate(cp, point, h)
        grad[i] = dp
    return grad

if __name__=='__main__':

    fig, ax = plt.subplots(subplot_kw={"projection":"3d"})

    res = 100
    x = np.linspace(-4,4,res)
    y = np.linspace(-4,4,res)

    X,Y = np.meshgrid(x,y)

    Z = f(X,Y)

    surf = ax.plot_surface(X,Y,Z, cmap=cm.cool)
    fig.colorbar(surf)
    plt.show()

    level_map = np.linspace(np.min(Z), np.max(Z), res)
    plt.contourf(X,Y,Z, levels=level_map, cmap=cm.cool)
    plt.colorbar()
    plt.title("Gradient descent")

    point = np.random.rand(2) * 8 - 4 #We muplitply by 8 minus 4 because we want the point to be between -4 and 4
    plt.plot(point[0], point[1], 'o', c='k')

    h = 0.01
    learning_rate = 0.01

    gradient(point, h)

    for i in  range(1000):
        point = point - learning_rate * gradient(point, h)
        if (i % 10 == 0): #This is only to plot every 10 ranges
            plt.plot(point[0], point[1], 'o', c='r') #Same here

    plt.plot(point[0], point[1], 'o', c='w')
    plt.show()