Introduction

I'm studying a lot because I want to use data assimilation for research. Lorenz96 model often used in data assimilation benchmarks https://journals.ametsoc.org/jas/article/55/3/399/24564/Optimal-Sites-for-Supplementary-Weather I wrote the code to solve the problem in Julia. I'm not sure if this is slow or fast, so I rewrote it in Python for the time being and compared it.

Lorenz 96 model

equation

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

However, $ j = 1 \ dots N $. This time N = 40.

If you initially set $ X_j = F $, you will get a stationary solution. This time, solve it as $ X_ {20} = 1.01F, ; F = 8 $. When F = 8, chaotic behavior can be seen.

This is solved with Runge-Kutta with 4 steps and 4th order accuracy.

Julia version

The version of Julia is 1.5.1. I ran it in Jupyter lab.

using LinearAlgebra

##Parameters
##Behavior changes depending on the size of F
F = 8
N = 40
dt = 0.01
tend = 146
nstep = Int32(tend/dt)

# dx/dt=f(x) 
#Lorentz 96 equation
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
    
    #Periodic boundary
    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

#Solve L96 with 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

#Calculate the true value in all steps
function progress(deltat)
    xn = zeros(Float64, N, nstep)
    t = zeros(Float64, nstep)
    
    #X = fill(F,N)+randn(N)
    X = fill(Float64(F), N)
    #The initial disturbance is j=0 of F value at 20 points.1%Gives the value of.
    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_Calculation of true
@time xn, t = progress(dt)

Python version

Python is 3.7.3. Python calculates with numpy installed with pip. Numpy in Anaconda is faster, but it's a test for the time being, so I don't care.

import numpy as np
from copy import copy

##Parameters
##Behavior changes depending on the size of F
F = 8
#Number of analysis data points (dimension)
N = 40
dt = 0.01
#1 year calculation
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

#Solve L96 with 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

#Calculate the true value in all steps
def progress(deltat):
    xn = np.zeros((N, nstep))
    t = np.zeros((nstep))
    
    X = float(F)*np.ones(N)
    #X = fill(Float64(F), N)
    #The initial disturbance is j=0 of F value at point 19 (considering 0 start).1%Gives the value of.
    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

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

Calculation time

Julia 0.73 s Python 2.94 s

And Julia was faster. I did it on a MacBook Pro with a 2.3 GHz Quad-Core Intel Core i5.

I haven't really thought about optimizing Python, so I think it will be faster. I changed the order of the indexes of array, but it didn't get faster. I can do JIT compilation with Python, but I haven't done it, so I put it on hold this time.

For the time being, Julia's calculation speed did not seem to be slow. By the way, I also wrote the code for data assimilation by Extended Kalman Filter, but Julia was faster.

Bonus: Drawing

Code to plot the calculation result in Julia

Hoffmeler diagram

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")

time change

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")

Bonus: Data storage

Save as HDF5. Save data every 5 steps in the latter half of the calculation result.

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

Reed looks like this

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

Solving the Lorenz 96 model with Julia and Python