RCAL-895: allow updating source catalog with tweaked WCS when running ELP

CHANGES.rst

@@ -21,7 +21,6 @@ resample
 
 pipeline
 --------
-
 - Added ``suffix`` to the spec of ExposurePipeline with a
   default value of ``cal``. Removed explicit setting of ``suffix``
   so that it can be passed as an argument to ``strun``. [#1378]

changes/1373.general.rst

@@ -0,0 +1 @@
+Update source catalog file with tweaked coordinates.

romancal/pipeline/exposure_pipeline.py

@@ -71,6 +71,8 @@
 
         # make sure source_catalog returns the updated datamodel
         self.source_catalog.return_updated_model = True
+        # make sure we update source catalog coordinates afer running TweakRegStep
+        self.tweakreg.update_source_catalog_coordinates = True
 
         log.info("Starting Roman exposure calibration pipeline ...")
         if isinstance(input, str):

romancal/tweakreg/tests/test_tweakreg.py

@@ -1,6 +1,8 @@
 import copy
+import filecmp
 import json
 import os
+import shutil
 from io import StringIO
 from pathlib import Path
 from typing import Tuple
@@ -18,6 +20,7 @@
 from gwcs import coordinate_frames as cf
 from gwcs import wcs
 from gwcs.geometry import CartesianToSpherical, SphericalToCartesian
+from numpy.random import default_rng
 from roman_datamodels import datamodels as rdm
 from roman_datamodels import maker_utils
 
@@ -1185,3 +1188,183 @@ def test_parse_catfile_raises_error_on_invalid_content(tmp_path, catfile_line_co
         trs._parse_catfile(catfile)
 
     assert type(exec_info.value) == ValueError
+
+
+def test_update_source_catalog_coordinates(tmp_path, base_image):
+    """Test that TweakReg updates the catalog coordinates with the tweaked WCS."""
+
+    os.chdir(tmp_path)
+
+    img = base_image(shift_1=1000, shift_2=1000)
+    add_tweakreg_catalog_attribute(tmp_path, img, catalog_filename="img_1")
+
+    tweakreg = trs.TweakRegStep()
+
+    # create SourceCatalogModel
+    source_catalog_model = setup_source_catalog_model(img)
+
+    # save SourceCatalogModel
+    tweakreg.save_model(
+        source_catalog_model,
+        output_file="img_1.asdf",
+        suffix="cat",
+        force=True,
+    )
+
+    # update tweakreg catalog name
+    img.meta.source_detection.tweakreg_catalog_name = "img_1_cat.asdf"
+
+    # run TweakRegStep
+    res = trs.TweakRegStep.call([img])
+
+    # tweak the current WCS using TweakRegStep and save the updated cat file
+    with res:
+        dm = res.borrow(0)
+        assert dm.meta.source_detection.tweakreg_catalog_name == "img_1_cat.asdf"
+        tweakreg.update_catalog_coordinates(
+            dm.meta.source_detection.tweakreg_catalog_name, dm.meta.wcs
+        )
+        res.shelve(dm, 0)
+
+    # read in saved catalog coords
+    cat = rdm.open("img_1_cat.asdf")
+    cat_ra_centroid = cat.source_catalog["ra_centroid"]
+    cat_dec_centroid = cat.source_catalog["dec_centroid"]
+    cat_ra_psf = cat.source_catalog["ra_psf"]
+    cat_dec_psf = cat.source_catalog["dec_psf"]
+
+    # calculate world coords using tweaked WCS
+    expected_centroid = img.meta.wcs(
+        cat.source_catalog["xcentroid"], cat.source_catalog["ycentroid"]
+    )
+    expected_psf = img.meta.wcs(
+        cat.source_catalog["x_psf"], cat.source_catalog["y_psf"]
+    )
+
+    # compare coordinates (make sure tweaked WCS was applied to cat file coords)
+    np.testing.assert_array_equal(cat_ra_centroid, expected_centroid[0])
+    np.testing.assert_array_equal(cat_dec_centroid, expected_centroid[1])
+    np.testing.assert_array_equal(cat_ra_psf, expected_psf[0])
+    np.testing.assert_array_equal(cat_dec_psf, expected_psf[1])
+
+
+def test_source_catalog_coordinates_have_changed(tmp_path, base_image):
+    """Test that the original catalog file content is different from the updated file."""
+
+    os.chdir(tmp_path)
+
+    img = base_image(shift_1=1000, shift_2=1000)
+    add_tweakreg_catalog_attribute(tmp_path, img, catalog_filename="img_1")
+
+    tweakreg = trs.TweakRegStep()
+
+    # create SourceCatalogModel
+    source_catalog_model = setup_source_catalog_model(img)
+
+    # save SourceCatalogModel
+    tweakreg.save_model(
+        source_catalog_model,
+        output_file="img_1.asdf",
+        suffix="cat",
+        force=True,
+    )
+    # save original data
+    shutil.copy("img_1_cat.asdf", "img_1_cat_original.asdf")
+
+    # update tweakreg catalog name
+    img.meta.source_detection.tweakreg_catalog_name = "img_1_cat.asdf"
+
+    # run TweakRegStep
+    res = trs.TweakRegStep.call([img])
+
+    # tweak the current WCS using TweakRegStep and save the updated cat file
+    with res:
+        dm = res.borrow(0)
+        assert dm.meta.source_detection.tweakreg_catalog_name == "img_1_cat.asdf"
+        tweakreg.update_catalog_coordinates(
+            dm.meta.source_detection.tweakreg_catalog_name, dm.meta.wcs
+        )
+        res.shelve(dm, 0)
+
+    cat_original = rdm.open("img_1_cat_original.asdf")
+    cat_updated = rdm.open("img_1_cat.asdf")
+
+    # compare contents
+    assert not filecmp.cmp("img_1_cat.asdf", "img_1_cat_original.asdf")
+    # compare coordinates using rtol=1e-07 and atol=0 (default)
+    np.testing.assert_allclose(
+        cat_original.source_catalog["ra_centroid"],
+        cat_updated.source_catalog["ra_centroid"],
+    )
+    np.testing.assert_allclose(
+        cat_original.source_catalog["dec_centroid"],
+        cat_updated.source_catalog["dec_centroid"],
+    )
+    np.testing.assert_allclose(
+        cat_original.source_catalog["ra_psf"],
+        cat_updated.source_catalog["ra_psf"],
+    )
+    np.testing.assert_allclose(
+        cat_original.source_catalog["dec_psf"],
+        cat_updated.source_catalog["dec_psf"],
+    )
+
+
+def setup_source_catalog_model(img):
+    """
+    Set up the source catalog model.
+
+    Notes
+    -----
+    This function reads the source catalog from a file, renames columns to match
+    expected names, adds mock PSF coordinates, applies random shifts to the centroid
+    and PSF coordinates, and calculates the world coordinates for the centroids.
+    """
+    cat_model = rdm.SourceCatalogModel
+    source_catalog_model = maker_utils.mk_datamodel(cat_model)
+    # this will be the output filename
+    source_catalog_model.meta.filename = "img_1.asdf"
+
+    # read in the mock table
+    source_catalog = Table.read("img_1", format="ascii.ecsv")
+    # rename columns to match expected column names
+    source_catalog.rename_columns(["x", "y"], ["xcentroid", "ycentroid"])
+    # add mock PSF coordinates
+    source_catalog["x_psf"] = source_catalog["xcentroid"]
+    source_catalog["y_psf"] = source_catalog["ycentroid"]
+
+    # add random fraction of a pixel shifts to the centroid coordinates
+    rng = default_rng()
+    shift_x = rng.uniform(-0.5, 0.5, size=len(source_catalog))
+    shift_y = rng.uniform(-0.5, 0.5, size=len(source_catalog))
+    source_catalog["xcentroid"] += shift_x
+    source_catalog["ycentroid"] += shift_y
+    # add random fraction of a pixel shifts to the PSF coordinates
+    shift_x = rng.uniform(-0.5, 0.5, size=len(source_catalog))
+    shift_y = rng.uniform(-0.5, 0.5, size=len(source_catalog))
+    source_catalog["x_psf"] += shift_x
+    source_catalog["y_psf"] += shift_y
+
+    # calculate centroid world coordinates
+    centroid = img.meta.wcs(
+        source_catalog["xcentroid"],
+        source_catalog["ycentroid"],
+    )
+    # calculate PSF world coordinates
+    psf = img.meta.wcs(
+        source_catalog["x_psf"],
+        source_catalog["y_psf"],
+    )
+    # add world coordinates to catalog
+    source_catalog["ra_centroid"], source_catalog["dec_centroid"] = centroid
+    source_catalog["ra_psf"], source_catalog["dec_psf"] = psf
+    # add units
+    source_catalog["ra_centroid"].unit = u.deg
+    source_catalog["dec_centroid"].unit = u.deg
+    source_catalog["ra_psf"].unit = u.deg
+    source_catalog["dec_psf"].unit = u.deg
+
+    # add source catalog to SourceCatalogModel
+    source_catalog_model.source_catalog = source_catalog
+
+    return source_catalog_model

romancal/tweakreg/tweakreg_step.py

@@ -162,41 +162,7 @@
                     images.shelve(image_model)
                     return image_model
 
-                if hasattr(image_model.meta, "source_detection"):
-                    is_tweakreg_catalog_present = hasattr(
-                        image_model.meta.source_detection, "tweakreg_catalog"
-                    )
-                    is_tweakreg_catalog_name_present = hasattr(
-                        image_model.meta.source_detection, "tweakreg_catalog_name"
-                    )
-                    if is_tweakreg_catalog_present:
-                        # read catalog from structured array
-                        catalog = Table(
-                            np.asarray(
-                                image_model.meta.source_detection.tweakreg_catalog
-                            )
-                        )
-                    elif is_tweakreg_catalog_name_present:
-                        catalog = self.read_catalog(
-                            image_model.meta.source_detection.tweakreg_catalog_name
-                        )
-                    else:
-                        images.shelve(image_model, i, modify=False)
-                        raise AttributeError(
-                            "Attribute 'meta.source_detection.tweakreg_catalog' is missing."
-                            "Please either run SourceDetectionStep or provide a"
-                            "custom source catalog."
-                        )
-                    # remove 4D numpy array from meta.source_detection
-                    if is_tweakreg_catalog_present:
-                        del image_model.meta.source_detection["tweakreg_catalog"]
-                else:
-                    images.shelve(image_model, i, modify=False)
-                    raise AttributeError(
-                        "Attribute 'meta.source_detection' is missing."
-                        "Please either run SourceDetectionStep or provide a"
-                        "custom source catalog."
-                    )
+                catalog = self.get_tweakreg_catalog(images, i, image_model)
 
                 for axis in ["x", "y"]:
                     if axis not in catalog.colnames:
@@ -495,10 +461,97 @@
                         del image_model.meta["wcs_fit_results"][k]
 
                     image_model.meta.wcs = imcat.wcs
+
+                    if getattr(self, "update_source_catalog_coordinates", False):
+                        self.update_catalog_coordinates(
+                            image_model.meta.source_detection.tweakreg_catalog_name,
+                            image_model.meta.wcs,
+                        )
                 images.shelve(image_model, i)
 
         return images
 
+    def get_tweakreg_catalog(self, images, i, image_model):
+        if hasattr(image_model.meta, "source_detection"):
+            is_tweakreg_catalog_present = hasattr(
+                image_model.meta.source_detection, "tweakreg_catalog"
+            )
+            is_tweakreg_catalog_name_present = hasattr(
+                image_model.meta.source_detection, "tweakreg_catalog_name"
+            )
+            if is_tweakreg_catalog_present:
+                # read catalog from structured array
+                catalog = Table(
+                    np.asarray(image_model.meta.source_detection.tweakreg_catalog)
+                )
+            elif is_tweakreg_catalog_name_present:
+                catalog = self.read_catalog(
+                    image_model.meta.source_detection.tweakreg_catalog_name
+                )
+            else:
+                images.shelve(image_model, i, modify=False)
+                raise AttributeError(
+                    "Attribute 'meta.source_detection.tweakreg_catalog' is missing."
+                    "Please either run SourceDetectionStep or provide a"
+                    "custom source catalog."
+                )
+                # remove 4D numpy array from meta.source_detection
+            if is_tweakreg_catalog_present:
+                del image_model.meta.source_detection["tweakreg_catalog"]
+        else:
+            images.shelve(image_model, i, modify=False)
+            raise AttributeError(
+                "Attribute 'meta.source_detection' is missing."
+                "Please either run SourceDetectionStep or provide a"
+                "custom source catalog."
+            )
+
+        return catalog
+
+    def update_catalog_coordinates(self, tweakreg_catalog_name, tweaked_wcs):
+        """
+        Update the source catalog coordinates using the tweaked WCS.
+
+        Parameters
+        ----------
+        tweakreg_catalog_name : str
+            The name of the TweakReg catalog file produced by `SourceCatalog`.
+        tweaked_wcs : `gwcs.wcs.WCS`
+            The tweaked World Coordinate System (WCS) object.
+
+        Returns
+        -------
+        None
+        """
+        # read in cat file
+        with rdm.open(tweakreg_catalog_name) as source_catalog_model:
+            # get catalog
+            catalog = source_catalog_model.source_catalog
+
+            # define mapping between pixel and world coordinates
+            colname_mapping = {
+                ("xcentroid", "ycentroid"): ("ra_centroid", "dec_centroid"),
+                ("x_psf", "y_psf"): ("ra_psf", "dec_psf"),
+            }
+
+            for k, v in colname_mapping.items():
+                # get column names
+                x_colname, y_colname = k
+                ra_colname, dec_colname = v
+
+                # calculate new coordinates using tweaked WCS and update catalog coordinates
+                catalog[ra_colname], catalog[dec_colname] = tweaked_wcs(
+                    catalog[x_colname], catalog[y_colname]
+                )
+
+            # save updated catalog (overwrite cat file)
+            self.save_model(
+                source_catalog_model,
+                output_file=source_catalog_model.meta.filename,
+                suffix="cat",
+                force=True,
+            )
+
     def read_catalog(self, catalog_name):
         if catalog_name.endswith("asdf"):
             with rdm.open(catalog_name) as source_catalog_model: