Add export_symplectic_model for simulation engine integration

[johannes-spies]

Summary

Introduces export_symplectic_model(), a utility that wraps a FlashMD model and a symplectic correction model into a single AtomisticModel that can be saved and loaded by LAMMPS via fix metatomic. This makes symplectic FlashMD available in simulation engines without any engine-side specialization — just load the exported .pt file like a regular FlashMD model.

How it works

The exported model runs a full symplectic step in its forward:

1. Calls the FlashMD model for an initial guess of the next positions/momenta
2. Refines it via an Anderson-accelerated midpoint fixed-point iteration, using the symplectic correction model at each midpoint estimate

Neighbor lists from the input system are reused at midpoint positions (valid as long as positional changes stay within the neighbor list skin — the same assumption LAMMPS already makes).

TorchScript compatibility

Several constraints arise from needing the module to be TorchScript-serializable:

• AtomisticModel wrappers cannot be nested; the inner .module is stored directly
• NeighborListOptions are not serializable; stored as three parallel primitive lists
• Labels/TensorBlock/TensorMap cannot be constructed fresh in the forward graph; output metadata is reused from the inner model instead
• The Anderson loop is inlined (TorchScript cannot take arbitrary Callable arguments); only the linear algebra step is shared via _anderson_update from fpi.py

Changes to fpi.py

Extracts _anderson_update() as a @torch.jit.script-decorated function containing the core Anderson step linear algebra. anderson_solver() now delegates to it, eliminating duplication between the Python-facing solver and the TorchScript forward pass.

Usage

from metatomic.torch import load_atomistic_model
from flashmd import export_symplectic_model

flashmd_model = load_atomistic_model("flashmd.pt")
symplectic_model = load_atomistic_model("symplectic-correction.pt")
exported = export_symplectic_model(flashmd_model, symplectic_model)
exported.save("symplectic-flashmd.pt")

Then in LAMMPS:

fix 0 all metatomic symplectic-flashmd.pt types 13

Notes

• metatrain dependency updated to git main with a TODO to re-pin once experimental.flashmd_symplectic lands in a PyPI release (same approach as Add Al example for symplectic integrator #44)
• This PR should be merged after Add symplectic integrator #43