ちょいめも

物理/Python/Cの雑記帳

フィッティング2次元

2次元フィッティングのメモ

"""
# 第一引数:2次元分布データ
# 第二引数:xデータ
# 第三引数:yデータ
"""
import os
import sys

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import Ellipse
from scipy.optimize import curve_fit

# 定義
C393 = 1.000  # 確率長円 39%
C500 = 1.177  # 確率長円 50%
C900 = 2.146  # 確率長円 90%
C950 = 2.448  # 確率長円 95%
C990 = 3.035  # 確率長円 99%


def fit_2d_normal_distribution(XY, a, rho, sigma_x, sigma_y, mu_x, mu_y, c):
    x, y = XY[0:2]
    z = (
        a
        * (1 / (2 * np.pi * sigma_x * sigma_y * np.sqrt(1 - rho**2)))
        * np.exp(
            -1
            / (2 * (1 - rho**2))
            * (
                ((x - mu_x) / sigma_x) ** 2
                + ((y - mu_y) / sigma_y) ** 2
                - 2 * rho * ((x - mu_x) * (y - mu_y)) / (sigma_x * sigma_y)
            )
        )
        + c
    )
    return z.ravel()  # 1次元に変換する


argvs = sys.argv
argc = len(argvs)

if argc != 4:
    print("arg1 : data\n")
    print("arg2 : x\n")
    print("arg3 : y\n")
    input("Press any key to exit\n")
    sys.exit()

# データ読み込み
data = np.loadtxt(argvs[1], delimiter=",")
data_x = np.loadtxt(argvs[2], delimiter=",")
data_y = np.loadtxt(argvs[3], delimiter=",")
X, Y = np.meshgrid(data_x, data_y)

# initial guesses
parameter_initial = (1000, 0.5, 8, 6, 0.5, -0.5, 0)
param, cov = curve_fit(fit_2d_normal_distribution, (X, Y), data.ravel(), p0=parameter_initial)
print(f"a       : {param[0]}")
print(f"rho     : {param[1]}")
print(f"sigma_x : {param[2]}")
print(f"sigma_y : {param[3]}")
print(f"mu_x    : {param[4]}")
print(f"mu_y    : {param[5]}")
print(f"c       : {param[6]}")

a = param[0]
rho = param[1]
sigma_x = param[2]
sigma_y = param[3]
mu_x = param[4]
mu_y = param[5]
c = param[6]

# 対角化後の分散共分散行列のsigma_u, siguma_v計算
t1 = sigma_x**2 + sigma_y**2
t2 = sigma_x**2 - sigma_y**2
t3 = 4 * (rho * sigma_x * sigma_y) ** 2  # 4*sigma_xy^2 = 4*(rho*sigma_x*sigma_y)^2
sigma_u = np.sqrt((t1 + np.sqrt(t2**2 + t3)) / 2)
sigma_v = np.sqrt((t1 - np.sqrt(t2**2 + t3)) / 2)
deg = np.rad2deg(np.arctan((sigma_u**2 - sigma_x**2) / np.sqrt(t3)))
sigma_u95 = 2 * C950 * sigma_u
sigma_v95 = 2 * C950 * sigma_v
print("・楕円長")
print(f"sigma_u(95%) : {sigma_u95}")
print(f"sigma_v(95%) : {sigma_v95}")
print(f"theta(deg)   : {deg}")

# fitting result
fitting_data = fit_2d_normal_distribution(
    (X, Y), param[0], param[1], param[2], param[3], param[4], param[5], param[6]
).reshape(data.shape[0], data.shape[1])

# plot graph
fig, axes = plt.subplots(nrows=2, ncols=2, figsize=(10, 8), sharex=False)
im = axes[0, 0].pcolormesh(X, Y, data, shading="auto", cmap="jet")
axes[0, 0].set_title("input data")
axes[0, 0].set_xlabel("x")
axes[0, 0].set_ylabel("y")
fig.colorbar(im, ax=axes[0, 0])

im = axes[0, 1].pcolormesh(X, Y, fitting_data, shading="auto", cmap="jet")
axes[0, 1].set_title("fit data")
axes[0, 1].set_xlabel("x")
axes[0, 1].set_ylabel("y")
fig.colorbar(im, ax=axes[0, 1])

im = axes[1, 0].pcolormesh(X, Y, data - fitting_data, shading="auto", cmap="jet")
axes[1, 0].set_title("diff(input - fit)")
axes[1, 0].set_xlabel("x")
axes[1, 0].set_ylabel("y")
fig.colorbar(im, ax=axes[1, 0])

# pcolormeshは図形描画ができないのでimshowで描画
im = axes[1, 1].imshow(data, extent=(data_x[0], data_x[-1], data_y[0], data_y[-1]), cmap="Greys")  # xy座標設定
axes[1, 1].set_title("check sigma (vs input)")
axes[1, 1].set_xlabel("x")
axes[1, 1].set_ylabel("y")
axes[1, 1].add_patch(
    Ellipse(
        xy=(mu_x, mu_y),
        width=sigma_u95,
        height=sigma_v95,
        angle=-deg,
        edgecolor="red",
        fill=False,
        label=f"2sigma(95%)\n sigma_u={sigma_u95:.1f}\n sigma_v={sigma_v95:.1f}\ntheta={deg:.1f}deg",
    )
)  # add_patchinvert_yaxisで反転しないのでangleにマイナスを掛ける
axes[1, 1].invert_yaxis()  # imshowは反転するのでここで反転
fig.colorbar(im, ax=axes[1, 1])
axes[1, 1].legend()

fig.tight_layout()
plt.savefig(f"{os.path.splitext(argvs[0])[0]}_fit.png")
plt.show()