IMDv2 generic support implementation with tests

imdclient/IMDClient.py

@@ -59,6 +59,11 @@ class IMDClient:
         If True, the client will attempt to change the simulation engine's waiting behavior to
         non-blocking after the client disconnects. If False, the client will attempt to change it
         to blocking. If None, the client will not attempt to change the simulation engine's behavior.
+    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. 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.
     **kwargs : dict (optional)
         Additional keyword arguments to pass to the :class:`BaseIMDProducer` and :class:`IMDFrameBuffer`
     """
@@ -71,6 +76,7 @@ def __init__(
         socket_bufsize=None,
         multithreaded=True,
         continue_after_disconnect=None,
+        transmission_rate=None,
         **kwargs,
     ):
 
@@ -120,6 +126,9 @@ def __init__(
 
         self._go()
 
+        if transmission_rate is not None and self._imdsinfo.version == 2:
+            self._trate(transmission_rate)
+
         if self._multithreaded:
             # Disconnect MUST occur. This covers typical cases (Python, IPython interpreter)
             signal.signal(signal.SIGINT, self.signal_handler)
@@ -293,13 +302,13 @@ def _await_IMD_handshake(self) -> IMDSessionInfo:
             sinfo = IMDSessionInfo(
                 version=ver,
                 endianness=end,
-                wrapped_coords=False,
                 time=False,
                 energies=True,
                 box=False,
                 positions=True,
                 velocities=False,
                 forces=False,
+                wrapped_coords=False,
             )
 
         elif ver == 3:
@@ -320,9 +329,17 @@ def _await_IMD_handshake(self) -> IMDSessionInfo:
 
         return sinfo
 
+    def _trate(self, rate):
+        """
+        Send a trate packet to the server to set transmission rate.
+        """
+        trate = create_header_bytes(IMDHeaderType.IMD_TRATE, rate)
+        self._conn.sendall(trate)
+        logger.debug("IMDClient: Sent transmission rate %s", rate)
+
     def _go(self):
         """
-        Send a go packet to the client to start the simulation
+        Send a go packet to the server to start the simulation
         and begin receiving data.
         """
         go = create_header_bytes(IMDHeaderType.IMD_GO, 0)
@@ -579,30 +596,33 @@ def _parse_imdframe(self):
         # Even if they are sent, energies might not be sent every frame
         # cache the last energies received
 
-        # Either receive energies + positions or just positions
+        # Consume any leading energy packets first, then handle positions.
         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:
             self._imdsinfo.energies = True
             self._read(self._energies)
             self._imdf.energies.update(
                 parse_energy_bytes(self._energies, self._imdsinfo.endianness)
             )
             self._prev_energies = self._imdf.energies
 
-            self._expect_header(
-                IMDHeaderType.IMD_FCOORDS, expected_value=self._n_atoms
-            )
-            self._read(self._positions)
-            np.copyto(
-                self._imdf.positions,
-                np.frombuffer(
-                    self._positions, dtype=f"{self._imdsinfo.endianness}f"
-                ).reshape((self._n_atoms, 3)),
+            leading_energies += 1
+
+            header = self._get_header()
+
+        if leading_energies == 0 or leading_energies > 1:
+            logger.warning(
+                f"IMDProducer: Received {leading_energies} leading IMDv2 energy packets before coordinates, energy values may be out of sync with coordinates"
             )
-        elif (
-            header.type == IMDHeaderType.IMD_FCOORDS
-            and header.length == self._n_atoms
-        ):
+
+        if header.type == IMDHeaderType.IMD_FCOORDS:
+            # check if the number of atoms is correct
+            if header.length != self._n_atoms:
+                raise RuntimeError(
+                    f"IMDProducer: Expected n_atoms value {self._n_atoms}, got {header.length}. "
+                    + "Ensure you are using the correct topology file."
+                )
             # If we received positions but no energies
             # use the last energies received
             if self._prev_energies is not None:
@@ -617,7 +637,9 @@ def _parse_imdframe(self):
                 ).reshape((self._n_atoms, 3)),
             )
         else:
-            raise RuntimeError("IMDProducer: Unexpected packet type or length")
+            raise RuntimeError(
+                f"IMDProducer: Unexpected packet type {header.type.name}"
+            )
 
     def _pause(self):
         logger.debug(

imdclient/data/gromacs/md/gromacs_v2_nst1.mdp

@@ -0,0 +1,51 @@
+title                   = PRODUCTION IN NPT
+ld-seed			= 1
+; Run parameters
+integrator              = md        ; leap-frog integrator
+nsteps                  = 100    ; 1 * 1000 = 1 ps
+dt                      = 0.001     ; 1 fs
+; Output control
+nstxout                 = 1      ; save coordinates every 1 fs
+nstvout                 = 1      ; save velocities every 1 fs
+nstfout                 = 1      ; save forces every 1 fs
+nstenergy               = 1      ; save energies every 1 fs
+nstlog                  = 10     
+; Center of mass (COM) motion
+nstcomm                 = 10        ; remove COM motion every 10 steps
+comm-mode               = Linear    ; remove only COM translation (liquids in PBC) 
+; Bond parameters
+continuation            = yes       ; first dynamics run
+constraint_algorithm    = lincs     ; holonomic constraints
+constraints             = all-bonds ; all bonds lengths are constrained
+lincs_iter              = 1         ; accuracy of LINCS
+lincs_order             = 4         ; also related to accuracy
+; Nonbonded settings
+cutoff-scheme           = Verlet    ; Buffered neighbor searching
+ns_type                 = grid      ; search neighboring grid cells
+nstlist                 = 10        ; 10 fs, largely irrelevant with Verlet
+rcoulomb                = 1.0       ; short-range electrostatic cutoff (in nm)
+rvdw                    = 1.0       ; short-range van der Waals cutoff (in nm)
+DispCorr                = EnerPres  ; account for cut-off vdW scheme
+; Electrostatics
+coulombtype             = PME       ; Particle Mesh Ewald for long-range electrostatics
+pme_order               = 4         ; cubic interpolation
+fourierspacing          = 0.12      ; grid spacing for FFT
+; Temperature coupling is on
+tcoupl                  = Nose-Hoover       ; good for production, after equilibration
+; we define separate thermostats for the solute and solvent (need to adapt)
+; see default groups defined by Gromacs for your system or define your own (make_ndx)
+tc-grps                 = Protein  SOL      ; the separate groups for the thermostats
+tau-t                   = 1.0      1.0      ; time constants for thermostats (ps)
+ref-t                   = 300      300      ; reference temperature for thermostats (K)
+; Pressure coupling is off
+pcoupl                  = Parrinello-Rahman ; good for production, after equilibration
+tau-p                   = 2.0               ; time constant for barostat (ps)
+compressibility         = 4.5e-5            ; compressibility (1/bar) set to water at ~300K
+ref-p                   = 1.0               ; reference pressure for barostat (bar)
+; Periodic boundary conditions
+pbc                     = xyz       ; 3-D PBC
+; Velocity generation
+gen_vel                 = no
+IMD-group               = System
+IMD-nst                 = 1
+IMD-version             = 2

imdclient/data/gromacs/md/gromacs_v2_nst8.mdp

@@ -0,0 +1,51 @@
+title                   = PRODUCTION IN NPT
+ld-seed                 = 1
+; Run parameters
+integrator              = md        ; leap-frog integrator
+nsteps                  = 100    ; 1 * 1000 = 1 ps
+dt                      = 0.001     ; 1 fs
+; Output control
+nstxout                 = 8      ; save coordinates every 1 fs
+nstvout                 = 8      ; save velocities every 1 fs
+nstfout                 = 8
+nstenergy               = 8      ; save energies every 1 fs
+nstlog                  = 10      ; update log file every 1 ps
+; Center of mass (COM) motion
+nstcomm                 = 10        ; remove COM motion every 10 steps
+comm-mode               = Linear    ; remove only COM translation (liquids in PBC) 
+; Bond parameters
+continuation            = yes       ; first dynamics run
+constraint_algorithm    = lincs     ; holonomic constraints
+constraints             = all-bonds ; all bonds lengths are constrained
+lincs_iter              = 1         ; accuracy of LINCS
+lincs_order             = 4         ; also related to accuracy
+; Nonbonded settings
+cutoff-scheme           = Verlet    ; Buffered neighbor searching
+ns_type                 = grid      ; search neighboring grid cells
+nstlist                 = 10        ; 10 fs, largely irrelevant with Verlet
+rcoulomb                = 1.0       ; short-range electrostatic cutoff (in nm)
+rvdw                    = 1.0       ; short-range van der Waals cutoff (in nm)
+DispCorr                = EnerPres  ; account for cut-off vdW scheme
+; Electrostatics
+coulombtype             = PME       ; Particle Mesh Ewald for long-range electrostatics
+pme_order               = 4         ; cubic interpolation
+fourierspacing          = 0.12      ; grid spacing for FFT
+; Temperature coupling is on
+tcoupl                  = Nose-Hoover       ; good for production, after equilibration
+; we define separate thermostats for the solute and solvent (need to adapt)
+; see default groups defined by Gromacs for your system or define your own (make_ndx)
+tc-grps                 = Protein  SOL      ; the separate groups for the thermostats
+tau-t                   = 1.0      1.0      ; time constants for thermostats (ps)
+ref-t                   = 300      300      ; reference temperature for thermostats (K)
+; Pressure coupling is off
+pcoupl                  = Parrinello-Rahman ; good for production, after equilibration
+tau-p                   = 2.0               ; time constant for barostat (ps)
+compressibility         = 4.5e-5            ; compressibility (1/bar) set to water at ~300K
+ref-p                   = 1.0               ; reference pressure for barostat (bar)
+; Periodic boundary conditions
+pbc                     = xyz       ; 3-D PBC
+; Velocity generation
+gen_vel                 = no
+IMD-group               = System
+IMD-nst                 = 8
+IMD-version             = 2

imdclient/data/lammps/md/lammps_v2_nst_1.in

@@ -0,0 +1,71 @@
+## Setup
+units metal 
+boundary p p p #Specify periodic boundary condition are needed in all three faces
+atom_style atomic #What style of atoms is to be used in the simulation
+log   logfile.txt #Write the log file to this text file. All thermodynamic information applicable to the entire system
+
+## Create Box
+#Refers to an abstract geometric region of space. units box refers to the fact that the size of the box is specified in the units as given in the units command. 
+# The name "forbox" refers to the region ID so that you can refer to it somewhere else in this input script.
+region forbox block 0 45.8 0 45.8 0 45.8 units box
+create_box 1 forbox 
+# Since we have given fcc as lattice type no need to mention basis for this
+lattice fcc 4.58 
+
+## Create atoms & define interactions
+# basis arg defines which atoms are created based on their lattice position (all are atom type 1)
+create_atoms 1 region forbox  basis 1 1 basis 2 1 basis 3 1 basis 4 1 units box
+# Mass of atom type 1 is 39.48 [mass units grams/mole]
+mass 1 39.948 
+# lj potential describes potential energy between two atoms as function of the dist between them
+# don't apply lj interactions beyond cutoff dist
+pair_style lj/cut 10 
+# The coefficient of the lj potential for the interactions of atom type 1 with atom type 1
+pair_coeff  1 1 0.01006418 3.3952 
+
+## Create atom group for argon atoms
+group ar type 1 #Group all the argon types (argon type is of type 1). All atoms of type 1 are in group with the name 'ar'
+
+
+## Write initial configuration
+dump dump_1 all custom 1 dump_initial_config.dump id type x y z ix iy iz vx vy vz
+
+
+## Perform energy minimization
+run 1
+# Stop dumping to this file
+undump dump_1 
+# Minimize the energy using a conjugate gradient step.
+minimize 1e-25 1e-19 10000 10000 
+print "Finished Minimizing"
+variable ener equal pe
+
+## Output the topology after minimization
+write_data topology_after_min.data
+
+## Prepare MD simulation
+timestep  0.001
+# Set the velocities of all the atoms so that the temperature of the system
+# is 300K. Make the distribution Gaussian.
+velocity all create 300 102939 dist gaussian mom yes rot yes
+# this is equlibration process.
+fix 1 all nve 
+
+# Create source of truth trajectory
+dump h5md1 all h5md 1 lammps_trj.h5md position
+dump_modify h5md1 unwrap no
+
+## IMD settings
+# https://docs.lammps.org/fix_imd.html
+fix 2 all imd 8888 version 2 nowait off trate 1
+
+## Run MD sim
+run 100
+
+# Stop dumping information to the dump file.
+undump h5md1 
+
+# Unfix the NVE. Additional lines if any will assume that this fix is off.
+unfix 1
+
+#End

imdclient/data/lammps/md/lammps_v2_nst_8.in

@@ -0,0 +1,71 @@
+## Setup
+units metal 
+boundary p p p #Specify periodic boundary condition are needed in all three faces
+atom_style atomic #What style of atoms is to be used in the simulation
+log   logfile.txt #Write the log file to this text file. All thermodynamic information applicable to the entire system
+
+## Create Box
+#Refers to an abstract geometric region of space. units box refers to the fact that the size of the box is specified in the units as given in the units command. 
+# The name "forbox" refers to the region ID so that you can refer to it somewhere else in this input script.
+region forbox block 0 45.8 0 45.8 0 45.8 units box
+create_box 1 forbox 
+# Since we have given fcc as lattice type no need to mention basis for this
+lattice fcc 4.58 
+
+## Create atoms & define interactions
+# basis arg defines which atoms are created based on their lattice position (all are atom type 1)
+create_atoms 1 region forbox  basis 1 1 basis 2 1 basis 3 1 basis 4 1 units box
+# Mass of atom type 1 is 39.48 [mass units grams/mole]
+mass 1 39.948 
+# lj potential describes potential energy between two atoms as function of the dist between them
+# don't apply lj interactions beyond cutoff dist
+pair_style lj/cut 10 
+# The coefficient of the lj potential for the interactions of atom type 1 with atom type 1
+pair_coeff  1 1 0.01006418 3.3952 
+
+## Create atom group for argon atoms
+group ar type 1 #Group all the argon types (argon type is of type 1). All atoms of type 1 are in group with the name 'ar'
+
+
+## Write initial configuration
+dump dump_1 all custom 1 dump_initial_config.dump id type x y z ix iy iz vx vy vz
+
+
+## Perform energy minimization
+run 1
+# Stop dumping to this file
+undump dump_1 
+# Minimize the energy using a conjugate gradient step.
+minimize 1e-25 1e-19 10000 10000 
+print "Finished Minimizing"
+variable ener equal pe
+
+## Output the topology after minimization
+write_data topology_after_min.data
+
+## Prepare MD simulation
+timestep  0.001
+# Set the velocities of all the atoms so that the temperature of the system
+# is 300K. Make the distribution Gaussian.
+velocity all create 300 102939 dist gaussian mom yes rot yes
+# this is equlibration process.
+fix 1 all nve 
+
+# Create source of truth trajectory
+dump h5md1 all h5md 8 lammps_trj.h5md position
+dump_modify h5md1 unwrap no
+
+## IMD settings
+# https://docs.lammps.org/fix_imd.html
+fix 2 all imd 8888 version 2 nowait off trate 8
+
+## Run MD sim
+run 100
+
+# Stop dumping information to the dump file.
+undump h5md1 
+
+# Unfix the NVE. Additional lines if any will assume that this fix is off.
+unfix 1
+
+#End

imdclient/data/lammps/md/lammps_v3_nst_1.in

@@ -57,7 +57,7 @@ dump_modify h5md1 unwrap no
 
 ## IMD settings
 # https://docs.lammps.org/fix_imd.html
-fix 2 all imd 8888 version 3 unwrap off nowait off
+fix 2 all imd 8888 version 3 unwrap off nowait off trate 1 time on box on coordinates on velocities on forces on
 
 ## Run MD sim
 run 100

imdclient/data/lammps/md/lammps_v3_nst_8.in

@@ -57,7 +57,7 @@ dump_modify h5md1 unwrap no
 
 ## IMD settings
 # https://docs.lammps.org/fix_imd.html
-fix 2 all imd 8888 version 3 unwrap off nowait off trate 8
+fix 2 all imd 8888 version 3 unwrap off nowait off trate 8 time on box on coordinates on velocities on forces on
 
 ## Run MD sim
 run 100