Implement 'gather_pairs', 'random_init'; add tests

Author: dburov190

examples/L96m/L96m.jl

@@ -31,7 +31,7 @@ variables are fast.
 
 Input:
  - `z`   : vector of size (K + K*J)
- - `p`   : parameters
+ - `p`   : parameters (L96m struct)
  - `t`   : time (not used here since L96m is autonomous)
 
 Output:
@@ -89,7 +89,7 @@ Both `rhs` and `x` are vectors of size p.K.
 
 Input:
  - `x`   : vector of size K
- - `p`   : parameters
+ - `p`   : parameters (L96m struct)
  - `t`   : time (not used here since L96m is autonomous)
 
 Output:
@@ -121,7 +121,7 @@ Both `rhs` and `x` are vectors of size p.K.
 
 Input:
  - `x`   : vector of size K
- - `p`   : parameters
+ - `p`   : parameters (L96m struct)
  - `t`   : time (not used here since L96m is autonomous)
 
 Output:
@@ -177,4 +177,55 @@ function set_G0(p::L96m, slope::Real)
   set_G0(p, slope = slope)
 end
 
+"""
+Gather (xk, Yk) pairs that are further used to train a GP regressor
+
+Input:
+ - `p`     : parameters (L96m struct)
+ - `sol`   : time series of a solution; time steps are in the 2nd dimension
+             (usually, just `sol` output from a time-stepper)
+
+Output:
+ - `pairs` : a 2-d array of size (N, 2) containing (xk, Yk) pairs, where N is
+             the 2nd dimension in `sol` (number of time steps)
+
+"""
+function gather_pairs(p::L96m, sol)
+  N = size(sol, 2)
+  pairs = Array{Float64, 2}(undef, p.K * N, 2)
+  for n in 1:N
+    pairs[p.K * (n-1) + 1 : p.K * n, 1] = sol[1:p.K, n]
+    pairs[p.K * (n-1) + 1 : p.K * n, 2] = compute_Yk(p, sol[:,n])
+  end
+  return pairs
+end
+
+"""
+Returns a randomly initialized array that can be used as an IC to ODE solver
+
+The returned array is of size `p.K + p.K * p.J`.
+The first `p.K` variables are slow and are drawn randomly ~ U[-5; 10]; each of
+the fast variables corresponding to the same slow variable is set to the value
+of that slow variable.
+
+For example, if K == 2 and J = 3, the returned array will be
+  [ rand1, rand2, rand1, rand1, rand1, rand2, rand2, rand2 ]
+
+Input:
+ - `p`     : parameters (L96m struct)
+
+Output:
+ - `z00`   : array of size `p.K + p.K * p.J` with random values
+"""
+function random_init(p::L96m)
+  z00 = Array{Float64}(undef, p.K + p.K * p.J)
+
+  z00[1:p.K] .= rand(p.K) * 15 .- 5
+  for k_ in 1:p.K
+    z00[p.K+1 + (k_-1)*p.J : p.K + k_*p.J] .= z00[k_]
+  end
+
+  return z00
+end
+
 

examples/L96m/main.jl

@@ -143,13 +143,7 @@ print_var(parameters)
 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)
-
-# 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
+z00 = random_init(l96)
 
 ################################################################################
 # main section #################################################################

test/L96m/data/dp5_pairs.npy

@@ -17,6 +17,7 @@ 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 regressed_test = NPZ.npzread(joinpath(data_dir, "regressed.npy"))
+const dp5_pairs = NPZ.npzread(joinpath(data_dir, "dp5_pairs.npy"))
 const dp5_test = NPZ.npzread(joinpath(data_dir, "dp5_full.npy"))
 const tp8_test = NPZ.npzread(joinpath(data_dir, "tsitpap8_full.npy"))
 const bal_test = NPZ.npzread(joinpath(data_dir, "dp5_bal.npy"))
@@ -59,6 +60,27 @@ const bal_test = NPZ.npzread(joinpath(data_dir, "dp5_bal.npy"))
     @test -0.5 * v1 == l96.G(v1)
     @test -0.5 * v2 == l96.G(v2)
   end
+
+  m1 = [1:9.0  1:9.0]
+  z00 = zeros(l96.K + l96.K * l96.J)
+  z00[1:l96.K] .= 1:l96.K
+  for k_ in 1:l96.K
+    z00[l96.K + 1 + (k_-1)*l96.J : l96.K + k_*l96.J] .= z00[k_]
+  end
+  @test isapprox(gather_pairs(l96, z00), m1)
+
+  z00 = random_init(l96)
+  x00 = @view(z00[1:l96.K])
+  y00 = @view(z00[l96.K+1:end])
+  @test size(z00,1) == l96.K + l96.K * l96.J
+  @test size(z00,1) == length(z00)
+  @test all(-5.0 .< x00 .< 10.0)
+  @test all(x00 .== reshape(y00, l96.J, l96.K)')
+
+  m2 = [x00  x00]
+  @test isapprox(gather_pairs(l96, z00), m2)
+
+  @test isapprox(gather_pairs(l96, dp5_test), dp5_pairs)
 end
 println("")