Merge pull request #4 from climate-machine/lorenz96

Add Lorenz 96 model to examples

Author: dburov190

Manifest.toml

@@ -4,8 +4,10 @@ authors = ["Simon Byrne <[email protected]>"]
 version = "0.1.0"
 
 [deps]
+DifferentialEquations = "0c46a032-eb83-5123-abaf-570d42b7fbaa"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+NPZ = "15e1cf62-19b3-5cfa-8e77-841668bca605"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

Project.toml

@@ -0,0 +1,129 @@
+using Parameters # lets you have defaults for fields
+
+@with_kw mutable struct L96m
+  """
+  Lorenz '96 multiscale
+
+  Parameters:
+  - `K`   : number of slow variables
+  - `J`   : number of fast variables per slow variable
+  - `hx`  : coupling term (in equations for slow variables)
+  - `hy`  : coupling term (in equations for fast variables)
+  - `F`   : forcing term for slow variables
+  - `eps` : scale separation term
+
+  Other:
+  - `G` : functional closure for slow variables (usually a GPR-closure)
+  """
+  K::Int = 9
+  J::Int = 8
+  hx::Float64 = -0.8
+  hy::Float64 = 1
+  F::Float64 = 10
+  eps::Float64 = 2^(-7)
+  G = nothing
+end
+
+function full(rhs::Array{<:Real,1}, z::Array{<:Real,1}, p::L96m, t)
+  """
+  Compute full RHS of the Lorenz '96 multiscale system.
+  The convention is that the first K variables are slow, while the rest K*J
+  variables are fast.
+
+  Input:
+  - `z`   : vector of size (K + K*J)
+  - `p`   : parameters
+  - `t`   : time (not used here since L96m is autonomous)
+
+  Output:
+  - `rhs` : RHS computed at `z`
+
+  """
+
+  K = p.K
+  J = p.J
+  x = @view(z[1:K])
+  y = @view(z[K+1:end])
+
+  ### slow variables subsystem ###
+  # compute Yk averages
+  Yk = compute_Yk(p, z)
+
+  # 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[1:K] .+= p.F
+
+  # add coupling w/ fast variables via averages
+  rhs[1:K] .+= p.hx * Yk
+
+  ### fast variables subsystem ###
+  # three boundary cases
+  rhs[K+1]   = -y[2]   * (y[3] - y[end]  ) - y[1]
+  rhs[end-1] = -y[end] * (y[1] - y[end-2]) - y[end-1]
+  rhs[end]   = -y[1]   * (y[2] - y[end-1]) - y[end]
+
+  # general case
+  rhs[K+2:end-2] = -y[3:end-1] .* (y[4:end] - y[1:end-3]) - y[2:end-2]
+
+  # add coupling w/ slow variables
+  for k in 1:K
+    rhs[K+1 + (k-1)*J : K + k*J] .+= p.hy * x[k]
+  end
+
+  # divide by epsilon
+  rhs[K+1:end] ./= p.eps
+
+  return rhs
+end
+
+function balanced(rhs::Array{<:Real,1}, x::Array{<:Real,1}, p::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 p.K.
+
+  Input:
+  - `x`   : vector of size K
+  - `p`   : parameters
+  - `t`   : time (not used here since L96m is autonomous)
+
+  Output:
+  - `rhs` : balanced RHS computed at `x`
+
+  """
+
+  K = p.K
+
+  # three boundary cases
+  rhs[1] = -x[K]   * (x[K-1] - x[2]) - (1 - p.hx*p.hy) * x[1]
+  rhs[2] = -x[1]   * (x[K]   - x[3]) - (1 - p.hx*p.hy) * x[2]
+  rhs[K] = -x[K-1] * (x[K-2] - x[1]) - (1 - p.hx*p.hy) * x[K]
+
+  # general case
+  rhs[3:K-1] = -x[2:K-2] .* (x[1:K-3] - x[4:K]) - (1 - p.hx*p.hy) * x[3:K-1]
+
+  # add forcing
+  rhs .+= p.F
+
+  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
+  by J to get averages
+  """
+  return dropdims(
+      sum( reshape(z[p.K+1:end], p.J, p.K), dims = 1 ),
+      dims = 1
+  ) / p.J
+end
+
+

examples/L96m/L96m.jl

@@ -0,0 +1,121 @@
+#!/usr/bin/julia --
+
+import DifferentialEquations
+import PyPlot
+const DE = DifferentialEquations
+const plt = PyPlot
+include("L96m.jl")
+
+
+################################################################################
+# constants section ############################################################
+################################################################################
+const RPAD = 42
+const LPAD_INTEGER = 7
+const LPAD_FLOAT = 13
+const SOLVER = DE.Tsit5()
+
+const T = 4 # integration time
+const T_compile = 1e-10 # force JIT compilation
+const T_conv = 100 # converging integration time
+#const T_learn = 15 # time to gather training data for GP
+#const T_hist = 10000 # time to gather histogram statistics
+
+# integrator parameters
+const dtmax = 1e-3 # maximum step size
+#const tau = 1e-3 # maximum step size for histogram statistics
+#const dt_conv = 0.01 # maximum step size for converging to attractor
+const reltol = 1e-3 # relative tolerance
+const abstol = 1e-6 # absolute tolerance
+
+const k = 1 # index of the slow variable to save etc.
+const j = 1 # index of the fast variable to save/plot etc.
+
+const hx = -0.8
+
+################################################################################
+# IC section ###################################################################
+################################################################################
+l96 = L96m(hx = hx, J = 8)
+
+z00 = Array{Float64}(undef, l96.K + l96.K * l96.J)
+
+# method 1
+z00[1:l96.K] .= rand(l96.K) * 15 .- 5
+for k_ in 1:l96.K
+  z00[l96.K+1 + (k_-1)*l96.J : l96.K + k_*l96.J] .= z00[k_]
+end
+
+################################################################################
+# main section #################################################################
+################################################################################
+
+# force compilation of functions used in numerical integration
+print(rpad("(JIT compilation)", RPAD))
+elapsed_jit = @elapsed begin
+  pb_jit = DE.ODEProblem(full, z00, (0.0, T_compile), l96)
+  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)
+end
+println(" " ^ (LPAD_INTEGER + 6),
+        "\t\telapsed:", lpad(elapsed_jit, LPAD_FLOAT))
+
+# full L96m integration (converging to attractor)
+print(rpad("(full, converging)", RPAD))
+elapsed_conv = @elapsed begin
+  pb_conv = DE.ODEProblem(full, z00, (0.0, T_conv), l96)
+  sol_conv = DE.solve(pb_conv,
+                      SOLVER,
+                      reltol = reltol,
+                      abstol = abstol,
+                      dtmax = dtmax
+                     )
+end
+println("steps:", lpad(length(sol_conv.t), LPAD_INTEGER),
+        "\t\telapsed:", lpad(elapsed_conv, LPAD_FLOAT))
+z0 = sol_conv[:,end]
+
+# full L96m integration
+print(rpad("(full)", RPAD))
+elapsed_dns = @elapsed begin
+  pb_dns = DE.ODEProblem(full, z0, (0.0, T), l96)
+  sol_dns = DE.solve(pb_dns,
+                     SOLVER,
+                     reltol = reltol,
+                     abstol = abstol,
+                     dtmax = dtmax
+                    )
+end
+println("steps:", lpad(length(sol_dns.t), LPAD_INTEGER),
+        "\t\telapsed:", lpad(elapsed_dns, LPAD_FLOAT))
+
+# balanced L96m integration
+print(rpad("(balanced)", RPAD))
+elapsed_bal = @elapsed begin
+  pb_bal = DE.ODEProblem(balanced, z0[1:l96.K], (0.0, T), l96)
+  sol_bal = DE.solve(pb_bal,
+                     SOLVER,
+                     reltol = reltol,
+                     abstol = abstol,
+                     dtmax = dtmax
+                    )
+end
+println("steps:", lpad(length(sol_bal.t), LPAD_INTEGER),
+        "\t\telapsed:", lpad(elapsed_bal, LPAD_FLOAT))
+
+################################################################################
+# plot section #################################################################
+################################################################################
+# plot DNS
+plt.plot(sol_dns.t, sol_dns[k,:], label = "DNS")
+plt.plot(sol_dns.t, sol_dns[l96.K + (k-1)*l96.J + j,:],
+    lw = 0.6, alpha = 0.6, color="gray")
+
+# plot balanced
+plt.plot(sol_bal.t, sol_bal[k,:], label = "balanced")
+
+plt.legend()
+plt.show()
+
+