Skip to content

Commit

Permalink
julia format
Browse files Browse the repository at this point in the history
  • Loading branch information
@odunbar
odunbar committed Jun 13, 2023
1 parent 81e6491 commit 62a6a02
Show file tree
Hide file tree
Showing 7 changed files with 208 additions and 116 deletions.
11 changes: 8 additions & 3 deletions examples/Emulator/RandomFeature/scalar_optimize_and_plot_RF.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -126,9 +126,14 @@ end

# setup random features
n_features = 180
optimizer_options= Dict("n_iteration" => 20, "prior_in_scale" =>0.1, "verbose" => true) #"Max" iterations (may do less)

srfi = ScalarRandomFeatureInterface(n_features, p, diagonalize_input = diagonalize_input, optimizer_options=optimizer_options)
optimizer_options = Dict("n_iteration" => 20, "prior_in_scale" => 0.1, "verbose" => true) #"Max" iterations (may do less)

srfi = ScalarRandomFeatureInterface(
n_features,
p,
diagonalize_input = diagonalize_input,
optimizer_options = optimizer_options,
)
emulator = Emulator(srfi, iopairs, obs_noise_cov = Σ, normalize_inputs = true)
println("build RF with $n training points and $(n_features) random features.")

Expand Down
26 changes: 20 additions & 6 deletions examples/Emulator/RandomFeature/vector_optimize_and_plot_RF.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -133,20 +133,34 @@ function main()

# setup random features
n_features = 200
if case ∈ ["svd-diag","svd-nondiag"]
optimizer_options= Dict("n_iteration" => 20, "prior_in_scale" => 0.01, "prior_out_scale" => 1.0, "verbose" => true) #"Max" iterations (may do less)
if case ∈ ["svd-diag", "svd-nondiag"]
optimizer_options =
Dict("n_iteration" => 20, "prior_in_scale" => 0.01, "prior_out_scale" => 1.0, "verbose" => true) #"Max" iterations (may do less)
else # without svd
optimizer_options= Dict("n_iteration" => 20, "prior_in_scale" => 0.01, "prior_out_scale" => 0.01, "verbose" => true) #"Max" iterations (may do less)
optimizer_options =
Dict("n_iteration" => 20, "prior_in_scale" => 0.01, "prior_out_scale" => 0.01, "verbose" => true) #"Max" iterations (may do less)
end

if case == "svd-diag"
vrfi = VectorRandomFeatureInterface(n_features, p, d, diagonalize_output = true, optimizer_options = optimizer_options)
vrfi = VectorRandomFeatureInterface(
n_features,
p,
d,
diagonalize_output = true,
optimizer_options = optimizer_options,
)
emulator = Emulator(vrfi, iopairs, obs_noise_cov = Σ, normalize_inputs = true)
elseif case == "svd-nondiag"
vrfi = VectorRandomFeatureInterface(n_features, p, d, optimizer_options = optimizer_options)
emulator = Emulator(vrfi, iopairs, obs_noise_cov = Σ, normalize_inputs = true)
elseif case == "nosvd-diag"
vrfi = VectorRandomFeatureInterface(n_features, p, d, diagonalize_output = true, optimizer_options = optimizer_options)
vrfi = VectorRandomFeatureInterface(
n_features,
p,
d,
diagonalize_output = true,
optimizer_options = optimizer_options,
)
emulator = Emulator(vrfi, iopairs, obs_noise_cov = Σ, normalize_inputs = true, decorrelate = false)
elseif case == "nosvd-nondiag"
vrfi = VectorRandomFeatureInterface(n_features, p, d, optimizer_options = optimizer_options)
Expand Down
115 changes: 61 additions & 54 deletions examples/Lorenz/calibrate.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@ function main()
# 4 is a linear fit over a batch of length Ts_days
stats_type = 5
# 5 is the mean over batches of length Ts_days:


###
### Define the (true) parameters
Expand Down Expand Up @@ -113,7 +113,7 @@ function main()
Fp = rand(Normal(0.0, 0.01), N)
kmax = 1
# Prescribe variance or use a number of forward passes to define true interval variability
var_prescribe = true
var_prescribe = false
# Use CES to learn ω?
ω_fixed = true

Expand Down Expand Up @@ -143,8 +143,8 @@ function main()
n_samples = 100
yt = zeros(length(gt), n_samples)
noise_sd = 0.25
# Γy = noise_sd^2 * convert(Array, Diagonal(gt))
Γy = noise_sd^2 * convert(Array, Diagonal(ones(length(gt))*mean(gt)))
# Γy = noise_sd^2 * convert(Array, Diagonal(gt))
Γy = noise_sd^2 * convert(Array, Diagonal(ones(length(gt)) * mean(gt)))
μ = zeros(length(gt))
# Add noise
for i in 1:n_samples
Expand All @@ -154,40 +154,40 @@ function main()

else
println("Using truth values to compute covariance")
n_samples = 50
n_samples = 100
nthreads = Threads.nthreads()
yt = zeros(nthreads, length(gt), n_samples)
Threads.@threads for i in 1:n_samples

#=
lorenz_settings_local = GModel.LSettings(
dynamics,
stats_type,
t_start + T * (i - 1),
T,
Ts,
Tfit,
Fp,
N,
dt,
t_start + T * (i - 1) + T,
kmax,
ω_fixed,
ω_true,
)=#
#=
lorenz_settings_local = GModel.LSettings(
dynamics,
stats_type,
t_start + T * (i - 1),
T,
Ts,
Tfit,
Fp,
N,
dt,
t_start + T * (i - 1) + T,
kmax,
ω_fixed,
ω_true,
)=#
lorenz_settings_local = lorenz_settings
tid = Threads.threadid()
yt[tid, :, i] = GModel.run_G_ensemble(params_true, lorenz_settings_local)
end
yt = dropdims(sum(yt, dims = 1), dims = 1)

# add a little extra variance for regularization here
noise_sd = 0.05
Γy_diag = noise_sd.^2 * convert(Array, I(size(yt,1)))
μ = zeros(size(yt,1))
noise_sd = 0.1
Γy_diag = noise_sd .^ 2 * convert(Array, I(size(yt, 1)))
μ = zeros(size(yt, 1))
# Add noise
yt += rand(MvNormal(μ, Γy_diag), size(yt,2))
yt += rand(MvNormal(μ, Γy_diag), size(yt, 2))

# Variance of truth data
#Γy = convert(Array, Diagonal(dropdims(mean((yt.-mean(yt,dims=1)).^2,dims=1),dims=1)))
# Covariance of truth data
Expand All @@ -197,10 +197,10 @@ function main()
end



# Construct observation object
truth = Observations.Observation(yt, Γy, data_names)
truth_sample = yt[:,end]
truth_sample = yt[:, end]
###
### Calibrate: Ensemble Kalman Inversion
###
Expand All @@ -209,29 +209,36 @@ function main()
lorenz_settings_G = lorenz_settings # initialize to truth settings

# EKP parameters
N_ens = 50 # number of ensemble members
N_ens = 30 # number of ensemble members
N_iter = 20 # number of EKI iterations
# initial parameters: N_params x N_ens
initial_params = EKP.construct_initial_ensemble(rng, priors, N_ens)

ekiobj = EKP.EnsembleKalmanProcess(initial_params, truth_sample, truth.obs_noise_cov, EKP.Inversion(), scheduler = EKP.DataMisfitController(), verbose=true)
ekiobj = EKP.EnsembleKalmanProcess(
initial_params,
truth_sample,
truth.obs_noise_cov,
EKP.Inversion(),
scheduler = EKP.DataMisfitController(),
verbose = true,
)

# EKI iterations
println("EKP inversion error:")
err = zeros(N_iter)
final_iter=[N_iter]
err = zeros(N_iter)
final_iter = [N_iter]
for i in 1:N_iter
params_i = EKP.get_ϕ_final(priors, ekiobj) # the `ϕ` indicates that the `params_i` are in the constrained space
g_ens = GModel.run_G_ensemble(params_i, lorenz_settings_G)
terminated = EKP.update_ensemble!(ekiobj, g_ens)
if !isnothing(terminated)
final_iter = i-1 # final update was previous iteration
final_iter = i - 1 # final update was previous iteration
break
end
err[i] = EKP.get_error(ekiobj)[end] #mean((params_true - mean(params_i,dims=2)).^2)
println("Iteration: " * string(i) * ", Error: " * string(err[i]))
end
N_iter = final_iter[1] #in case it terminated early
N_iter = final_iter[1] #in case it terminated early

# EKI results: Has the ensemble collapsed toward the truth?
println("True parameters: ")
Expand Down Expand Up @@ -261,27 +268,27 @@ N_iter = final_iter[1] #in case it terminated early
params_true, #constrained here, as these are in a physically constrained space (unlike the u inputs),
)

# plots in constrained space
gr(dpi = 300, size = (400,400))
ϕ_init = EKP.get_ϕ(priors, ekiobj, 1)
p = plot(title = "EKI iterations", xlims = extrema(ϕ_init[1, :]), xaxis = "F", yaxis = "A", ylims = extrema(ϕ_init[2, :]))
for i in 1:N_iter
ϕ_i = EKP.get_ϕ(priors, ekiobj, i)
scatter!(p, ϕ_i[1, :], ϕ_i[2, :], color = :grey, label = false)
end
ϕ_fin = EKP.get_ϕ_final(priors, ekiobj)
scatter!(p, ϕ_fin[1, :], ϕ_fin[2, :], color = :magenta, label = false)

vline!(
p,
[params_true[1]],
linestyle = :dash,
linecolor = :red,
label = false,
# plots in constrained space
gr(dpi = 300, size = (400, 400))
ϕ_init = EKP.get_ϕ(priors, ekiobj, 1)
p = plot(
title = "EKI iterations",
xlims = extrema(ϕ_init[1, :]),
xaxis = "F",
yaxis = "A",
ylims = extrema(ϕ_init[2, :]),
)
hline!(p, [params_true[2]], linestyle = :dash, linecolor = :red, label=false)
savefig(p, joinpath(figure_save_directory,"calibration.png"))
savefig(p, joinpath(figure_save_directory,"calibration.pdf"))
for i in 1:N_iter
ϕ_i = EKP.get_ϕ(priors, ekiobj, i)
scatter!(p, ϕ_i[1, :], ϕ_i[2, :], color = :grey, label = false)
end
ϕ_fin = EKP.get_ϕ_final(priors, ekiobj)
scatter!(p, ϕ_fin[1, :], ϕ_fin[2, :], color = :magenta, label = false)

vline!(p, [params_true[1]], linestyle = :dash, linecolor = :red, label = false)
hline!(p, [params_true[2]], linestyle = :dash, linecolor = :red, label = false)
savefig(p, joinpath(figure_save_directory, "calibration.png"))
savefig(p, joinpath(figure_save_directory, "calibration.pdf"))
end

main()
35 changes: 14 additions & 21 deletions examples/Lorenz/emulate_sample.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,35 +45,26 @@ function main()
"RF-vector-nosvd-nondiag",
]

# I hope one day we can use just the in-built priors.
case_priors =[
nothing,

]
#### CHOOSE YOUR CASE:
mask = 1:7
prior_scalings = [
[0.0,0.0],
[1e-3, 0.0],
[1e-3, 0.0],
[1e-2, 1e-2],
[1e-2, 1e-2],
[1e-2, 1e-2],
[1e-2, 1e-2],
]
for (case,scaling) in zip(cases[mask], prior_scalings[mask])
mask = 2:7
# One day we can use good heuristics here
# Currently set so that learnt hyperparameters stays relatively far inside the prior. (use "verbose" => true in optimizer overrides to see this info)
prior_scalings = [[0.0, 0.0], [1e-2, 0.0], [1e-1, 0.0], [1e-2, 1e-2], [1e-2, 1e-1], [1e-2, 1e-2], [1e-2, 1e-2]]
for (case, scaling) in zip(cases[mask], prior_scalings[mask])

#case = cases[7]
println("case: ", case)
max_iter = 6 # number of EKP iterations to use data from is at most this
overrides = Dict(
"verbose" => true,
"train_fraction" => 0.8,
"n_iteration" => 40,
"scheduler" => DataMisfitController(terminate_at=20),
"scheduler" => DataMisfitController(),
"prior_in_scale" => scaling[1],
"prior_out_scale" => scaling[2],
)
# we do not want termination, as our priors have relatively little interpretation

# Should be loaded:
τc = 5.0
F_true = 8.0 # Mean F
Expand Down Expand Up @@ -108,7 +99,7 @@ function main()
truth_params = transform_constrained_to_unconstrained(priors, truth_params_constrained)
Γy = ekiobj.obs_noise_cov
println(Γy)

n_params = length(truth_params) # "input dim"
output_dim = size(Γy, 1)
###
Expand Down Expand Up @@ -140,7 +131,7 @@ function main()
# do we want to assume that the outputs are decorrelated in the machine-learning problem?
diagonalize_output = case ∈ ["RF-vector-svd-diag", "RF-vector-nosvd-diag"] ? true : false
n_features = 300

mlt = VectorRandomFeatureInterface(
n_features,
n_params,
Expand All @@ -162,8 +153,10 @@ function main()
println(norm_factor)

# Get training points from the EKP iteration number in the second input term
N_iter = length(get_u(ekiobj)) - 1 # number of paired iterations taken from EKP
N_iter = min(max_iter, length(get_u(ekiobj)) - 1) # number of paired iterations taken from EKP

input_output_pairs = Utilities.get_training_points(ekiobj, N_iter - 1) # 1:N-1

input_output_pairs_test = Utilities.get_training_points(ekiobj, N_iter:N_iter) # just "next" iteration used for testing
# Save data
@save joinpath(data_save_directory, "input_output_pairs.jld2") input_output_pairs
Expand Down
Loading

0 comments on commit 62a6a02

Please sign in to comment.