Pass additional metadata to WorkUnit

src/kbmod/image_collection.py

@@ -867,14 +867,16 @@ def toImageStack(self):
         layeredImages = [img for std in self._standardizers for img in std.toLayeredImage()]
         return ImageStack(layeredImages)
 
-    def toWorkUnit(self, search_config=None, **kwargs):
+    def toWorkUnit(self, search_config=None, extra_meta=None, **kwargs):
         """Return an `~kbmod.WorkUnit` object for processing with
         KBMOD.
 
         Parameters
         ----------
         search_config : `~kbmod.SearchConfiguration` or None, optional
             Search configuration. Default ``None``.
+        extra_meta : `list` or None, optional
+            Extra meta data columns to add to the WorkUnit.
 
         Returns
         -------
@@ -884,11 +886,45 @@ def toWorkUnit(self, search_config=None, **kwargs):
         from .work_unit import WorkUnit
 
         logger.info("Building WorkUnit from ImageCollection")
+
+        # Create a storage location for additional meta data to save.
+        # Not all of this may be present in from the ImageCollection,
+        # in which case we will append None.
+        meta_data_vals = {
+            "visit": [],
+            "filter": [],
+        }
+        if extra_meta is not None:
+            for col in extra_meta:
+                meta_data_vals[col] = []
+
+        # Extract data from each standardizer and each LayeredImage within
+        # that standardizer.
         layeredImages = []
         for std in self.get_standardizers(**kwargs):
+            num_added = 0
             for img in std["std"].toLayeredImage():
                 layeredImages.append(img)
+                num_added += 1
+
+            # Get each meta data value from the standardizer so it can be
+            # passed to the WorkUnit. Use the same value for all images from
+            # this standardizer.
+            meta_data = std["std"].standardizeMetadata()
+            for col in meta_data_vals.keys():
+                value = meta_data.get(col, None)
+                meta_data_vals[col].extend([value] * num_added)
+
+        # Create the basic WorkUnit from the ImageStack.
         imgstack = ImageStack(layeredImages)
         if None not in self.wcs:
-            return WorkUnit(imgstack, search_config, per_image_wcs=list(self.wcs))
-        return WorkUnit(imgstack, search_config)
+            work = WorkUnit(imgstack, search_config, per_image_wcs=list(self.wcs))
+        else:
+            work = WorkUnit(imgstack, search_config)
+
+        # Add any extra metadata to the WorkUnit. This may include a column of
+        # None values if the information was not available to the standardizer.
+        for col in meta_data_vals.keys():
+            work.add_org_img_meta_data(col, meta_data_vals[col])
+
+        return work

src/kbmod/standardizers/fits_standardizers/kbmodv05.py

@@ -141,14 +141,15 @@ def __init__(self, location=None, hdulist=None, config=None, **kwargs):
         ]
 
     def translateHeader(self):
-        """Returns the following metadata, read from the primary header, as a
-        dictionary:
+        """Returns at least the following metadata, read from the primary header,
+         as a dictionary:
 
         ======== ========== ===================================================
         Key      Header Key Description
         ======== ========== ===================================================
         mjd      DATE-AVG   Decimal MJD timestamp of the middle of the exposure
         filter   FILTER     Filter band
+        visit    EXPID      Exposure ID
         ======== ========== ===================================================
         """
         # this is the 1 mandatory piece of metadata we need to extract
@@ -159,6 +160,7 @@ def translateHeader(self):
 
         # these are all optional things
         standardizedHeader["filter"] = self.primary["FILTER"]
+        standardizedHeader["visit"] = self.primary["EXPID"]
 
         # If no observatory information is given, default to the Deccam data
         # (Cerro Tololo Inter-American Observatory).

tests/test_imagecollection.py

@@ -155,7 +155,15 @@ def test_workunit(self):
 
         data = self.fitsFactory.get_n(3, spoof_data=True)
         ic = ImageCollection.fromTargets(data)
-        wu = ic.toWorkUnit(search_config=SearchConfiguration())
+        wu = ic.toWorkUnit(search_config=SearchConfiguration(), extra_meta=["FILTER"])
+        self.assertEqual(len(wu), 3)
+
+        # We can retrieve the meta data from the WorkUnit.
+        filter_info = wu.get_constituent_meta("FILTER")
+        self.assertEqual(len(filter_info), 3)
+        self.assertIsNotNone(filter_info[0])
+
+        # We can write the whole work unit to a file.
         with tempfile.TemporaryDirectory() as dir_name:
             wu.to_fits(f"{dir_name}/test.fits")