Python plot of a piecewise defined surface

Question

I am trying to make a 3d plot of a surface that is defined in different ways for different regions. As an example, take f(x,y) that is defined as 1 if x > y and as x^2 if x <= y.

I defined f with logical operators, and tried to plot it with the "plot_surface" function, evaluating it in a grid. Unfortunately, I got an error saying that "the truth value of an array with more than one element is ambiguous".

Do you know any way of solving this?

Possible duplicate of [ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()](http://stackoverflow.com/questions/10062954/valueerror-the-truth-value-of-an-array-with-more-than-one-element-is-ambiguous) — m00am, Jun 29 '16 at 09:25
Possible duplicate of [piecewise function with 3d plot](http://stackoverflow.com/questions/22430429/piecewise-function-with-3d-plot) — Serenity, Jun 29 '16 at 09:43
I read the entries and I would say that the specific problems that are treated there do not solve my issue... — Ignatius, Jun 29 '16 at 10:22

score 0 · Answer 1 · answered Jun 30 '16 at 16:14

Taking from the link posted by Serenity you need to define f using np.piecewise

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import axes3d, Axes3D 

num_steps = 500
x_arr = np.linspace(0,100, num_steps)
y_arr = np.linspace(0,100, num_steps)

def zfunc(x, y):
    return np.piecewise(x, [x>y, x<=y], [lambda x: 1, lambda x: x**2])

x,y = np.meshgrid(x_arr, y_arr)
z =zfunc(x,y)

fig=plt.figure()
ax=fig.add_subplot(1,1,1,projection='3d')
ax.set_xlabel('X axis')
ax.set_ylabel('Y axis')

ax.plot_surface(x,y,z,cmap='viridis') #cmap to make it easier to see 
ax.view_init(30, 70)
plt.show()

Giving you this plot:

Python plot of a piecewise defined surface

1 Answers1

Linked