python/numpy.piecewise 場合分けのある関数を扱う

pythonで条件分岐のある関数を扱ったとき詰まったのでメモ。
kumpeishiraishi.hatenablog.com

↑のブログで用いられている関数を拝借します。

これを例えば図示したいとき、

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()
とのエラーが出ます。引数がnumpy.arrayなのがダメみたいです。
定義域ごとに別々に関数を定義してもいいんですが、その関数を使って最少二乗フィッティングやろうと思ったときなど、のちのち困ることになります。

解決策として一番短くかけるのは、numpy.piecewiseという場合分け

import numpy as np
import matplotlib.pyplot as plt
def multi_func(x):
    condlist = [x <=1 , (x >1) & (x <= 2), x >2]
    funclist = [lambda a: a**2, lambda a: 2-a,lambda a: a**2-3*a+2]
    return np.piecewise(x, condlist, funclist)

x=np.linspace(0,3,100)
y=multi_func(x)
plt.plot(x,y)
plt.show()

(x >1) & (x <= 2)のところでカッコを忘れがちなので要注意だ！