Modified run_gmx_cmd and moved compare_MDPs from utils to gmx_parser

ensemble_md/tests/test_analyze_traj.py

@@ -11,22 +11,24 @@
 Unit tests for the module analyze_traj.py.
 """
 import os
+import numpy as np
 from ensemble_md.analysis import analyze_traj
 
 current_path = os.path.dirname(os.path.abspath(__file__))
 input_path = os.path.join(current_path, "data")
 
 
 def test_extract_state_traj():
-    traj = analyze_traj.extract_state_traj(os.path.join(input_path, 'dhdl/dhdl_0.xvg'))
+    traj, t = analyze_traj.extract_state_traj(os.path.join(input_path, 'dhdl/dhdl_0.xvg'))
     state_list = [
         0, 0, 3, 1, 4, 4, 5, 4, 5, 5,
         4, 4, 5, 4, 2, 4, 5, 2, 1, 2,
         3, 1, 2, 4, 1, 0, 2, 4, 3, 2,
         1, 3, 3, 4, 2, 3, 1, 1, 0, 1,
         2, 3, 1, 0, 1, 4, 3, 1, 3, 2, 5]
+    t_true = [0.02 * i + 3 for i in range(len(state_list))]
     assert traj == state_list
-
+    assert np.allclose(t, t_true)
 
 def test_stitch_trajs():
     pass

ensemble_md/utils/gmx_parser.py

@@ -15,6 +15,7 @@
 import six
 import logging
 import warnings
+from itertools import combinations
 from collections import OrderedDict as odict
 
 from ensemble_md.utils import utils
@@ -383,3 +384,43 @@ def write(self, filename=None, skipempty=False):
                         mdp.write("{k!s} = {v!s}\n".format(**vars()))
                     else:
                         mdp.write("{} = {}\n".format(k, " ".join(map(str, v))))
+
+
+def compare_MDPs(mdp_list):
+    """
+    Given a list of MDP files, identify the parameters for which not all MDP
+    files have the same values. Note that this function is not aware of the default
+    values of GROMACS parameters. (Currently, this function is not used in the
+    workflow adopted in :code:`run_EEXE.py` but it might be useful in some places,
+    so we decided to keep it.)
+
+    Returns
+    -------
+    diff_params : list
+        The list of parameters differing between the input MDP files.
+    """
+    compare_list = list(combinations(mdp_list, r=2))
+    diff_params = []
+    for i in range(len(compare_list)):
+        params_1 = MDP(compare_list[i][0])
+        params_2 = MDP(compare_list[i][1])
+
+        mdp_1 = odict([(k.replace('-', '_'), v) if type(v) is str else (k.replace('-', '_'), v.replace('-', '_')) for k, v in params_1.items()])  # noqa: E501
+        mdp_2 = odict([(k.replace('-', '_'), v) if type(v) is str else (k.replace('-', '_'), v.replace('-', '_')) for k, v in params_2.items()])  # noqa: E501
+
+        # First figure out the union set of the parameters and exclude blanks and comments
+        all_params = set(list(mdp_1.keys()) + list(mdp_2.keys()))
+        all_params = [p for p in all_params if not p.startswith(('B', 'C'))]
+
+        for p in all_params:
+            if p in diff_params:
+                pass  # already in the list, no need to compare again
+            else:
+                if p not in mdp_1 or p not in mdp_2:
+                    diff_params.append(p)
+                else:
+                    # the parameter is in both MDP files
+                    if mdp_1[p] != mdp_2[p]:
+                        diff_params.append(p)
+
+    return diff_params

ensemble_md/utils/utils.py

@@ -17,9 +17,6 @@
 import subprocess
 import collections
 import numpy as np
-from itertools import combinations
-from collections import OrderedDict
-from ensemble_md.utils import gmx_parser
 
 
 class Logger:
@@ -72,15 +69,17 @@ def flush(self):
         pass
 
 
-def run_gmx_cmd(arguments):
+def run_gmx_cmd(arguments, prompt_input=None):
     """
-    Run a GROMACS command as a subprocess
+    Runs a GROMACS command as a subprocess
 
     Parameters
     ----------
     arguments : list
         A list of arguments that compose of the GROMACS command to run, e.g.
         :code:`['gmx', 'mdrun', '-deffnm', 'sys']`.
+    prompt_input : str or None
+        The input to be passed to the GROMACS command when it prompts for input.
 
     Returns
     -------
@@ -93,54 +92,14 @@ def run_gmx_cmd(arguments):
 
     """
     try:
-        result = subprocess.run(arguments, capture_output=True, text=True, check=True)
+        result = subprocess.run(arguments, capture_output=True, text=True, input=prompt_input, check=True)
         return_code, stdout, stderr = result.returncode, result.stdout, None
     except subprocess.CalledProcessError as e:
         return_code, stdout, stderr = e.returncode, None, e.stderr
 
     return return_code, stdout, stderr
 
 
-def compare_MDPs(mdp_list):
-    """
-    Given a list of MDP files, identify the parameters for which not all MDP
-    files have the same values. Note that this function is not aware of the default
-    values of GROMACS parameters. (Currently, this function is not used in the
-    workflow adopted in :code:`run_EEXE.py` but it might be useful in some places,
-    so we decided to keep it.)
-
-    Returns
-    -------
-    diff_params : list
-        The list of parameters differing between the input MDP files.
-    """
-    compare_list = list(combinations(mdp_list, r=2))
-    diff_params = []
-    for i in range(len(compare_list)):
-        params_1 = gmx_parser.MDP(compare_list[i][0])
-        params_2 = gmx_parser.MDP(compare_list[i][1])
-
-        mdp_1 = OrderedDict([(k.replace('-', '_'), v) if type(v) is str else (k.replace('-', '_'), v.replace('-', '_')) for k, v in params_1.items()])  # noqa: E501
-        mdp_2 = OrderedDict([(k.replace('-', '_'), v) if type(v) is str else (k.replace('-', '_'), v.replace('-', '_')) for k, v in params_2.items()])  # noqa: E501
-
-        # First figure out the union set of the parameters and exclude blanks and comments
-        all_params = set(list(mdp_1.keys()) + list(mdp_2.keys()))
-        all_params = [p for p in all_params if not p.startswith(('B', 'C'))]
-
-        for p in all_params:
-            if p in diff_params:
-                pass  # already in the list, no need to compare again
-            else:
-                if p not in mdp_1 or p not in mdp_2:
-                    diff_params.append(p)
-                else:
-                    # the parameter is in both MDP files
-                    if mdp_1[p] != mdp_2[p]:
-                        diff_params.append(p)
-
-    return diff_params
-
-
 def format_time(t):
     """
     Converts time in seconds to the "most readable" format.