Search code examples
pythonmatplotlibplotgraphing

Is there a way in Matplotlib to plot a function of two variables, y = y(x₀, x₁)?

I have a function that takes two inputs as arguments and returns an output based on the two inputs.

The two inputs, r and E range between 3-14 and 0.05-0.75 respectively and will be used to create some values for inputs based on some intervals.

Each set of the two inputs (r and E) will generate an output. I want to plot the function based on that data using Python's Matplotlib library.

I have used matplotlib for graphing an input x with output y, but I can't think of a way to plot a function whose output depends equally on two inputs.

Solution

  • The most critical point in what you want to do is to understand that Matplotlib requires that the data is organized in a 2D grid, consequently it is convenient that you compute the data exploiting Numpy's ability to operate on multidimensional data structures.

    1. compute two 1D arrays with the coordinates of the grid — in the example the grid points are equally spaced but the grids can be also irregular,
    2. from the two 1D arrays, using numpy.meshgrid, compute two 2D arrays containing, the first the x₀ coordinates in all the grid points, the second the x₁ coordinates in all the grid points.
    3. using the vector-enabled Numpy's mathematical functions, compute the values of y on the grid at once, using the 2D arrays as the x₀ and x₁,
    4. plot the contour lines, in black, using a relatively small number of levels to avoid a crowded figure,
    5. plot a color-mapped representation of y, this time we can use a larger number of levels to have a continuous effect,
    6. add the labels on the contour lines (note that we use the object returned by contour) and add a colorbar that explains the values in the figure (note that we use the object returned by contourf):

    And this is the simple code that implements the lengthy explanation above:

    In [35]: import numpy as np
    ...: import matplotlib.pyplot as plt
    ...: 
    ...: X0 = np.linspace(3, 14, 111)
    ...: X1 = np.linspace(0.05, 0.75, 71)
    ...: x0, x1 = np.meshgrid(X0, X1)
    ...: y = np.cos(x0-10*x1)*(x0-40*(x1-0.2)**2)
    ...: 
    ...: ctr = plt.contour( x0, x1, y, levels=4, colors='k')
    ...: fil = plt.contourf(x0, x1, y, levels=60)
    ...: plt.clabel(ctr)
    ...: plt.colorbar(fil)
    ...: plt.show()

    plot of a function of two variables using contours and colormap

    Another possibility is to make a 3D plot, using the plot_surface method. The plot_surface method is only available for Axes that have been instantiated using the keyword argument projection='3d', so we cannot use the convenience plt namespace as in the previous example…

    Everything we discussed about the grid is still valid, the signature in terms of requested arguments is the same, …(x0, x1, y, **kwd_args), the example is hence straightforward if you followed the previous one, but it uses a couple of tricks

    1. when you plot a surface on your computer, you can rotate the graph at will to look at the data from different points of view, when you publish the figure in a traditional way (e.g., on Stack Overflow) you lose this interactivity. This answer introduced me to a technique that alleviates the problem.
    2. if you want to place a colormap aside of the multiple subplots, it can be difficult to place it correctly, difficult unless you have found this answer…

    Again, the code and its product

    In [41]: fig, axes = plt.subplots(2,2,figsize=(8,6),
    ...:                          constrained_layout=True,
    ...:                          subplot_kw=dict(projection='3d'))
    ...: surfs = [ax.view_init(elev=10.+i*10, azim=25+i*25) or
    ...:          ax.plot_surface(x0, x1, y, cmap='viridis')
    ...:              for i, ax in enumerate(axes.flat)]
    ...: fig.colorbar(surfs[0], ax=axes.flat, aspect=50);

In [42]:

    surface plots of a function of two variables