Implement 'balanced'; add a unit test for that

Author: dburov190

examples/L96m/L96m.jl

@@ -83,6 +83,38 @@ function full(rhs::Array{<:Real,1}, z::Array{<:Real,1}, _s::L96m, t)
   return rhs
 end
 
+function balanced(rhs::Array{<:Real,1}, x::Array{<:Real,1}, _s::L96m, t)
+  """
+  Compute balanced RHS of the Lorenz '96 multiscale system; i.e. only slow
+  variables with the linear closure.
+  Both `rhs` and `x` are vectors of size _s.K.
+
+  Input:
+  - `x`   : vector of size K
+  - `_s`  : parameters
+  - `t`   : time (not used here since L96m is autonomous)
+
+  Output:
+  - `rhs` : balanced RHS computed at `x`
+
+  """
+
+  K = _s.K
+
+  # three boundary cases
+  rhs[1] = -x[K]   * (x[K-1] - x[2]) - (1 - _s.hx*_s.hy) * x[1]
+  rhs[2] = -x[1]   * (x[K]   - x[3]) - (1 - _s.hx*_s.hy) * x[2]
+  rhs[K] = -x[K-1] * (x[K-2] - x[1]) - (1 - _s.hx*_s.hy) * x[K]
+
+  # general case
+  rhs[3:K-1] = -x[2:K-2] .* (x[1:K-3] - x[4:K]) - (1 - _s.hx*_s.hy) * x[3:K-1]
+
+  # add forcing
+  rhs .+= _s.F
+
+  return rhs
+end
+
 function compute_Yk(_s::L96m, z::Array{<:Real,1})
   """
   Reshape a vector of y_{j,k} into a matrix, then sum along one dim and divide

test/L96m/data/balanced.npy

@@ -15,6 +15,7 @@ const data_dir = joinpath(@__DIR__, "data")
 const z0 = NPZ.npzread(joinpath(data_dir, "ic.npy"))
 const Yk_test = NPZ.npzread(joinpath(data_dir, "Yk.npy"))
 const full_test = NPZ.npzread(joinpath(data_dir, "full.npy"))
+const balanced_test = NPZ.npzread(joinpath(data_dir, "balanced.npy"))
 const dp5_test = NPZ.npzread(joinpath(data_dir, "dp5.npy"))
 const tp8_test = NPZ.npzread(joinpath(data_dir, "tsitpap8.npy"))
 
@@ -23,10 +24,13 @@ const tp8_test = NPZ.npzread(joinpath(data_dir, "tsitpap8.npy"))
 ################################################################################
 Yk_comp = compute_Yk(l96, z0)
 full_comp = similar(z0)
-full(full_comp, z0, l96, 0)
+full(full_comp, z0, l96, 0.0)
+balanced_comp = similar(z0[1:l96.K])
+balanced(balanced_comp, z0[1:l96.K], l96, 0.0)
 @testset "unit testing" begin
   @test isapprox(Yk_test, Yk_comp, atol=1e-15)
   @test isapprox(full_test, full_comp, atol=1e-15)
+  @test isapprox(balanced_test, balanced_comp, atol=1e-15)
 end
 println("")