IMDv2 generic support implementation with tests by amruthesht · Pull Request #141 · Becksteinlab/imdclient

amruthesht · 2026-06-01T16:15:10Z

This PR provides generic IMDv2 support for handling various v2 streaming scenarios in imdclient

Changes made in this Pull Request:

IMDv2 schema as described in IMDv2 support and tests for receiving positions #140
intergation tests for IMDv2 simulationengine implementationss
unit test for IMDClientv2
box wrapping for TRR and IMD based checking when using v2 or IMD-unwrap=True in session info
v2 trate sending via IMDClient object
Added v2 tests addressing Add tests for IMDv2 streams #82

Note: IMDv2 in GROMACS seems to expects transmission rate from the user via the client, as discussed in #143
It ignores IMD-nst in the *.mdp file. Thus the IMDClient object was modified to send IMD_TRATE right after IMD_GO as set in the object initiation when using v2.

PR Checklist

Tests?
Docs?
CHANGELOG updated?
Issue raised/referenced?

codecov · 2026-06-01T16:16:52Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 86.58%. Comparing base (0769db1) to head (7e3ba9a).

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Copilot

Pull request overview

This PR extends imdclient to handle more IMDv2 streaming patterns (notably variable energy-packet ordering) and adds/updates integration tests and simulation inputs to validate IMDv2 behavior across supported engines.

Changes:

Update IMDv2 frame parsing to tolerate multiple leading energy packets before coordinates and to reuse cached energies when energies are omitted.
Add IMDv2 integration test coverage alongside existing IMDv3 tests for NAMD, GROMACS, and LAMMPS.
Add new IMDv2 engine input/config files and rename/refactor test datafile constants to separate v2/v3 variants.

Reviewed changes

Copilot reviewed 15 out of 15 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
imdclient/IMDClient.py	Adjust IMDv2 producer parsing to handle multiple/optional energy packets before coordinates.
imdclient/tests/base.py	Refactor integration test base class and add an IMDv2-specific integration test class.
imdclient/tests/test_namd.py	Split shared NAMD fixtures and add IMDv2 + IMDv3 integration test classes with version-specific inputs.
imdclient/tests/test_lammps.py	Split shared LAMMPS fixtures and add IMDv2 + IMDv3 integration test classes with version-specific inputs.
imdclient/tests/test_gromacs.py	Split shared GROMACS fixtures and add IMDv2 + IMDv3 integration test classes with version-specific inputs.
imdclient/tests/minimalreader.py	Handle missing box data when converting to MDAnalysis dimensions during stream processing.
imdclient/tests/datafiles.py	Rename/split datafile constants to explicitly distinguish v2 vs v3 simulation inputs.
imdclient/data/namd/md/namd_v2_nst_1.namd	Add NAMD IMDv2 test configuration (NST=1).
imdclient/data/namd/md/namd_v2_nst_8.namd	Add NAMD IMDv2 test configuration (NST=8).
imdclient/data/lammps/md/lammps_v2_nst_1.in	Add LAMMPS IMDv2 test input (NST=1).
imdclient/data/lammps/md/lammps_v2_nst_8.in	Add LAMMPS IMDv2 test input (NST=8).
imdclient/data/lammps/md/lammps_v3_nst_1.in	Expand LAMMPS IMDv3 fix options to explicitly enable time/box/coords/vel/forces.
imdclient/data/lammps/md/lammps_v3_nst_8.in	Expand LAMMPS IMDv3 fix options to explicitly enable time/box/coords/vel/forces.
imdclient/data/gromacs/md/gromacs_v2_nst1.mdp	Add GROMACS IMDv2 test MDP (NST=1).
imdclient/data/gromacs/md/gromacs_v2_nst8.mdp	Add GROMACS IMDv2 test MDP (NST=8).

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

orbeckst

As discussed in meeting, wrap the GROMACS streamed coordinates into the unit cell by

extracting the ts.dimensions from the saved trajectory
create two on-the-fly transformations
- set_box_dimensions
- wrap
use the transformations with the in-memory universe that you build with the numpy array of the IMD data

This should generate an in-memory trajectory of wrapped coordinates.

- v2 unexpected packets case - trate sending only in v2

amruthesht · 2026-06-20T00:36:01Z

@orbeckst - This PR is finally done after a bit of back and forth trying to figure out issues with using v2 in GROMACS.
Please let me know if you still see any outstanding issues.

Thank you to you and @HeydenLabASU for the discussion and suggestions on solving this issue.

ljwoods2 · 2026-06-27T15:23:29Z

-            assert_allclose(
-                universe.trajectory[i].triclinic_dimensions, imdf.box
-            )
+            if imdsinfo.time:


Drive-by comment, feel free to ignore if I'm off base

I'm guessing you move this test to the base class to avoid code duplication, but the if statements on imdsinfo.<attribute> could mask real issues, like if in v3, the time packet was somehow not saved from the stream. I'd suggest just having two version of the test, each without these if statements, even though it's more code

orbeckst

I only managed a partial review (code, not tests). The fact that it works is great!

Do we use the the same IMDClient for v3 and v2? Do any of the changes affect v3?

My primary comments so far are just all related to docs.

We need to document the v2 capabilities. Can you create a protocol page for the v2 format as we understand it (then we have a v2 and a v3 page). On the IMDClient docs page we should also have a section on v2 that lists any shortcomings; some out of our control (no dimensions, no v/f), some our implementation (energies can be out of sync as I understand it).

Add an entry to CHANGELOG. This will have to be 0.3.0.

orbeckst · 2026-06-27T14:58:43Z

+        IMD transmission rate to be set after client send go signal. 
+        This parameter is only set to the server when using IMDv2. Default behavior is to not set the transmission rate. 
+        This parameter is useful when running GROMACS with IMDv2, where transmission rate can only be set via the client. 
+        IMDv2 implementations in LAMMPS and NAMD support setting the transmission rate via relevant input file parameters.


Add an definition of what "IMD transmission rate" is, e.g., a which integrator step IMD packages are sent. Include a link to https://imdclient.readthedocs.io/en/latest/protocol_v3.html#transmission-rate

IMD :ref:`transmission rate <transmission-rate>` to be set after client send go signal.

State what the default value for IMDv2 in GROMACS is.

Does GROMACS default to transmission_rate = 1? I.e., is it safe for users to leave it as None or will it hang their simulations?

You could add a note box

.. warning:: When using GROMACS with IMDv2, users should set `transmission_rate` explicitly. Otherwise the GROMACS default of 1 will be chosen. This default setting will send every integrator time step and may impact performance substantially.

And add an indication in which version this option was added

transmission_rate : int, optional [``None``] IMD transmission rate to be set after client send go signal. This parameter is only set to the server when using IMDv2. Defaul... ... IMDv2 implementations in LAMMPS and NAMD support setting the transmission rate via relevant input file parameters. .. versionadded:: 0.3.0

Check that it looks right when the docs are rendered.

orbeckst · 2026-06-27T15:02:38Z

+        This parameter is useful when running GROMACS with IMDv2, where transmission rate can only be set via the client. 
+        IMDv2 implementations in LAMMPS and NAMD support setting the transmission rate via relevant input file parameters.
    **kwargs : dict (optional)
        Additional keyword arguments to pass to the :class:`BaseIMDProducer` and :class:`IMDFrameBuffer`


Add a

.. versionchanged:: 0.3.0 Added capability to process IMDv2 protocol streams.

(Note: need two blank lines above)

orbeckst · 2026-06-27T15:04:53Z

+    def _trate(self, rate):
+        """
+        Send a trate packet to the server to set transmission rate.
+        """


Suggested change

"""

.. versionadded:: 0.3.0

"""

orbeckst · 2026-06-27T15:17:26Z

        header = self._get_header()
-        if header.type == IMDHeaderType.IMD_ENERGIES and header.length == 1:
+        leading_energies = 0
+        while header.type == IMDHeaderType.IMD_ENERGIES and header.length == 1:


When do you get multiple energy packets in sequence? Only in v2?

Add a note that energy handling in v2 can miss something/report out of sync.

orbeckst · 2026-06-27T15:19:05Z

-        # Either receive energies + positions or just positions
+        # Consume any leading energy packets first, then handle positions.


Does this change how v3 streams are processed?

ljwoods2 · 2026-06-27T15:31:19Z


        self._go()

+        if transmission_rate is not None and self._imdsinfo.version == 2:


If version == 3 and transmission rate is set, issue a warning that the kwarg is ignored, for least surprise to user

(I know the docs already say this, so if you think that's sufficient, feel free to ignore , just thinking of principle of least surprise)

ljwoods2 · 2026-06-27T15:33:08Z

        :class:`IMDFrameBuffer` will be filled with as many :class:`IMDFrame` fit in `buffer_size` bytes [``10MB``]
    timeout : int, optional
        Timeout for the socket in seconds [``5``]
    continue_after_disconnect : bool, optional [``None``]


make note that this is v3-only

ljwoods2 · 2026-06-27T15:33:48Z

        self._conn = self._connect_to_server(host, port, socket_bufsize)
        self._imdsinfo = self._await_IMD_handshake()
        self._multithreaded = multithreaded
        self._continue_after_disconnect = continue_after_disconnect


if v2, fail out if continue after disconnect is anything except None

ljwoods2 · 2026-06-27T15:34:45Z

        self._conn.sendall(go)
        logger.debug("IMDClient: Sent go packet to server")

        if self._continue_after_disconnect is not None:


to be extra safe:

Suggested change

if self._continue_after_disconnect is not None and self._imdsinfo.version == 3:

amruthesht added 4 commits May 28, 2026 15:18

IMDv2 test input files

9de16b1

IMDv2 schema for reading stream

f2010f8

IMDv2 integration tests

4918016

IMDv2 leading energies warning updated

aedc2f3

amruthesht requested a review from Copilot June 1, 2026 16:15

amruthesht self-assigned this Jun 1, 2026

Copilot started reviewing on behalf of amruthesht June 1, 2026 16:16 View session

amruthesht requested a review from orbeckst June 1, 2026 16:16

Copilot AI reviewed Jun 1, 2026

View reviewed changes

Comment thread imdclient/IMDClient.py Outdated

Comment thread imdclient/tests/test_lammps.py Outdated

Comment thread imdclient/tests/base.py Outdated

Comment thread imdclient/tests/minimalreader.py Outdated

Comment thread imdclient/IMDClient.py

amruthesht added 6 commits June 2, 2026 12:11

first frame set inside v2,v3 classes

54ea9dd

Remove unnecessary fixture requirement

9dccabe

Fix typo

4f075f9

Update v2 h5md output to just positions

2ee1660

Change IMD MD-engine base class

abc43b2

v2 trjconv mol-whole treatment

bd67a88

amruthesht force-pushed the imdv2-support branch 2 times, most recently from fca1842 to bd67a88 Compare June 5, 2026 19:35

Merge branch 'main' into imdv2-support

abf1f8f

orbeckst requested changes Jun 5, 2026

View reviewed changes

amruthesht added 10 commits June 5, 2026 18:37

v2 trjconv mol-whole treatment additional step

24c1031

v2 trjconv mol-whole treatment correction

260a06b

v2 trjconv mol-whole container testing

8c2b6e9

v2 trjconv mol-whole post simulation container waiting

abebc6a

v2 trjconv whole-only post simulation container waiting

f39d6b3

Fix: wrap imd_u into box sintead of trjconv

f99320b

Fix: wrap true_u into box as well

2d27171

Fix: trate set by client in IMDv2

2c95c41

Add: IMDCLient tests for v2 and v3

40a2fba

black linter cleanup

cb77b43

amruthesht force-pushed the imdv2-support branch from e7a22a3 to f39d6b3 Compare June 19, 2026 07:11

amruthesht added 5 commits June 19, 2026 15:48

Fix: IMDSessionInfo missing time issue

edc7767

Fix: incorrect n_atoms Error raised in v2

8c00c7b

Add: test coverage

751cf39

- v2 unexpected packets case - trate sending only in v2

Clean: black reformat

79c2b56

Add: docstring for trate

c79f62f

amruthesht requested a review from orbeckst June 20, 2026 00:36

Merge branch 'main' into imdv2-support

7e3ba9a

amruthesht changed the title ~~IMDv2 generic support implmentation with tests~~ IMDv2 generic support implementation with tests Jun 26, 2026

ljwoods2 requested changes Jun 27, 2026

View reviewed changes

orbeckst requested changes Jun 27, 2026

View reviewed changes

ljwoods2 requested changes Jun 27, 2026

View reviewed changes

orbeckst mentioned this pull request Jun 27, 2026

quick 0.2.4. release #146

Open

2 tasks

-        """
+.. versionadded:: 0.3.0
+        """

		# Either receive energies + positions or just positions
		# Consume any leading energy packets first, then handle positions.


		self._go()

		if transmission_rate is not None and self._imdsinfo.version == 2:


	if self._continue_after_disconnect is not None and self._imdsinfo.version == 3:

Uh oh!

Conversation

amruthesht commented Jun 1, 2026 • edited by orbeckst Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR Checklist

Uh oh!

codecov Bot commented Jun 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

orbeckst left a comment

Choose a reason for hiding this comment

Uh oh!

amruthesht commented Jun 20, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

orbeckst left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

amruthesht commented Jun 1, 2026 •

edited by orbeckst

Loading

codecov Bot commented Jun 1, 2026 •

edited

Loading