Найти функцию с NumPy

Question

У меня есть пустой массив, и я хочу найти все индексы, которые проверяют определенное условие. Например, я хочу построить график функции Хевисайда;

import numpy as np
x=np.linspace(-5,5,11)
k_neg=x.find(x<0)
k_pos=x.find(x>=0)
y=np.zeros(len(x))
y(k_neg)=-1
y(k_pos)=1

Я не нахожу такую ​​функцию (как она существует в Matlab).

Примечание: моя настоящая проблема не в том, чтобы построить Heavyside, конечно;)

Answer 1

Как сказал Пол Панцер;

Похоже, вы ищете np.where

Что решило мою проблему.

Answer 2

Я бы сделал это в одной строке с numpy:

import numpy as np
x = np.linspace(-5,5,11)
y = ((x>=0)*2)-1