Additional SkyDriver Updates (#119)

Co-authored-by: github-actions <[email protected]>

Author: ric-evans

skymap_scanner/utils/scan_result.py

@@ -1,6 +1,9 @@
 """For encapsulating the results of an event scan in a single instance."""
 
+# fmt: off
+# pylint: skip-file
 
+import dataclasses as dc
 import io
 import itertools as it
 import json
@@ -19,6 +22,7 @@
 from matplotlib import pyplot as plt
 from matplotlib import text
 
+from .event_tools import EventMetadata
 from .pixelreco import NSidesDict, PixelReco
 
 ###############################################################################
@@ -47,7 +51,6 @@ class InvalidPixelValueError(Exception):
 ###############################################################################
 # MAIN CLASS
 
-# fmt: off
 class ScanResult:
     """This class parses a nsides_dict and stores the relevant numeric result
     of the scan. Ideally it should serve as the basic data structure for
@@ -217,6 +220,21 @@ def has_metadata(self) -> bool:
                     return False
         return True
 
+    def get_event_metadata(self) -> EventMetadata:
+        """Get the EventMetadata portion of the result's metadata."""
+        if self.has_metadata():
+            first_metadata = self.result[list(self.result.keys())[0]].dtype.metadata
+            return EventMetadata(
+                first_metadata['run_id'],
+                first_metadata['event_id'],
+                first_metadata['event_type'],
+                first_metadata['mjd'],
+                first_metadata.get('is_real', True),  # assume real event
+            )
+        else:
+            self.logger.warning("Metadata doesn't seem to exist and will not be used for plotting.")
+            return EventMetadata(0, 0, '', 0, False)
+
     def is_close(
         self,
         other: "ScanResult",
@@ -315,12 +333,17 @@ def get_nside_string(self) -> str:
         # TODO: possibly better to use integer values as keys in self.result
         return "_".join([str(nside) for nside in self.nsides])
 
-    def get_filename(self, event_id: str, extension: str, output_dir: Union[str, Path, None] = None) -> Path:
+    def get_filename(
+        self,
+        event_metadata: EventMetadata,
+        extension: str,
+        output_dir: Union[str, Path, None] = None
+    ) -> Path:
         """Make a filepath for writing representations of `self` to disk."""
         if not extension.startswith('.'):
             extension = '.' + extension
 
-        filename = Path(f"{event_id}_{self.get_nside_string()}{extension}")
+        filename = Path(f"{str(event_metadata)}_{self.get_nside_string()}{extension}")
         if output_dir is not None:
             filename = output_dir / Path(filename)
         return filename
@@ -330,13 +353,21 @@ def get_filename(self, event_id: str, extension: str, output_dir: Union[str, Pat
     """
 
     @classmethod
-    def from_nsides_dict(cls, nsides_dict: NSidesDict, **kwargs) -> "ScanResult":
+    def from_nsides_dict(
+        cls,
+        nsides_dict: NSidesDict,
+        event_metadata: Optional[EventMetadata] = None
+    ) -> "ScanResult":
         """Factory method for nsides_dict."""
         logger = logging.getLogger(__name__)
 
+        event_metadata_dict = {}
+        if event_metadata:
+            event_metadata_dict = dc.asdict(event_metadata)
+
         result = dict()
         for nside, pixel_dict in nsides_dict.items():
-            _dtype = np.dtype(cls.PIXEL_TYPE, metadata=dict(nside=nside, **kwargs))
+            _dtype = np.dtype(cls.PIXEL_TYPE, metadata=dict(nside=nside, **event_metadata_dict))
             nside_pixel_values = np.zeros(len(pixel_dict), dtype=_dtype)
             logger.info(f"nside {nside} has {len(pixel_dict)} pixels / {12 * nside**2} total.")
 
@@ -384,9 +415,13 @@ def read_npz(cls, filename: Union[str, Path]) -> "ScanResult":
                 result[key] = np.array(list(npz[key]), dtype=_dtype)
         return cls(result=result)
 
-    def to_npz(self, event_id: str, output_dir: Union[str, Path, None] = None) -> Path:
+    def to_npz(
+        self,
+        event_metadata: EventMetadata,
+        output_dir: Union[str, Path, None] = None,
+    ) -> Path:
         """Save to .npz file."""
-        filename = self.get_filename(event_id, '.npz', output_dir)
+        filename = self.get_filename(event_metadata, '.npz', output_dir)
 
         try:
             metadata_dtype = np.dtype(
@@ -412,9 +447,13 @@ def read_json(cls, filename: Union[str, Path]) -> "ScanResult":
             pydict = json.load(f)
         return cls.deserialize(pydict)
 
-    def to_json(self, event_id: str, output_dir: Union[str, Path, None] = None) -> Path:
+    def to_json(
+        self,
+        event_metadata: EventMetadata,
+        output_dir: Union[str, Path, None] = None
+    ) -> Path:
         """Save to .json file."""
-        filename = self.get_filename(event_id, '.json', output_dir)
+        filename = self.get_filename(event_metadata, '.json', output_dir)
         pydict = self.serialize()
         with open(filename, 'w') as f:
             json.dump(pydict, f, indent=4)
@@ -550,7 +589,6 @@ def create_plot(self,
                     log_func=None,
                     upload_func=None,
                     final_channels=None):
-        from .icetrayless import create_event_id_string
         from .plotting_tools import DecFormatter, RaFormatter
 
         if log_func is None:
@@ -576,15 +614,9 @@ def upload_func(file_buffer, name, title):
             if "nside-" not in k:
                 raise RuntimeError("\"nside\" not in result file..")
 
-        if self.has_metadata():
-            run_id, event_id, event_type, mjd = [
-                self.result[k].dtype.metadata[_] for _ in "run_id event_id event_type mjd".split()]
-        else:
-            self.logger.warn(f"Metadata doesn't seem to exist and will not be used for plotting.")
-            run_id, event_id, event_type, mjd = [0]*4
-        unique_id = f'{create_event_id_string(run_id, event_id, event_type)}_{self.get_nside_string()}'
-
-        plot_title = f"Run: {run_id} Event {event_id}: Type: {event_type} MJD: {mjd}"
+        event_metadata = self.get_event_metadata()
+        unique_id = f'{str(event_metadata)}_{self.get_nside_string()}'
+        plot_title = f"Run: {event_metadata.run_id} Event {event_metadata.event_id}: Type: {event_metadata.event_type} MJD: {event_metadata.mjd}"
 
         plot_filename = f"{unique_id}.{'plot_zoomed.' if dozoom else ''}pdf"
         print(f"saving plot to {plot_filename}")
@@ -814,7 +846,6 @@ def create_plot_zoomed(self,
                            plot_4fgl=False,
                            final_channels=None):
         """Uses healpy to plot a map."""
-        from .icetrayless import create_event_id_string
         from .plotting_tools import format_fits_header, hp_ticklabels, plot_catalog
 
         if log_func is None:
@@ -850,15 +881,9 @@ def bounding_box(ra, dec, theta, phi):
             if "nside-" not in k:
                 raise RuntimeError("\"nside\" not in result file..")
 
-        if self.has_metadata():
-            run_id, event_id, event_type, mjd = [
-                self.result[k].dtype.metadata[_] for _ in "run_id event_id event_type mjd".split()]
-        else:
-            self.logger.warn(f"Metadata doesn't seem to exist and will not be used for plotting.")
-            run_id, event_id, event_type, mjd = [0]*4
-        unique_id = f'{create_event_id_string(run_id, event_id, event_type)}_{self.get_nside_string()}'
-
-        plot_title = f"Run: {run_id} Event {event_id}: Type: {event_type} MJD: {mjd}"
+        event_metadata = self.get_event_metadata()
+        unique_id = f'{str(event_metadata)}_{self.get_nside_string()}'
+        plot_title = f"Run: {event_metadata.run_id} Event {event_metadata.event_id}: Type: {event_metadata.event_type} MJD: {event_metadata.mjd}"
 
         nsides = self.nsides
         print(f"available nsides: {nsides}")
@@ -1139,9 +1164,13 @@ def area(vs):
                 log_func("Memory Error prevented contours from being written")
 
         uncertainty = [(ra_minus, ra_plus), (dec_minus, dec_plus)]
-        fits_header = format_fits_header((run_id, event_id, event_type), 0,
-            np.degrees(minRA), np.degrees(minDec), uncertainty,
-           )
+        fits_header = format_fits_header(
+            (event_metadata.run_id, event_metadata.event_id, event_metadata.event_type),
+            0,
+            np.degrees(minRA),
+            np.degrees(minDec),
+            uncertainty,
+        )
         mmap_nside = healpy.get_nside(master_map)
 
         # Pixel numbers as is gives a map that is reflected