Kumpei Shiraishi's blog

物理、プログラミング、クライミングに関する雑記

matplotlibでグラフの場合分け

やりたいこと

matplotlibで、場合分けしたグラフを描きたい。

この記事では参考ページに倣って、

{
f(x) = \left\{ \begin{array}{ll}
x^2 & (0 \leq x \leq 1)\\
2-x & (1 < x \leq 2)\\
x^2 - 3x +2 & (2 < x \leq 3)\\
\end{array} \right.
}

を描く。

ダメな例

import numpy as np
import matplotlib.pyplot as plt

def f(x):
    if 0<=x and x<=1:
        return x**2
    elif 1<x and x<=2:
        return 2-x
    elif 2<x and x<=3:
        return x**2-3*x+2

x = np.arange(0,3,0.001)
y = f(x)
plt.plot(x,y)

plt.show()

これを実行すると

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

というエラーになる。

関数f(x)の定義部分でxの大きさを判定しているのに、x = np.arange(0,3,0.001)というnumpyの配列を渡しているのが原因らしい。

上手くいく例

import numpy as np
import matplotlib.pyplot as plt

def f1(x):
    return x**2
def f2(x):
    return 2-x
def f3(x):
    return x**2-3*x+2

x1 = np.arange(0,1,0.001)
y1 = f1(x1)
plt.plot(x1,y1,"b")

x2 = np.arange(1,2,0.001)
y2 = f2(x2)
plt.plot(x2,y2,"b")

x3 = np.arange(2,3,0.001)
y3 = f3(x3)
plt.plot(x3,y3,"b")

plt.show()

明示的に別々の関数を定義すると上手くいく。複数の関数をmatplotlibで描画するには、単純にplotを複数回呼べばよい。グラフの色が別々だと格好悪いので、"b"を渡して全部青色にした。

そうして得られたグラフは以下。

Case in matplotlib

参考