How to evaluate and plot a 2D function in python ?

To evaluate a two-variable function in python such as for example

\begin{equation}
f: (x_1,x_2) \rightarrow x_1 * \exp^{-(x_1^2+x_2^2)}
\end{equation}

a solution is to use the numpy function meshgrid.

Using meshgrid

Example

from pylab import figure, cm

import matplotlib.pyplot as plt
import numpy as np


def f(x1,x2):
    return x1 * np.exp(-(x1**2+x2**2))


x1_min = -2.0
x1_max = 2.0
x2_min = -2.0
x2_max = 2.0

x1, x2 = np.meshgrid(np.arange(x1_min,x1_max, 0.1), np.arange(x2_min,x2_max, 0.1))

y = f(x1,x2)

Plot the function using imshow from matplotlib

To visulaize the results, a quick solution is to use matplotlib with imshow:

plt.imshow(y,extent=[x1_min,x1_max,x2_min,x2_max], cmap=cm.jet, origin='lower')

plt.colorbar()

plt.title("How to evaluate a 2D function using a python grid?" , fontsize=8)

plt.savefig("evaluate_2d_function_using_meshgrid_03.png", bbox_inches='tight')

plt.show()

We can also plot x1:

plt.imshow(x1, origin='lower', cmap=cm.jet)

plt.title("How to evaluate a 2D function using a python grid?" , fontsize=8)

plt.colorbar()

plt.savefig("evaluate_2d_function_using_meshgrid_01.png", bbox_inches='tight')

plt.show()

plt.close()

or x2

plt.imshow(x2, origin='lower', cmap=cm.jet)

plt.title("How to evaluate a 2D function using a python grid?" , fontsize=8)

plt.colorbar()

plt.savefig("evaluate_2d_function_using_meshgrid_02.png", bbox_inches='tight')

plt.show()

plt.close()

Plot with the matplotlib contour function

You can also use contour

plt.contour(x1,x2,y,extent=[x1_min,x1_max,x2_min,x2_max], cmap=cm.jet, origin='lower')

plt.colorbar()

plt.title("How to evaluate a 2D function using a python grid?" , fontsize=8)

plt.savefig("evaluate_2d_function_using_meshgrid_04.png", bbox_inches='tight')

Plot a 3D function

Note: To visualize in 3D, there are several solutions, we can for example use the contour3D function:

from mpl_toolkits.mplot3d import Axes3D

fig = plt.figure()

ax = fig.add_subplot(111, projection='3d')

ax.contour3D(x1,x2,y, 100,cmap=cm.jet)

ax.view_init(60, 35)

plt.title("How to evaluate a 2D function using a python grid?" , fontsize=8)

plt.savefig("evaluate_2d_function_using_meshgrid_05.png", bbox_inches='tight')

plt.close()

Source code

from pylab import figure, cm

import matplotlib.pyplot as plt
import numpy as np


def f(x1,x2):
    return x1 * np.exp(-(x1**2+x2**2))


x1_min = -2.0
x1_max = 2.0
x2_min = -2.0
x2_max = 2.0

x1, x2 = np.meshgrid(np.arange(x1_min,x1_max, 0.1), np.arange(x2_min,x2_max, 0.1))


plt.imshow(x1, origin='lower', cmap=cm.jet)

plt.title("How to evaluate a 2D function using a python grid?" , fontsize=8)

plt.colorbar()

plt.savefig("evaluate_2d_function_using_meshgrid_01.png", bbox_inches='tight')

#plt.show()

plt.close()


plt.imshow(x2, origin='lower', cmap=cm.jet)

plt.title("How to evaluate a 2D function using a python grid?" , fontsize=8)

plt.colorbar()

plt.savefig("evaluate_2d_function_using_meshgrid_02.png", bbox_inches='tight')

#plt.show()

plt.close()


y = f(x1,x2)

plt.imshow(y,extent=[x1_min,x1_max,x2_min,x2_max], cmap=cm.jet, origin='lower')

plt.colorbar()

plt.title("How to evaluate a 2D function using a python grid?" , fontsize=8)

plt.savefig("evaluate_2d_function_using_meshgrid_03.png", bbox_inches='tight')

#plt.show()

plt.close()


plt.contour(x1,x2,y,extent=[x1_min,x1_max,x2_min,x2_max], cmap=cm.jet, origin='lower')

plt.colorbar()

plt.title("How to evaluate a 2D function using a python grid?" , fontsize=8)

plt.savefig("evaluate_2d_function_using_meshgrid_04.png", bbox_inches='tight')

plt.close()



from mpl_toolkits.mplot3d import Axes3D

fig = plt.figure()

ax = fig.add_subplot(111, projection='3d')

ax.contour3D(x1,x2,y, 100,cmap=cm.jet)

ax.view_init(60, 35)

plt.title("How to evaluate a 2D function using a python grid?" , fontsize=8)

plt.savefig("evaluate_2d_function_using_meshgrid_05.png", bbox_inches='tight')

plt.close()