Contour plots using a function

Question

I am stuck with the following problem;

I have to do a contour plot of a function of two variables chi2(X,Y) in a certain point (chi=2.3). I have tried many ways but I still do not know how to do it. For instance, I have done the following:

  import numpy as np
  import matplotlib.pyplot as plt
  from function import chi2  # Two variable function imported from another file

  x = np.arange(-2,2,0.02)
  y = np.arange(0,1.5,0.02)

  X,Y= np.meshgrid(x,y)

  Z = chi2(X,Y)

  plt.figure()
  CS = plt.contour(X,Y,Z)
  plt.clabel(CS,inline=2.3, fontsize=10)
  plt.show()

The error obtained is that the Supplied function does not return a valid float. But I can print values and work using that function, so I don't know what is exactly the problem.

Thank you in advance.

Bernat.

Answer 1

If you want to use a generic function, the key is to vectorized it first.

for example

import numpy as np
import matplotlib.pyplot as plt

def f(x, y):
    return 5 * x ** 2 + 2 * y + 8

v_func = np.vectorize(f)    # major key!

x, y = np.meshgrid(np.linspace(-5, 5, 100),
                   np.linspace(-5, 5, 100))

fig, ax = plt.subplots(1)
ax.contour(x, y, v_func(x, y))
plt.show()

you should get something like this