From 12248f029353a30e488b7e02969cc798dae395b0 Mon Sep 17 00:00:00 2001
From: Lyubov Yamshchikova <YamLyubov@gmail.com>
Date: Fri, 16 Feb 2024 15:02:21 +0300
Subject: [PATCH] Use HistoryReader and from_saved_histories when possible

---
 examples/molecule_search/experiment.py | 12 +-----------
 test/unit/adaptive/test_mab_agent.py   |  4 ++--
 2 files changed, 3 insertions(+), 13 deletions(-)

diff --git a/examples/molecule_search/experiment.py b/examples/molecule_search/experiment.py
index 79cea2bbe..324593e3b 100644
--- a/examples/molecule_search/experiment.py
+++ b/examples/molecule_search/experiment.py
@@ -209,18 +209,8 @@ def run_experiment(optimizer_setup: Callable,
 
 def plot_experiment_comparison(experiment_ids: Sequence[str], metric_id: int = 0, results_dir='./results'):
     root = Path(results_dir)
-    histories = {}
-    for exp_name in experiment_ids:
-        trials = []
-        for history_filename in os.listdir(root / exp_name):
-            if history_filename.startswith('history'):
-                history = OptHistory.load(root / exp_name / history_filename)
-                trials.append(history)
-        histories[exp_name] = trials
-        print(f'Loaded {len(trials)} trial histories for experiment: {exp_name}')
     # Visualize
-    MultipleFitnessLines.from_histories(histories).visualize(metric_id=metric_id)
-    return histories
+    MultipleFitnessLines.from_saved_histories(experiment_ids, root).visualize(metric_id=metric_id)
 
 
 if __name__ == '__main__':
diff --git a/test/unit/adaptive/test_mab_agent.py b/test/unit/adaptive/test_mab_agent.py
index e808a3bca..68fa699e7 100644
--- a/test/unit/adaptive/test_mab_agent.py
+++ b/test/unit/adaptive/test_mab_agent.py
@@ -12,7 +12,7 @@
                           (os.path.join(Path(__file__).parent), '0_mab.pkl')])
 def test_save_mab(path_to_save, file_name):
     """ Tests if MAB is saved with specifying file_nam and without. """
-    mab = MultiArmedBanditAgent(actions=[0, 1, 2],
+    mab = MultiArmedBanditAgent(actions=['0', '1', '2'],
                                 n_jobs=1,
                                 path_to_save=path_to_save)
     mab.save()
@@ -25,7 +25,7 @@ def test_load_mab():
     file_name = 'test_mab.pkl'
     path_to_load = os.path.join(Path(__file__).parent, file_name)
     # save mab to load it later
-    mab = MultiArmedBanditAgent(actions=[0, 1, 2],
+    mab = MultiArmedBanditAgent(actions=['0', '1', '2'],
                                 n_jobs=1,
                                 path_to_save=path_to_load)
     mab.save()