diff --git a/src/kbmod/image_collection.py b/src/kbmod/image_collection.py index dce44354d..00e711720 100644 --- a/src/kbmod/image_collection.py +++ b/src/kbmod/image_collection.py @@ -885,10 +885,18 @@ def toWorkUnit(self, search_config=None, **kwargs): logger.info("Building WorkUnit from ImageCollection") layeredImages = [] + locations = [] for std in self.get_standardizers(**kwargs): + location = std["std"].location for img in std["std"].toLayeredImage(): layeredImages.append(img) + locations.append(location) 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) + per_img_wcs = list(self.wcs) if None not in self.wcs else None + + return WorkUnit( + imgstack, + search_config, + per_image_wcs=per_img_wcs, + constituent_images=locations, + ) diff --git a/tests/test_imagecollection.py b/tests/test_imagecollection.py index d320e5d0a..9ff0bd688 100644 --- a/tests/test_imagecollection.py +++ b/tests/test_imagecollection.py @@ -156,6 +156,11 @@ def test_workunit(self): data = self.fitsFactory.get_n(3, spoof_data=True) ic = ImageCollection.fromTargets(data) wu = ic.toWorkUnit(search_config=SearchConfiguration()) + self.assertEqual(len(wu), 3) + for loc in wu.constituent_images: + self.assertEqual(loc, ":memory:") + + # Write out the WorkUnit. with tempfile.TemporaryDirectory() as dir_name: wu.to_fits(f"{dir_name}/test.fits")