diff --git a/CHANGELOG.md b/CHANGELOG.md index 06fa4ce..0b2a3f5 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/src/asyncmd/gromacs/mdengine.py b/src/asyncmd/gromacs/mdengine.py index d15797c..9fd6761 100644 --- a/src/asyncmd/gromacs/mdengine.py +++ b/src/asyncmd/gromacs/mdengine.py @@ -508,8 +508,8 @@ async def _0step_md(self, conf_in: Trajectory, conf_out_name: str, *, 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 @@ async def _0step_md(self, conf_in: Trajectory, conf_out_name: str, *, # 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,7 +537,7 @@ async def _0step_md(self, conf_in: Trajectory, conf_out_name: str, *, 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 @@ -563,7 +563,7 @@ async def _0step_md(self, conf_in: Trajectory, conf_out_name: str, *, # 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, ) @@ -571,17 +571,17 @@ async def _0step_md(self, conf_in: Trajectory, conf_out_name: str, *, # 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: diff --git a/src/asyncmd/trajectory/propagate.py b/src/asyncmd/trajectory/propagate.py index 26f04a3..022a6ea 100644 --- a/src/asyncmd/trajectory/propagate.py +++ b/src/asyncmd/trajectory/propagate.py @@ -253,6 +253,7 @@ async def remove_parts(self, workdir: str, deffnm: str, *, folder=workdir, deffnm=deffnm, file_ending=ending.lower(), + engine=self.engine_cls, ) # make sure we dont miss anything because we have different # capitalization @@ -261,6 +262,7 @@ async def remove_parts(self, workdir: str, deffnm: str, *, 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 diff --git a/src/asyncmd/utils.py b/src/asyncmd/utils.py index 5db178f..632e16d 100644 --- a/src/asyncmd/utils.py +++ b/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: diff --git a/tests/helper_classes.py b/tests/helper_classes.py index def8c49..6d4d7a2 100644 --- a/tests/helper_classes.py +++ b/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 diff --git a/tests/test_utils.py b/tests/test_utils.py index 5ed2ecd..e5a6345 100644 --- a/tests/test_utils.py +++ b/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