Merge branch 'dev' into feature/1623-typehints

src/primaite/utils/__init__.py

@@ -1 +1,2 @@
+# Crown Owned Copyright (C) Dstl 2023. DEFCON 703. Shared in confidence.
 """Utilities for PrimAITE."""

src/primaite/utils/package_data.py

@@ -1,4 +1,4 @@
-# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
+# Crown Owned Copyright (C) Dstl 2023. DEFCON 703. Shared in confidence.
 import os
 from pathlib import Path
 from typing import TYPE_CHECKING

src/primaite/utils/session_metadata_parser.py

@@ -0,0 +1,59 @@
+# Crown Owned Copyright (C) Dstl 2023. DEFCON 703. Shared in confidence.
+import json
+from pathlib import Path
+from typing import Union
+
+import yaml
+
+from primaite import getLogger
+
+_LOGGER = getLogger(__name__)
+
+
+def parse_session_metadata(session_path: Union[Path, str], dict_only=False):
+    """
+    Loads a session metadata from the given directory path.
+
+    :param session_path: Directory where the session metadata file is in
+    :param dict_only: If dict_only is true, the function will only return the dict contents of session metadata
+
+    :return: Dictionary which has all the session metadata contents
+    :rtype: Dict
+
+    :return: Path where the YAML copy of the training config is dumped into
+    :rtype: str
+    :return: Path where the YAML copy of the laydown config is dumped into
+    :rtype: str
+    """
+    if not isinstance(session_path, Path):
+        session_path = Path(session_path)
+
+    if not session_path.exists():
+        # Session path does not exist
+        msg = f"Failed to load PrimAITE Session, path does not exist: {session_path}"
+        _LOGGER.error(msg)
+        raise FileNotFoundError(msg)
+
+    # Unpack the session_metadata.json file
+    md_file = session_path / "session_metadata.json"
+    with open(md_file, "r") as file:
+        md_dict = json.load(file)
+
+    # if dict only, return dict without doing anything else
+    if dict_only:
+        return md_dict
+
+    # Create a temp directory and dump the training and lay down
+    # configs into it
+    temp_dir = session_path / ".temp"
+    temp_dir.mkdir(exist_ok=True)
+
+    temp_tc = temp_dir / "tc.yaml"
+    with open(temp_tc, "w") as file:
+        yaml.dump(md_dict["env"]["training_config"], file)
+
+    temp_ldc = temp_dir / "ldc.yaml"
+    with open(temp_ldc, "w") as file:
+        yaml.dump(md_dict["env"]["lay_down_config"], file)
+
+    return [md_dict, temp_tc, temp_ldc]

src/primaite/utils/session_output_reader.py

@@ -1,5 +1,6 @@
+# Crown Owned Copyright (C) Dstl 2023. DEFCON 703. Shared in confidence.
 from pathlib import Path
-from typing import Dict, Union
+from typing import Any, Dict, Tuple, Union
 
 # Using polars as it's faster than Pandas; it will speed things up when
 # files get big!
@@ -13,8 +14,33 @@ def av_rewards_dict(av_rewards_csv_file: Union[str, Path]) -> Dict[int, float]:
     The dictionary keys are the episode number, and the values are the mean reward that episode.
 
     :param av_rewards_csv_file: The average rewards per episode csv file path.
-    :return: The average rewards per episode cdv as a dict.
+    :return: The average rewards per episode csv as a dict.
     """
-    df = pl.read_csv(av_rewards_csv_file).to_dict()
+    df_dict = pl.read_csv(av_rewards_csv_file).to_dict()
 
-    return {v: df["Average Reward"][i] for i, v in enumerate(df["Episode"])}
+    return {v: df_dict["Average Reward"][i] for i, v in enumerate(df_dict["Episode"])}
+
+
+def all_transactions_dict(all_transactions_csv_file: Union[str, Path]) -> Dict[Tuple[int, int], Dict[str, Any]]:
+    """
+    Read an all transactions csv file and return as a dict.
+
+    The dict keys are a tuple with the structure (episode, step). The dict
+    values are the remaining columns as a dict.
+
+    :param all_transactions_csv_file: The all transactions csv file path.
+    :return: The all transactions csv file as a dict.
+    """
+    df_dict = pl.read_csv(all_transactions_csv_file).to_dict()
+    new_dict = {}
+
+    episodes = df_dict["Episode"]
+    steps = df_dict["Step"]
+    keys = list(df_dict.keys())
+
+    for i in range(len(episodes)):
+        key = (episodes[i], steps[i])
+        value_dict = {key: df_dict[key][i] for key in keys if key not in ["Episode", "Step"]}
+        new_dict[key] = value_dict
+
+    return new_dict

src/primaite/utils/session_output_writer.py

@@ -1,3 +1,4 @@
+# Crown Owned Copyright (C) Dstl 2023. DEFCON 703. Shared in confidence.
 import csv
 from logging import Logger
 from typing import Final, List, Tuple, TYPE_CHECKING, Union