[Brean0]

Summary

momentum_step and position_step document support for a timestep tensor shaped [n_systems], but multiply that vector directly by [n_atoms,3] quantities. Two systems produce a shape error. With exactly three systems, broadcasting silently applies the timesteps to x, y, and z coordinates instead of to each system. For unit-velocity free particles and dt=[0.1,0.2,0.3], every atom moves by [0.1,0.2,0.3] regardless of system membership.

This change maps vector timesteps through state.system_idx and adds a trailing coordinate axis in both update functions. Each system then advances all coordinates by its own timestep. Scalar timesteps keep the existing path.

Regression tests cover two and three systems, with equal and unequal atom counts, zero forces, and unit velocities through the complete NVE step, using a separated-particle, zero-epsilon Lennard-Jones model. Direct momentum-step tests use constant nonzero forces so the momentum mapping is checked independently. Both tests include scalar-timestep controls.

Validation

All checks used Python 3.12.9 on macOS CPU, with the isolated checkout imported through PYTHONPATH and a temporary Warp cache.

• Before the fix: the initial four vector-timestep cases failed, reproducing the two-system shape errors and three-system coordinate scaling; all four initial scalar controls passed.
• Final python -m pytest tests/test_integrators.py -q: 47 passed, including 16 scalar/vector timestep cases across equal and unequal atom counts.
• Final repository-wide python -m pytest: 1019 passed, 208 skipped, exit 0. Each final run includes all regression additions. Three PyTables UnclosedFileWarning messages followed the summary during trajectory-file cleanup.
• Ruff 0.17.1 lint and format checks on both changed files: passed.
• git diff --check: passed.

The production change is limited to the timestep mapping in the two core MD updates. The new tests cover NVE and direct momentum propagation; they do not establish vector-timestep support throughout every thermostat or barostat. GPU execution and skipped optional-model tests were not exercised.

The F004 evidence report includes the original reproduction, scalar controls, and independently executed native ASE VelocityVerlet comparisons.

Checklist

• Docstrings have been added in the Google docstring format.
• Ruff lint and format checks pass.
• Regression tests have been added for the fix.