ошибка при использовании np.meshgrid для построения трехмерного графика из-за петель

У меня есть функция вида:

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

эта функция, когда я пытаюсь с заданными значениями x и y, работает хорошо.

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

Я хочу создать трехмерный контурный график, используя matplotlib. Пробовал с


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');

Мне пришлось создать сетку сетки для 3D-графика. Я получаю сообщение об ошибке, когда пытаюсь сделать это из-за цикла в моей функции. я получаю сообщение об ошибке

ValueError: setting an array element with a sequence.

Как построить трехмерные графики, если моя функция имеет цикл?


person cvg    schedule 22.05.2019    source источник
comment
Попробуйте использовать форму (10, len(x), len(y)) (u = np.zeros((10, len(x), len(y)))) для матрицы нулей в вашей функции f. Более позднее использование pyplot также вызывает некоторые ошибки после этого.   -  person Endor 8    schedule 22.05.2019
comment
Спасибо @ Endor8 Это единственный способ? потому что моя исходная функция, которая немного сложна, также имеет некоторые сравнения. Отредактировал вопрос, чтобы он соответствовал этому случаю.   -  person cvg    schedule 22.05.2019
comment
Добро пожаловать :) Я видел здесь, что вы можете использовать атрибут dtype, позволяющий вашей матрице хранить любой объект (например, u = np.zeros(10, dtype=object)), который будет делать то, что вы хотите. Однако наилучшая производительность будет достигнута при использовании только функций numpy. Когда вы сравниваете массив numpy с числом, вы получаете массив логических значений, который нельзя использовать непосредственно в операторе if. Возможно, вы захотите использовать затем np.any или np.all для результата (if np.any(u[i] < 10):), чтобы проверить, являются ли какие-либо или все элементы истинными.   -  person Endor 8    schedule 22.05.2019


Ответы (1)


Вам нужно 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]])

(Я проверил с imshow, но contour3D тоже подойдет.)

np.vectorize возьмет функцию, которая принимает скалярные (не массивные) аргументы и волшебным образом перебирает массивы. Это номинально эквивалентно:

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

но быстрее: после удаления print в 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)

(Мне пришлось удалить отпечатки по времени, потому что печать идет очень медленно.)

person Ahmed Fasih    schedule 22.05.2019