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 90.09%. Comparing base (b243331) to head (406d8bb).

🚀 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.

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.

Modified and added based on IMDv2 behavior in GROMACS, LAMMPS, and NAMD

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: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?

No, this is only implemented under IMDProducerV2

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)

orbeckst · 2026-07-13T20:27:55Z

I released 0.2.4 #146 so please create the 0.3.0 section in the CHANGELOG when you add your entry.

    @pytest.fixture()
-    def first_frame(self):
-        return 0
+    def imd_nst(self, mdp):


+        return request.param
+
+    @pytest.fixture()
+    def first_frame(self, inp):


+        return request.param
+
+    @pytest.fixture()
+    def first_frame(self, inp):


-    # def match_string(self):
-    #     return "IMD: Will wait until I have a connection and IMD_GO orders."
+
+class TestIMDv3Gromacs(IMDGromacsTest, IMDv3IntegrationTest):


+        return request.param
+
+
+class TestIMDv2Gromacs(IMDGromacsTest, IMDv2IntegrationTest):


        yield u
+
+
+class TestIMDv3Lammps(IMDLammpsTest, IMDv3IntegrationTest):


+            return 0
+
+
+class TestIMDv2Lammps(IMDLammpsTest, IMDv2IntegrationTest):


+        return 0
+
+
+class TestIMDv3NAMD(IMDNAMDTest, IMDv3IntegrationTest):


            )
+
+
+class TestIMDv2NAMD(IMDNAMDTest, IMDv2IntegrationTest):


- added ValueError raising in `__init__` - tests to check error - updated API doc

orbeckst · 2026-09-04T22:02:59Z

Can you please update to sync it with the base branch? Thanks!

Will try to look at it over the weekend.

orbeckst

Add entry to CHANGELOG and write somewhere in the docs explicitly that we can also work with v2. Mention why this is not great for analysis.

orbeckst · 2026-09-07T02:06:50Z

+gen_vel                 = no
+IMD-group               = System
+IMD-nst                 = 8
+IMD-version             = 2


This will only work with our patched version of GROMACS that also speaks v3. I think you should leave it out.

Added new docker image support for vanilla gromacs versions as a PR. once mereged these tests would tests imdv2 support for imdclient using appropriate version of all 3 MD engines (GROMACS, LAMMPS and NAMD)

orbeckst · 2026-09-18T21:02:03Z

As discussed: use a vanilla version of GROMACS (eg 2025 from conda-forge) to test v2, too.

orbeckst · 2026-10-13T21:06:15Z

The actual MD engine tests are failing, eg https://github.com/Becksteinlab/imdclient/actions/runs/37972189712/job/113961342160?pr=141 with

E           ConnectionRefusedError: IMDClient: Connection to localhost:56093 refused

What's the issue?

orbeckst · 2026-10-13T21:19:22Z

From discussion: first need new docker images Becksteinlab/streaming-md-docker#32

orbeckst · 2026-11-05T23:55:06Z

We have new images (I fixed the deployment workflow, see also Becksteinlab/streaming-md-docker#34 ). If this works here now, then we should be good.

orbeckst · 2026-11-06T00:18:08Z

@amruthesht the MD engine tests fail with multiple Connection to localhost:...... refused errors. This occurs for GROMACS, NAMD, and LAMMPS. Eg all GROMACS tests failed

=========================== short test summary info ============================
FAILED imdclient/tests/test_gromacs.py::TestIMDv3Gromacs::test_continue_after_disconnect[/home/runner/work/imdclient/imdclient/imdclient/data/gromacs/md/gromacs_v3_nst1.mdp] - ConnectionRefusedError: IMDClient: Connection to localhost:56499 refused
FAILED imdclient/tests/test_gromacs.py::TestIMDv3Gromacs::test_continue_after_disconnect[/home/runner/work/imdclient/imdclient/imdclient/data/gromacs/md/gromacs_v3_nst8.mdp] - ConnectionRefusedError: IMDClient: Connection to localhost:60221 refused
FAILED imdclient/tests/test_gromacs.py::TestIMDv3Gromacs::test_wait_after_disconnect[/home/runner/work/imdclient/imdclient/imdclient/data/gromacs/md/gromacs_v3_nst1.mdp] - ConnectionRefusedError: IMDClient: Connection to localhost:37825 refused
FAILED imdclient/tests/test_gromacs.py::TestIMDv3Gromacs::test_wait_after_disconnect[/home/runner/work/imdclient/imdclient/imdclient/data/gromacs/md/gromacs_v3_nst8.mdp] - ConnectionRefusedError: IMDClient: Connection to localhost:34553 refused
ERROR imdclient/tests/test_gromacs.py::TestIMDv3Gromacs::test_compare_imd_to_true_traj[/home/runner/work/imdclient/imdclient/imdclient/data/gromacs/md/gromacs_v3_nst1.mdp] - ConnectionRefusedError: IMDClient: Connection to localhost:35123 refused
ERROR imdclient/tests/test_gromacs.py::TestIMDv3Gromacs::test_compare_imd_to_true_traj[/home/runner/work/imdclient/imdclient/imdclient/data/gromacs/md/gromacs_v3_nst8.mdp] - ConnectionRefusedError: IMDClient: Connection to localhost:47581 refused
ERROR imdclient/tests/test_gromacs.py::TestIMDv2Gromacs::test_compare_imd_to_true_traj[/home/runner/work/imdclient/imdclient/imdclient/data/gromacs/md/gromacs_v2_nst1.mdp] - ConnectionRefusedError: IMDClient: Connection to localhost:46023 refused
ERROR imdclient/tests/test_gromacs.py::TestIMDv2Gromacs::test_compare_imd_to_true_traj[/home/runner/work/imdclient/imdclient/imdclient/data/gromacs/md/gromacs_v2_nst8.mdp] - ConnectionRefusedError: IMDClient: Connection to localhost:55955 refused
============== 4 failed, 1 warning, 4 errors in 252.62s (0:04:12) ==============

and LAMMPS

=========================== short test summary info ============================
ERROR imdclient/tests/test_lammps.py::TestIMDv2Lammps::test_compare_imd_to_true_traj[/home/runner/work/imdclient/imdclient/imdclient/data/lammps/md/lammps_v2_nst_1.in] - ConnectionRefusedError: IMDClient: Connection to localhost:35983 refused
ERROR imdclient/tests/test_lammps.py::TestIMDv2Lammps::test_compare_imd_to_true_traj[/home/runner/work/imdclient/imdclient/imdclient/data/lammps/md/lammps_v2_nst_8.in] - ConnectionRefusedError: IMDClient: Connection to localhost:33241 refused
============== 6 passed, 1 warning, 2 errors in 437.71s (0:07:17) ==============

and NAMD

=========================== short test summary info ============================
FAILED imdclient/tests/test_namd.py::TestIMDv3NAMD::test_wait_after_disconnect[/home/runner/work/imdclient/imdclient/imdclient/data/namd/md/namd_v3_nst_1.namd] - ConnectionRefusedError: IMDClient: Connection to localhost:44577 refused
FAILED imdclient/tests/test_namd.py::TestIMDv3NAMD::test_compare_imd_to_true_traj_forces[/home/runner/work/imdclient/imdclient/imdclient/data/namd/md/namd_v3_nst_1.namd] - AssertionError: Arrays are not almost equal.
FAILED imdclient/tests/test_namd.py::TestIMDv3NAMD::test_compare_imd_to_true_traj_forces[/home/runner/work/imdclient/imdclient/imdclient/data/namd/md/namd_v3_nst_8.namd] - AssertionError: Arrays are not almost equal.
ERROR imdclient/tests/test_namd.py::TestIMDv2NAMD::test_compare_imd_to_true_traj[/home/runner/work/imdclient/imdclient/imdclient/data/namd/md/namd_v2_nst_1.namd] - ConnectionRefusedError: IMDClient: Connection to localhost:48575 refused
ERROR imdclient/tests/test_namd.py::TestIMDv2NAMD::test_compare_imd_to_true_traj[/home/runner/work/imdclient/imdclient/imdclient/data/namd/md/namd_v2_nst_8.namd] - ConnectionRefusedError: IMDClient: Connection to localhost:48851 refused
======== 3 failed, 7 passed, 19 warnings, 2 errors in 561.34s (0:09:21) ========

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 Resolved

Comment thread imdclient/tests/test_lammps.py Resolved

Comment thread imdclient/tests/base.py Resolved

Comment thread imdclient/tests/minimalreader.py Resolved

Comment thread imdclient/IMDClient.py Resolved

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

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

Closed

2 tasks done

Merge branch 'main' into imdv2-support

77a0af3

ljwoods2 mentioned this pull request Jul 19, 2026

Modify disconnect method for proper socket shutdown #149

Merged

3 of 4 tasks

Merge 'main' into imdv2-support

90628a7

github-code-quality Bot found potential problems Jul 20, 2026

View reviewed changes

Merge branch 'main' into imdv2-support

9f521ea

orbeckst mentioned this pull request Jul 23, 2026

Add UserGuide documentation for IMDReader MDAnalysis/UserGuide#430

Open

This was linked to issues Jul 24, 2026

Add tests for IMDv2 streams #82

Open

IMDv2 in GROMACS requires client to set trate #143

Open

orbeckst mentioned this pull request Jul 24, 2026

IMDv2 in GROMACS requires client to set trate #143

Open

amruthesht added 4 commits September 4, 2026 11:15

Add: continue_after_disconnect only for v2

74ae4ce

- added ValueError raising in `__init__` - tests to check error - updated API doc

Add: trate related doc, warning and description

a94c60f

Add: trate option warning for v3

47dc7f6

Format: black

78fda98

amruthesht requested review from ljwoods2 and orbeckst September 4, 2026 19:18

Merge branch 'main' into imdv2-support

daded54

orbeckst requested changes Sep 7, 2026

View reviewed changes

Change: use version appropriate (imdv2/v3) MD engines for tests

9e9fb28

Merge branch 'main' into imdv2-support

a388a85

Merge branch 'main' into imdv2-support

406d8bb

		# 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:

		return request.param


		class TestIMDv2Gromacs(IMDGromacsTest, IMDv2IntegrationTest):

		yield u


		class TestIMDv3Lammps(IMDLammpsTest, IMDv3IntegrationTest):

		return 0


		class TestIMDv2Lammps(IMDLammpsTest, IMDv2IntegrationTest):

		return 0


		class TestIMDv3NAMD(IMDNAMDTest, IMDv3IntegrationTest):

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!

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!

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!

Uh oh!

Uh oh!

Uh oh!

orbeckst commented Jul 13, 2026

Uh oh!

orbeckst commented Sep 4, 2026

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!

orbeckst commented Sep 18, 2026

Uh oh!

orbeckst commented Oct 13, 2026

Uh oh!

orbeckst commented Oct 13, 2026

Uh oh!

orbeckst commented Nov 5, 2026

Uh oh!

orbeckst commented Nov 6, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

amruthesht commented Jun 1, 2026 •

edited by orbeckst

Loading

codecov Bot commented Jun 1, 2026 •

edited

Loading