Skip to content

Add FlashMD energy rescaling to fix metatomic#51

Draft
johannes-spies wants to merge 1 commit into
metatomicfrom
flashmd-energy-rescaling
Draft

Add FlashMD energy rescaling to fix metatomic#51
johannes-spies wants to merge 1 commit into
metatomicfrom
flashmd-energy-rescaling

Conversation

@johannes-spies

Copy link
Copy Markdown

Implements the energy-conservation enforcement from App. C of the FlashMD paper (arXiv:2505.19350): with rescale_energy on, the predicted momenta are rescaled after each step by alpha = sqrt(1 - (E'-E)/K') so the total energy matches its pre-step value.

The potential energy U is taken from a pair_style metatomic energy model on top of the fix, read every step via an internal compute pe (mirroring how fix npt obtains the virial). The rescale is applied in post_force, i.e. after the FlashMD map and before the thermostat half-step, matching i-PI's operator splitting.

Verified on a 108-atom fcc Al NVE example: total energy drifts ~30 meV/atom without rescaling and is flat to machine precision with it, consistent with an i-PI reference and with an independent recomputation of the energy from dumped frames. Example inputs added under examples/PACKAGES/metatomic.

fix metatomic/kk overrides initial_integrate/post_force, so the host-side rescale does not run there; it now errors out instead of silently doing nothing. On-device Kokkos support is left as follow-up.

Implements the energy-conservation enforcement from App. C of the FlashMD
paper (arXiv:2505.19350): with `rescale_energy on`, the predicted momenta
are rescaled after each step by alpha = sqrt(1 - (E'-E)/K') so the total
energy matches its pre-step value.

The potential energy U is taken from a `pair_style metatomic` energy model
on top of the fix, read every step via an internal `compute pe` (mirroring
how fix npt obtains the virial). The rescale is applied in post_force, i.e.
after the FlashMD map and before the thermostat half-step, matching i-PI's
operator splitting.

Verified on a 108-atom fcc Al NVE example: total energy drifts ~30 meV/atom
without rescaling and is flat to machine precision with it, consistent with
an i-PI reference and with an independent recomputation of the energy from
dumped frames. Example inputs added under examples/PACKAGES/metatomic.

fix metatomic/kk overrides initial_integrate/post_force, so the host-side
rescale does not run there; it now errors out instead of silently doing
nothing. On-device Kokkos support is left as follow-up.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
@Luthaf

Luthaf commented Jun 24, 2026

Copy link
Copy Markdown
Member

I don't love doing this inside the engine. Would it make sense/be possible to implement this by having the FlashMD model request the energy as an extra input, and do the rescaling internally?

@johannes-spies

Copy link
Copy Markdown
Author

Hm, can the models have state/remember the last energy? Because that's what's needed here.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants