F(x)=2x (x<5)、F(x)=x**2 (5<x<30) みたいにF(x)が変化するような関数をプロットしたい

Question

xの条件によって、F(x)が変化する関数をグラフにプロットしたい。

import matplotlib.pyplot as plt
from scipy.integrate import quad
import numpy as np
import math

x = np.linspace(-100, 100, 10000)

def A(x):
    if x < 5.0:
        return 2*x 
    elif 5.0 < x < 30:
        return x**2
    elif 30 < x:
        return -3*x

fig = plt.figure()
ax = fig.add_subplot(1,1,1)
ax.grid()
ax.plot(x,A(x))
ax.set_xlim(-100, 100)
plt.show()

エラー
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-2-f89226c37b9c> in <module>
     17 ax = fig.add_subplot(1,1,1)
     18 ax.grid()
---> 19 ax.plot(x,A(x))
     20 ax.set_xlim(-100, 100)
     21 plt.show()

<ipython-input-2-f89226c37b9c> in A(x)
      7 
      8 def A(x):
----> 9     if x < 5.0:
     10         return 2*x
     11     elif 5.0 < x < 30:

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

Answer 1

3つの解決方法を示します。

内包表記を使う

内包表記を使ってループして下さい。

y = np.array([A(xx) for xx in x])
...
ax.plot(x, y)

条件分岐にnp.where()を使う

よく見ると条件分岐が完全ではない(x==5やx==30の場合にNoneが返る)ため完全に同一とは言えませんが、それを除けばnp.where()でif文を代用することができます。

def A(x):
    return np.where(x < 5.0, 2 * x, np.where(x < 30, x ** 2, -3 * x))

関数をユニバーサル関数化する

vA = np.frompyfunc(A, 1, 1)
...
ax.plot(x, vA(x))

もしくは

vA = np.vectorize(A)
...
ax.plot(x, vA(x))

参考URL

• numpyのarrayの要素をforで回さないで済むようにする方法チェックリスト
• Python高速化実験～map関数とか～

Answer 2

if x < 5.0: ここで、xがarrayなのでエラー出ているみたいですね。
下記のように一旦yをつくってみるのはどうでしょうか。（もう少しスマートなやり方がありそうな気もしますが）

import matplotlib.pyplot as plt
from scipy.integrate import quad
import numpy as np
import math

x = np.linspace(-100, 100, 10000)

def A(x):
    if x < 5.0:
        return 2*x 
    elif 5.0 < x < 30:
        return x**2
    elif 30 < x:
        return -3*x

y = [A(i) for i in x]
fig = plt.figure()
ax = fig.add_subplot(1,1,1)
ax.grid()
ax.plot(x,y)
ax.set_xlim(-100, 100)
plt.show()