Access to neighborlist after JAX-MD backend finishes

[em819]

@pabloferz This sounds like a very basic functionality that is probably implemented, but after looking through the code I still fail to find it.

I intend to run my MDs with jax-md+PySAGES. After the enhanced sampling simulation finishes, i.e after pysages.run(), I need to have access to the updated neighborlists and (in case of an NPT simulation) to the box-parameters. However, the Results class does not seem to contain these objects.

Is there a way to get the updated neighborlist after PySAGES+jax-md finishes?

[em819]

Update: I figured out that the snapshots-object of the results contains most of the information, like positions, velocities and masses, from which I can also recompute the neighbor lists.

The last piece, that is missing to reconstruct the NVTNoseHooverState from jax-md is the chain-information including the chain-positions and chain-momentum.

In that sense I would like to update my question: Can I access the thermostat parameters after PySAGES returns control to jax-md?