O3.7.1 Momentum-accelerated EKP

[odunbar]

- [ ] https://github.com/CliMA/EnsembleKalmanProcesses.jl/issues/323
- [ ] https://github.com/CliMA/EnsembleKalmanProcesses.jl/issues/340
- [ ] https://github.com/CliMA/EnsembleKalmanProcesses.jl/issues/339
- [ ] https://github.com/CliMA/EnsembleKalmanProcesses.jl/issues/333
- [ ] add documentation for nonlinear examples.
- [ ] Check Implementation for Sparse EKI and unit tests
- [ ] Implementation of EKS acceleration based on Kalman-Wasserstein gradient flow
- [ ] Failure Handler compatability

Offline we have been experimenting with a few momentum-based accelerations of EKI with good success. Integration of such approaches into the code for application with general EKP-type methods is detailed below. In general we have identified a rough template of what such accelerators will look like:

# type for flow accelerators
abstract type Accelerator

# e.g. the Nesterov type contains the parameter and a history of states
Nesterov <: Accelerator 
    momentum_parameter
    state_history
end

# it lives in the general EKP object
struct EnsembleKalmanProcess{..., A<:Accelerator}
    ...
    accelerator::A
end

# the state history is intialized with the initial ensemble
function initial_ensemble(...)
    ...
    ekp.accelerator.state_history = initial_ensemble 
end

# the update involves updating the particles with the usual method, along with the augmented state
function update_ensemble(g)
    new_ens= update_ensemble!(ekp.u[end],g,ekp.process)  # update to perform an approximate gradient step 
    new_momentum_ens = accelerator_update(ekp.accelerator, new_ens) # shift new ensemble by accelerator
    append!(ekp.u, new_momentum_ens)
    ekp.accelerator.state_history = new_ens
    ...
return new_momentum_ens
end

It is also important that the saved pairs are the momentum-shifted ensemble members and the evaluations of them (i.e. v and G(v) remain consistent). This works because accelerators typically still involve the calculation of a typical gradient-based update, and are simply followed by a perturbation of the particles by momentum.

Hopefully UKI,EKI,Transform-EKI,SEKI, (maybe even EKS, which also represents a gradient flow) can be all be updated with the same interface.

Happy to iterate a bit here.

[sydneyvernon]

Note on timestepping in momentum-accelerated EKP methods: Our goal is to implement momentum-inspired accelerators without changing existing update methods, as detailed above. EKI and "momentum-accelerated" EKI apply virtually the same update, but scale by dt and dt^2 respectively. It is possible that this distinction won't cause a problem in individual applications, since we have optimized timesteppers in place. That said I will need to verify that the currently-implemented EKP timestep (related to the observation error covariance) makes sense in the accelerated context.