LoginSignup
3
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

More than 3 years have passed since last update.

@ishigaki

Lorenz96モデルをJuliaとPythonで解く

Last updated at Posted at 2020-10-03

はじめに

データ同化を研究に使いたくて色々勉強してます。
データ同化のベンチマークでよく使われるLorenz96モデル https://journals.ametsoc.org/jas/article/55/3/399/24564/Optimal-Sites-for-Supplementary-Weather
を解くコードをJuliaで書いてました。これが遅いのか早いのかはっきりしないので，とりあえずPythonで書き直して比較しました。

Lorenz96モデル

方程式

$$
\frac{dX_j}{dt} = (X_{j+1} - X_{j-2})X_{j-1} - X_j +F
$$

ただし，$j= 1\dots N$。今回はN=40。

初期に$X_j=F$とすると，定常解が得られます。今回は$X_{20}=1.001F, ; F=8$として解きます。F=8だとカオス的な挙動が見られます。

これを4段4次精度のルンゲクッタで解きます。

Julia版

Juliaのバージョンは1.5.1です。Jupyter labで実行しました。

using LinearAlgebra

## パラメータ
## Fの大きさで挙動が変わる
F = 8
N = 40
dt = 0.01
tend = 146
nstep = Int32(tend/dt)

# dx/dt=f(x) 
# Lorentz96の方程式
function f(x)
    f = fill(0.0, N)
    for i=3:N-1
        f[i] = (x[i+1]-x[i-2])x[i-1] - x[i] + F
    end
    
    # 周期境界
    f[1] = (x[2]-x[N-1])x[N] - x[1] + F
    f[2] = (x[3]-x[N])x[1] - x[2] + F
    f[N] = (x[1]-x[N-2])x[N-1] - x[N] + F
    
    return f
end

# L96をRK4で解く
function RK4(xold, dt)
    k1 = f(xold)
    k2 = f(xold + k1*dt/2.)
    k3 = f(xold + k2*dt/2.)
    k4 = f(xold + k3*dt)
    
    xnew = xold + dt/6.0*(k1 + 2.0k2 + 2.0k3 + k4)
end

# 真値を全ステップ計算
function progress(deltat)
    xn = zeros(Float64, N, nstep)
    t = zeros(Float64, nstep)
    
    #X = fill(F,N)+randn(N)
    X = fill(Float64(F), N)
    # 初期擾乱はj=20の点にFの値の0.1%の値を与える。
    X[20] = 1.001*F
    
    for j=1:nstep
        X = RK4(X, deltat)
        for i=1:N
            xn[i, j] = X[i]
        end
        t[j] = dt*j
    end
    
    return xn, t
end

# x_trueの計算
@time xn, t = progress(dt)

Python版

Pythonは3.7.3です。
Pythonはpipでインストールしたnumpyで計算します。Anacondaでいれたnumpyのほうが早いですが，とりあえずテストなので，気にしないことにします。

import numpy as np
from copy import copy

## パラメータ
## Fの大きさで挙動が変わる
F = 8
# 解析データ点の個数（次元）
N = 40
dt = 0.01
# 1年分計算
tend=146
nstep = int(tend/dt)

def f(x):
    f = np.zeros((N))
    for i in range(2, N-1):
        f[i] = (x[i+1]-x[i-2])*x[i-1] - x[i] + F
    
    f[0] = (x[1]-x[N-2])*x[N-1] - x[0] + F
    f[1] = (x[2]-x[N-1])*x[0] - x[1] + F
    f[N-1] = (x[0]-x[N-3])*x[N-2] - x[N-1] + F
    
    return f

# L96をRK4で解く
def Model(xold, dt):
    k1 = f(xold)
    k2 = f(xold + k1*dt/2.)
    k3 = f(xold + k2*dt/2.)
    k4 = f(xold + k3*dt)
    
    xnew = xold + dt/6.0*(k1 + 2.0*k2 + 2.0*k3 + k4)
    return xnew

# 真値を全ステップ計算
def progress(deltat):
    xn = np.zeros((N, nstep))
    t = np.zeros((nstep))
    
    X = float(F)*np.ones(N)
    #X = fill(Float64(F), N)
    # 初期擾乱はj=19（0始まりを考慮）の点にFの値の0.1%の値を与える。
    X[19] = 1.001*F
    
    for i in range(N):
        xn[i, 0] = copy(X[i])
    
    for j in range(1, nstep):
        X = Model(X, deltat)
        xn[:, j] = X[:]
        #print(X[20], X[1], deltat)
        t[j] = dt*j

    return xn, t

# 実行
start = time.time()
xn, t = progress(dt)
elapsed_time = time.time() - start
print ("elapsed_time:{0}".format(elapsed_time) + "[sec]")

計算時間

Julia 0.73 s
Python 2.94 s

でJuliaのほうが早かったです。2.3 GHz Quad-Core Intel Core i5積んだMacBook Proでやりました。

Pythonの最適化はほとんど考えてないので，もっと早くなると思います。arrayのインデックスの順番変えたりはしましたが，早くなりませんでした。PythonでもJITコンパイルできたりしますが，やったことないので，今回は保留しました。

とりあえず，Juliaの計算速度は遅くはなさそうでした。ちなみにExtended Kalman Filterによるデータ同化のコードも書いてみましたが，Juliaのほうが早かったです。

おまけ：作図

Juliaで計算結果をプロットするコード

ホフメラー図

using PyPlot

fig = plt.figure(figsize=(6, 5.))
ax1 = fig.add_subplot(111)

ax1.set_title("true")
ax1.set_xlabel("X_i")
ax1.set_ylabel("time step")

image1 = ax1.contourf(xn'[1:5000,:], )
cb1 = fig.colorbar(image1, ax=ax1, label="X")

時間変化

using PyPlot

fig = plt.figure(figsize=(5, 8.))
ax1 = fig.add_subplot(211)
nhalf = Int(nstep/2)
nstart = nhalf
nend = nstep
ax1.plot(t[nstart:nend], xn[1, nstart:nend], label="x1")
ax1.plot(t[nstart:nend], xn[2, nstart:nend], label="x2")
ax1.plot(t[nstart:nend], xn[3, nstart:nend], label="x3")

おまけ：データ保存

HDF5で保存。計算結果の後半，5ステップおきのデータを保存。

using HDF5

h5open("L96_true_obs.h5", "w") do file
    write(file, "Xn_true", xn[:, nhalf:5:nstep])
    write(file, "t", t[nhalf:5:nstep])
end

リードはこんな感じ

file = h5open("L96_true_obs.h5", "r") 
Xn_true = read(file, "Xn_true")
t = read(file, "t")
close(file)

おまけ：Fortran

Fortranでも書いてみました。久しぶりに書いたけど，Fotranめんどう。

program main
implicit none

double precision :: F, dt, tend
integer :: N, nstep, i
real(8) :: xn(40, 14600), t(14600)

F = 8.0
N = 40
dt = 0.01
tend = 146
nstep = tend/dt

write(*, *) "F=", F

call progress(dt, N, nstep, F, xn, t)
write(*,*) "xn(:, 1)="
write(*,*) xn(1:N,1)
write(*,*) "xn(:, 2)="
write(*,*) xn(1:N,2)
write(*,*) "xn(:, nstep)=", xn(1:N,nstep)

!do i=1, N
!   write(1,*) xn(i,1)
!   write(10,*) xn(i,10)
!   !write(5,*) xn(i,5)
!   write(100,*) xn(i,nstep)
!enddo

!do i=1, nstep
!   write(200,*) xn(1:N, i)
!enddo

end program main

!! Lorenz96の式
subroutine func(x, N, F, g)
  implicit none
  integer :: i, N
  real(8) :: g(N), x(N), F

  do i=3, N-1
     g(i) = (x(i+1)-x(i-2))*x(i-1) - x(i) + F
  enddo

  g(1) = (x(2)-x(N-1))*x(N) - x(1) + F
  g(2) = (x(3)-x(N))*x(1) - x(2) + F
  g(N) = (x(1)-x(N-2))*x(N-1) - x(N) + F
  
  return 

end subroutine func


subroutine RK4(xold, N, dt, F, xnew)
  implicit none
  integer :: i, N
  real(8) :: k1(N), k2(N), k3(N), k4(N), xold(N), xnew(N), dt, F
  real(8) :: temp(N)

  do i=1, N
     temp(i) = xold(i)
  enddo
    
  call func(temp, N, F, k1)

  do i=1, N
     temp(i) = xold(i) + k1(i)*dt/2.
  enddo
  
  call func(temp, N, F, k2)

  do i=1, N
     temp(i) = xold(i) + k2(i)*dt/2.
  enddo
  
  call func(temp, N, F, k3)

  do i=1, N
     temp(i) = xold(i) + k3(i)*dt
  enddo
  
  call func(temp, N, F, k4)

  do i=1, N
     xnew(i)=xold(i) + dt/6.0*(k1(i)+2.0*k2(i)+2.0*k3(i)+k4(i))
  enddo
  
  return
end subroutine RK4

subroutine progress(dt, N, nstep, F, xn, t)
  implicit none
  integer :: N, nstep, i, j
  real(8) :: dt, F, xn(N, nstep), t(nstep), x(N), temp(N)


  do i=1, N
     x(i) = F
  enddo
  x(20) = 1.001*F

  do i=1,N
     temp(i) = x(i)
  enddo
     
  do j=1, nstep
     call RK4(temp, N, dt, F, temp)

     do i=1, N
        xn(i, j) = temp(i)
     enddo
     t(j) = dt*j

  enddo
  
  return
end subroutine progress

Homebrewでいれたgfortranでやったところ，
0.045s でした。

上のJuliaのコードはこれに比べるとだいぶ遅いですが，黒木さんによって最適化されたコード　https://gist.github.com/genkuroki/f0426a7f3772c804ace0e6875445ea49
ではJuliaで5msで計算できたとのことです。
上記のコードはグローバル変数使ってるので，そこを直すと早くなるはず。
（これからやります）

3
1
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
3
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?