imdclient 0.1.2__py3-none-any.whl → 0.1.4__py3-none-any.whl

imdclient/{IMDREADER.py → IMD.py}

@@ -50,6 +50,7 @@ class IMDReader(StreamReaderBase):
     ):
         super(IMDReader, self).__init__(filename, **kwargs)
 
+        self._imdclient = None
         logger.debug("IMDReader initializing")
 
         if n_atoms is None:
@@ -83,9 +84,8 @@ class IMDReader(StreamReaderBase):
 
         try:
             imdf = self._imdclient.get_imdframe()
-        except EOFError:
-            # Not strictly necessary, but for clarity
-            raise StopIteration
+        except EOFError as e:
+            raise e
 
         self._frame = frame
         self._load_imdframe_into_ts(imdf)
@@ -126,6 +126,7 @@ class IMDReader(StreamReaderBase):
     def close(self):
         """Gracefully shut down the reader. Stops the producer thread."""
         logger.debug("IMDReader close() called")
-        self._imdclient.stop()
+        if self._imdclient is not None:
+            self._imdclient.stop()
         # NOTE: removeme after testing
         logger.debug("IMDReader shut down gracefully.")

imdclient/IMDClient.py

@@ -25,6 +25,7 @@ import time
 import numpy as np
 from typing import Union, Dict
 import signal
+import atexit
 
 logger = logging.getLogger(__name__)
 
@@ -42,7 +43,13 @@ class IMDClient:
     socket_bufsize : int, (optional)
         Size of the socket buffer in bytes. Default is to use the system default
     buffer_size : int (optional)
-        IMDFramebuffer will be filled with as many :class:`IMDFrame` fit in `buffer_size` [``10MB``]
+        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``]
+        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.
     **kwargs : dict (optional)
         Additional keyword arguments to pass to the :class:`BaseIMDProducer` and :class:`IMDFrameBuffer`
     """
@@ -54,6 +61,7 @@ class IMDClient:
         n_atoms,
         socket_bufsize=None,
         multithreaded=True,
+        continue_after_disconnect=None,
         **kwargs,
     ):
 
@@ -61,6 +69,7 @@ class IMDClient:
         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 self._multithreaded:
             self._buf = IMDFrameBuffer(
@@ -68,8 +77,10 @@ class IMDClient:
                 n_atoms,
                 **kwargs,
             )
+            self._error_queue = queue.Queue()
         else:
             self._buf = None
+            self._error_queue = None
         if self._imdsinfo.version == 2:
             self._producer = IMDProducerV2(
                 self._conn,
@@ -77,6 +88,7 @@ class IMDClient:
                 self._imdsinfo,
                 n_atoms,
                 multithreaded,
+                self._error_queue,
                 **kwargs,
             )
         elif self._imdsinfo.version == 3:
@@ -86,23 +98,60 @@ class IMDClient:
                 self._imdsinfo,
                 n_atoms,
                 multithreaded,
+                self._error_queue,
                 **kwargs,
             )
 
         self._go()
 
         if self._multithreaded:
+            # Disconnect MUST occur. This covers typical cases (Python, IPython interpreter)
             signal.signal(signal.SIGINT, self.signal_handler)
+            signal.signal(signal.SIGTERM, self.signal_handler)
+
+            # Disconnect and socket shutdown MUST occur. This covers Jupyter use
+            # since in jupyter, the signal handler is reset to the default
+            # by pre- and post- hooks
+            # https://stackoverflow.com/questions/70841648/jupyter-reverts-signal-handler-to-default-when-running-next-cell
+            try:
+                import IPython
+            except ImportError:
+                has_ipython = False
+            else:
+                has_ipython = True
+
+            if has_ipython:
+                try:
+                    from IPython import get_ipython
+
+                    if get_ipython() is not None:
+                        kernel = get_ipython().kernel
+                        kernel.pre_handler_hook = lambda: None
+                        kernel.post_handler_hook = lambda: None
+                        logger.debug("Running in Jupyter")
+                except NameError:
+                    logger.debug("Running in non-jupyter IPython environment")
+
+            # Final case: error is raised outside of IMDClient code
+            logger.debug("Registering atexit")
+            atexit.register(self.stop)
+
             self._producer.start()
 
     def signal_handler(self, sig, frame):
         """Catch SIGINT to allow clean shutdown on CTRL+C
         This also ensures that main thread execution doesn't get stuck
         waiting in buf.pop_full_imdframe()"""
+        logger.debug("Intercepted signal")
         self.stop()
+        logger.debug("Shutdown success")
 
     def get_imdframe(self):
         """
+        Returns
+        -------
+        IMDFrame
+            The next frame from the IMD server
         Raises
         ------
         EOFError
@@ -116,6 +165,9 @@ class IMDClient:
                 # and doesn't need to be notified
                 self._disconnect()
                 self._stopped = True
+
+                if self._error_queue.qsize():
+                    raise EOFError(f"{self._error_queue.get()}")
                 raise EOFError
         else:
             try:
@@ -125,14 +177,23 @@ class IMDClient:
                 raise EOFError
 
     def get_imdsessioninfo(self):
+        """
+        Returns
+        -------
+        IMDSessionInfo
+            Information about the IMD session
+        """
         return self._imdsinfo
 
     def stop(self):
+        """
+        Stop the client and close the connection
+        """
         if self._multithreaded:
             if not self._stopped:
-                self._buf.notify_consumer_finished()
-                self._disconnect()
                 self._stopped = True
+                self._disconnect()
+                self._buf.notify_consumer_finished()
         else:
             self._disconnect()
 
@@ -239,6 +300,17 @@ class IMDClient:
         self._conn.sendall(go)
         logger.debug("IMDClient: Sent go packet to server")
 
+        if self._continue_after_disconnect is not None:
+            wait_behavior = (int)(not self._continue_after_disconnect)
+            wait_packet = create_header_bytes(
+                IMDHeaderType.IMD_WAIT, wait_behavior
+            )
+            self._conn.sendall(wait_packet)
+            logger.debug(
+                "IMDClient: Attempted to change wait behavior to %s", 
+                not self._continue_after_disconnect 
+            )
+
     def _disconnect(self):
         # MUST disconnect before stopping execution
         # if simulation already ended, this method will do nothing
@@ -254,6 +326,13 @@ class IMDClient:
         finally:
             self._conn.close()
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.stop()
+        return False
+
 
 class BaseIMDProducer(threading.Thread):
     """
@@ -269,11 +348,14 @@ class BaseIMDProducer(threading.Thread):
         Information about the IMD session
     n_atoms : int
         Number of atoms in the simulation
-    multithreaded : bool, optional
+    multithreaded : bool
         If True, socket interaction will occur in a separate thread &
         frames will be buffered. Single-threaded, blocking IMDClient
-        should only be used in testing [[``True``]]
-
+        should only be used in testing
+    error_queue: queue.Queue
+        Queue to hold errors produced by the producer thread
+    timeout : int, optional
+        Timeout for the socket in seconds [``5``]
     """
 
     def __init__(
@@ -282,7 +364,8 @@ class BaseIMDProducer(threading.Thread):
         buffer,
         sinfo,
         n_atoms,
-        multithreaded=True,
+        multithreaded,
+        error_queue,
         timeout=5,
         **kwargs,
     ):
@@ -291,6 +374,7 @@ class BaseIMDProducer(threading.Thread):
         self._imdsinfo = sinfo
         self._paused = False
 
+        self.error_queue = error_queue
         # Timeout for first frame should be longer
         # than rest of frames
         self._timeout = timeout
@@ -385,6 +469,7 @@ class BaseIMDProducer(threading.Thread):
             logger.debug("IMDProducer: Simulation ended normally, cleaning up")
         except Exception as e:
             logger.debug("IMDProducer: An unexpected error occurred: %s", e)
+            self.error_queue.put(e)
         finally:
             logger.debug("IMDProducer: Stopping run loop")
             # Tell consumer not to expect more frames to be added
@@ -400,9 +485,19 @@ class BaseIMDProducer(threading.Thread):
             )
         # Sometimes we do not care what the value is
         if expected_value is not None and header.length != expected_value:
-            raise RuntimeError(
-                f"IMDProducer: Expected header value {expected_value}, got {header.length}"
-            )
+            if expected_type in [
+                IMDHeaderType.IMD_FCOORDS,
+                IMDHeaderType.IMD_VELOCITIES,
+                IMDHeaderType.IMD_FORCES,
+            ]:
+                raise RuntimeError(
+                    f"IMDProducer: Expected n_atoms value {expected_value}, got {header.length}. "
+                    + "Ensure you are using the correct topology file."
+                )
+            else:
+                raise RuntimeError(
+                    f"IMDProducer: Expected header value {expected_value}, got {header.length}"
+                )
 
     def _get_header(self):
         self._read(self._header)
@@ -422,9 +517,18 @@ class BaseIMDProducer(threading.Thread):
 
 
 class IMDProducerV2(BaseIMDProducer):
-    def __init__(self, conn, buffer, sinfo, n_atoms, multithreaded, **kwargs):
+    def __init__(
+        self,
+        conn,
+        buffer,
+        sinfo,
+        n_atoms,
+        multithreaded,
+        error_queue,
+        **kwargs,
+    ):
         super(IMDProducerV2, self).__init__(
-            conn, buffer, sinfo, n_atoms, multithreaded, **kwargs
+            conn, buffer, sinfo, n_atoms, multithreaded, error_queue, **kwargs
         )
 
         self._energies = bytearray(IMDENERGYPACKETLENGTH)
@@ -517,6 +621,7 @@ class IMDProducerV3(BaseIMDProducer):
         sinfo,
         n_atoms,
         multithreaded,
+        error_queue,
         **kwargs,
     ):
         super(IMDProducerV3, self).__init__(
@@ -525,6 +630,7 @@ class IMDProducerV3(BaseIMDProducer):
             sinfo,
             n_atoms,
             multithreaded,
+            error_queue,
             **kwargs,
         )
         # The body of an x/v/f packet should contain
@@ -633,9 +739,6 @@ class IMDProducerV3(BaseIMDProducer):
                 ).reshape((self._n_atoms, 3)),
             )
 
-    def __del__(self):
-        logger.debug("IMDProducer: I am being deleted")
-
 
 class IMDFrameBuffer:
     """

imdclient/IMDProtocol.py

@@ -34,6 +34,7 @@ class IMDHeaderType(Enum):
     IMD_BOX = 13
     IMD_VELOCITIES = 14
     IMD_FORCES = 15
+    IMD_WAIT = 16
 
 
 def parse_energy_bytes(data, endianness):

imdclient/data/gromacs/md/gromacs_v3_nst1.mdp

@@ -1,5 +1,5 @@
 title                   = PRODUCTION IN NPT
-ld-seed                 = 1
+ld-seed			= 1
 ; Run parameters
 integrator              = md        ; leap-frog integrator
 nsteps                  = 100    ; 1 * 1000 = 1 ps
@@ -7,9 +7,9 @@ dt                      = 0.001     ; 1 fs
 ; Output control
 nstxout                 = 1      ; save coordinates every 1 fs
 nstvout                 = 1      ; save velocities every 1 fs
-nstfout                 = 1
+nstfout                 = 1      ; save forces every 1 fs
 nstenergy               = 1      ; save energies every 1 fs
-nstlog                  = 10      ; update log file every 1 ps
+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) 

imdclient/data/gromacs/md/gromacs_v3_nst8.mdp

@@ -0,0 +1,58 @@
+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             = 3
+IMD-time                = yes
+IMD-box                 = yes
+IMD-coords              = yes
+IMD-unwrap              = no
+IMD-vels                = yes
+IMD-forces              = yes
+IMD-energies            = no

imdclient/data/lammps/md/{lammps_v3.in → lammps_v3_nst_1.in}

@@ -52,8 +52,8 @@ velocity all create 300 102939 dist gaussian mom yes rot yes
 fix 1 all nve 
 
 # Create source of truth trajectory
-# dump h5md1 all h5md 1 lammps_trj.h5md position velocity force box yes
-# dump_modify h5md1 unwrap no
+dump h5md1 all h5md 1 lammps_trj.h5md position velocity force box yes
+dump_modify h5md1 unwrap no
 
 ## IMD settings
 # https://docs.lammps.org/fix_imd.html
@@ -63,7 +63,7 @@ fix 2 all imd 8888 version 3 unwrap off nowait off
 run 100
 
 # Stop dumping information to the dump file.
-# undump h5md1 
+undump h5md1 
 
 # Unfix the NVE. Additional lines if any will assume that this fix is off.
 unfix 1

imdclient/data/lammps/md/lammps_v3_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 velocity force box yes
+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
+
+## 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/namd/md/{namd_v3.namd → namd_v3_nst_1.namd}

@@ -17,8 +17,20 @@ switchdist	7.0
 cutoff		8.0
 pairlistdist 	9.0
 
-dcdfile	alanin.dcd
-dcdfreq	1
+# Add box dimensions 
+cellBasisVector1  32.76  0.0   0.0
+cellBasisVector2   0.0  31.66  0.0
+cellBasisVector3   0.0   0.0  32.89
+
+DCDfile	alanin.dcd
+DCDfreq	1
+DCDUnitCell    yes
+velDcdFile	alanin.vel.dcd
+velDcdFreq	1
+forceDcdFile	alanin.force.dcd
+forceDcdFreq	1
+XSTFile     alanin.xst
+xstFreq		1
 
 #restartname	alanin.restart
 #restartfreq	10
@@ -40,8 +52,8 @@ IMDwait		on
 IMDversion     3
 IMDsendPositions        yes
 IMDsendEnergies     yes
-#IMDsendTime        yes
-#IMDsendBoxDimensions       yes
+IMDsendTime        yes
+IMDsendBoxDimensions       yes
 IMDsendVelocities       yes
-#IMDsendForces      yes
+IMDsendForces      yes
 IMDwrapPositions       yes

imdclient/data/namd/md/namd_v3_nst_8.namd

@@ -0,0 +1,59 @@
+# This is a test namd configuration file
+
+timestep	0.5
+numsteps	100
+structure	alanin.psf
+parameters	alanin.params
+coordinates	alanin.pdb
+exclude		scaled1-4
+1-4scaling	0.4
+outputname 	output[myReplica]
+margin		1.0
+stepspercycle   3
+temperature	0
+
+switching 	on
+switchdist	7.0
+cutoff		8.0
+pairlistdist 	9.0
+
+# Add box dimensions 
+cellBasisVector1  32.76  0.0   0.0
+cellBasisVector2   0.0  31.66  0.0
+cellBasisVector3   0.0   0.0  32.89
+
+DCDfile	alanin.dcd
+DCDfreq	8
+DCDUnitCell    yes
+velDcdFile	alanin.vel.dcd
+velDcdFreq	8
+forceDcdFile	alanin.force.dcd
+forceDcdFreq	8
+XSTFile     alanin.xst
+xstFreq		8
+
+#restartname	alanin.restart
+#restartfreq	10
+
+#langevin	on
+#langevinTemp	300.0
+#langevincol	O
+
+#constraints	on
+
+#fma		on
+
+seed		12345
+
+IMDon		yes
+IMDport		8888
+IMDfreq		8
+IMDwait		on
+IMDversion     3
+IMDsendPositions        yes
+IMDsendEnergies     yes
+IMDsendTime        yes
+IMDsendBoxDimensions       yes
+IMDsendVelocities       yes
+IMDsendForces      yes
+IMDwrapPositions       yes