Refactor integrators.py to make it more general. (#589)

* Refactor integrators.py to make it more general.
Also add momentum update based on Esh dynamics

Co-authored-by: Reuben Cohn-Gordon <reubenharry@gmail.com>

* Refactor to use integrator generation functions

* Additional refactoring

Also add test for esh momentum update.

Co-authored-by: Reuben Cohn-Gordon <reubenharry@gmail.com>

* Minor clean up.

* Use standard JAX ops

* Adding a test for energy preservation.

Co-authored-by: Reuben Cohn-Gordon <reubenharry@gmail.com>

* fix formatting

---------

Co-authored-by: Reuben Cohn-Gordon <reubenharry@gmail.com>