Updated accelerators with variable timestep, add localization case to…

… examples (#348) * Accelerators remove docstring format removed ratio of timsteps new example cases, plotting etc added ConstantNesterov, renamed ConstantStep to FirstOrder ConstantStep -> FirstOrder const -> firstorder in legends * bugfix off-by-one error * format * add GC * format
CliMA · Nov 16, 2023 · a7126a8 · a7126a8
1 parent 3ae43fb
commit a7126a8
Show file tree

Hide file tree

Showing 5 changed files with 248 additions and 101 deletions.
diff --git a/examples/Accelerators/DarcyAccelerated/darcy_accelerated.jl b/examples/Accelerators/DarcyAccelerated/darcy_accelerated.jl
@@ -44,8 +44,23 @@ function main()
     seed = 100234
     rng = Random.MersenneTwister(seed)
 
-    USE_SCHEDULER = false
-    scheduler_def = DataMisfitController(on_terminate = "continue")
+    cases = ["const", "dmc", "dmc-loc-small-ens"]
+    case = cases[3]
+
+    @info "running case $case"
+    if case == "const"
+        scheduler = DefaultScheduler()
+        N_ens = 52 # dofs+2  
+        localization_method = EKP.Localizers.NoLocalization()
+    elseif case == "dmc"
+        scheduler = DataMisfitController(terminate_at = 1e4)
+        N_ens = 52  # dofs+2
+        localization_method = EKP.Localizers.NoLocalization()
+    elseif case == "dmc-loc-small-ens"
+        scheduler = DataMisfitController(terminate_at = 1e4)
+        N_ens = 10
+        localization_method = EKP.Localizers.SEC(1.0, 0.01)
+    end
 
     # Define the spatial domain and discretization 
     dim = 2
@@ -57,7 +72,6 @@ function main()
     smoothness = 1.0
     corr_length = 0.25
     dofs = 50
-
     grf = GRF.GaussianRandomField(
         GRF.CovarianceFunction(dim, GRF.Matern(smoothness, corr_length)),
         GRF.KarhunenLoeve(dofs),
@@ -91,19 +105,16 @@ function main()
     # We define some algorithm parameters, here we take ensemble members larger than the dimension of the parameter space
     N_ens = dofs + 2    # number of ensemble members
     N_iter = 20         # number of EKI iterations
-    N_trials = 5       # number of trials 
+    N_trials = 10       # number of trials 
+    @info "obtaining statistics over $N_trials trials"
 
     errs = zeros(N_trials, N_iter)
     errs_acc = zeros(N_trials, N_iter)
     errs_acc_cs = zeros(N_trials, N_iter)
 
-    for trial in 1:N_trials
-        if USE_SCHEDULER
-            scheduler = scheduler_def
-        else
-            scheduler = DefaultScheduler(0.1)
-        end
+    for (idx, trial) in enumerate(1:N_trials)
 
+        @info "computing trial $idx"
         # We sample the initial ensemble from the prior, and create three EKP objects to 
         # perform EKI algorithm using three different acceleration methods.
         initial_params = construct_initial_ensemble(rng, prior, N_ens)
@@ -113,6 +124,7 @@ function main()
             obs_noise_cov,
             Inversion(),
             scheduler = deepcopy(scheduler),
+            localization_method = deepcopy(localization_method),
         )
         ekiobj_acc = EKP.EnsembleKalmanProcess(
             initial_params,
@@ -121,14 +133,16 @@ function main()
             Inversion(),
             accelerator = NesterovAccelerator(),
             scheduler = deepcopy(scheduler),
+            localization_method = deepcopy(localization_method),
         )
         ekiobj_acc_cs = EKP.EnsembleKalmanProcess(
             initial_params,
             truth_sample,
             obs_noise_cov,
             Inversion(),
-            accelerator = ConstantStepNesterovAccelerator(),
+            accelerator = FirstOrderNesterovAccelerator(),
             scheduler = deepcopy(scheduler),
+            localization_method = deepcopy(localization_method),
         )
 
         # Run EKI algorithm, recording parameter error after each iteration.
@@ -156,14 +170,16 @@ function main()
             errs_acc_cs[trial, :] = err_acc_cs
         end
 
+        GC.gc()
+
     end
 
 
-    # compare recorded convergences with default, Nesterov, and Constant-Step Nesterov accelerators
+    # compare recorded convergences with default, Nesterov, and First-Order Nesterov accelerators
     gr(legend = true)
     conv_plot = plot(1:N_iter, mean((errs), dims = 1)[:], label = "No acceleration", color = "black")
     plot!(1:N_iter, mean((errs_acc), dims = 1)[:], label = "Nesterov", color = "blue")
-    plot!(1:N_iter, mean((errs_acc_cs), dims = 1)[:], label = "Nesterov, Constant Step", color = "red")
+    plot!(1:N_iter, mean((errs_acc_cs), dims = 1)[:], label = "Nesterov, FirstOrder", color = "red")
     # error bars
     plot!(
         1:N_iter,
@@ -211,7 +227,7 @@ function main()
     title!("EKI convergence on Darcy IP")
     xlabel!("Iteration")
     ylabel!("log(Error)")
-    savefig(conv_plot, joinpath(fig_save_directory, "darcy_conv_comparison.png"))
+    savefig(conv_plot, joinpath(fig_save_directory, case * "_darcy_conv_comparison.png"))
 end
 
 main()
diff --git a/examples/Accelerators/ExpSinAccelerated/exp_sinusoid_accelerated.jl b/examples/Accelerators/ExpSinAccelerated/exp_sinusoid_accelerated.jl
@@ -17,7 +17,10 @@ using Distributions, Plots
 using EnsembleKalmanProcesses
 using EnsembleKalmanProcesses.ParameterDistributions
 const EKP = EnsembleKalmanProcesses
-fig_save_directory = @__DIR__
+fig_save_directory = joinpath(@__DIR__, "output")
+if !isdir(fig_save_directory)
+    mkdir(fig_save_directory)
+end
 
 ## Setting up the model and data for our inverse problem
 dt = 0.01
@@ -40,8 +43,24 @@ function G(u)
 end
 
 function main()
-    USE_SCHEDULER = false
-    scheduler_def = DataMisfitController(on_terminate = "continue")
+
+    cases = ["const", "dmc", "dmc-loc-small-ens"]
+    case = cases[3]
+
+    @info "running case $case"
+    if case == "const"
+        scheduler = DefaultScheduler()
+        N_ens = 10
+        localization_method = EKP.Localizers.NoLocalization()
+    elseif case == "dmc"
+        scheduler = DataMisfitController(terminate_at = 1e4)
+        N_ens = 10
+        localization_method = EKP.Localizers.NoLocalization()
+    elseif case == "dmc-loc-small-ens"
+        scheduler = DataMisfitController(terminate_at = 1e4)
+        N_ens = 4
+        localization_method = EKP.Localizers.SEC(1.0, 0.01)
+    end
 
     dim_output = 2
     Γ = 0.01 * I
@@ -58,10 +77,9 @@ function main()
 
     # To compare the two EKI methods, we will average over several trials, 
     # allowing the methods to run with different initial ensembles and noise samples.
-    N_ens = 10
-    N_iterations = 15
+    N_iterations = 20
     N_trials = 50
-
+    @info "obtaining statistics over $N_trials trials"
     # Define cost function to compare convergences. We use a logarithmic cost function 
     # to best interpret exponential model. Note we do not explicitly penalize distance from the prior here.
     function cost(theta)
@@ -80,35 +98,34 @@ function main()
         # to compare convergence.
         initial_ensemble = EKP.construct_initial_ensemble(rng, prior, N_ens)
 
-        if USE_SCHEDULER
-            scheduler = deepcopy(scheduler_def)
-            scheduler_acc = deepcopy(scheduler_def)
-            scheduler_acc_cs = deepcopy(scheduler_def)
-        else
-            scheduler = DefaultScheduler()
-            scheduler_acc = DefaultScheduler()
-            scheduler_acc_cs = DefaultScheduler()
-        end
-
-        ensemble_kalman_process =
-            EKP.EnsembleKalmanProcess(initial_ensemble, y, Γ, Inversion(); rng = rng, scheduler = scheduler)
+        ensemble_kalman_process = EKP.EnsembleKalmanProcess(
+            initial_ensemble,
+            y,
+            Γ,
+            Inversion();
+            rng = rng,
+            scheduler = deepcopy(scheduler),
+            localization_method = deepcopy(localization_method),
+        )
         ensemble_kalman_process_acc = EKP.EnsembleKalmanProcess(
             initial_ensemble,
             y,
             Γ,
             Inversion();
             accelerator = NesterovAccelerator(),
             rng = rng,
-            scheduler = scheduler_acc,
+            scheduler = deepcopy(scheduler),
+            localization_method = deepcopy(localization_method),
         )
         ensemble_kalman_process_acc_cs = EKP.EnsembleKalmanProcess(
             initial_ensemble,
             y,
             Γ,
             Inversion();
-            accelerator = ConstantStepNesterovAccelerator(),
+            accelerator = FirstOrderNesterovAccelerator(),
             rng = rng,
-            scheduler = scheduler_acc_cs,
+            scheduler = deepcopy(scheduler),
+            localization_method = deepcopy(localization_method),
         )
 
         global convs = zeros(N_iterations)
@@ -141,10 +158,10 @@ function main()
         all_convs_acc_cs[trial, :] = convs_acc_cs
     end
 
-    gr(size = (400, 400), legend = true)
+    gr(size = (600, 600), legend = true)
     p = plot(1:N_iterations, mean(all_convs, dims = 1)[:], color = :black, label = "No acceleration")
     plot!(1:N_iterations, mean(all_convs_acc, dims = 1)[:], color = :blue, label = "Nesterov")
-    plot!(1:N_iterations, mean(all_convs_acc_cs, dims = 1)[:], color = :red, label = "Nesterov, Constant Step")
+    plot!(1:N_iterations, mean(all_convs_acc_cs, dims = 1)[:], color = :red, label = "Nesterov, FirstOrder")
     # error bars
     plot!(
         1:N_iterations,
@@ -194,7 +211,7 @@ function main()
     ylabel!("log(Cost)")
     title!("EKI convergence on exp sin IP")
 
-    savefig(p, joinpath(fig_save_directory, "exp_sin.png"))
+    savefig(p, joinpath(fig_save_directory, case * "_exp_sin.png"))
 end
 
 main()
diff --git a/examples/Accelerators/LorenzAccelerated/lorenz_accelerated.jl b/examples/Accelerators/LorenzAccelerated/lorenz_accelerated.jl
@@ -10,7 +10,10 @@ using EnsembleKalmanProcesses.Localizers
 import EnsembleKalmanProcesses: construct_mean, construct_cov, construct_sigma_ensemble
 const EKP = EnsembleKalmanProcesses
 
-fig_save_directory = @__DIR__
+fig_save_directory = joinpath(@__DIR__, "output")
+if !isdir(fig_save_directory)
+    mkdir(fig_save_directory)
+end
 
 function rk4(f::F, y0::Array{Float64, 1}, t0::Float64, t1::Float64, h::Float64; inplace::Bool = true) where {F}
     y = y0
@@ -53,8 +56,24 @@ end
 
 function main()
     D = 20
-    USE_SCHEDULER = false
-    scheduler_def = DataMisfitController(on_terminate = "continue")
+
+    cases = ["const", "dmc", "dmc-loc-small-ens"]
+    case = cases[3]
+
+    @info "running case $case"
+    if case == "const"
+        scheduler = DefaultScheduler()
+        N_ens = 20
+        localization_method = EKP.Localizers.NoLocalization()
+    elseif case == "dmc"
+        scheduler = DataMisfitController(terminate_at = 1e4)
+        N_ens = 20
+        localization_method = EKP.Localizers.NoLocalization()
+    elseif case == "dmc-loc-small-ens"
+        scheduler = DataMisfitController(terminate_at = 1e4)
+        N_ens = 5
+        localization_method = EKP.Localizers.SEC(1.0, 0.01)
+    end
 
     lorenz96_sys = (t, u) -> lorenz96(t, u, Dict("N" => D))
 
@@ -77,45 +96,48 @@ function main()
     end
     prior = combine_distributions(priors)
 
-    N_ens = 20
+
     N_iter = 20
     N_trials = 50
+    @info "obtaining statistics over $N_trials trials"
+
     errs = zeros(N_trials, N_iter)
     errs_acc = zeros(N_trials, N_iter)
     errs_acc_cs = zeros(N_trials, N_iter)
 
     for trial in 1:N_trials
         initial_ensemble = EKP.construct_initial_ensemble(rng, prior, N_ens)
 
-        if USE_SCHEDULER
-            scheduler_van = deepcopy(scheduler_def)
-            scheduler_acc = deepcopy(scheduler_def)
-            scheduler_acc_cs = deepcopy(scheduler_def)
-        else
-            scheduler_van = DefaultScheduler()
-            scheduler_acc = DefaultScheduler()
-            scheduler_acc_cs = DefaultScheduler()
-        end
+
         # We create 3 EKP Inversion objects to compare acceleration.
-        ekiobj_vanilla =
-            EKP.EnsembleKalmanProcess(initial_ensemble, y, Γ, Inversion(); rng = rng, scheduler = scheduler_van)
+        ekiobj_vanilla = EKP.EnsembleKalmanProcess(
+            initial_ensemble,
+            y,
+            Γ,
+            Inversion();
+            rng = rng,
+            scheduler = deepcopy(scheduler),
+            localization_method = deepcopy(localization_method),
+        )
         ekiobj_acc = EKP.EnsembleKalmanProcess(
             initial_ensemble,
             y,
             Γ,
             Inversion();
             rng = rng,
             accelerator = NesterovAccelerator(),
-            scheduler = scheduler_acc,
+            scheduler = deepcopy(scheduler),
+            localization_method = deepcopy(localization_method),
         )
         ekiobj_acc_cs = EKP.EnsembleKalmanProcess(
             initial_ensemble,
             y,
             Γ,
             Inversion();
             rng = rng,
-            accelerator = ConstantStepNesterovAccelerator(),
-            scheduler = scheduler_acc_cs,
+            accelerator = FirstOrderNesterovAccelerator(),
+            scheduler = deepcopy(scheduler),
+            localization_method = deepcopy(localization_method),
         )
 
         err = zeros(N_iter)
@@ -138,8 +160,9 @@ function main()
 
     convplot = plot(1:(N_iter), mean(errs, dims = 1)[:], color = :black, label = "No acceleration")
     plot!(1:(N_iter), mean(errs_acc, dims = 1)[:], color = :blue, label = "Nesterov")
-    plot!(1:(N_iter), mean(errs_acc_cs, dims = 1)[:], color = :red, label = "Nesterov, Constant Step")
+    plot!(1:(N_iter), mean(errs_acc_cs, dims = 1)[:], color = :red, label = "Nesterov, FirstOrder")
     title!("EKI convergence on Lorenz96 IP, N_trials=" * string(N_trials))
+    xlabel!("Iteration")
 
     # error bars
     plot!(
@@ -185,7 +208,7 @@ function main()
         label = "",
     )
 
-    savefig(convplot, joinpath(fig_save_directory, "lorenz96.png"))
+    savefig(convplot, joinpath(fig_save_directory, case * "_lorenz96.png"))
 end
 
 main()