diff --git a/src/epyscan/__init__.py b/src/epyscan/__init__.py index 2c03625..7284c8d 100644 --- a/src/epyscan/__init__.py +++ b/src/epyscan/__init__.py @@ -1,3 +1,4 @@ +import re from copy import deepcopy from pathlib import Path from typing import Union @@ -7,8 +8,63 @@ from scipy.stats import qmc +def allocate_next_run_number(root: Path): + """ + Return the next available run number in the given root directory. + + This function searches all existing run directories and returns the + highest run number plus 1, or 0 if no runs exist. Only the lowest-level + run directories are considered (e.g., 'run_0', 'run_50'). + + Parameters + ---------- + root : Path + The root directory to search for existing run directories. + + Returns + ------- + int + The next available run number, or 0 if none exist. + """ + run_nums = [ + int(path.name[4:]) + for path in root.rglob("*") + if re.fullmatch(r"run_\d+", path.name) + ] + + return max(run_nums, default=-1) + 1 + + def rundir_hierarchy(root: Path, run_num: int) -> Path: - """Create nested directory structure for a run""" + """ + Create the nested directory structure for a given run number. + + Campaign runs are organised into a hierarchy to avoid putting too many + files in a single directory. The structure follows `rundir_hierarchy`: + + ./run_0_1000000/run_0_10000/run_0_100/run_[0-99] + + As the run numbers increase, higher-level directories are split further: + + ./run_0_1000000/run_0_10000/run_100_200/run_[100-199] + ./run_0_1000000/run_10000_20000/run_10000_10100/run_[10000-10100] + + The hierarchy ensures that only the lowest-level directories (e.g., 'run_0', + 'run_50') contain individual simulation runs, while higher-level directories + manage grouping of runs in batches. + + Parameters + ---------- + root : Path + The root directory where the run hierarchy should be created. + run_num : int + The run number for which the directories will be created. + + Returns + ------- + Path + The full path to the newly created run directory. + """ def level_dir(exponent: int) -> Path: level = 100**exponent @@ -52,10 +108,12 @@ class Campaign: Base template deck as a Python dict (for example, as created by `epydeck`) root: Path to root run directory + append: + Flag to decide whether to append new runs to the end of an existing + campaign (if True) or overwrite it (if False) Examples -------- - >>> with open("template.deck") as f: template = epydeck.load(f) >>> campaign = Campaign(template, "grid_root") @@ -63,8 +121,8 @@ class Campaign: """ - def __init__(self, template: dict, root: Union[str, Path]): - self._counter = 0 + def __init__(self, template: dict, root: Union[str, Path], append: bool = False): + self._counter = allocate_next_run_number(root) if append else 0 self.template = template self.root = Path(root) diff --git a/tests/test_epyscan.py b/tests/test_epyscan.py index b226b07..bc0295d 100644 --- a/tests/test_epyscan.py +++ b/tests/test_epyscan.py @@ -1,3 +1,5 @@ +from pathlib import Path + import epydeck import numpy as np @@ -143,3 +145,66 @@ def test_campaign(tmp_path): actual_case_deck = epydeck.load(f) assert actual_case_deck == expected_case_deck + + +def test_load_existing_campaign(tmp_path): + # Simulate a pre-existing campaign by making fake paths + base_path = tmp_path / "run_0_1000000/run_0_10000/run_0_100" + + existing_campaign_paths = [ + base_path / "run_0", + base_path / "run_1", + base_path / "run_2", + base_path / "run_3", + base_path / "run_4", + ] + + for path in existing_campaign_paths: + Path(path).mkdir(parents=True, exist_ok=True) + + template = {"block": {"var4": 1.23}, "other_block": {"var5": True}} + campaign = epyscan.Campaign(template, tmp_path, append=True) + + new_path = campaign.setup_case({"block": {"var4": 1.24}}) + assert new_path == (base_path / "run_5") + assert new_path.exists() + + +def test_override_existing_campaign(tmp_path): + base_path = tmp_path / "run_0_1000000/run_0_10000/run_0_100" + + existing_campaign_paths = [ + base_path / "run_0", + base_path / "run_1", + base_path / "run_2", + base_path / "run_3", + base_path / "run_4", + ] + + for path in existing_campaign_paths: + Path(path).mkdir(parents=True, exist_ok=True) + + template = {"block": {"var4": 1.23}, "other_block": {"var5": True}} + campaign = epyscan.Campaign(template, tmp_path) + + new_path = campaign.setup_case({"block": {"var4": 1.24}}) + assert new_path == (base_path / "run_0") + assert new_path.exists() + + +def test_retrieve_existing_run_count(tmp_path): + base_path = tmp_path / "run_0_1000000/run_0_10000/run_0_100" + (base_path / "run_0").mkdir(parents=True) + (base_path / "run_1").mkdir() + (base_path / "run_50").mkdir() + (base_path / "run_100").mkdir() + + assert epyscan.allocate_next_run_number(tmp_path) == 101 + + +def test_retrieve_existing_run_count_edge_cases(tmp_path): + (tmp_path / "run_10").mkdir() + (tmp_path / "run_20_30").mkdir() + (tmp_path / "not_a_run").mkdir() + + assert epyscan.allocate_next_run_number(tmp_path) == 11