Add symplectic integrator#43
Conversation
Pass model=(flashmd_model, (symplectic_model, config_dict)) to use the fixed-point midpoint-rule integrator instead of plain FlashMD. The plain single-model API is completely unchanged. Works in both i-PI and ASE (VelocityVerlet, Bussi, Langevin). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
The `._nl.algorithm = "cell_list"` workaround guards against a vesin CUDA bug with triclinic cells (vesin#157), but accesses a private attribute. It was not present in the original PR#32 design, and verified to be unnecessary on CPU (the only target for now). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…per" This reverts commit 7c1aa52.
frostedoyster
left a comment
frostedoyster
left a comment
There was a problem hiding this comment.
Looks good. Can we also support (direct_model, symplectic_model) without the symplectic config? It should be trivial. We also need tests. For that, we will need to upload the symplectic models to the HF repository. We also need some documentation
Allow passing the symplectic model directly as the second tuple element, without wrapping it in a (model, config) pair. Defaults to empty config. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
VelocityVerlet assumed the symplectic model was always passed as (model, config); passing just a plain model crashed with TypeError. Also adds pytest filterwarnings for torch.jit and metatomic 0.1.12+ deprecations (renamed output names/quantities) that now surface with newer dependency versions. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Extend get_pretrained() with symplectic=True to download and export the symplectic FlashMD checkpoint from HuggingFace. Add tests covering both the plain-model and (model, config) tuple API. Metatrain is pinned to the git HEAD for now since experimental.flashmd_symplectic is not yet in a released version. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Pin metatrain to 2026.2.1 in pyproject.toml. In get_pretrained, check the mtt export return code and give a helpful error if the experimental.flashmd_symplectic architecture is missing. Override metatrain with a pinned metatensor/metatrain commit in CI so the symplectic tests can actually run the export. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
warp-lang>=1.14.0 deprecated warp.config.quiet, which is used by nvalchemi-toolkit-ops==0.3.0. Suppress it so pytest's filterwarnings error mode doesn't turn it into a failure. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
The existing filter only matched APICLaunchParamRecord; warp-lang 1.14.0 triggers the same warning for other ctypes structures (e.g. APICLaunchPtrLocation) on Python 3.14. Use .* to suppress all of them. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
frostedoyster
left a comment
frostedoyster
left a comment
There was a problem hiding this comment.
Almost good. By the way, are we checking the time step of the symplectic model against that of the FlashMD model at any point?
| echo "metatrain @ git+https://github.com/metatensor/metatrain.git@870fae2a253ab33b20c52b958a16eac4640d840f" > /tmp/metatrain-override.txt | ||
| tox -e tests | ||
| env: | ||
| UV_OVERRIDE: /tmp/metatrain-override.txt |
There was a problem hiding this comment.
Can you try to coordinate with Pol on the metatrain release so we don't have to use this?
There was a problem hiding this comment.
I asked him, and he told me that they're fighting with the CI and can't merge enough PRs right now. Unfortunately, a release there does not seem very likely. Therefore, I thought that this "hack" is fine for now. We can revert the override as soon as there is a release of metatrain.
| from flashmd.fpi import anderson_solver | ||
|
|
||
|
|
||
| class SymplecticStepper: |
There was a problem hiding this comment.
Do you think a better design would be to inherit this from FlashMDStepper?
There was a problem hiding this comment.
What would we gain from that? I don't see a real benefit, it just makes the classes more coupled.
Validate that the symplectic correction model's timestep matches the FlashMD model's in SymplecticStepper, covering both the i-PI and ASE paths. Add a test pairing mismatched-timestep models. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
|
Yes, you're right about verifying that the time steps match. I added that in 0aa9a58. |
2 participants
Summary
Adds a fixed-point midpoint-rule symplectic integrator with minimal changes to the existing API.
src/flashmd/fpi.py: Anderson-accelerated fixed-point solver (ported from Add Symplectic Integrator #32).src/flashmd/symplectic_stepper.py:SymplecticStepper— wraps aFlashMDStepperfor the initial guess and refines it using a separate symplectic correction model.ipi.pyandase/velocity_verlet.py: tuple model detection. The plain single-model API is completely unchanged.Usage
Pass a tuple instead of a plain model wherever a model is expected:
The
configdict keys are all optional and passed straight toanderson_solver:tol,max_iter,m,beta,lambda_reg.This supersedes #32, which is being closed in favour of this smaller focused change. The example will follow in a separate PR on top of this one.