More api consistency

CHANGELOG.md

@@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- all `asyncmd.utils` methods now require a MDEngine (class) to dispatch to the correct engine submodule methods
 - GmxEngine `apply_constraints` and `generate_velocities` methods: rename `wdir` argument to `workdir` to make it consistent with `prepare` and `prepare_from_files` (also add the `workdir` argument to the MDEngine ABC).
 
 ### Fixed

src/asyncmd/gromacs/mdengine.py

@@ -211,7 +211,7 @@
                                            top_file=top_file,
                                            )
         self._engine_state = _GmxEngineState()
         # TODO: store a hash/the file contents for gro, top, ndx to check against
         #       when we load from storage/restart? if we do this, do it in the property!
         self.gro_file = gro_file
         self.top_file = top_file
@@ -508,8 +508,8 @@
                                           k=6,
                                           )
                            )
-        swdir = os.path.join(wdir, run_name)
-        await aiofiles.os.mkdir(swdir)
+        wdir = os.path.join(wdir, run_name)
+        await aiofiles.os.mkdir(wdir)
         constraints_mdp = copy.deepcopy(self.mdp)
         constraints_mdp["continuation"] = "no" if constraints else "yes"
         constraints_mdp["gen-vel"] = "yes" if generate_velocities else "no"
@@ -523,12 +523,12 @@
             # make sure we have draw a new/different random number for gen-vel
             constraints_mdp["gen-seed"] = -1
         constraints_mdp["nsteps"] = 0
-        await self._run_grompp(workdir=swdir, deffnm=run_name,
+        await self._run_grompp(workdir=wdir, deffnm=run_name,
                                trr_in=conf_in.trajectory_files[0],
-                               tpr_out=os.path.join(swdir, f"{run_name}.tpr"),
+                               tpr_out=os.path.join(wdir, f"{run_name}.tpr"),
                                mdp_obj=constraints_mdp)
-        cmd_str = self._mdrun_cmd(tpr=os.path.join(swdir, f"{run_name}.tpr"),
-                                  workdir=swdir,
+        cmd_str = self._mdrun_cmd(tpr=os.path.join(wdir, f"{run_name}.tpr"),
+                                  workdir=wdir,
                                   deffnm=run_name)
         logger.debug("About to execute gmx mdrun command for constraints and"
                      "/or velocity generation: %s",
@@ -537,9 +537,9 @@
         stdout = bytes()
         await self._acquire_resources_gmx_mdrun()
         mdrun_proc = await self._start_gmx_mdrun(
-                        cmd_str=cmd_str, workdir=swdir,
+                        cmd_str=cmd_str, workdir=wdir,
                         run_name=run_name,
                         # TODO: we hardcode that the 0step MD runs can not be longer than 15 min
                         # (but i think this should be fine for randomizing velocities and/or
                         #  applying constraints?!)
                         walltime=0.25,
@@ -563,25 +563,25 @@
             # the FrameExtractor (i.e. MDAnalysis) handle any potential conversions
             engine_traj = Trajectory(
                             trajectory_files=os.path.join(
-                                swdir, f"{run_name}{self._num_suffix(1)}.trr"
+                                wdir, f"{run_name}{self._num_suffix(1)}.trr"
                                 ),
                             structure_file=conf_in.structure_file,
                             )
             extractor = NoModificationFrameExtractor()
             # Note: we use extract (and not extract_async) because otherwise
             #       it can happen in super-rare circumstances that the Trajectory
             #       we just instantiated is "replaced" by a Trajectory with the
-            #       same hash but a different filename/path, then the extraction
+            #       same hash but a different filename/path. If then in addition
+            #       this trajectory is removed before extracting, the extraction
             #       fails. If we dont await this can not happen since we do not
             #       give up control in between.
-            out_traj = extractor.extract(outfile=conf_out_name,
-                                         traj_in=engine_traj,
-                                         idx=len(engine_traj) - 1,
-                                         )
-            return out_traj
+            return extractor.extract(outfile=conf_out_name,
+                                     traj_in=engine_traj,
+                                     idx=len(engine_traj) - 1,
+                                     )
         finally:
-            await self._cleanup_gmx_mdrun(workdir=swdir, run_name=run_name)
-            shutil.rmtree(swdir)  # remove the whole directory we used as wdir
+            await self._cleanup_gmx_mdrun(workdir=wdir, run_name=run_name)
+            shutil.rmtree(wdir)  # remove the whole directory we used as wdir
 
     async def prepare(self, starting_configuration: Trajectory | None | str,
                       workdir: str, deffnm: str) -> None:
@@ -965,7 +965,7 @@
             cmd += f" -n {ndx_file}"
         if trr_in is not None:
             # input trr is optional
             # TODO /NOTE: currently we do not pass '-time', i.e. we just use the
             #            gmx default frame selection: last frame from trr
             trr_in = os.path.relpath(trr_in, start=workdir)
             cmd += f" -t {trr_in}"
@@ -996,7 +996,7 @@
         #       however gromacs -deffnm is deprecated (and buggy),
         #       so we just make our own 'deffnm', i.e. we name all files the same
         #       except for the ending but do so explicitly
         # TODO /FIXME: we dont specify the names for e.g. pull outputfiles,
         #             so they will have their default names and will collide
         #             when running multiple engines in the same folder!
         cmd = f"{self.mdrun_executable} -noappend -s {tpr}"
@@ -1026,7 +1026,7 @@
     # purposes) be used as a drop-in replacement. Therefore we only need to
     # reimplement `_start_gmx_mdrun()`, `_acquire_resources_gmx_mdrun()` and
     # `_cleanup_gmx_mdrun()` to have a working SlurmGmxEngine.
     # TODO: use SLURM also for grompp?! (would it make stuff faster?)
     #       I (hejung) think probably not by much because we already use
     #       asyncios subprocess for grompp (i.e. do it asynchronous) and grompp
     #       will most likely not take much resources on the login (local) node

src/asyncmd/trajectory/propagate.py

@@ -253,6 +253,7 @@
                                                 folder=workdir,
                                                 deffnm=deffnm,
                                                 file_ending=ending.lower(),
+                                                engine=self.engine_cls,
                                                        )
             # make sure we dont miss anything because we have different
             # capitalization
@@ -261,12 +262,13 @@
                                                     folder=workdir,
                                                     deffnm=deffnm,
                                                     file_ending=ending.upper(),
+                                                    engine=self.engine_cls,
                                                            )
             await asyncio.gather(*(aiofiles.os.unlink(f)
                                    for f in parts_to_remove
                                    )
                                  )
             # TODO: address the note below?
             # NOTE: this is a bit hacky: we just try to remove the offset and
             #       lock files for every file we remove (since we do not know
             #       if the file we remove is a trajectory [and therefore

src/asyncmd/utils.py

@@ -20,6 +20,8 @@
 It also includes various functions to retrieve or ensure important parameters from
 MDConfig/MDEngine combinations, such as nstout_from_mdconfig and ensure_mdconfig_options.
 """
+import logging
+
 from .mdengine import MDEngine
 from .mdconfig import MDConfig
 from .trajectory.trajectory import Trajectory
@@ -28,7 +30,11 @@
 from .gromacs import mdconfig as gmx_config
 
 
-async def get_all_traj_parts(folder: str, deffnm: str, engine: MDEngine) -> list[Trajectory]:
+logger = logging.getLogger(__name__)
+
+
+async def get_all_traj_parts(folder: str, deffnm: str, engine: MDEngine | type[MDEngine],
+                             ) -> list[Trajectory]:
     """
     List all trajectories in folder by given engine class with given deffnm.
 
@@ -37,10 +43,12 @@ async def get_all_traj_parts(folder: str, deffnm: str, engine: MDEngine) -> list
     folder : str
         Absolute or relative path to a folder.
     deffnm : str
-        deffnm used by the engines simulation run from which we want the trajs.
-    engine : MDEngine
-        The engine that produced the trajectories
-        (or one from the same class and with similar init args)
+        deffnm used by the engines simulation run from which we want the trajectories.
+    engine : MDEngine | type[MDEngine]
+        The engine that produced the trajectories (or one from the same class
+        and with similar init args). Note that it is also possible to pass an
+        uninitialized engine class, but then the default trajectory output type
+        will be returned.
 
     Returns
     -------
@@ -52,7 +60,17 @@ async def get_all_traj_parts(folder: str, deffnm: str, engine: MDEngine) -> list
     ValueError
         Raised when the engine class is unknown.
     """
-    if isinstance(engine, (gmx_engine.GmxEngine, gmx_engine.SlurmGmxEngine)):
+    # test for uninitialized engine classes, we warn but return the default traj type
+    if isinstance(engine, type) and issubclass(engine, MDEngine):
+        logger.warning("Engine %s is not initialized, i.e. it is an engine class. "
+                       "Returning the default output trajectory type for this "
+                       "engine class.", engine)
+    if (
+        isinstance(engine, (gmx_engine.GmxEngine, gmx_engine.SlurmGmxEngine))
+        or (isinstance(engine, type)  # check that it is a type otherwise issubclass might not work
+            and issubclass(engine, (gmx_engine.GmxEngine, gmx_engine.SlurmGmxEngine))
+            )
+    ):
         return await gmx_utils.get_all_traj_parts(folder=folder, deffnm=deffnm,
                                                   traj_type=engine.output_traj_type,
                                                   )
@@ -61,6 +79,7 @@ async def get_all_traj_parts(folder: str, deffnm: str, engine: MDEngine) -> list
 
 
 async def get_all_file_parts(folder: str, deffnm: str, file_ending: str,
+                             engine: MDEngine | type[MDEngine],
                              ) -> list[str]:
     """
     Find and return all files with given ending produced by a `MDEngine`.
@@ -75,16 +94,24 @@ async def get_all_file_parts(folder: str, deffnm: str, file_ending: str,
         deffnm (prefix of filenames) used in the simulation.
     file_ending : str
         File ending of the requested filetype (with or without preceding ".").
+    engine : MDEngine | type[MDEngine]
+        The engine or engine class that produced the file parts.
 
     Returns
     -------
     list[str]
         Ordered list of filepaths for files with given ending.
     """
-    # TODO: we just use the function from the gromacs engines for now, i.e. we
-    #       assume that the filename scheme will be the same for other engines
-    return await gmx_utils.get_all_file_parts(folder=folder, deffnm=deffnm,
-                                              file_ending=file_ending)
+    if (
+        isinstance(engine, (gmx_engine.GmxEngine, gmx_engine.SlurmGmxEngine))
+        or (isinstance(engine, type)  # check that it is a type otherwise issubclass might not work
+            and issubclass(engine, (gmx_engine.GmxEngine, gmx_engine.SlurmGmxEngine))
+            )
+    ):
+        return await gmx_utils.get_all_file_parts(folder=folder, deffnm=deffnm,
+                                                  file_ending=file_ending)
+    raise ValueError(f"Engine {engine} is not a known MDEngine (class)."
+                     + " Maybe someone just forgot to add the function?")
 
 
 def nstout_from_mdconfig(mdconfig: MDConfig, output_traj_type: str) -> int:

tests/helper_classes.py

@@ -28,13 +28,14 @@ class NoOpMDEngine(MDEngine):
     current_trajectory = None
     output_traj_type = "TEST"
     steps_done = 0
-    async def apply_constraints(self, conf_in: Trajectory, conf_out_name: str) -> Trajectory:
+    async def apply_constraints(self, conf_in: Trajectory, conf_out_name: str,
+                                *, workdir: str = ".") -> Trajectory:
         pass
     async def prepare(self, starting_configuration: Trajectory, workdir: str, deffnm: str) -> None:
         pass
     async def prepare_from_files(self, workdir: str, deffnm: str) -> None:
         pass
-    async def run_walltime(self, walltime: float) -> Trajectory:
+    async def run_walltime(self, walltime: float, max_steps: int | None = None) -> Trajectory:
         pass
     async def run_steps(self, nsteps: int, steps_per_part: bool = False) -> Trajectory:
         pass

tests/test_utils.py

@@ -19,9 +19,12 @@
 i.e. no tests for, e.g., functions from asyncmd.gromacs.utils.
 """
 import pytest
+import logging
 
 from conftest import NoOpMDEngine, NoOpMDConfig
-from asyncmd.utils import (get_all_traj_parts, nstout_from_mdconfig,
+from asyncmd.utils import (get_all_traj_parts,
+                           get_all_file_parts,
+                           nstout_from_mdconfig,
                            ensure_mdconfig_options,
                            )
 
@@ -34,6 +37,13 @@ async def test_get_all_traj_parts(self):
                                      engine=NoOpMDEngine(),
                                      )
 
+    @pytest.mark.asyncio
+    async def test_get_all_file_parts(self):
+        with pytest.raises(ValueError):
+            await get_all_file_parts(folder="test", deffnm="test",
+                                     file_ending=".test", engine=NoOpMDEngine(),
+                                     )
+
     def test_nstout_from_mdconfig(self):
         with pytest.raises(ValueError):
             nstout_from_mdconfig(mdconfig=NoOpMDConfig(),
@@ -42,3 +52,19 @@ def test_nstout_from_mdconfig(self):
     def test_ensure_mdconfig_options(self):
         with pytest.raises(ValueError):
             ensure_mdconfig_options(mdconfig=NoOpMDConfig())
+
+
+class Test_warn_for_default_value_from_engine_class:
+    @pytest.mark.asyncio
+    async def test_get_all_traj_parts(self, caplog):
+        with pytest.raises(ValueError):
+            with caplog.at_level(logging.WARNING):
+                await get_all_traj_parts(folder="test", deffnm="test",
+                                         # this time we use an uninitialized
+                                         # engine class so we get the warning
+                                         # (and then fail after)
+                                         engine=NoOpMDEngine,
+                                         )
+        warn_text = f"Engine {NoOpMDEngine} is not initialized, i.e. it is an engine class. "
+        warn_text += "Returning the default output trajectory type for this engine class."
+        assert warn_text in caplog.text