Implement 'regressed' RHS & G0 closure

Author: dburov190

examples/L96m/L96m.jl

@@ -18,8 +18,8 @@ using Parameters # lets you have defaults for fields
   K::Int = 9
   J::Int = 8
   hx::Float64 = -0.8
-  hy::Float64 = 1
-  F::Float64 = 10
+  hy::Float64 = 1.0
+  F::Float64 = 10.0
   eps::Float64 = 2^(-7)
   G = nothing
 end
@@ -115,6 +115,42 @@ function balanced(rhs::Array{<:Real,1}, x::Array{<:Real,1}, p::L96m, t)
   return rhs
 end
 
+function regressed(rhs::Array{<:Real,1}, x::Array{<:Real,1}, p::L96m, t)
+  """
+  Compute slow-variable closed RHS of the Lorenz '96 Multiscale system;
+  i.e. only slow variables with some closure instead of Yk.
+  Closure is taken from p.G.
+  Both `rhs` and `x` are vectors of size p.K.
+
+  Input:
+  - `x`   : vector of size K
+  - `p`   : parameters
+  - `t`   : time (not used here since L96m is autonomous)
+
+  Output:
+  - `rhs` : regressed RHS computed at `x`
+
+  """
+
+  K = p.K
+
+  # three boundary cases
+  rhs[1] = -x[K]   * (x[K-1] - x[2]) - x[1]
+  rhs[2] = -x[1]   * (x[K]   - x[3]) - x[2]
+  rhs[K] = -x[K-1] * (x[K-2] - x[1]) - x[K]
+
+  # general case
+  rhs[3:K-1] = -x[2:K-2] .* (x[1:K-3] - x[4:K]) - x[3:K-1]
+
+  # add forcing
+  rhs .+= p.F
+
+  # add closure
+  rhs .+= p.hx * p.G(x)
+
+  return rhs
+end
+
 function compute_Yk(p::L96m, z::Array{<:Real,1})
   """
   Reshape a vector of y_{j,k} into a matrix, then sum along one dim and divide
@@ -126,4 +162,22 @@ function compute_Yk(p::L96m, z::Array{<:Real,1})
   ) / p.J
 end
 
+function set_G0(p::L96m; slope = nothing)
+  """
+  Set the closure `p.G` to a linear one with slope `slope`.
+  If unspecified, slope is equal to `p.hy`.
+  """
+  if (slope == nothing) || (!isa(slope, Real))
+    slope = p.hy
+  end
+  p.G = x -> slope * x
+end
+
+function set_G0(p::L96m, slope::Real)
+  """
+  Wrapper for set_G0(p::L96m; slope = nothing).
+  """
+  set_G0(p, slope = slope)
+end
+
 

examples/L96m/main.jl

@@ -37,6 +37,7 @@ const hx = -0.8
 # IC section ###################################################################
 ################################################################################
 l96 = L96m(hx = hx, J = 8)
+set_G0(l96) # set the linear closure (essentially, as in balanced)
 
 z00 = Array{Float64}(undef, l96.K + l96.K * l96.J)
 
@@ -57,6 +58,8 @@ elapsed_jit = @elapsed begin
   DE.solve(pb_jit, SOLVER, reltol = reltol, abstol = abstol, dtmax = dtmax)
   pb_jit = DE.ODEProblem(balanced, z00[1:l96.K], (0.0, T_compile), l96)
   DE.solve(pb_jit, SOLVER, reltol = reltol, abstol = abstol, dtmax = dtmax)
+  pb_jit = DE.ODEProblem(regressed, z00[1:l96.K], (0.0, T_compile), l96)
+  DE.solve(pb_jit, SOLVER, reltol = reltol, abstol = abstol, dtmax = dtmax)
 end
 println(" " ^ (LPAD_INTEGER + 6),
         "\t\telapsed:", lpad(elapsed_jit, LPAD_FLOAT))
@@ -104,6 +107,20 @@ end
 println("steps:", lpad(length(sol_bal.t), LPAD_INTEGER),
         "\t\telapsed:", lpad(elapsed_bal, LPAD_FLOAT))
 
+# regressed L96m integration
+print(rpad("(regressed)", RPAD))
+elapsed_reg = @elapsed begin
+  pb_reg = DE.ODEProblem(regressed, z0[1:l96.K], (0.0, T), l96)
+  sol_reg = DE.solve(pb_reg,
+                     SOLVER,
+                     reltol = reltol,
+                     abstol = abstol,
+                     dtmax = dtmax
+                    )
+end
+println("steps:", lpad(length(sol_reg.t), LPAD_INTEGER),
+        "\t\telapsed:", lpad(elapsed_reg, LPAD_FLOAT))
+
 ################################################################################
 # plot section #################################################################
 ################################################################################
@@ -115,6 +132,9 @@ plt.plot(sol_dns.t, sol_dns[l96.K + (k-1)*l96.J + j,:],
 # plot balanced
 plt.plot(sol_bal.t, sol_bal[k,:], label = "balanced")
 
+# plot regressed
+plt.plot(sol_reg.t, sol_reg[k,:], label = "regressed")
+
 plt.legend()
 plt.show()