Pythonのmatplotlibで自分で作った関数を適用したいけれど何故か動かないです

Question

Pythonを初めて間もない初心者です。

matplotlibで自分で作った関数を適用したいと思い、以下のように書きました。

import matplotlib.pyplot as plt
import numpy as np

def hoge(x):
	if x >= 0:
		return 1
	else:
		return 0

x = np.arange(-5, 5, 1);
y = hoge(x)

plt.stem(x, y,'-.')
plt.show()

xが０以上の時、yが常に１というように出力したいのですが、エラーで実行できません。

def hoge(x):
  x**2

とすると動作するのですが・・・

お力お貸しいただけると幸いです

Answer 1

いくつかやり方があります。

愚直な方法

x = np.arange(-5, 5, 1);
y = [hoge(n) for n in x]

これでyに [0, 0, 0, 0, 0, 1, 1, 1, 1, 1] が代入されます。

愚直な方法は速度を考慮していないので、xのデータ量が多くなると遅くなりそうです。

f = np.vectorize(hoge)
y = f(x)

これでyに array([0, 0, 0, 0, 0, 1, 1, 1, 1, 1]) が代入されます。
前の方法より速いことが期待されます。

Comments

[KotaroYamame]

丁寧にありがとうございます
大変わかりやすかったです
他にもやりたい事があるのですが、教えていただいた内容で解決出来そうです

Answer 2

エラーとなる原因は hoge()の引数x にベクトル(np.array)が渡るからです。

ベクトルの各要素に対して三項演算子っぽい処理が行える np.where() を使うという手もあります。

def hoge(x):
    return np.where(x>=0, 1, 0)

x = np.arange(-5, 5, 1);
y = hoge(x)

https://docs.scipy.org/doc/numpy-1.12.0/reference/generated/numpy.where.html

Comments

[KotaroYamame]

引数xは数値かと思っていました
大変勉強になりました
ありがとうございます

Answer 3

pythonではbool型を使って数値計算すると、
Trueは1、Falseは0として計算されます。
従って

def hoge(x):
    return (x >= 0) + 0

y = hoge(x)

で望みの値が得られます。

Comments

[KotaroYamame]

ありがとうございます
Pythonはそういう性質があるのですね
勉強になります

[shimizukawa]

こちらに説明がありますので参考にどうぞ
https://docs.python.jp/3/reference/datamodel.html#the-standard-type-hierarchy
True と 1 は完全に同じ、ではない、ことに注意してください。
````
>>> int(True)
1
>>> True == 1
True
>>> True + True
2
>>> True is 1
False
>>> str(True)
'True'
```