pymc3でのモデル関数が条件分岐を含む場合の書き方

pymcで推定を行う際に、モデル関数がifを含む条件分岐がある場合についての書き方についてです。numpyとは少し書き方が違います。
teratailで回答した内容
pymc3でのモデル関数が条件分岐を含む場合の書き方
を再度編集したものです。
動作環境は、Win10、python3.5、Anaconda、Jupyterを利用。

今回の条件分岐関数
条件分岐を含むモデル関数
uの位置で、1次関数として立ち上がる関数です。

def model_line(x, u, nor, bg):
    ud = x - u 
    if ud <= 0:
        f =  bg
    else:
        f = nor*ud + bg
    return f

importモジュール
前の関数を実行する前に以下のモジュールをインポートしておいてください。

import pymc3 as pm
import numpy as np
import scipy.stats as stats
import matplotlib.pyplot as plt
# import seaborn as sns
plt.style.use('seaborn-darkgrid')
np.set_printoptions(precision=2)
%matplotlib inline

条件分岐を含む関数の作図

#条件分岐を含む関数の作図
xx = np.arange(3, 6.1, 0.1)

#引数はfloatで書く、np.vectorizeしています。
vec_model_line = np.vectorize(model_line)
y_data = vec_model_line(xx,4.5,1.0,1.0)

plt.plot(xx,y_data)
plt.show()

グラフ表示

ノイズを付加

#ノイズを付加
noise=np.random.normal(1, 0.1, len(xx))
yy_data=y_data+noise
plt.plot(xx,yy_data)
plt.show()

ノイズを付加したグラフ
4.5付近で立ち上がっています。

モデル関数の書き換え
事後分布を求める前に、model_line関数をpymc用に書き換えます。
pymc3のswich関数を用います。引数の1番目は条件、2番目は真の条件での実行、3番目は、偽の条件での実行を入力します。
参考：PyMC3を使って変化点検出のベイズ推論をする

def model_line2(x, u, nor, bg):

    ud = x - u 
    f1=bg
    f2=nor*ud + bg
    
    return pm.math.switch(ud<=0,f1,f2)

事後分布推定

# x, 
# u, gauss
# nor guass
# bg gauss
with pm.Model() as model:
    u= pm.Normal('u', mu=4.0, sd=1.0)
    nor = pm.Normal('nor', mu=1, sd=1)    
    bg = pm.Normal('bg', mu=1, sd=1)
    epsilon = pm.HalfCauchy('epsilon', 1)
    
    mu =model_line2(xx,u,nor,bg)

    y_pred = pm.Normal('y_pred', mu=mu, sd=epsilon, observed=yy_data)

    trace_model = pm.sample(2000, njobs=1)

pm.traceplot(trace_model)
plt.tight_layout()

plt.figure()

計算結果

最初に設定したu=4.5,nor=1.0,bg=1.0にたいして、ノイズを足したものと比較すると、大体良い一致をしています。bgについては、ノイズを足すときに平均値を1にしているため2に近い値になっています。