Merge pull request #32 from aurora-multiphysics/31-improve-file-manag…

…ement-by-design-evaluators

31 improve file management by design evaluators

examples/example_1_simple_monoblock.py

@@ -25,7 +25,7 @@
 # In general, the user will set their own paths and pass them where required.
 EXAMPLES_DIR = SLEDO_ROOT / "examples"
 INPUT_FILE = EXAMPLES_DIR / "input_files" / "simple_monoblock_thermomech.i"
-WORKING_DIR = EXAMPLES_DIR / "results" / "example_1"
+WORKING_DIR = EXAMPLES_DIR / "results"
 PICKLE_FILEPATH = WORKING_DIR / "example_1_optimiser.pickle"
 
 if __name__ == "__main__":
@@ -39,7 +39,6 @@
     design_evaluator = MooseHerderDesignEvaluator(
         metrics,
         INPUT_FILE,  # The base input file to be modified per design iteration.
-        working_dir=WORKING_DIR,  # Directory to store generated files.
         config_path=MOOSE_CONFIG_FILE,  # Contains required MOOSE paths.
     )
 
@@ -58,7 +57,7 @@
         design_evaluator,
         search_space,
         max_total_trials=20,
-        name="simple-monoblock-optimiser",
+        name="example_1",
         data_dir=WORKING_DIR,
     )
 

examples/example_2_monoblock.py

@@ -25,7 +25,7 @@
 # In general, the user will set their own paths and pass them where required.
 EXAMPLES_DIR = SLEDO_ROOT / "examples"
 INPUT_FILE = EXAMPLES_DIR / "input_files" / "monoblock_thermomech.i"
-WORKING_DIR = EXAMPLES_DIR / "results" / "example_2"
+WORKING_DIR = EXAMPLES_DIR / "results"
 PICKLE_FILEPATH = WORKING_DIR / "example_2_optimiser.pickle"
 
 if __name__ == "__main__":
@@ -39,7 +39,6 @@
     design_evaluator = MooseHerderDesignEvaluator(
         metrics,
         INPUT_FILE,  # The base input file to be modified per design iteration.
-        working_dir=WORKING_DIR,  # Directory to store generated files.
         config_path=MOOSE_CONFIG_FILE,  # Contains required MOOSE paths.
     )
 
@@ -60,7 +59,7 @@
         design_evaluator,
         search_space,
         max_total_trials=20,
-        name="monoblock-optimiser",
+        name="example_2",
         data_dir=WORKING_DIR,
     )
 

examples/example_3_catbird_monoblock.py

@@ -38,7 +38,7 @@
 # Set the paths required for this example.
 # In general, the user will set their own paths and pass them where required.
 EXAMPLES_DIR = SLEDO_ROOT / "examples"
-WORKING_DIR = EXAMPLES_DIR / "results" / "example_3"
+WORKING_DIR = EXAMPLES_DIR / "results"
 FACTORY_CONFIG_PATH = WORKING_DIR / "factory_config.json"
 INPUT_FILE_PATH = WORKING_DIR / "trial.i"
 PICKLE_FILEPATH = WORKING_DIR / "example_3_optimiser.pickle"
@@ -76,7 +76,6 @@
         design_evaluator = CatBirdMooseHerderDesignEvaluator(
             metrics,
             model,
-            working_dir=WORKING_DIR,  # Directory to store generated files.
             config_path=MOOSE_CONFIG_FILE,  # Contains required MOOSE paths.
         )
 
@@ -97,7 +96,7 @@
             design_evaluator,
             search_space,
             max_total_trials=20,
-            name="catbird-monoblock-optimiser",
+            name="example_3",
             data_dir=WORKING_DIR,
         )
 

src/sledo/design_evaluator.py

@@ -107,7 +107,6 @@ def __init__(
         self,
         metrics: list[str],
         base_input_file: Path | str,
-        working_dir: Path | str = Path.cwd(),
         config_path: Path | str = MOOSE_CONFIG_FILE,
         run_options: dict = {
             "n_tasks": 1,
@@ -127,9 +126,6 @@ def __init__(
         base_input_file : Path | str
             Path to the base MOOSE input file (.i) to use as the basis for
             generating modified files. This file will not be modified.
-        working_dir : Path | str, optional
-            Path to the working directory to use for storing modified MOOSE
-            input files (.i) and running MOOSE, by default Path.cwd().
         config_path : Path | str, optional
             Path to the config file containing the required paths to run MOOSE,
             by default 'moose_config.json' in the sledo root folder.
@@ -139,7 +135,6 @@ def __init__(
         """
         self._metrics = metrics
         self.base_input_file = Path(base_input_file)
-        self.working_dir = Path(working_dir)
         self.config_path = Path(config_path)
         self.run_options = run_options
 
@@ -170,7 +165,7 @@ def evaluate_design(self, parameters: dict, timestep: int = -1) -> dict:
         # Generate input file.
         trial_filepath = generate_modified_input_file(
             self.base_input_file,
-            self.working_dir / "trial.i",
+            Path.cwd() / "trial.i",
             parameters,
         )
         # Run simulation.
@@ -199,7 +194,6 @@ def __init__(
         self,
         metrics: list[str],
         model: MooseModel,
-        working_dir: Path | str = Path.cwd(),
         config_path: Path | str = MOOSE_CONFIG_FILE,
         run_options: dict = {
             "n_tasks": 1,
@@ -219,9 +213,6 @@ def __init__(
         model : MooseModel
             A catbird MooseModel capable of updating parameters and writing a
             MOOSE input file.
-        working_dir : Path | str, optional
-            Path to the working directory to use for storing modified MOOSE
-            input files (.i) and running MOOSE, by default Path.cwd().
         config_path : Path | str, optional
             Path to the config file containing the required paths to run MOOSE,
             by default 'moose_config.json' in the sledo root folder.
@@ -231,7 +222,6 @@ def __init__(
         """
         self._metrics = metrics
         self._model = model
-        self.working_dir = Path(working_dir)
         self.config_path = Path(config_path)
         self.run_options = run_options
 
@@ -333,7 +323,7 @@ def evaluate_design(self, parameters: dict, timestep: int = -1) -> dict:
         """
         # Generate input file.
         trial_filepath = self.generate_input_file(
-            self.working_dir / "trial.i",
+            Path.cwd() / "trial.i",
             parameters,
         )
         # Run simulation.

src/sledo/optimiser.py

@@ -4,6 +4,7 @@
 (c) Copyright UKAEA 2023-2024.
 """
 
+import os
 from pathlib import Path
 import dill
 
@@ -88,27 +89,43 @@ def __init__(
         if not self.data_dir.exists():
             self.data_dir.mkdir()
 
-        # Add a step to the evaluation function reporting the results of each
-        # design evaluation to the ray tune optimiser.
-        self.evaluation_function = lambda results_dict: train.report(
-            self.design_evaluator.evaluate_design(results_dict)
-        )
+        # Workaround to a bug which causes ray to save to both the passed
+        # storage directory and the default (~/ray-results).
+        os.environ['TUNE_RESULT_DIR'] = str(self.data_dir)
 
         # Instantiate ray tune Tuner object.
         self.tuner = tune.Tuner(
-            self.evaluation_function,
+            self.trial,
             tune_config=tune.TuneConfig(
                 mode=mode,
                 metric=self.metrics[0],
                 search_alg=self.search_alg,
                 num_samples=max_total_trials,
             ),
             run_config=train.RunConfig(
+                storage_path=self.data_dir,
                 name=self.name,
+                log_to_file=True,
             ),
             param_space=search_space,
         )
 
+    def trial(self, parameters: dict):
+        """Trial evaluation function to be passed to the tuner object.
+
+        Calls the evaluate_design method of the passed DesignEvaluator subclass
+        with the parameters for a given trial and reports the result to
+        the optimiser via train.report().
+
+        Note: users should not need to call this method directly.
+
+        Parameters
+        ----------
+        parameters : dict
+            Dictionary of parameters describing the design to be evaluated.
+        """
+        train.report(self.design_evaluator.evaluate_design(parameters))
+
     def run_optimisation(self) -> ResultGrid:
         """Run the optimisation loop and return the results.