Search code examples
pythonmatplotlibmachine-learninglinear-regressiongradient-descent

How to illustrate a 3D graph of gradient descent using python matplotlib?


I am having trouble with plotting a 3d graph for gradient descent using python's matplotlib. The commented code in the gradient_descent function was what I tried but doesn't work. I would appreciate any solutions, including solutions that uses other libraries. The graph below is an example of what I wanted to plot.

enter image description here

Here's the code:

import numpy as np
import matplotlib.pyplot as plt


def scatter_plot(x, y, x_title, y_title, g_title):
    plt.plot(x, y, 'bo')
    plt.xlabel(x_title)
    plt.ylabel(y_title)
    plt.title(g_title)


def plot_hypothesis(x, y, hyp, x_title, y_title, g_title):
    scatter_plot(x, y, x_title, y_title, g_title)
    plt.plot(x, hyp, '--')
    plt.show()


def gradient_descent(x, y, iterations, r):
    t0 = t1 = 0                                                                 # t0 = y-intercept, t1 = gradient
    m = len(x)                                                                  # Number of training examples
    h = 0                                                                       # Initialize 0 to the hypothesis
    cost = 0
    print("Learning Rate = ", r)
    print('Number of Iterations = ', iterations)
    for i in range(iterations):
        h = t0 + (t1 * x)                                                       # Set hypothesis
        cost = (1/(2 * m)) * sum([val**2 for val in (h - y)])                   # Calculate cost
        t0 = t0 - r * (1 / m) * sum(h - y)                                      # Partial derivative of t0 and update t0
        t1 = t1 - r * (1 / m) * sum((h - y) * x)                                # Partial derivative of t1 and update t1
        print("i={}, cost={}, t0={}, t1={}".format(i, cost, t0, t1))
    plot_hypothesis(x, y, h, 'year', 'life expectancy', 'Malaysian Males Life Expectancy At Birth')
    # fig = plt.figure()
    # ax = fig.add_subplot(111, projection='3d')
    # ax.plot_trisurf(cost, t1, t0, colot='None', alpha=0.5)
    # ax.set_xlabel('J(\u03B80,\u03B81)')
    # ax.set_ylabel('\u03B81')
    # ax.set_zlabel('\u03B80')

# main()
# x = year
x = np.array([1966, 1967, 1968, 1969, 1970, 1971, 1972, 1973, 1974, 1975, 1976, 1977, 1978, 1979, 1980, 1981, 1982, 1983, 1984, 1985, 1986, 1987, 1988, 1989, 1990, 1991, 1992, 1993, 1994, 1995, 1996, 1997, 1998, 1999, 2000, 2001, 2002, 2003, 2004, 2005, 2006, 2007, 2008, 2009, 2010, 2011, 2012, 2013, 2014, 2015])
# y = life expectancy at birth
y = np.array([63.1, 63.5, 63.3, 63.8, 61.6, 62.6, 62.8, 63.2, 63.6, 64.3, 64.7, 65.3, 65.6, 65.8, 66.4, 66.9, 67.1, 67.1, 67.2, 67.7, 68.2, 68.5, 68.7, 68.8, 68.9, 69.2, 69.4, 69.6, 69.6, 69.5, 69.5, 69.7, 69.5, 69.7, 70, 70.6, 70.7, 70.8, 71.1, 71.4, 71.6, 71.6, 71.6, 71.7, 71.9, 72.1, 72.2, 72.4, 72.5, 72.5])

scatter_plot(x, y, 'year', 'life expectancy', 'Malaysian Males Life Expectancy At Birth')
plt.show()

gradient_descent(x, y, 100, 0.0000001)  

                                    

Solution

  • In order to create a 3D surface a very basic code is something like the following:

    import matplotlib.pyplot as plt
    from mpl_toolkits.mplot3d import Axes3D
    from matplotlib import cm
    import numpy as np
    
    fig = plt.figure()
    ax = fig.gca(projection='3d')   # Create the axes
    
    # Data
    X = np.linspace(-8, 8, 100)
    Y = np.linspace(-4, 4, 100)
    X, Y = np.meshgrid(X, Y)
    Z = X**2 + Y**2
    
    # Plot the 3d surface
    surface = ax.plot_surface(X, Y, Z,
                              cmap=cm.coolwarm,
                              rstride = 2,
                              cstride = 2)
    
    # Set some labels
    ax.set_xlabel('x-axis')
    ax.set_ylabel('y-axis')
    ax.set_zlabel('z-axis')
    
    plt.show()
    

    Resulting this:

    However, to create a 3D surface for gradient descent as you want, you should consider again which data you need to plot it. You need for example a list of all thetas and costs. Based on how plot_surface works try to figure out what data you need and modify your gradient_descent function accordingly. Also take a look at this implementation.