Source code for surface

# -*- coding: utf-8 -*-
r"""Plot a surface in three-dimensions.
"""
import numpy as np
import matplotlib.pyplot as plt




def surface(z, x_range, y_range):
    r"""Plot a surface in three-dimensions.


    :param z: The function that represents the surface.

        If you have the equation of the figure in the form
            .. math::
                F(x, y, z) = 0,

        You should first re-write it into the form

            .. math::
                z(x, y) = 0.

    :param x_range: The range of :math:`x`-coordinates.
    :param y_range: The range of :math:`y`-coordinates.

    """
    ax = plt.figure().add_subplot(projection='3d')
    x_linspace = np.linspace(x_range[0], x_range[1], 100)
    y_linspace = np.linspace(y_range[0], y_range[1], 100)
    meshgrid = np.meshgrid(x_linspace, y_linspace)
    z = z(*meshgrid)
    ax.plot_surface(*meshgrid, z, cmap="coolwarm", linewidth=0)
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z')
    plt.show()



if __name__ == '__main__':
    def z(x, y):
        return np.sin(2 * np.pi * x) * np.cos(2 * np.pi * y)
    surface(z, [0, 1], [0, 1])