Merge pull request #278 from lsst/tickets/DM-39168

DM-39168: Add proper motion/parallax correction to astrometry metrics

Author: cmsaunders

pipelines/matchedVisitQualityCore.yaml

@@ -88,6 +88,8 @@ tasks:
     config:
       connections.sourceCatalogs: preSourceTable_visit
       connections.outputName: matchedPreVisitCore
+      # Proper motion catalogs are not available for `preSourceTables`
+      applyAstrometricCorrections: false
       atools.stellarAstrometricSelfRepeatabilityRA: AstrometricRepeatability
       atools.stellarAstrometricSelfRepeatabilityRA.coordinate: 'RA'
       atools.stellarAstrometricSelfRepeatabilityDec: AstrometricRepeatability

pipelines/visitQualityCore.yaml

@@ -39,6 +39,8 @@ tasks:
       atools.photoDiffPsfScatterVisitPlot: TargetRefCatDeltaPsfScatterVisitPlot
       atools.photoDiffCModelScatterVisitPlot: TargetRefCatDeltaAp09ScatterVisitPlot
       python: from lsst.analysis.tools.atools import *
+  objectEpochTable:
+    class: lsst.analysis.tools.tasks.ObjectEpochTableTask
   sourceObjectMatch:
     class: lsst.analysis.tools.tasks.SourceObjectTableAnalysisTask
     config:

python/lsst/analysis/tools/tasks/associatedSourcesTractAnalysis.py

@@ -22,12 +22,18 @@
 
 __all__ = ("AssociatedSourcesTractAnalysisConfig", "AssociatedSourcesTractAnalysisTask")
 
+import astropy.time
+import astropy.units as u
+import lsst.pex.config as pexConfig
 import numpy as np
-from astropy.table import join, vstack
+import pandas as pd
+from astropy.table import Table, join, vstack
+from lsst.drp.tasks.gbdesAstrometricFit import calculate_apparent_motion
 from lsst.geom import Box2D
 from lsst.pipe.base import NoWorkFound
 from lsst.pipe.base import connectionTypes as ct
 from lsst.skymap import BaseSkyMap
+from smatch import Matcher
 
 from ..interfaces import AnalysisBaseConfig, AnalysisBaseConnections, AnalysisPipelineTask
 
@@ -70,11 +76,59 @@ class AssociatedSourcesTractAnalysisConnections(
         dimensions=("instrument",),
         isCalibration=True,
     )
+    astrometricCorrectionCatalogs = ct.Input(
+        doc="Catalog containing proper motion and parallax.",
+        name="gbdesAstrometricFit_starCatalog",
+        storageClass="ArrowNumpyDict",
+        dimensions=("instrument", "skymap", "tract", "physical_filter"),
+        multiple=True,
+        deferLoad=True,
+    )
+    visitTable = ct.Input(
+        doc="Catalog containing visit information.",
+        name="visitTable",
+        storageClass="DataFrame",
+        dimensions=("instrument",),
+    )
+
+    def __init__(self, *, config=None):
+        super().__init__(config=config)
+
+        if not config.applyAstrometricCorrections:
+            self.inputs.remove("astrometricCorrectionCatalogs")
+            self.inputs.remove("visitTable")
 
 
 class AssociatedSourcesTractAnalysisConfig(
     AnalysisBaseConfig, pipelineConnections=AssociatedSourcesTractAnalysisConnections
 ):
+    applyAstrometricCorrections = pexConfig.Field(
+        dtype=bool,
+        default=True,
+        doc="Apply proper motion and parallax corrections to source positions.",
+    )
+    matchingRadius = pexConfig.Field(
+        dtype=float,
+        default=0.2,
+        doc=(
+            "Radius in mas with which to match the mean positions of the sources with the positions in the"
+            " astrometricCorrectionCatalogs."
+        ),
+    )
+    astrometricCorrectionParameters = pexConfig.DictField(
+        keytype=str,
+        itemtype=str,
+        # TODO: DM-45845 Update default names when catalog gets updated.
+        default={
+            "ra": "starX",
+            "dec": "starY",
+            "pmRA": "starPMx",
+            "pmDec": "starPMy",
+            "parallax": "starParallax",
+        },
+        doc="Column names for position and motion parameters in the astrometric correction catalogs.",
+    )
+
     def setDefaults(self):
         super().setDefaults()
 
@@ -91,28 +145,39 @@ def getBoxWcs(skymap, tract):
         tractBox = tractInfo.getBBox()
         return tractBox, wcs
 
-    @classmethod
-    def callback(cls, inputs, dataId):
+    def callback(self, inputs, dataId):
         """Callback function to be used with reconstructor."""
-        return cls.prepareAssociatedSources(
+        return self.prepareAssociatedSources(
             inputs["skyMap"],
             dataId["tract"],
             inputs["sourceCatalogs"],
             inputs["associatedSources"],
+            inputs["astrometricCorrectionCatalogs"],
+            inputs["visitTable"],
         )
 
-    @classmethod
-    def prepareAssociatedSources(cls, skymap, tract, sourceCatalogs, associatedSources):
+    def prepareAssociatedSources(
+        self,
+        skymap,
+        tract,
+        sourceCatalogs,
+        associatedSources,
+        astrometricCorrectionCatalogs=None,
+        visitTable=None,
+    ):
         """Concatenate source catalogs and join on associated object index."""
 
         # Keep only sources with associations
         sourceCatalogStack = vstack(sourceCatalogs)
         dataJoined = join(sourceCatalogStack, associatedSources, keys="sourceId", join_type="inner")
 
+        if astrometricCorrectionCatalogs is not None:
+            self.applyAstrometricCorrections(dataJoined, astrometricCorrectionCatalogs, visitTable)
+
         # Determine which sources are contained in tract
         ra = np.radians(dataJoined["coord_ra"])
         dec = np.radians(dataJoined["coord_dec"])
-        box, wcs = cls.getBoxWcs(skymap, tract)
+        box, wcs = self.getBoxWcs(skymap, tract)
         box = Box2D(box)
         x, y = wcs.skyToPixelArray(ra, dec)
         boxSelection = box.contains(x, y)
@@ -123,6 +188,77 @@ def prepareAssociatedSources(cls, skymap, tract, sourceCatalogs, associatedSourc
 
         return dataFiltered
 
+    def applyAstrometricCorrections(self, dataJoined, astrometricCorrectionCatalogs, visitTable):
+        """Use proper motion/parallax catalogs to shift positions to median
+        epoch of the visits.
+
+        Parameters
+        ----------
+        dataJoined : `astropy.table.Table`
+            Table containing source positions, which will be modified in place.
+        astrometricCorrectionCatalogs: `dict` [`pd.DataFrame`]
+            Dictionary keyed by band with proper motion and parallax catalogs.
+        visitTable : `pd.DataFrame`
+            Table containing the MJDs of the visits.
+        """
+        for band in np.unique(dataJoined["band"]):
+            bandInd = dataJoined["band"] == band
+            bandSources = dataJoined[bandInd]
+            # Add key for sorting below.
+            bandSources["__index__"] = np.arange(len(bandSources))
+            bandSourcesDf = bandSources.to_pandas()
+            meanRAs = bandSourcesDf.groupby("obj_index")["coord_ra"].aggregate("mean")
+            meanDecs = bandSourcesDf.groupby("obj_index")["coord_dec"].aggregate("mean")
+
+            bandPMs = astrometricCorrectionCatalogs[band]
+            with Matcher(meanRAs, meanDecs) as m:
+                idx, i1, i2, d = m.query_radius(
+                    bandPMs[self.config.astrometricCorrectionParameters["ra"]],
+                    bandPMs[self.config.astrometricCorrectionParameters["dec"]],
+                    (self.config.matchingRadius * u.mas).to(u.degree),
+                    return_indices=True,
+                )
+
+            catRAs = np.zeros_like(meanRAs)
+            catDecs = np.zeros_like(meanRAs)
+            pmRAs = np.zeros_like(meanRAs)
+            pmDecs = np.zeros_like(meanRAs)
+            parallaxes = np.zeros(len(meanRAs))
+            catRAs[i1] = bandPMs[self.config.astrometricCorrectionParameters["ra"]][i2]
+            catDecs[i1] = bandPMs[self.config.astrometricCorrectionParameters["dec"]][i2]
+            pmRAs[i1] = bandPMs[self.config.astrometricCorrectionParameters["pmRA"]][i2]
+            pmDecs[i1] = bandPMs[self.config.astrometricCorrectionParameters["pmDec"]][i2]
+            parallaxes[i1] = bandPMs[self.config.astrometricCorrectionParameters["parallax"]][i2]
+
+            pmDf = Table(
+                {
+                    "ra": catRAs * u.degree,
+                    "dec": catDecs * u.degree,
+                    "pmRA": pmRAs * u.mas / u.yr,
+                    "pmDec": pmDecs * u.mas / u.yr,
+                    "parallax": parallaxes * u.mas,
+                    "obj_index": meanRAs.index,
+                }
+            )
+
+            dataWithPM = join(bandSources, pmDf, keys="obj_index", join_type="left")
+
+            visits = bandSourcesDf["visit"].unique()
+            mjds = [visitTable.loc[visit]["expMidptMJD"] for visit in visits]
+            mjdTable = Table(
+                [astropy.time.Time(mjds, format="mjd", scale="tai"), visits], names=["MJD", "visit"]
+            )
+            dataWithMJD = join(dataWithPM, mjdTable, keys="visit", join_type="left")
+            # After astropy 7.0, it should be possible to use "keep_order=True"
+            # in the join and avoid sorting.
+            dataWithMJD.sort("__index__")
+            medianMJD = astropy.time.Time(np.median(mjds), format="mjd", scale="tai")
+
+            raCorrection, decCorrection = calculate_apparent_motion(dataWithMJD, medianMJD)
+
+            dataJoined["coord_ra"][bandInd] = dataWithMJD["coord_ra"] - raCorrection.value
+            dataJoined["coord_dec"][bandInd] = dataWithMJD["coord_dec"] - decCorrection.value
+
     def runQuantum(self, butlerQC, inputRefs, outputRefs):
         inputs = butlerQC.get(inputRefs)
 
@@ -135,6 +271,18 @@ def runQuantum(self, butlerQC, inputRefs, outputRefs):
             sourceCatalogs.append(self.loadData(handle, names))
         inputs["sourceCatalogs"] = sourceCatalogs
 
+        if self.config.applyAstrometricCorrections:
+            astrometricCorrections = {}
+            for pmCatRef in inputs["astrometricCorrectionCatalogs"]:
+                pmCat = pmCatRef.get(
+                    parameters={"columns": self.config.astrometricCorrectionParameters.values()}
+                )
+                astrometricCorrections[pmCatRef.dataId["band"]] = pd.DataFrame(pmCat)
+            inputs["astrometricCorrectionCatalogs"] = astrometricCorrections
+        else:
+            inputs["astrometricCorrectionCatalogs"] = None
+            inputs["visitTable"] = None
+
         dataId = butlerQC.quantum.dataId
         plotInfo = self.parsePlotInfo(inputs, dataId, connectionName="associatedSources")