Matplotlib | GUIで2領域を選択してフィッティング

2022年12月22日2023年1月10日

この記事で分かること

Matplotlibを用いてGUI上でフィッティングを実施。
選択範囲を選ぶと結果をリアルタイムにプロットできます。
key_press_eventを使って、２つの領域（ROI）を選択してフィッティング。

はじめに

以前、下記の記事でmatplotlibを用いてマウス操作でフィット範囲（ROI）を選ぶpythonスクリプトを作成しました。

今回はROIを２か所使うように改造したパターンをご紹介します。バックグラウンド処理のようにデータ全体ではなくピークの前後にフィット範囲を取りたいときなどに有用です。

私の調べた限りmatplotlibは複数のROIをセット出来ないようなので（可能なら教えて下さい！）、fig.canvas.mpl_connectを利用して二か所を選択します。

コード例

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.widgets import SpanSelector
from lmfit import Model

def gaussian(x, amp, cen, wid):
    return (amp / (np.sqrt(2*np.pi) * wid)) * np.exp(-(x-cen)**2 / (2*wid**2))

def fit_gaussian(x, data):
    mod = Model(gaussian) 
    pars = mod.make_params(amp=1, cen=1, wid=1)
    result = mod.fit(data, pars, x=x)
    return result

class Fit:
    def __init__(self, ax, df):
        self.ax = ax or plt.gca()
        self.df = df
        self.flag = 1
        self.xx = [0,0,0,0]

    def fitting(self):
        x0, x1, x2, x3 = self.xx

        df_ROI = self.df.query(f'{x0}<x<{x1} or {x2}<x<{x3}')
        xx,yy = df_ROI.x, df_ROI.y

        result = fit_gaussian(xx, yy)
        print("-------------Fitting result------------------")
        pars =  list(result.best_values.values())

        fitted.set_data(self.df.x, gaussian(self.df.x, *pars) )
        residual.set_data(self.df.x, self.df.y - gaussian(self.df.x, *pars) )
        
        ROI0.set_data(self.df.query(f'{x0}<x<{x1}').x, self.df.query(f'{x0}<x<{x1}').y )
        ROI1.set_data(self.df.query(f'{x2}<x<{x3}').x, self.df.query(f'{x2}<x<{x3}').y )

        fig.canvas.draw()
        fig.canvas.flush_events()   

    def select_callback(self, aa, bb):
        if(self.flag==1):
            self.xx[0], self.xx[1] = aa, bb
            
        elif(self.flag==-1):
            self.xx[2], self.xx[3] = aa, bb

        self.fitting()

    def onclick(self, event):
        self.flag *= -1
        print("flag:", self.flag)

# デモデータ
data = np.random.randn(10000)
histo,bins = np.histogram(data,range=(-5,5),bins=100,density=True)
x=bins[1:]
df = pd.DataFrame(data=np.stack([x,histo]).T, columns=['x', 'y'])

# プロット
fig, ax = plt.subplots()
ax.plot(df.x, df.y, ".", c="gray", label="data")

# あとで更新する
fitted, = ax.plot([],[], c="orange", label = "fitted")
residual, = ax.plot([], [], drawstyle = "steps-post", label = "residual")
ROI0, = ax.plot([],[], c="b", linewidth = 0, marker =".", label = "ROI0")
ROI1, = ax.plot([],[], c="g", linewidth = 0,marker =".",  label = "ROI1")
ax.legend()
fit = Fit(ax, df)

fig.canvas.mpl_connect('key_press_event', fit.onclick)

# x軸を選ぶ
span = SpanSelector(
    ax,
    fit.select_callback,
    "horizontal",
    useblit=True,
    props=dict(alpha=0.2, facecolor="tab:blue"),
    interactive=True,
    drag_from_anywhere=True
)

plt.show()

使用方法概略

マウスでフィットROIを設定
ROI2つめ切り替えはキーを何かヒット！

fitted, = ax.plot([],[], c="orange", label = "fitted")
residual, = ax.plot([], [], drawstyle = "steps-post", label = "residual")
ROI0, = ax.plot([],[], c="b", linewidth = 0, marker =".", label = "ROI0")
ROI1, = ax.plot([],[], c="g", linewidth = 0,marker =".",  label = "ROI1")
ax.legend()
fit = Fit(ax, df)

fig.canvas.mpl_connect('key_press_event', fit.onclick)

最後の行にeventとfigを結びつけるコードが入っています。
key_press_eventなので、キーを何か押せば対応してonclick関数が呼ばれます。
fitクラスの中で定義されているonclickは呼ばれるたびにflagを1か-1にスイッチします。

    def onclick(self, event):
        self.flag *= -1
        print("flag:", self.flag)

マウス操作のcallback関数で、そのflagの値を参照して1と-1の場合で選んだ範囲を分けて覚えておきます。

    def select_callback(self, aa, bb):
        if(self.flag==1):
            self.xx[0], self.xx[1] = aa, bb
            
        elif(self.flag==-1):
            self.xx[2], self.xx[3] = aa, bb

選んだ範囲はリストself.xxに入れて、その値を参照してself.df.query()でフィットする範囲だけを持つデータdf_ROIを作ります。フィットする流れは前回と同じですね。

  def fitting(self):
        x0, x1, x2, x3 = self.xx

        df_ROI = self.df.query(f'{x0}<x<{x1} or {x2}<x<{x3}')
        xx,yy = df_ROI.x, df_ROI.y

        result = fit_gaussian(xx, yy)
        print("-------------Fitting result------------------")
        pars =  list(result.best_values.values())

        fitted.set_data(self.df.x, gaussian(self.df.x, *pars) )
        residual.set_data(self.df.x, self.df.y - gaussian(self.df.x, *pars) )
        
        ROI0.set_data(self.df.query(f'{x0}<x<{x1}').x, self.df.query(f'{x0}<x<{x1}').y )
        ROI1.set_data(self.df.query(f'{x2}<x<{x3}').x, self.df.query(f'{x2}<x<{x3}').y )