emceeで関数フィッティングをする

この記事は CAMPHOR- Advent Calendar の17日目の記事です。

 

初めまして。今年からCAMPHOR-の運営メンバーをやっております、&esです。あんですとお読みください。

 

さて、本日はPythonのライブラリ、emceeを使った関数フィッティングを紹介したいと思います。

 

皆さん人生一度は、任意のデータに任意の関数をフィッティングしたいと思ったことがあるんじゃないでしょうか。そんな時に役立つのが、emceeです。

emcee以外にも関数フィッティングをやってくれるツールは色々あると思います。Excelでも簡単な関数はフィッティングしてくれますし、PythonだったらSciPyとかもやってくれますね。ただ、emceeはパラメータが多くなっても極値で止まってしまうことなく、ちゃんと更新していってくれるため、より良いフィッティングができる(はず)という利点があります。その代わりに計算コストが高いので時間がかかってしまいますが、普通に便利なので使ってみても良いと思います。

emceeとは

emceeはPythonマルコフ連鎖モンテカルロ法(以下MCMC)を扱うためのライブラリです。MCMCというのは、本来

求める確率分布を均衡分布として持つマルコフ連鎖を作成することによって確率分布のサンプリングを行う種々のアルゴリズムの総称

引用元:マルコフ連鎖モンテカルロ法 - Wikipedia

だそうですが、関数フィッティングもできてしまうんですね。しかもフィッティングの誤差まで求まってしまいます。どうしてできてしまうのかはこの記事では詳しく解説しません。私もよく知りません。むしろ教えてください。

emceeの公式ドキュメントにもFitting a model to dataというのがあるのでそれを見てもらえたら良いんですけど、この記事ではもっと簡単に動けばいいの精神で使ってみたいと思います。

emceeを使う

とりあえずemceeとNumPyをpypiなりcondaなりでインストールしておいてください。

適当なデータを用意する

まずはいい感じのデータを用意しましょう。モデルがいい感じの関数で、いい感じにばらついていて、いい感じに誤差が乗っていると良いですね。今回はNumPyを使って適当な感じで作ってみます。random.seedを同じにすれば皆さんも同じデータを作ることができます。

import numpy as np

np.random.seed(0)

N = 40
x = np.linspace(-3, 3, N)
y = 0.5 * x ** 3 + 2 * x ** 2 + 3 * np.sin(3 * x) + 0.5 * np.random.randn(N)

\displaystyle y=\frac{1}{2}x^3 + 2x^2 + 3\sin3x\ \ \ (-3\leq x\leq3)

の関数に正規分布の乱数を載せたものですね。ちなみにこんな感じです。

f:id:and_es:20211216232211p:plain

必要な関数などを定義する

以下の関数をフィッティングすることにしましょう。

ax^3+bx^2+c\sin dx

a, b, c, dがフィッティングパラメータですね。

まず対数尤度関数を定義します。MCMCによるフィッティングは基本的に対数尤度を最大化することで行います。

尤度ってなんぞ、という方のためにほんの少しだけ解説すると、簡単に言えば、ある関数(フィッティングする関数)からそのデータが生成される確率のことです。データがその関数の周りで、正規分布に従って生成されるとすれば、

\displaystyle \prod_i\frac{1}{\sqrt{2\pi\sigma^2}}\exp{\left(-\frac{(x_i-f(x_i))^2}{2\sigma^2}\right)}

となりますね。ただ、積の計算はコストが高いので、対数をとってあげます。

\displaystyle \sum_i\left(\ln\frac{1}{\sqrt{2\pi\sigma^2}} -\frac{(x_i-f(x_i))^2}{2\sigma^2}\right)

そうするといい感じに和の形になってくれます。さらに第1項はフィッティングする関数と無関係なので無視してあげれば、

\displaystyle -\frac{1}{2}\sum_i\left(\frac{x_i-f(x_i)}{\sigma}\right)^2

を最大化すれば良いということになりますね。

したがって、対数尤度関数は、

def lnlike(theta, x, y):
    sigma = 0.5
    a, b, c, d = theta
    model = a *x ** 3 + b * x ** 2 + c * np.sin(d * x)
    chi2 = - np.sum(((y - model) / sigma) ** 2) / 2
    return chi2

こんな感じに定義できます。

あとはこんな感じに必要なものを定義しておきます。

import emcee

init = np.array([2, 2, 2, 2])
pos = init + 0.5 * np.random.randn(32, 4)
nwalkers, ndim = pos.shape
sampler = emcee.EnsembleSampler(nwalkers, ndim, lnlike, args=(x, y))

initはパラメータの初期値です。理論値とかを入れると良いですが、今回は特に物理的なモデルでも何でもないので、適当にすべて2にしておきます。

posでどの値からパラメータを更新するか、と、パラメータの初期化と更新を何回繰り返すか、を定義しています。np.random.radn()の1つ目の値が初期化と更新を何回繰り返すかで、2つ目の値にはパラメータの個数を入れてください。今回は、initを中心に正規分布で値をばらけさせています。違う値からパラメータを更新させることで、より大局的に尤度が最大となるパラメータを探すことができます。

nwalkersとndimはそれぞれパラメータの初期化と更新を繰り返す回数と、パラメータの個数ですね。posのshapeなので当たり前ですが。

samplerでemceeのサンプラーを定義します。3つ目に対数尤度関数を、4つ目に対数尤度関数に渡すデータを与えています。

MCMCでフィッティングする

あとはMCMCを走らせるだけです。

sampler.run_mcmc(pos, 10000, progress=True)

こんな感じですね。2つ目の引数はパラメータの更新を行う回数です。今回は10000回更新してみます。あまり少なすぎるとemceeに怒られますが、あまり多すぎても時間がかかってしょうがないのでいい感じに変えてみてください。

progress=Trueとすることで進捗を表示することができます。tqdmがないと怒られたような気がしますね。

フィッティングパラメータの値を取り出す

flat_samples = sampler.get_chain(discard=100, thin=15, flat=True)
theta = [np.percentile(flat_samples[:, i], [16, 50, 84])[0] for i in range(flat_samples.shape[1])]

こんな感じで値を取り出せます。

ちなみにこんな感じになりました。

print(theta)
[0.47993708089393267, 2.0188253746545453, 2.77740977885968, 2.960187497302918]

まあいいんじゃないでしょうか。グラフにするとこんな感じです。

f:id:and_es:20211217015904p:plain

割ときれいに合ってますね。

あとがき

パラメータに制限をかけていなかったり、プログラムがあまりにも適当だったり、正直言ってこの記事はガバガバです。ただ、emceeを使って複雑な関数のフィッティングができるんだということを知っていただければ幸いです。ぜひ公式ドキュメントなどを読んで試してみてください。それではまたっ。