Pythonで場合分けされたグラフを描こう！

結論

np.vectorizeを使え！

詳しく

場合分けされたグラフを描く記事があまり見当たらなかったので、書いてみました。

次のような関数のグラフを描くことを目標にします(ただし、$-2 \leq x \leq 4$)。

f(x) = 

\left\{ \begin{array}{lll}
  \mathrm{e}^{-|x|}\cos(50x) & (x \geq 0) \\
  x + 1 & (-1 \leq x < 0) \\
  \cos(\frac{13\pi}{3 + x}) & (otherwise)
\end{array} \right.

import numpy as np

x = np.linspace(-2, 4, 1000)
def myFunc(x):
  if x >= 0:
    return np.exp(-np.abs(x)) * np.cos(50*x)
  elif -1 <= x and x < 0:
    return x + 1
  else:
    return np.cos(13*np.pi/(3+x))

npmyFunc = np.vectorize(myFunc)

fig, ax = plt.subplots()
ax.plot(x, npmyFunc(x))
plt.show()

関数をベクトル化することで、ndarrayを引数として受け取れるようになっています。

for文で関数に値をつっこみ、新しいリストに入れてプロットするやり方もできますが、こっちの方が簡潔に書けると思います。