WIP: adding momentum-inspired accelerators to EKP

[sydneyvernon]

Purpose

Closes #323

Content

---

• I have read and checked the items on the review checklist.

Current results from nonlinear inversion unit tests:

Accelerator: DefaultAccelerator Process: Inversion
┌ Info: Convergence:
│   cost_initial = 55.61433558387613
└   cost_final = 8.644475010222061
Accelerator: NesterovAccelerator Process: Inversion
┌ Info: Convergence:
│   cost_initial = 55.32972049471293
└   cost_final = 12.14289870253091
Accelerator: DefaultAccelerator Process: Sampler
┌ Info: Convergence:
│   cost_initial = 56.12630362217311
└   cost_final = 13.797245580534
Accelerator: NesterovAccelerator Process: Sampler
┌ Info: Convergence:
│   cost_initial = 58.07591023338068
└   cost_final = 7.951552079718973
Accelerator: DefaultAccelerator Process: TransformInversion
┌ Info: Convergence:
│   cost_initial = 58.242226993212626
└   cost_final = 5.25456484944421
Accelerator: NesterovAccelerator Process: TransformInversion
┌ Info: Convergence:
│   cost_initial = 58.64590869107448
└   cost_final = 4.183960610134597

[odunbar]

Looks great so far - just adding some hints as you go.

As you may have seen the Unscented version is a bit trickier to parse, so we can look at how to get an interface working here later.

[sydneyvernon]

Last commit shows an initial version (not yet functional) of what NesterovAccelerator could look like. Notes:

• incorporating additional states: I am trying to be as noninvasive as possible and keep my contributions contained in Accelerators.jl. In the current framework, I think I would have to change existing (widely used) functions get_u or update_ensemble if I wanted to implement momentum EKI exactly as we derived it, because momentum EKI follows 2 states (v and u), and evaluates G on v but returns u as the ensemble value. However, I don't see any problem with accepting v as the "ensemble value" and storing information on u for acceleration purposes only. I looked at the convergence of the v state in my smaller project/examples, and it seems to work just fine (?). Only issue is that I will have to be mindful of potentially shifting my iteration index.
• timestep: again hoping to keep things contained to Accelerators.jl: EKI and momentum EKI apply virtually the same update, but scale by dt and dt^2 respectively. For now I haven't addressed this since for individual applications, if we have optimized timesteppers in place, I don't think this is significant(?). I think the issue may come when comparing methods, though.

[odunbar]

First, if you have addressed my comments above, feel free to "Resolve conversation" on them

On your points

• (1) is indeed what is described in O3.7.1 Momentum-accelerated EKP #321, that is, the ensemble update produces new_ens (the old "stored state") but this now gets moved by the accelerator before being stored as ekp.u. So following this structure should produce the desired result
• (2) excellent point about the timestep. I think it would be good to be consistent, can you add this into O3.7.1 Momentum-accelerated EKP #321 as a comment, rather than dealing with it in this PR.

[eviatarbach]

Yes, good point about the timestep. Note also that the timestep is implemented in EKP as a rescaling of the observation error covariance; please check the consistency of this with the timestep in momentum EKI.

[odunbar]

EDIT: Moved to a new issues for a different PR

Some notes regarding UKI. It saves only the mean and cov of the ensemble, in process.u_mean and process.uu_cov
A possible strategy (I'd be happy to implement if you are tight on time etc.)

• Make these returned, not saved, (as we did with other method ensembles). So for example UKI will return u_mean, uu_cov, and also the sigma points. Note the whole triple return values will get stored in "u"
• in the current update_state!(ekp,u) we define a new update method for UKI typed ekp. This shifts the mean, and contracts the covariance appropriately as with the acceleration. It then saves them in the process object, and redraws the new ensemble ("sigma points") and saves these.
• you will also need to add a new update method for default accelerator with UKI typed ekp too. This just saves the mean,cov and sigma ensemble

[odunbar]

Hi Sydney a couple of comments to help us to the finish line on this one

[odunbar]

To be safe, I would create a new process object for each test. There is no guarantee that they aren't modified during the calls to EKP updates.

[odunbar]

LGTM! Thanks for all the work!

[eviatarbach]

This looks great Sydney! Thank you for all the hard work!

It almost looks ready to go. Could you please just verify that TransformInversion works with acceleration, and add a test for it?

I would also prefer to add a warning that states that acceleration for EKS is experimental, although still allowing the user to do it. What do you think @odunbar?

[odunbar]

So we did add a unit test for ETKI (see L116 in the added runtests file) which seems to run and converge, so i think we are covered there. I think we can defer the actual acceleration experiment for the next PR (with proper examples).

For EKS, this is a good point. Sydney, could you add a warning for when someone creates the accelerator and the EKP object is of type Sampler? Just saying that it is an experimental feature and may affect behaviour

[eviatarbach]

LGTM!

[odunbar]

bors r+

[bors]

Build succeeded!

The publicly hosted instance of bors-ng is deprecated and will go away soon.

If you want to self-host your own instance, instructions are here.
For more help, visit the forum.

If you want to switch to GitHub's built-in merge queue, visit their help page.

• docs-build
• format
• test_success