3

I have a function which is of the form : 

def f(x, y):
    total = 0
    u = np.zeros(10)
    for i in range(0,10):
        u[i] = x * i + y* i
        if u[i] < 10:
            print('do something')
    total = total + u[i]        
    return total

this function when i try with a given x and y values works well. 

f(3,4)
Out[49]: 63.0

I want to create a 3d contour plot using matplotlib. Tried with 


x = np.linspace(-6, 6, 30)
y = np.linspace(-6, 6, 30)

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

fig = plt.figure()
ax = plt.axes(projection='3d')
ax.contour3D(X, Y, Z, 50, cmap='binary')
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('z');

I had to create a mesh grid for 3d plot. I get an error when I try with this because of the loop in my function. I get an error

ValueError: setting an array element with a sequence.

How to plot 3d graphs if my function has a loop?

chink
  • 1,505
  • 3
  • 28
  • 70
  • 3
    Try using shape `(10, len(x), len(y))` (`u = np.zeros((10, len(x), len(y)))`) for your zeros matrix on your `f` function. The later pyplot usage also throws some errors after that. – Endor 8 May 22 '19 at 06:26
  • Thanks @Endor8 Is this the only way? because my original function which is bit complex has some comparisons also. Edited the question to match that case. – chink May 22 '19 at 09:29
  • You're welcome :) I've seen [here](/questions/21803256/error-setting-an-array-element-with-a-sequence-python-numpy) that you can use the dtype attribute to allow your matrix to store any object (like `u = np.zeros(10, dtype=object)`) which would do what you want. However, best performance would be achieved using only numpy features. When you compare a numpy array with a number you get an array of bools, which can't be used directly in an `if` statement. You might want to use then `np.any` or `np.all` on the result (`if np.any(u[i] < 10):`) to check if any or all of the elements are True. – Endor 8 May 22 '19 at 19:43

1 Answers1

3

You need np.vectorize:

# same setup as above, then
Z = np.vectorize(f)(X, Y)

import pylab as plt
plt.imshow(Z, extent=[x[0], x[-1], y[0], y[-1]])

(I checked with imshow but contour3D will work too.)

np.vectorize will take a function that accepts scalar (non-array) arguments and magically loops over arrays. It's nominally equivalent to:

Z2 = np.array([f(xx, yy) for xx in x for yy in y]).reshape(X.shape)
print(np.abs(Z - Z2).max()) # should print 0

but faster: after I remove the print in f:

In [47]: %timeit Z = np.vectorize(f)(X, Y)
6 ms ± 339 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [48]: %timeit Z2 = np.array([f(xx, yy) for xx in x for yy in y]).reshape(X.shape)
13.7 ms ± 310 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

(I had to remove the prints for timing because printing is Very Slow.)

Ahmed Fasih
  • 6,458
  • 7
  • 54
  • 95