diff --git a/janus_core/calculations/md.py b/janus_core/calculations/md.py index d3ea513f..006bce4c 100644 --- a/janus_core/calculations/md.py +++ b/janus_core/calculations/md.py @@ -808,7 +808,18 @@ def _prepare_restart(self) -> None: restart_stem = self._restart_stem # Use restart_stem.name otherwise T300.0 etc. counts as extension - poss_restarts = restart_stem.parent.glob(f"{restart_stem.name}*.extxyz") + poss_restarts = list( + restart_stem.parent.glob(f"{restart_stem.name}*.extxyz") + ) + + # If no restarts found, check potential last results directory + if not poss_restarts: + results_dirs = sorted(restart_stem.parent.parent.glob("janus_results*")) + if results_dirs: + poss_restarts = results_dirs[-1].glob( + f"{restart_stem.name}*.extxyz" + ) + try: last_restart = max(poss_restarts, key=getmtime) diff --git a/janus_core/helpers/utils.py b/janus_core/helpers/utils.py index 41456b64..432d714c 100644 --- a/janus_core/helpers/utils.py +++ b/janus_core/helpers/utils.py @@ -7,6 +7,7 @@ from io import StringIO from logging import Logger from pathlib import Path +import re from typing import Any, Literal, TextIO from ase import Atoms @@ -110,8 +111,23 @@ def _get_default_prefix( struct_name = struct.get_chemical_formula() # Set default output directory - prefix = Path("./janus_results") / struct_name + results_dir = Path("./janus_results") + results_dirs = Path(".").glob("janus_results_*") + + if results_dirs: + idx = ( + int(match.group(0)) + for dir in results_dirs + if (match := re.search("([0-9]+)$", str(dir))) + ) + new_idx = max(idx, default=-1) + 1 + + if new_idx: + results_dir = Path(f"./janus_results_{new_idx}") + elif results_dir.exists(): + results_dir = Path("./janus_results_1") + prefix = results_dir / struct_name return "-".join((str(prefix), *filter(None, additional))) def _build_filename( diff --git a/tests/test_filenamemixin.py b/tests/test_filenamemixin.py index c3a01955..02ccbabd 100644 --- a/tests/test_filenamemixin.py +++ b/tests/test_filenamemixin.py @@ -8,6 +8,7 @@ import pytest from janus_core.helpers.utils import FileNameMixin +from tests.utils import chdir DATA_PATH = Path(__file__).parent / "data" STRUCT = read(DATA_PATH / "benzene.xyz") @@ -125,3 +126,35 @@ def test_file_name_mixin_build(mixin_params, file_args, file_kwargs, file_name): assert ( file_mix.build_filename(*file_args, **file_kwargs) == Path(file_name).absolute() ) + + +@pytest.mark.parametrize( + "existing_dirs, expected_new_dir", + ( + ((), "janus_results"), + (("janus_results",), "janus_results_1"), + (("janus_results", "janus_results_1"), "janus_results_2"), + (("janus_results", "janus_results_1", "janus_results_2"), "janus_results_3"), + # Check non-sequential output directories + (("janus_results_1",), "janus_results_2"), + (("janus_results_2",), "janus_results_3"), + (("janus_results_1", "janus_results_2"), "janus_results_3"), + (("janus_results", "janus_results_2"), "janus_results_3"), + # Check non-default output directories + (("janus_results1"), "janus_results"), + (("janus_results", "janus_results1"), "janus_results_1"), + (("janus_result"), "janus_results"), + (("janus_results_abc"), "janus_results"), + ), +) +def test_new_results_dir(tmp_path, existing_dirs, expected_new_dir): + """Test the creation of new results directories.""" + with chdir(tmp_path): + for dir in existing_dirs: + Path(dir).mkdir(parents=True, exist_ok=True) + + file_mix = DummyFileHandler(STRUCT, None, None) + assert ( + file_mix.file_prefix.as_posix() + == (Path(expected_new_dir) / "C6H6").as_posix() + )