How do I create a contourplot with a custom function?

Viewed 41

I have seen many examples online of create a contourplot as follows

import numpy as np
xlist = np.linspace(-3.0, 3.0, 3)
ylist = np.linspace(-3.0, 3.0, 4)
X, Y = np.meshgrid(xlist, ylist)
Z = np.sqrt(X**2 + Y**2)
cp = plt.contourf(X, Y, Z)
plt.colorbar(cp)

ax.set_title('Contour Plot')
ax.set_xlabel('x (cm)')
ax.set_ylabel('y (cm)')
plt.show()

now this is quite straightforward because the sqrt function in numpy will happily accept meshgrid inputs and evaluate on all points of interest. What I don't understand however is how I would do this if I have a function of my own. For example say I have a much more complex function, lambda nllh x1,x2 : my_function(x1,x2) which could be arbitrarily complicated but takes in two scalars and itself returns a scalar.

When I try to implement this with my function in the same way (i.e. the same code above except with the line Z=nllh(X,Y)), I get an error that the input array dimensions don't make any sense, as this function is not designed to take a meshgrid as an input. How can i rectify this issue or make python understand that the function Z=nllh(X,Y) needs to be evaluated on the meshgrid value pairs independently instead of the meshgrids themselves?

Thanks

1 Answers

Answered by Tadhg McDonald-Jensen in a comment. Create a new function with np.vectorize(nllh)

Related