From a2862f3286bcfefb1a71e86c6009d1f5d9bd86ce Mon Sep 17 00:00:00 2001
From: Robert David Stein <rdstein@caltech.edu>
Date: Thu, 28 Mar 2024 15:55:49 -0700
Subject: [PATCH] Add cache paths and extra loading option (#394)

---
 nuztf/base_scanner.py | 54 +++++++++++++++++++++++++++++++++++++++----
 pyproject.toml        |  2 +-
 2 files changed, 50 insertions(+), 6 deletions(-)

diff --git a/nuztf/base_scanner.py b/nuztf/base_scanner.py
index 09d04617..e56afa72 100644
--- a/nuztf/base_scanner.py
+++ b/nuztf/base_scanner.py
@@ -318,6 +318,41 @@ def query_ampel(
 
         return query_res
 
+    def get_initial_cache_path(self) -> Path:
+        """
+        Get the path to the initial cache file
+
+        :return: Cache file path
+        """
+        return self.get_cache_dir().joinpath("initial_stage.json")
+
+    def get_final_cache_path(self) -> Path:
+        """
+        Get the path to the final cache file
+
+        :return: Cache file path
+        """
+        return self.get_cache_dir().joinpath("final_stage.json")
+
+    def load_from_cache(self):
+        """
+        Load the cache from the final cache file
+
+        :return: None
+        """
+        if not self.get_final_cache_path().exists():
+            err = (
+                f"Final cache file {self.get_final_cache_path()} does not exist. "
+                f"Run scan_area first."
+            )
+            self.logger.error(err)
+            raise FileNotFoundError(err)
+
+        with open(self.get_final_cache_path(), "r") as infile:
+            results = json.load(infile)
+
+        self.add_results(results)
+
     def scan_area(
         self,
         t_min=None,
@@ -329,9 +364,7 @@ def scan_area(
         """
         query_res = self.query_ampel(t_min=t_min, t_max=t_max)
 
-        cache_file_initial_stage = self.get_cache_dir().joinpath("initial_stage.json")
-
-        with open(cache_file_initial_stage, "w") as outfile:
+        with open(self.get_initial_cache_path(), "w") as outfile:
             json.dump(query_res, outfile)
 
         ztf_ids_first_stage = []
@@ -350,13 +383,24 @@ def scan_area(
 
         results = self.ampel_object_search(ztf_ids=ztf_ids_first_stage)
 
+        with open(self.get_final_cache_path(), "w") as outfile:
+            json.dump(results, outfile)
+
+        self.add_results(results)
+        self.create_candidate_summary()
+
+    def add_results(self, results):
+        """
+        Add the results to the cache and create a summary
+
+        :param results: Results from the AMPEL API
+        :return: None
+        """
         for res in results:
             self.add_res_to_cache(res)
 
         self.logger.info(f"Found {len(self.cache)} candidates")
 
-        self.create_candidate_summary()
-
     def filter_f_no_prv(self, res):
         raise NotImplementedError
 
diff --git a/pyproject.toml b/pyproject.toml
index 4b2b5917..8182f015 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nuztf"
-version = "2.7.2"
+version = "2.8.0"
 description = "Package for multi-messenger correlation searches with ZTF"
 authors = ["Robert Stein <rdstein@caltech.edu>", "Simeon Reusch <simeon.reusch@desy.de>", "Jannis Necker <jannis.necker@desy.de>"]
 repository = "https://github.com/desy-multimessenger/nuztf"