surface.py 文件源码

def surface_bounds(f, bounds, cmap=cm.magma_r, resolution=[100, 100], type='2D'):
    """
    Plot a function over a grid
    :param f: function to plot in Z
    :param bounds:
    :param cmap: cmap
    :param resolution:
    :param type: '2D' or '3D'
    :return:
    """
    resolution = np.array(resolution)
    assert bounds.get_n_dim() == 2, 'Bounds are not 2D'

    x = np.linspace(start=bounds.get_min(0), stop=bounds.get_max(0), num=resolution[0])
    y = np.linspace(start=bounds.get_min(1), stop=bounds.get_max(1), num=resolution[1])
    X, Y = np.meshgrid(x, y)
    Z = f(np.vstack((X.flatten(), Y.flatten())).T).reshape(X.shape)  # Evaluate grid on the function
    surface(X=X, Y=Y, Z=Z, type=type, cmap=cmap)


# def surface_grid(f, x, y)