Plotting the 3D conditional function in Python

Question

I need to make a 3D plot with the condition that in area of the grid where (x-(M/2))^2+(y-(N/2))^2 < r it draws the paraboloid, in the other grid points there is a flat plane.

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

M=70
N=70
r=(N/2)**2

def f(x, y):
    if (x-(M/2))**2+(y-(N/2))**2<r:
        return (x-(M/2))**2/122.5+(y-(N/2))**2/122.5
    return 10.0

x = np.linspace(0, 70, M)
y = np.linspace(0, 70, N)
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')
plt.show()

But there is some kind of mistake with this condition, so I get:

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

Please, help me to make this conditional work.

Is `[1, 5, 0]` greater or smaller than `3`? Read on `numpy.where`. — ImportanceOfBeingErnest, Nov 06 '19 at 13:58

score 0 · Accepted Answer · answered Nov 06 '19 at 14:08

Your error comes from the fact that you compare an array (the result of (x-(M/2))**2+(y-(N/2))**2 againt a scalar value r = 1225.

This result is an array of boolean, which cannot be used as is as a condition (see this answer).

As advised by @ImportanceOfBeingErnest, you should use np.where in this case.

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

M=70
N=70
r=(N/2)**2

def f(x, y):
    computed = (x-(M/2))**2+(y-(N/2))**2
    return np.where(computed < r, r / 122.5, 10.0)

x = np.linspace(0, 70, M)
y = np.linspace(0, 70, N)
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')
plt.show()

Plotting the 3D conditional function in Python

1 Answers1