Lorentz方程式をETKF (Ensemble Transform Kalman Filter) でデータ同化

今回はETKFでデータ同化します

EKF (Extend Kalman Filter)でやったのが下記。わりと見てもらえたので，ETKF版も作りました。

コード

やってることが前回とほぼおなじなので，コードだけ。mainのなかでETKFの計算をしています。ローカライゼーションは適用していません。

固有値求めるところで，

Symmetric(I - dY'*inv(dY*dY' + (m-1)*R)*dY)

としていますが，こうしないと複素数の固有値が出ることがあったので，対称行列であることを担保するためにやっています。

using PyPlot
using LinearAlgebra
using HDF5
using Random

function Lorentz63(x::Vector{Float64})
    σ::Float64 = 10
    ρ::Float64 = 28
    β::Float64 = 8/3
    
    Lorentz63 = zeros(3)
    Lorentz63[1] = σ*(x[2]-x[1])
    Lorentz63[2] = x[1]*(ρ - x[3]) - x[2]
    Lorentz63[3] = x[1]*x[2] - β*x[3]
    
    return Lorentz63
end

function RK4(f, dt, x)
    k1 = f(x)
    k2 = f(x + k1*dt*0.5)
    k3 = f(x + k2*dt*0.5)
    k4 = f(x + k3*dt)
    x .= x .+ dt .*(k1 .+ 2k2 .+ 2k3 .+ k4) ./6.
end

function makeobs(y, xtrue, dstep, nmax)
    Random.seed!(10)
    for i=1:nmax
        if mod(i, dstep) == 0
            y[:,i] .= xtrue[:,i] .+ randn()
        end
    end
    
end

function main(m, α)
    nmax = 10000
    dstep = 10
    dt = 0.01
    
    x0 = zeros(3, nmax+1)
    x1 = zeros(3, nmax+1)
    x0[:,1] = [1.0, 0., 0.]
    x1[:,1] = [1.00001, 0., 0.]
    t = zeros(nmax+1)
    
    for i=1:nmax
        x0[:, i+1] .= RK4(Lorentz63, dt, x0[:,i]) # true
        x1[:, i+1] .= RK4(Lorentz63, dt, x1[:,i]) # solution with error
        t[i+1] = dt*i
    end
    #println(x)
    
    y = zeros(3, nmax +1)     
    makeobs(y, x0, dstep, nmax)
    #println(y)
    
    p = 3
    N = 3
    #m = 15 # member number
    Xa = zeros(Float64, N, m)
    Xf = zeros(Float64, N, m)
    xfm = zeros(Float64, N, nmax+1)
    xam = zeros(Float64, N, nmax+1)
    dXf = zeros(Float64, N, m)
    dXa = zeros(Float64, N, m)
    dY = zeros(Float64, p, m)
    
    R = Matrix{Float64}(I, p, p)
    K = zeros(Float64, N, N)
    H = Matrix{Float64}(I, p, N)
    #α = 1.3
    E = Matrix{Float64}(I, N, N)
    
   
    # initialize Xa
    for i = 1:m
        Xa[:,i] = [1.0 + 0.1randn(), 0., 0.] 
        
    end
    println(Xa[:,1:div(m,2)])
    xam[:, 1] = mean(Xa, dims=2)
    println("xam_1=", xam[:,1])
    
    for i=1:nmax
        # forecast
        for l=1:m
            Xf[:, l] .= RK4(Lorentz63, dt, Xa[:,l])
        end
        
        xfm[:, i+1] = mean(Xf, dims=2)
        for l=1:m
            dXf[:,l] = Xf[:,l] - xfm[:,i+1]
        end
        
        if mod(i+1, dstep) == 0
            #inflation 
            dXf = α * dXf
            dY = H*dXf
            K = dXf*dY'*inv(dY*dY' + (m-1)*R)
            xam[:,i+1] = xfm[:,i+1] + K*(y[:,i+1] - H*xfm[:,i+1])
            
            val, vec = eigen( Symmetric(I - dY'*inv(dY*dY' + (m-1)*R)*dY) )
            #println(val)
            T = vec*sqrt( Diagonal(val) ) * vec'
            #println("T=", size(T))
            
            dXa = dXf*T
            
            for l = 1:m
                Xa[:,l] = xam[:,i+1] + dXa[:,l]
            end
            
        else
            xam[:,i+1] .= xfm[:,i+1]
            Xa[:,:] .= Xf[:,:]
        end
        #Xa[:,i+1] = Xf[:,i+1]
        #Pa[:,:, i+1] = Pa[:,:, i]
    end
    
    plot_start=1000
    plot_end=9500
    fig = plt.figure(figsize=(9, 7.))
    ax1 = fig.add_subplot(211)
    ax1.plot(t[plot_start:plot_end], x0[1,plot_start:plot_end], 
        label="x1: true")
    ax1.plot(t[plot_start:plot_end], xam[1,plot_start:plot_end], 
        label="x1: DA",)
    ax1.plot(t[plot_start:plot_end], x1[1,plot_start:plot_end], 
        label="x1: no DA")
    ax2 = fig.add_subplot(212)
    ax2.plot(t[plot_start:plot_end], x0[3,plot_start:plot_end], 
        label="x3: true")
    ax2.plot(t[plot_start:plot_end], xam[3,plot_start:plot_end], 
        label="x3: DA",)
    ax2.plot(t[plot_start:plot_end], x1[3,plot_start:plot_end], 
        label="x3: no DA")
    
    ax1.set_xlabel("time")
    ax1.set_ylabel("x1")
    ax1.legend()
    ax2.set_xlabel("time")
    ax2.set_ylabel("x3")
    ax2.legend()
    plt.savefig("true_ETKF.jpeg")
    
    
    rms_a = zeros(nmax+1)
    rms_o = zeros(nmax+1)
    rms_x1 = zeros(nmax+1)
    for i=1:nmax
        rms_a[i] = norm(xam[:,i]-x0[:,i])/sqrt(N)
        rms_o[i] = norm(y[:,i]-x0[:,i])/sqrt(N)
        rms_x1[i] = norm(x1[:,i]-x0[:,i])/sqrt(N)
    end
    
    fig = plt.figure()
    ax1 = fig.add_subplot(111)
    ax1.set_ylim(-0.5,5)
    ax1.plot(t[plot_start:dstep:plot_end], rms_o[plot_start:dstep:plot_end], 
        label="RMSE obs.", "o", markersize=2.5)
    ax1.plot(t[plot_start:plot_end], rms_a[plot_start:plot_end], 
        label="RMSE ana.")
    ax1.plot(t[plot_start:plot_end], rms_x1[plot_start:plot_end], 
        label="RMSE sol. w/ error")
  
    ax1.legend()
    plt.savefig("RMSE_ETKF_m="*string(m)*"_alpha="*string(α)*".jpeg")
    
    ax = plt.figure().add_subplot(projection="3d")
    ax.plot(x0[1,plot_start:plot_end], x0[2,plot_start:plot_end], x0[3,plot_start:plot_end], )
    ax.plot(xam[1,plot_start:plot_end], xam[2,plot_start:plot_end], xam[3,plot_start:plot_end], )

    return x0, x1, xam, t
end

実行は下記。アンサンブルメンバー数mとインフレーションファクターのαを引数で与えるようにしました。

@time main(10, 1.0)

m=10で0.3sくらいです。無駄にあたらしく変数を生成したりしますが，そこそこ速いです。