WIP: adding momentum-inspired accelerators to EKP

Project.toml

@@ -40,4 +40,4 @@ StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["StableRNGs", "Test", "Plots"]
+test = ["StableRNGs", "Test", "Plots"]

src/Accelerators.jl

@@ -0,0 +1,94 @@
+# included in EnsembleKalmanProcess.jl
+
+export DefaultAccelerator, NesterovAccelerator
+export update_state!, set_initial_acceleration!
+
+"""
+$(TYPEDEF)
+
+Default accelerator provides no acceleration, runs traditional EKI
+"""
+struct DefaultAccelerator <: Accelerator end
+
+"""
+$(TYPEDEF)
+
+Accelerator that adapts Nesterov's momentum method for EKI. 
+Stores a previous state value u_prev for computational purposes (note this is distinct from state returned as "ensemble value")
+
+$(TYPEDFIELDS)
+"""
+mutable struct NesterovAccelerator{FT <: AbstractFloat} <: Accelerator
+    r::FT
+    u_prev::Any
+end
+
+function NesterovAccelerator(r = 3.0, initial = Float64[])
+    return NesterovAccelerator(r, initial)
+end
+
+
+"""
+Sets u_prev to the initial parameter values
+"""
+function set_ICs!(accelerator::NesterovAccelerator{FT}, u::MA) where {FT <: AbstractFloat, MA <: AbstractMatrix{FT}}
+    accelerator.u_prev = u
+end
+
+
+"""
+Performs traditional state update with no momentum. 
+"""
+function update_state!(
+    ekp::EnsembleKalmanProcess{FT, IT, P, LRS, DefaultAccelerator},
+    u::MA,
+) where {FT <: AbstractFloat, IT <: Int, P <: Process, LRS <: LearningRateScheduler, MA <: AbstractMatrix{FT}}
+    push!(ekp.u, DataContainer(u, data_are_columns = true))
+end
+
+"""
+Performs state update with modified Nesterov momentum approach.
+"""
+function update_state!(
+    ekp::EnsembleKalmanProcess{FT, IT, P, LRS, NesterovAccelerator{FT}},
+    u::MA,
+) where {FT <: AbstractFloat, IT <: Int, P <: Process, LRS <: LearningRateScheduler, MA <: AbstractMatrix{FT}}
+    ## update "v" state:
+    k = get_N_iterations(ekp) + 2
+    v = u .+ (1 - ekp.accelerator.r / k) * (u .- ekp.accelerator.u_prev)
+
+    ## update "u" state: 
+    ekp.accelerator.u_prev = u
+
+    ## push "v" state to EKP object
+    push!(ekp.u, DataContainer(v, data_are_columns = true))
+end
+
+
+"""
+State update method for UKI with no acceleration.
+The Accelerator framework has not yet been integrated with UKI process;
+UKI tracks its own states, so this method is empty.
+"""
+function update_state!(
+    ekp::EnsembleKalmanProcess{FT, IT, P, LRS, DefaultAccelerator},
+    u::MA,
+) where {FT <: AbstractFloat, IT <: Int, P <: Unscented, LRS <: LearningRateScheduler, MA <: AbstractMatrix{FT}}
+
+end
+
+"""
+Placeholder state update method for UKI with Nesterov Accelerator.
+The Accelerator framework has not yet been integrated with UKI process, so this
+method throws an error.
+"""
+function update_state!(
+    ekp::EnsembleKalmanProcess{FT, IT, P, LRS, NesterovAccelerator{FT}},
+    u::MA,
+) where {FT <: AbstractFloat, IT <: Int, P <: Unscented, LRS <: LearningRateScheduler, MA <: AbstractMatrix{FT}}
+    throw(
+        ArgumentError(
+            "option `accelerator = NesterovAccelerator` is not implemented for UKI, please use `DefaultAccelerator`",
+        ),
+    )
+end

src/EnsembleKalmanInversion.jl

@@ -138,17 +138,17 @@ function update_ensemble!(
 
     u = fh.failsafe_update(ekp, u, g, y, scaled_obs_noise_cov, failed_ens)
 
-    # store new parameters (and model outputs)
-    push!(ekp.u, DataContainer(u, data_are_columns = true))
     push!(ekp.g, DataContainer(g, data_are_columns = true))
 
     # Store error
     compute_error!(ekp)
 
     # Diagnostics
-    cov_new = cov(get_u_final(ekp), dims = 2)
+    cov_new = cov(u, dims = 2)
 
     if ekp.verbose
         @info "Covariance-weighted error: $(get_error(ekp)[end])\nCovariance trace: $(tr(cov_new))\nCovariance trace ratio (current/previous): $(tr(cov_new)/tr(cov_init))"
     end
+
+    return u
 end

src/EnsembleKalmanProcess.jl

@@ -13,7 +13,7 @@
 export get_u_prior, get_u_final, get_g_final, get_ϕ_final
 export get_N_iterations, get_error, get_cov_blocks
 export get_u_mean, get_u_cov, get_g_mean, get_ϕ_mean
-export get_u_mean_final, get_u_cov_prior, get_u_cov_final, get_g_mean_final, get_ϕ_mean_final
+export get_u_mean_final, get_u_cov_prior, get_u_cov_final, get_g_mean_final, get_ϕ_mean_final, get_accelerator
 export compute_error!
 export update_ensemble!
 export sample_empirical_gaussian, split_indices_by_success
@@ -29,6 +29,9 @@
 # Failure handlers
 abstract type FailureHandlingMethod end
 
+# Accelerators
+abstract type Accelerator end
+
 
 
 "Failure handling method that ignores forward model failures"
@@ -104,7 +107,13 @@
 
 $(METHODLIST)
 """
-struct EnsembleKalmanProcess{FT <: AbstractFloat, IT <: Int, P <: Process, LRS <: LearningRateScheduler}
+struct EnsembleKalmanProcess{
+    FT <: AbstractFloat,
+    IT <: Int,
+    P <: Process,
+    LRS <: LearningRateScheduler,
+    ACC <: Accelerator,
+}
     "array of stores for parameters (`u`), each of size [`N_par × N_ens`]"
     u::Array{DataContainer{FT}}
     "vector of the observed vector size [`N_obs`]"
@@ -119,6 +128,8 @@
     err::Vector{FT}
     "Scheduler to calculate the timestep size in each EK iteration"
     scheduler::LRS
+    "accelerator object that informs EK update steps, stores additional state variables as needed"
+    accelerator::ACC
     "stored vector of timesteps used in each EK iteration"
     Δt::Vector{FT}
     "the particular EK process (`Inversion` or `Sampler` or `Unscented` or `TransformInversion` or `SparseInversion`)"
@@ -139,6 +150,7 @@
     obs_noise_cov::Union{AbstractMatrix{FT}, UniformScaling{FT}},
     process::P;
     scheduler::Union{Nothing, LRS} = nothing,
+    accelerator::Union{Nothing, ACC} = nothing,
     Δt = nothing,
     rng::AbstractRNG = Random.GLOBAL_RNG,
     failure_handler_method::FM = IgnoreFailures(),
@@ -147,6 +159,7 @@
 ) where {
     FT <: AbstractFloat,
     LRS <: LearningRateScheduler,
+    ACC <: Accelerator,
     P <: Process,
     FM <: FailureHandlingMethod,
     LM <: LocalizationMethod,
@@ -193,23 +206,39 @@
     # timestep store
     Δt = FT[]
 
+    # set up accelerator
+    if isnothing(accelerator)
+        acc = DefaultAccelerator()
+    else
+        acc = accelerator
+    end
+    AC = typeof(acc)
+
+    if AC <: NesterovAccelerator
+        set_ICs!(acc, params)
+        if P <: Sampler
+            @warn "Acceleration is experimental for Sampler processes and may affect convergence."
+        end
+    end
+
     # failure handler
     fh = FailureHandler(process, failure_handler_method)
     # localizer
     loc = Localizer(localization_method, N_par, N_obs, N_ens, FT)
 
     if verbose
-        @info "Initializing ensemble Kalman process of type $(nameof(typeof(process)))\nNumber of ensemble members: $(N_ens)\nLocalization: $(nameof(typeof(localization_method)))\nFailure handler: $(nameof(typeof(failure_handler_method)))\nScheduler: $(nameof(typeof(lrs)))"
+        @info "Initializing ensemble Kalman process of type $(nameof(typeof(process)))\nNumber of ensemble members: $(N_ens)\nLocalization: $(nameof(typeof(localization_method)))\nFailure handler: $(nameof(typeof(failure_handler_method)))\nScheduler: $(nameof(typeof(lrs)))\nAccelerator: $(nameof(typeof(acc)))"
     end
 
-    EnsembleKalmanProcess{FT, IT, P, RS}(
+    EnsembleKalmanProcess{FT, IT, P, RS, AC}(
         [init_params],
         obs_mean,
         obs_noise_cov,
         N_ens,
         g,
         err,
         lrs,
+        acc,
         Δt,
         process,
         rng,
@@ -222,7 +251,6 @@
 
 include("LearningRateSchedulers.jl")
 
-
 """
     get_u(ekp::EnsembleKalmanProcess, iteration::IT; return_array=true) where {IT <: Integer}
 
@@ -423,6 +451,14 @@
     return ekp.scheduler
 end
 
+"""
+    get_accelerator(ekp::EnsembleKalmanProcess)
+Return accelerator type of EnsembleKalmanProcess.
+"""
+function get_accelerator(ekp::EnsembleKalmanProcess)
+    return ekp.accelerator
+end
+
 
 """
     construct_initial_ensemble(
@@ -628,7 +664,8 @@
 
     terminate = calculate_timestep!(ekp, g, Δt_new)
     if isnothing(terminate)
-        update_ensemble!(ekp, g, get_process(ekp); ekp_kwargs...)
+        u = update_ensemble!(ekp, g, get_process(ekp); ekp_kwargs...)
+        update_state!(ekp, u)
         if s > 0.0
             multiplicative_inflation ? multiplicative_inflation!(ekp; s = s) : nothing
             additive_inflation ? additive_inflation!(ekp; use_prior_cov = use_prior_cov, s = s) : nothing
@@ -664,3 +701,6 @@
 export Gaussian_2d
 export construct_initial_ensemble, construct_mean, construct_cov
 include("UnscentedKalmanInversion.jl")
+
+# struct Accelerator
+include("Accelerators.jl")

src/EnsembleKalmanSampler.jl

@@ -139,17 +139,18 @@ function update_ensemble!(
     u = fh.failsafe_update(ekp, u_old, g, failed_ens)
 
     # store new parameters (and model outputs)
-    push!(ekp.u, DataContainer(u, data_are_columns = true))
     push!(ekp.g, DataContainer(g, data_are_columns = true))
     # u_old is N_ens × N_par, g is N_ens × N_obs,
     # but stored in data container with N_ens as the 2nd dim
 
     compute_error!(ekp)
 
     # Diagnostics
-    cov_new = get_u_cov_final(ekp)
+    cov_new = cov(u, dims = 2)
 
     if ekp.verbose
         @info "Covariance-weighted error: $(get_error(ekp)[end])\nCovariance trace: $(tr(cov_new))\nCovariance trace ratio (current/previous): $(tr(cov_new)/tr(cov_init))"
     end
+
+    return u
 end

src/EnsembleTransformKalmanInversion.jl

@@ -124,16 +124,17 @@ function update_ensemble!(
     u = fh.failsafe_update(ekp, u, g, y, scaled_obs_noise_cov, failed_ens)
 
     # store new parameters (and model outputs)
-    push!(ekp.u, DataContainer(u, data_are_columns = true))
     push!(ekp.g, DataContainer(g, data_are_columns = true))
 
     # Store error
     compute_error!(ekp)
 
     # Diagnostics
-    cov_new = cov(get_u_final(ekp), dims = 2)
+    cov_new = cov(u, dims = 2)
 
     if ekp.verbose
         @info "Covariance-weighted error: $(get_error(ekp)[end])\nCovariance trace: $(tr(cov_new))\nCovariance trace ratio (current/previous): $(tr(cov_new)/tr(cov_init))"
     end
+
+    return u
 end

src/SparseEnsembleKalmanInversion.jl

@@ -222,12 +222,13 @@ function update_ensemble!(
     u = fh.failsafe_update(ekp, u, g, y, scaled_obs_noise_cov, failed_ens)
 
     # store new parameters (and model outputs)
-    push!(ekp.u, DataContainer(u, data_are_columns = true))
     push!(ekp.g, DataContainer(g, data_are_columns = true))
 
     # Store error
     compute_error!(ekp)
 
     # Check convergence
     cov_new = cov(get_u_final(ekp), dims = 2)
+
+    return u
 end