plot-sample-landscape.py

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits import mplot3d

def plot_surface(domain, fn, title=None, grid_samples=100):
    x = np.linspace(domain[0][0], domain[0][1], grid_samples)
    y = np.linspace(domain[1][0], domain[1][1], grid_samples)
    
    X, Y = np.meshgrid(x, y)
    
    fn_vectorized = np.vectorize(fn)
    Z = fn_vectorized(X, Y)
    
    plt.figure(figsize=(20,10))
    ax = plt.axes(projection="3d")
    ax.plot_surface(X, Y, Z, rstride=1, cstride=1, 
                   cmap='terrain', edgecolor=None)
    ax.set(xlabel="x1", ylabel="x2", zlabel="Loss(x1, x2)", title=title);
    
def func(x, y):
    return -5*x*y*np.exp(-x**2-y**2)
  

domain = [(-2, 2), (-2, 2)] 
plot_surface(domain, func)