RCAL-828: add PSF bits to source catalog. (spacetelescope#1243)

Co-authored-by: zacharyburnett <zburnett@stsci.edu>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Brett <brettgraham@gmail.com>
Co-authored-by: D Davis <49163225+ddavis-stsci@users.noreply.github.com>

CHANGES.rst

@@ -14,11 +14,16 @@ general
 
 - Rename highlevelpipeline to mosaic pipeline [#1249]
 
+
+source_catalog
+--------------
+- Add PSF photometry capability. [#1243]
+
 dq_init
 -------
-
 -  Refactor DQInitStep to use the RampModel method of creating ramps. [#1258]
 
+
 0.15.1 (2024-05-15)
 ===================
 

docs/roman/source_catalog/arguments.rst

@@ -28,16 +28,19 @@ The ``source_catalog`` step uses the following user-settable arguments:
   energy value
 
 * ``--ci1_star_threshold``: A floating-point value of the threshold
-  compared to the concentration index calculated from aperture_ee1
-  and aperture_ee2 that is used to determine whether a source is a
-  star. Sources must meet the criteria of both ci1_star_threshold and
-  ci2_star_threshold to be considered a star.
+  compared to the concentration index calculated from ``aperture_ee1``
+  and ``aperture_ee2`` that is used to determine whether a source is a
+  star. Sources must meet the criteria of both ``ci1_star_threshold`` and
+  ``ci2_star_threshold`` to be considered a star.
 
 * ``--ci2_star_threshold``: A floating-point value of the threshold
-  compared to the concentration index calculated from aperture_ee2
-  and aperture_ee3 that is used to determine whether a source is a
-  star. Sources must meet the criteria of both ci1_star_threshold and
-  ci2_star_threshold to be considered a star.
+  compared to the concentration index calculated from ``aperture_ee2``
+  and ``aperture_ee3`` that is used to determine whether a source is a
+  star. Sources must meet the criteria of both ``ci1_star_threshold`` and
+  ``ci2_star_threshold`` to be considered a star.
 
 * ``--suffix``: A string value giving the file name suffix to use for
-  the output catalog file [default='cat']
+  the output catalog file (default is ``'cat'``).
+
+* ``--fit_psf``: A boolean value indicating whether to perform
+  PSF photometry (default is ``True``)

docs/roman/source_catalog/index.rst

@@ -9,5 +9,8 @@ Source Catalog
 
    main.rst
    arguments.rst
+   psf.rst
 
 .. automodapi:: romancal.source_catalog
+
+.. automodapi:: romancal.lib.psf

docs/roman/source_catalog/main.rst

@@ -60,6 +60,17 @@ Photometric errors are calculated from the resampled total-error
 array contained in the ``ERR`` (``model.err``) array. Note that this
 total-error array includes source Poisson noise.
 
+PSF Fitting
+-----------
+
+Star finding algorithms like `~photutils.detection.DAOStarFinder` provide
+approximate stellar centroids. More precise centroids may be inferred by
+fitting model PSFs to the observations. Setting the SourceDetectionStep's
+option `fit_psf` to True will generate model Roman PSFs with
+`WebbPSF <https://webbpsf.readthedocs.io/en/latest/roman.html>`_, and fit
+those models to each of the sources detected by
+`~photutils.detection.DAOStarFinder`. More details are in :doc:`psf`.
+
 Output Products
 ---------------
 
@@ -170,6 +181,36 @@ columns (assuming the default encircled energies of 30, 50, and 70):
 |                        | minimal bounding box of the source                 |
 +------------------------+----------------------------------------------------+
 
+
+If ``fit_psf=True``, the following columns will also be available:
+
++------------------------+----------------------------------------------------+
+| Column                 | Description                                        |
++========================+====================================================+
+| x_psf, x_psf_err       | X pixel value of the source and its associated     |
+|                        | error as determined by PSF fitting                 |
++------------------------+----------------------------------------------------+
+| y_psf, y_psf_err       | Y pixel value of the source and its associated     |
+|                        | error as determined by PSF fitting                 |
++------------------------+----------------------------------------------------+
+| flux_psf, flux_psf_err | Flux of the source and its associated error as     |
+|                        | determined by PSF fitting                          |
++------------------------+----------------------------------------------------+
+| flag_psf               | DQ flag of the resulting PSF fitting.              |
+|                        | Possible values are [1]_:                          |
+|                        |                                                    |
+|                        | - 1 : one or more pixels in the fitting region     |
+|                        |   were masked                                      |
+|                        | - 2 : the fit x and/or y position lies outside of  |
+|                        |   the input data                                   |
+|                        | - 4 : the fit flux is less than or equal to zero   |
+|                        | - 8 : the fitter may not have converged            |
+|                        | - 16 : the fitter parameter covariance matrix was  |
+|                        |   not returned                                     |
++------------------------+----------------------------------------------------+
+
+.. [1] See `PSFPhotometry <https://photutils.readthedocs.io/en/stable/api/photutils.psf.PSFPhotometry.html#photutils.psf.PSFPhotometry>`_ for more details.
+
 Note that pixel coordinates are 0 indexed, matching the Python 0-based
 indexing. That means pixel coordinate ``0`` is the center of the first
 pixel.

docs/roman/source_catalog/psf.rst

@@ -0,0 +1,29 @@
+PSF Fitting
+===========
+
+A few PSF fitting utilities are included to interface between observations
+within Roman datamodels and methods within dependencies that generate and
+fit PSF models to observations.
+
+Create PSF models
+-----------------
+
+`~romancal.lib.psf.create_gridded_psf_model` computes a gridded PSF model for
+a given detector using `~webbpsf.gridded_library.CreatePSFLibrary` from
+`WebbPSF <https://webbpsf.readthedocs.io/>`_. The defaults are chosen to
+balance more accurate PSF models with the cost of increased runtime. For
+further reading on the WebbPSF approach to ePSFs, see the WebbPSF docs on
+`Using PSF Grids <https://webbpsf.readthedocs.io/en/latest/psf_grids.html>`_.
+
+Fit model PSFs to an ImageModel
+-------------------------------
+
+Once PSF models are generated, you can fit those models to observations
+in an ImageModel with `~romancal.lib.psf.fit_psf_to_image_model` using
+`photutils <https://photutils.readthedocs.io/en/stable/psf.html>`_.
+By default the fits are done with `~photutils.psf.PSFPhotometry`, and
+crowded fields may benefit from using `~photutils.psf.IterativePSFPhotometry`.
+For neighboring sources that are near one another on the detector, grouping
+the sources and fitting their PSFs simultaneously improves the fit quality.
+Initial guesses for target centroids can be given or source
+detection can be performed with, e.g., `~photutils.detection.DAOStarFinder`.

docs/roman/source_detection/index.rst

@@ -12,5 +12,3 @@ Source Detection
    psf.rst
 
 .. automodapi:: romancal.source_detection
-
-.. automodapi:: romancal.lib.psf

romancal/lib/psf.py

@@ -327,10 +327,12 @@ def fit_psf_to_image_model(
     # at some point in the future.
 
     if dq is None:
-        if image_model is not None:
+        if image_model is not None and isinstance(image_model, ImageModel):
+            # L2 images have a dq array
             mask = dq_to_boolean_mask(image_model, ignore_flags=ignore_flags)
         else:
-            mask = None
+            # L3 images
+            mask = image_model.weight == 0
     else:
         mask = dq_to_boolean_mask(dq)
 
@@ -352,9 +354,9 @@ def fit_psf_to_image_model(
         # outside the detector boundaries:
         init_centroid_in_range = (
             (guesses["x_init"] > 0)
-            & (guesses["x_init"] < data.shape[0])
+            & (guesses["x_init"] < data.shape[1])
             & (guesses["y_init"] > 0)
-            & (guesses["y_init"] < data.shape[1])
+            & (guesses["y_init"] < data.shape[0])
         )
         guesses = guesses[init_centroid_in_range]
 

romancal/source_catalog/source_catalog.py

@@ -4,6 +4,8 @@
 
 import logging
 import warnings
+from pathlib import Path
+from typing import List
 
 import astropy.units as u
 import numpy as np
@@ -22,6 +24,7 @@
 from scipy.spatial import KDTree
 
 from romancal import __version__ as romancal_version
+from romancal.lib import psf
 
 from ._wcs_helpers import pixel_scale_angle_at_skycoord
 
@@ -84,6 +87,7 @@ def __init__(
         aperture_params,
         ci_star_thresholds,
         kernel_fwhm,
+        fit_psf,
         flux_unit="uJy",
     ):
 
@@ -96,6 +100,7 @@ def __init__(
         self.aperture_params = aperture_params
         self.kernel_sigma = kernel_fwhm * gaussian_fwhm_to_sigma
         self.flux_unit = flux_unit
+        self.fit_psf = fit_psf
 
         self.sb_unit = "MJy/sr"
         self.l2_unit = "DN/s"
@@ -113,6 +118,12 @@ def __init__(
         self.wcs = self.model.meta.wcs
         self.column_desc = {}
         self.meta = {}
+        if self.fit_psf:
+            # get the necessary columns for the PSF photometry table
+            # and its name mapping to the final catalog
+            self.psf_photometry_catalog_mapping = (
+                self.get_psf_photometry_catalog_colnames_mapping()
+            )
 
     @lazyproperty
     def _pixscale_angle(self):
@@ -1007,10 +1018,12 @@ def colnames(self):
         """
         The column name order for the final source catalog.
         """
-        colnames = self.segment_colnames[0:4]
+        colnames = self.segment_colnames[:4]
         colnames.extend(self.aperture_colnames)
         colnames.extend(self.extras_colnames)
         colnames.extend(self.segment_colnames[4:])
+        if self.fit_psf:
+            colnames.extend(self.psf_photometry_colnames)
         return colnames
 
     def _update_metadata(self):
@@ -1073,6 +1086,125 @@ def _split_skycoord(self, table):
 
         return table
 
+    @staticmethod
+    def get_psf_photometry_catalog_colnames_mapping() -> List:
+        """
+        Set the mapping between the PSF photometry table column names
+        and the final catalog column names.
+
+        Returns
+        -------
+            List of dictionaries containing old column names, new column names,
+            and descriptions for PSF photometry catalog.
+        """
+
+        return [
+            {
+                "old_name": "flags",
+                "new_name": "flag_psf",
+                "desc": "Data quality flags",
+            },
+            {
+                "old_name": "x_fit",
+                "new_name": "x_psf",
+                "desc": "X coordinate as determined by PSF fitting",
+            },
+            {
+                "old_name": "x_err",
+                "new_name": "x_psf_err",
+                "desc": "Error on X coordinate of PSF fitting",
+            },
+            {
+                "old_name": "y_fit",
+                "new_name": "y_psf",
+                "desc": "Y coordinate as determined by PSF fitting",
+            },
+            {
+                "old_name": "y_err",
+                "new_name": "y_psf_err",
+                "desc": "Error on Y coordinate of PSF fitting",
+            },
+            {
+                "old_name": "flux_fit",
+                "new_name": "flux_psf",
+                "desc": "Source flux as determined by PSF photometry",
+            },
+            {
+                "old_name": "flux_err",
+                "new_name": "flux_psf_err",
+                "desc": "Source flux error as determined by PSF photometry",
+            },
+        ]
+
+    @lazyproperty
+    def psf_photometry_colnames(self) -> List:
+        """
+        Update and return column descriptions for PSF photometry results.
+
+        This method updates the column descriptions with PSF photometry-related information
+        and returns a list of column names.
+
+        Returns
+        -------
+            List of column names for PSF photometry results.
+        """
+        desc = {x["new_name"]: x["desc"] for x in self.psf_photometry_catalog_mapping}
+
+        self.column_desc.update(desc)
+
+        return list(desc.keys())
+
+    def do_psf_photometry(self) -> None:
+        """
+        Perform PSF photometry by fitting PSF models to detected sources for refined astrometry.
+
+        This method constructs a gridded PSF model based on instrument and detector information.
+        It then fits the PSF model to the image model's sources to improve astrometric precision.
+
+        """
+        log.info("Constructing a gridded PSF model.")
+        if hasattr(self.model.meta, "instrument"):
+            # ImageModel (L2 datamodel)
+            filt = self.model.meta.instrument.optical_element
+            detector = self.model.meta.instrument.detector.replace("WFI", "SCA")
+        else:
+            # MosaicModel (L3 datamodel)
+            filt = self.model.meta.basic.optical_element
+            detector = "SCA02"
+        # prefix of the temporary FITS file that will contain the gridded PSF model
+        path_prefix = "tmp"
+        gridded_psf_model, _ = psf.create_gridded_psf_model(
+            path_prefix=path_prefix,
+            filt=filt,
+            detector=detector,
+            overwrite=True,
+            logging_level="ERROR",
+        )
+
+        log.info("Fitting a PSF model to sources for improved astrometric precision.")
+        psf_photometry_table, photometry = psf.fit_psf_to_image_model(
+            image_model=self.model,
+            psf_model=gridded_psf_model,
+            x_init=self.xcentroid,
+            y_init=self.ycentroid,
+            exclude_out_of_bounds=True,
+        )
+
+        # mapping between the columns of the PSF photometry table
+        # and the name that will be used in the final catalog
+        old_name_to_new_name_mapping = {
+            x["old_name"]: x["new_name"] for x in self.psf_photometry_catalog_mapping
+        }
+
+        # append PSF results to the class instance with the proper column name
+        for old_name, new_name in old_name_to_new_name_mapping.items():
+            setattr(self, new_name, psf_photometry_table[old_name])
+
+        # remove temporary file containing gridded_psf_model
+        filepath = Path().cwd().glob(f"{path_prefix}*{detector.lower()}*.fits")
+        for filename in filepath:
+            filename.unlink(missing_ok=True)
+
     @lazyproperty
     def catalog(self):
         """
@@ -1086,6 +1218,8 @@ def catalog(self):
         self.set_segment_properties()
         self.set_aperture_properties()
         self.set_ci_properties()
+        if self.fit_psf:
+            self.do_psf_photometry()
 
         catalog = QTable()
         for column in self.colnames: