Skip to content

Commit aa389fb

Browse files
ArneBouillonArneBouillon
committed
Separate spatial-dep Lorenz into its own file
1 parent 1dd5db1 commit aa389fb

2 files changed

Lines changed: 130 additions & 130 deletions

File tree

‎examples/Lorenz/calibrate_spatial_dep.jl‎

Lines changed: 1 addition & 130 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,3 @@
1-
include("GModel.jl") # Contains Lorenz 96 source code
2-
31
# Import modules
42
using Distributions # probability distributions and associated functions
53
using LinearAlgebra
@@ -16,134 +14,7 @@ using CalibrateEmulateSample.EnsembleKalmanProcesses.Localizers
1614

1715
const EKP = EnsembleKalmanProcesses
1816

19-
# G(θ) = H(Ψ(θ,x₀,t₀,t₁))
20-
# y = G(θ) + η
21-
22-
# This will change for different Lorenz simulators
23-
struct LorenzConfig{FT <: Real}
24-
"Length of a fixed integration timestep"
25-
dt::FT
26-
"Total duration of integration (T = N*dt)"
27-
T::FT
28-
end
29-
30-
# This will change for each ensemble member
31-
struct EnsembleMemberConfig{VV <: AbstractVector}
32-
"state-dependent-forcing"
33-
F::VV
34-
end
35-
36-
# This will change for different "Observations" of Lorenz
37-
struct ObservationConfig{FT1 <: Real, FT2 <: Real}
38-
"initial time to gather statistics (T_start = N_start*dt)"
39-
T_start::FT1
40-
"end time to gather statistics (T_end = N_end*dt)"
41-
T_end::FT2
42-
end
43-
#########################################################################
44-
############################ Model Functions ############################
45-
#########################################################################
46-
47-
# Forward pass of forward model
48-
# Inputs:
49-
# - params: structure with F (state-dependent-forcing vector)
50-
# - x0: initial condition vector
51-
# - config: config of forward run
52-
# - observation_config: config for observations
53-
54-
function lorenz_forward(
55-
params::EnsembleMemberConfig,
56-
x0::VV,
57-
config::LorenzConfig,
58-
observation_config::ObservationConfig,
59-
) where {VV <: AbstractVector}
60-
# run the Lorenz simulation
61-
xn = lorenz_solve(params, x0, config)
62-
# Get statistics
63-
gt = stats(xn, config, observation_config)
64-
return gt
65-
end
66-
67-
# Calculates statistics for forward model output
68-
# Inputs:
69-
# - xn: timeseries of states for length of simulation through Lorenz96
70-
# - config: config of forward run
71-
# - observation_config: config for observations
72-
function stats(xn::VorM, config::LorenzConfig, observation_config::ObservationConfig) where {VorM <: AbstractVecOrMat}
73-
T_start = observation_config.T_start
74-
T_end = observation_config.T_end
75-
dt = config.dt
76-
N_start = Int(ceil(T_start / dt))
77-
N_end = Int(ceil(T_end / dt))
78-
xn_stat = xn[:, N_start:N_end]
79-
N_state = size(xn_stat, 1)
80-
gt = zeros(2 * N_state)
81-
gt[1:N_state] = mean(xn_stat, dims = 2)
82-
gt[(N_state + 1):(2 * N_state)] = std(xn_stat, dims = 2)
83-
return gt
84-
end
85-
86-
# Forward pass of the Lorenz 96 model
87-
# Inputs:
88-
# - params: structure with F (state-dependent-forcing vector)
89-
# - x0: initial condition vector
90-
# - config: structure including dt (timestep Float64(1)) and T (total time Float64(1))
91-
function lorenz_solve(params::EnsembleMemberConfig, x0::VorM, config::LorenzConfig) where {VorM <: AbstractVecOrMat}
92-
# Initialize
93-
nstep = Int(ceil(config.T / config.dt))
94-
state_dim = size(x0, 1)
95-
xn = zeros(state_dim, nstep + 1)
96-
xn[:, 1] = x0
97-
98-
# March forward in time
99-
for j in 1:nstep
100-
xn[:, j + 1] = RK4(params, xn[:, j], config)
101-
end
102-
103-
return xn
104-
end
105-
106-
# Lorenz 96 system
107-
# f = dx/dt
108-
# Inputs:
109-
# - params: structure with F (state-dependent-forcing vector)
110-
# - x: current state
111-
function f(params::EnsembleMemberConfig, x::VV) where {VV <: AbstractVector}
112-
F = params.F
113-
N = length(x)
114-
f = zeros(N)
115-
# Loop over N positions
116-
for i in 3:(N - 1)
117-
f[i] = -x[i - 2] * x[i - 1] + x[i - 1] * x[i + 1] - x[i] + F[i]
118-
end
119-
# Periodic boundary conditions
120-
f[1] = -x[N - 1] * x[N] + x[N] * x[2] - x[1] + F[1]
121-
f[2] = -x[N] * x[1] + x[1] * x[3] - x[2] + F[2]
122-
f[N] = -x[N - 2] * x[N - 1] + x[N - 1] * x[1] - x[N] + F[N]
123-
# Output
124-
return f
125-
end
126-
127-
# RK4 solve
128-
# Inputs:
129-
# - params: structure with F (state-dependent-forcing vector)
130-
# - xold: current state
131-
# - config: structure including dt (timestep Float64(1)) and T (total time Float64(1))
132-
function RK4(params::EnsembleMemberConfig, xold::VorM, config::LorenzConfig) where {VorM <: AbstractVecOrMat}
133-
N = length(xold)
134-
dt = config.dt
135-
136-
# Predictor steps (note no time-dependence is needed here)
137-
k1 = f(params, xold)
138-
k2 = f(params, xold + k1 * dt / 2.0)
139-
k3 = f(params, xold + k2 * dt / 2.0)
140-
k4 = f(params, xold + k3 * dt)
141-
# Step
142-
xnew = xold + (dt / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4)
143-
# Output
144-
return xnew
145-
end
146-
17+
include("./lorenz_spatial_dep.jl")
14718

14819
########################################################################
14920
############################ Problem setup #############################

‎examples/Lorenz/lorenz_spatial_dep.jl‎

Lines changed: 129 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
1+
include("GModel.jl") # Contains Lorenz 96 source code
2+
3+
# G(θ) = H(Ψ(θ,x₀,t₀,t₁))
4+
# y = G(θ) + η
5+
6+
# This will change for different Lorenz simulators
7+
struct LorenzConfig{FT <: Real}
8+
"Length of a fixed integration timestep"
9+
dt::FT
10+
"Total duration of integration (T = N*dt)"
11+
T::FT
12+
end
13+
14+
# This will change for each ensemble member
15+
struct EnsembleMemberConfig{VV <: AbstractVector}
16+
"state-dependent-forcing"
17+
F::VV
18+
end
19+
20+
# This will change for different "Observations" of Lorenz
21+
struct ObservationConfig{FT1 <: Real, FT2 <: Real}
22+
"initial time to gather statistics (T_start = N_start*dt)"
23+
T_start::FT1
24+
"end time to gather statistics (T_end = N_end*dt)"
25+
T_end::FT2
26+
end
27+
#########################################################################
28+
############################ Model Functions ############################
29+
#########################################################################
30+
31+
# Forward pass of forward model
32+
# Inputs:
33+
# - params: structure with F (state-dependent-forcing vector)
34+
# - x0: initial condition vector
35+
# - config: config of forward run
36+
# - observation_config: config for observations
37+
38+
function lorenz_forward(
39+
params::EnsembleMemberConfig,
40+
x0::VV,
41+
config::LorenzConfig,
42+
observation_config::ObservationConfig,
43+
) where {VV <: AbstractVector}
44+
# run the Lorenz simulation
45+
xn = lorenz_solve(params, x0, config)
46+
# Get statistics
47+
gt = stats(xn, config, observation_config)
48+
return gt
49+
end
50+
51+
# Calculates statistics for forward model output
52+
# Inputs:
53+
# - xn: timeseries of states for length of simulation through Lorenz96
54+
# - config: config of forward run
55+
# - observation_config: config for observations
56+
function stats(xn::VorM, config::LorenzConfig, observation_config::ObservationConfig) where {VorM <: AbstractVecOrMat}
57+
T_start = observation_config.T_start
58+
T_end = observation_config.T_end
59+
dt = config.dt
60+
N_start = Int(ceil(T_start / dt))
61+
N_end = Int(ceil(T_end / dt))
62+
xn_stat = xn[:, N_start:N_end]
63+
N_state = size(xn_stat, 1)
64+
gt = zeros(2 * N_state)
65+
gt[1:N_state] = mean(xn_stat, dims = 2)
66+
gt[(N_state + 1):(2 * N_state)] = std(xn_stat, dims = 2)
67+
return gt
68+
end
69+
70+
# Forward pass of the Lorenz 96 model
71+
# Inputs:
72+
# - params: structure with F (state-dependent-forcing vector)
73+
# - x0: initial condition vector
74+
# - config: structure including dt (timestep Float64(1)) and T (total time Float64(1))
75+
function lorenz_solve(params::EnsembleMemberConfig, x0::VorM, config::LorenzConfig) where {VorM <: AbstractVecOrMat}
76+
# Initialize
77+
nstep = Int(ceil(config.T / config.dt))
78+
state_dim = size(x0, 1)
79+
xn = zeros(state_dim, nstep + 1)
80+
xn[:, 1] = x0
81+
82+
# March forward in time
83+
for j in 1:nstep
84+
xn[:, j + 1] = RK4(params, xn[:, j], config)
85+
end
86+
87+
return xn
88+
end
89+
90+
# Lorenz 96 system
91+
# f = dx/dt
92+
# Inputs:
93+
# - params: structure with F (state-dependent-forcing vector)
94+
# - x: current state
95+
function f(params::EnsembleMemberConfig, x::VV) where {VV <: AbstractVector}
96+
F = params.F
97+
N = length(x)
98+
f = zeros(N)
99+
# Loop over N positions
100+
for i in 3:(N - 1)
101+
f[i] = -x[i - 2] * x[i - 1] + x[i - 1] * x[i + 1] - x[i] + F[i]
102+
end
103+
# Periodic boundary conditions
104+
f[1] = -x[N - 1] * x[N] + x[N] * x[2] - x[1] + F[1]
105+
f[2] = -x[N] * x[1] + x[1] * x[3] - x[2] + F[2]
106+
f[N] = -x[N - 2] * x[N - 1] + x[N - 1] * x[1] - x[N] + F[N]
107+
# Output
108+
return f
109+
end
110+
111+
# RK4 solve
112+
# Inputs:
113+
# - params: structure with F (state-dependent-forcing vector)
114+
# - xold: current state
115+
# - config: structure including dt (timestep Float64(1)) and T (total time Float64(1))
116+
function RK4(params::EnsembleMemberConfig, xold::VorM, config::LorenzConfig) where {VorM <: AbstractVecOrMat}
117+
N = length(xold)
118+
dt = config.dt
119+
120+
# Predictor steps (note no time-dependence is needed here)
121+
k1 = f(params, xold)
122+
k2 = f(params, xold + k1 * dt / 2.0)
123+
k3 = f(params, xold + k2 * dt / 2.0)
124+
k4 = f(params, xold + k3 * dt)
125+
# Step
126+
xnew = xold + (dt / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4)
127+
# Output
128+
return xnew
129+
end

0 commit comments

Comments
 (0)