Fixes and optimizations to jax-md backend

[em819]

I have performed following fixes concerning the jax-md backend of PySAGES:

1. Changed the point in _step() at which the snapshot is updated. This addresses the corner case of performing a single MD step. In the previous implementation after a single MD step, the snapshot would still contain the original state (i.e the 0th state). Thus if one would repeatedly restart with single-step-runs, the trajectory of snapshots would not change. This is now fixed.
2. Extended the contents of the Snapshot class to also optionally contain thermostat information, which by jax-md is saved in state.chain.
3. Rewritten run-method in jax-md backend to optionally use jax.lax.fori_loop() if jit_compile=True. In addition, now the step and run function is defined/instantiated only once even if pysages.run is called multiple times in succession. If pysages.run is called repeatedly these functions are now reused and thus don't need to be jit-recompiled. Particularly for smaller systems with less than 10000 atoms, I have observed a speed-up of up to factor 10 compared to the original version.

[trunk-io]

Merging to main in this repository is managed by Trunk.

• To merge this pull request, check the box to the left or comment /trunk merge below.

[p-zubieta]

Hi @em819. Thank you interest and for the contribution to the repo! I'm sorry we've been slow to reply here. I'll try to review the code soon and provide any appropriate feedback.

[em819]

Hey @p-zubieta ,
Thanks for having a look, much appreciated.

[em819]

Hi there @p-zubieta . Just wanted to ask if I should take action in resolving the conflicts and merging or you prefer to do it on your end? Don't want to inappropriately interfere with the pull request.

[p-zubieta]

I was thinking on doing it on our end, cherry-picking your commits. We did some internal changes to allow for the inclusion of additional data into Snapshots a bit easier and want to move this PR on top of those.

[em819]

Sure thing, sounds good.