PythonでUnscented Kalman Filter (UKF)

以下の書籍の例をPythonで試しました.

【関連記事】
Pythonで拡張カルマンフィルタ - Notes_JP

Unscented Kalman Filter (UKF)

以下が参考になるかも.
UKF (Unscented Kalman Filter)って何?
UKF(Unscented Kalman Filter)とKFはどう違うの?(アイ・サイ問答教室)

numpyによる実装

def mapcols(dimf, f, x):
    """
    ベクトル値関数を列に作用させた結果を返す
    x = (x1, x2, ...) -> res = (f(x1), f(x2), ...)
    """
    res = np.empty((dimf, x.shape[1]))
    for j in range(x.shape[1]):
        res[:, j] = f(x[:, j].reshape(x.shape[0], 1)).reshape(res[:, j].shape)
    return res

def ut(dimf, f, xm, Pxx):
    """
    Unscented transform
    """
    n = xm.shape[0]
    kappa = 3 - n
    w0 = kappa / (n + kappa)
    wi = 1 / (2 * (n + kappa))
    wT = np.array([[w0] + [wi for _ in range(2 * n)]]).reshape(1, -1)
    W = np.diag([w0] + [wi for _ in range(2 * n)])
    L = np.linalg.cholesky(Pxx)
    # --- X ---
    X = xm.T.reshape((1, n))
    a = np.ones((n, 1)) @ xm.T.reshape((1, n))
    b = np.sqrt(n + kappa) * L.T
    X = np.concatenate([X, a + b], axis=0)
    X = np.concatenate([X, a - b], axis=0)
    Xd = X - xm.T.reshape((1, n))
    # --- Y ---
    Y = mapcols(dimf, f, X.T).T
    ym = (wT @ Y).reshape((-1, 1))
    Yd = Y - ym.T.reshape((1, dimf))
    # --- cov ---
    Pyy = Yd.T @ W @ Yd
    Pxy = Xd.T @ W @ Yd
    return ym, Pyy, Pxy

def ukf(f, h, B, Q, R, y, xhat, P):
    """
    x(k+1) = f(x(k)) + Bv(k)
    y(k) = h(x(k)) + w(k)
    """
    xhatm, Pm, _ = ut(len(xhat), f, xhat, P)
    Pm += Q * (B @ B.T)
    yhatm, Pyy, Pxy = ut(1, h, xhatm, Pm)
    G = Pxy / (Pyy + R)
    xhat_new = xhatm + G * (y - yhatm)
    P_new = Pm - G @ Pxy.T
    return xhat_new, P_new, G

def calc_timeseries_ukf(N, f, h, B, Q, R, y, x0, P_ukf):
    xhat_ukf = [x0]
    G0 = np.empty(x0.shape)
    G0[:] = np.nan
    G_ukf = [G0]
    for k in range(1, N):
        xhat_new, P_new, G = ukf(f, h, B, Q, R, y[k], xhat_ukf[-1], P_ukf)
        xhat_ukf.append(xhat_new)
        G_ukf.append(G)
        P_ukf = P_new
    xhat_ukf = np.array(xhat_ukf)
    G_ukf= np.array(G_ukf)
    return xhat_ukf, G_ukf

物体の落下運動(拡張カルマンフィルタとの比較)

拡張カルマンフィルタ(Pythonで拡張カルマンフィルタ - Notes_JP)と結果を比較します.
EKFとUKFの比較
EKFとUKFの比較

# --- パラメータ ---
rho = 1.23
g = 9.81
eta = 6e3
M = 3e4
a = 3e4

T = 0.5
EndTime = 30
N = int(EndTime / T + 1)
time = T * np.arange(0, N)
n = 3
R = 4e3
Q = 0
B = np.array([
    [0],
    [0],
    [0]
])

# --- 関数 ---
def f(x):
    x1, x2, x3 = x
    x1, x2, x3 = x1.item(0), x2.item(0), x3.item(0)
    res = np.array([
        [x1 + T * x2],
        [x2 + T * (0.5 * rho * np.exp(- x1 / eta) * (x2 ** 2) * x3 - g)],
        [x3]
    ])
    return res

def h(x):
    x1, x2, x3 = x
    x1, x2, x3 = x1.item(0), x2.item(0), x3.item(0)
    # res = np.array([
    #     [np.sqrt((M ** 2) + ((x1 - a) ** 2))]
    # ])
    res = np.sqrt((M ** 2) + ((x1 - a) ** 2))
    return res

def dfdx(x):
    x1, x2, x3 = x
    x1, x2, x3 = x1.item(0), x2.item(0), x3.item(0)
    res = np.array([
        [1, T, 0],
        [- T * 0.5 * (rho / eta) * np.exp(- x1 / eta) * (x2 ** 2) * x3,
         1 + T * rho * np.exp(- x1 / eta) * x2 * x3,
         T * 0.5 * rho * np.exp(- x1 / eta) * (x2 ** 2)
        ],
        [0, 0, 1]
    ])
    return res

def dhdx(x):
    x1, x2, x3 = x
    x1, x2, x3 = x1.item(0), x2.item(0), x3.item(0)
    res = np.array([
        [(x1 - a) / h(x), 0, 0]
    ])
    return res

# --- 真値, 測定値 ---
w = np.random.normal(loc = 0.0, scale = np.sqrt(R), size = N)
x = [np.array([
    [9e4], 
    [-6e3],
    [3e-3]
])]
y = [h(x[-1]) + w[0]]

for k in range(N - 1):
    x.append(f(x[-1]))
    y.append(h(x[-1]) + w[k + 1]) 
x = np.array(x)
y = np.array(y)

# --- EKF ---
P_ekf = np.array([
    [9e3, 0, 0],
    [0, 4e5, 0],
    [0, 0, 0.4]
])

xhat_ekf, G_ekf = calc_timeseries_ekf(N, f, h, dfdx, B, dhdx, Q, R, y, x[0], P_ekf)

# --- UKF ---
P_ukf = np.array([
    [9e3, 0, 0],
    [0, 4e5, 0],
    [0, 0, 0.4]
])

xhat_ukf, G_ukf = calc_timeseries_ukf(N, f, h, B, Q, R, y, x[0], P_ukf)

# --- プロット ---
fig_num = 3
fig = plt.figure(figsize = (10, 3 * fig_num), tight_layout = True)
axes = [fig.add_subplot(fig_num, 1, i + 1) for i in range(fig_num)]
axes2 = [ax.twinx() for ax in axes]
ylim = [(0, 1e5), (-1e4, 0), (-0.5, 0.5)]
scilimits = [4, 3, 0]
for i in range(fig_num):
    ax, ax2 = axes[i], axes2[i]
    ax.plot(time, x[:, i], label = 'true', ls = 'dotted')
    ax.plot(time, xhat_ekf[:, i], label = 'estimated (EKF)', ls = 'dashed', c = 'tab:orange')
    ax.plot(time, xhat_ukf[:, i], label = 'estimated (UKF)', ls = 'solid', c = 'tab:orange')
    ax.set_ylim(ylim[i])
    ax2.plot(time, G_ekf[:, i], label = 'gain (EKF)', ls = 'dashed', c = 'tab:gray')
    ax2.plot(time, G_ukf[:, i], label = 'gain (UKF)', ls = 'solid', c = 'tab:gray')
    if scilimits[i] != 0:
        ax.yaxis.set_major_formatter(ptick.ScalarFormatter(useMathText = True))
        ax.ticklabel_format(style="sci", axis="y", scilimits=(scilimits[i], scilimits[i]))
    ax.set_xlabel('time (s)')
    ax.set_ylabel(f'x{i + 1}')
    ax2.set_ylabel(f'g{i + 1}')
    lines, labels = ax.get_legend_handles_labels()
    li2, la2 = ax2.get_legend_handles_labels()
    lines += li2
    labels += la2
    ax2.legend(lines, labels, loc = 'lower right')
plt.show()    

計算メモ

上の実装の数式部分です.
\begin{aligned}
f:\mathbb{R}^{n} \to \mathbb{R}^{n}
\end{aligned}

$n$次元確率変数$\bm{x}$の平均値を$\bm{\mu}_{x}$,共分散行列を$P_{xx}$とする.$L=(\bm{l}_{1}, \bm{l}_{2}, \ldots, \bm{l}_{n})$を,$P_{xx} = L {}^{t}\!L$を満たす行列とする.

\begin{aligned}
X &=
\begin{pmatrix}
{}^{t}\!\bm{\mathcal{X}}_{0} \\
{}^{t}\!\bm{\mathcal{X}}_{1} \\
\vdots \\
\end{pmatrix} \\
\bm{\mathcal{X}}_{i} &=
\begin{cases}
\, \bm{\mu}_{x} & (i = 0) \\
\, \bm{\mu}_{x} + \sqrt{n + \kappa} \, \bm{l}_{i} & (i = 1,\ldots,n) \\
\, \bm{\mu}_{x} - \sqrt{n + \kappa} \, \bm{l}_{i} & (i = n + 1,\ldots,2n)
\end{cases}
\end{aligned}

$F(\bm{a}_{1},\bm{a}_{2},\ldots) = (f(\bm{a}_{1}),f(\bm{a}_{2}),\ldots)$として

\begin{aligned}
Y
&=
\begin{pmatrix}
{}^{t}\!\bm{\mathcal{Y}}_{0} \\
{}^{t}\!\bm{\mathcal{Y}}_{1} \\
\vdots \\
\end{pmatrix}
=
\begin{pmatrix}
{}^{t}\!f(\bm{\mathcal{X}}_{0}) \\
{}^{t}\!f(\bm{\mathcal{X}}_{1}) \\
\vdots \\
\end{pmatrix} \\
&= {}^{t}\![F({}^{t}\!X)]
\end{aligned}
とする.

\begin{aligned}
\bm{w}
&= {}^{t}\!(w_{0}, w_{1}, \ldots ,w_{2n}) \\
w_{i}
&=
\begin{cases}
\, \frac{\kappa}{n + \kappa} & (i = 0) \\
\, \frac{\kappa}{2(n + \kappa)} & (i = 1,\ldots 2n)
\end{cases}
\end{aligned}

\begin{aligned}
\bm{\mu}_{y}
&= \sum_{k = 0}^{2n} w_{k} \bm{\mathcal{Y}}_{k} \\
&= \Biggl(\sum_{k = 0}^{2n} w_{k} (\bm{\mathcal{Y}}_{k})_{i}\Biggr)_{i}
= \Biggl(\sum_{k = 0}^{2n} w_{k} (Y)_{ki} \Biggr)_{i} \\
&= {}^{t}\! \bm{w} Y
\end{aligned}

\begin{aligned}
P_{yy}
&= \sum_{k = 0}^{2n}
w_{k}
(\bm{\mathcal{Y}}_{k} - \bm{\mu}_{y})
{}^{t}\!(\bm{\mathcal{Y}}_{k} - \bm{\mu}_{y}) \\
&=\sum_{k = 0}^{2n} w_{k} ({}^{t}\!Y_{d})_{k} (Y_{d})_{k} \\
&={}^{t}\!Y_{d} W Y_{d}
\end{aligned}

\begin{aligned}
X_{d} &=
\begin{pmatrix}
{}^{t}\!(\bm{\mathcal{X}}_{0} - \bm{\mu}_{x}) \\
{}^{t}\!(\bm{\mathcal{X}}_{1} - \bm{\mu}_{x}) \\
\vdots \\
\end{pmatrix}
\end{aligned}
\begin{aligned}
Y_{d} &=
\begin{pmatrix}
{}^{t}\!(\bm{\mathcal{Y}}_{0} - \bm{\mu}_{y}) \\
{}^{t}\!(\bm{\mathcal{Y}}_{1} - \bm{\mu}_{y}) \\
\vdots \\
\end{pmatrix}
\end{aligned}

\begin{aligned}
W
&=
\begin{pmatrix}
w_{0} & & & &\\
& w_{1}& & &\\
& & \ddots & &\\
& & & & w_{2n}
\end{pmatrix}
\end{aligned}