0

I was trying to plot a function which was defined something as follows:

def f(x,y);
    if x==y:
        return x
    else:
        return x-y

Now there's nothing wrong with that syntactically, except that I can't plot it:

from numpy import mgrid
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
[X,Y] = mgrid[0:1:0,01,0:1:0.01]
fig = plt.figure()
ax = fig.gca(projection='3d')
surf = ax.plot_surface(X,Y,f(X,Y))

This brings up the error:

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

which is due to Python trying to apply a boolean function to an array. I tried to rewrite the function with

if np.equal(x,y)

as the test, but that didn't work either. I know I can create the array of f(X,Y) values "by hand" so to speak with various loops, but what I want to know is: is there a way of defining a function which includes a boolean condition, and which can then be easily plotted with matplotlib/pyplot?

Alasdair
  • 1,300
  • 4
  • 16
  • 28

1 Answers1

1

In the definition of f, if you expect that x and y can be arrays (which they are), you have to do something different than "if x==y:" - for example "if all(x==y):".

Additionally you may have a problem after that with what you mentioned - then use something like [f(x, y) for x, y in zip(X, Y)].

Or map(f, *zip(*zip(X, Y))), if you want to be more idiomatic, but less obvious.

Edit: Above I assume that indeed you expect f to accept arrays, and just guessed that all() may be what you need.

Petar Donchev
  • 399
  • 4
  • 9