Skip to content

Commit

Permalink
Pass entire struct up through ssoAssociation
Browse files Browse the repository at this point in the history
  • Loading branch information
Gerenjie committed Jan 10, 2025
1 parent f468c93 commit 7c860aa
Show file tree
Hide file tree
Showing 4 changed files with 85 additions and 52 deletions.
26 changes: 18 additions & 8 deletions python/lsst/ap/association/diaPipe.py
Original file line number Diff line number Diff line change
Expand Up @@ -128,6 +128,13 @@ class DiaPipelineConnections(
storageClass="DataFrame",
dimensions=("instrument", "visit", "detector"),
)
unassociatedSsObjects = connTypes.Output(
doc="Expected locations of an ssObject with no source",
name="ssUnassociatedObjects",
storageClass="ArrowAstropy",
dimensions=("instrument", "visit", "detector"),
)

diaForcedSources = connTypes.Output(
doc="Optional output storing the forced sources computed at the diaObject positions.",
name="{fakesType}{coaddName}Diff_diaForcedSrc",
Expand All @@ -154,6 +161,7 @@ def __init__(self, *, config=None):
self.inputs.remove("solarSystemObjectTable")
if (not config.doWriteAssociatedSources) or (not config.doSolarSystemAssociation):
self.outputs.remove("associatedSsSources")
self.outputs.remove("unassociatedSsObjects")

def adjustQuantum(self, inputs, outputs, label, dataId):
"""Override to make adjustments to `lsst.daf.butler.DatasetRef` objects
Expand Down Expand Up @@ -434,9 +442,10 @@ def run(self,
diaObjects = preloadedDiaObjects

# Associate DiaSources with DiaObjects
associatedDiaSources, newDiaObjects, associatedSsSources = self.associateDiaSources(
diaSourceTable, solarSystemObjectTable, diffIm, diaObjects
)

(
associatedDiaSources, newDiaObjects, associatedSsSources, unassociatedSsObjects
) = self.associateDiaSources(diaSourceTable, solarSystemObjectTable, diffIm, diaObjects)

# Merge associated diaSources
mergedDiaSourceHistory, mergedDiaObjects, updatedDiaObjectIds = self.mergeAssociatedCatalogs(
Expand Down Expand Up @@ -517,7 +526,8 @@ def run(self,
associatedDiaSources=associatedDiaSources,
diaForcedSources=diaForcedSources,
diaObjects=diaCalResult.diaObjectCat,
associatedSsSources=associatedSsSources
associatedSsSources=associatedSsSources,
unassociatedSsObjects=unassociatedSsObjects
)

def createNewDiaObjects(self, unAssocDiaSources):
Expand Down Expand Up @@ -598,15 +608,15 @@ def associateDiaSources(self, diaSourceTable, solarSystemObjectTable, diffIm, di
toAssociate.append(ssoAssocResult.ssoAssocDiaSources)
nTotalSsObjects = ssoAssocResult.nTotalSsObjects
nAssociatedSsObjects = ssoAssocResult.nAssociatedSsObjects
self.metadata['numTotalSolarSystemObjects'] = nTotalSsObjects
self.metadata['numAssociatedSsObjects'] = nAssociatedSsObjects
associatedSsSources = ssoAssocResult.ssSourceData
associatedSsSources = ssoAssocResult.associatedSsSources
unassociatedSsObjects = ssoAssocResult.unassociatedSsObjects
else:
# Create new DiaObjects from unassociated diaSources.
createResults = self.createNewDiaObjects(assocResults.unAssocDiaSources)
nTotalSsObjects = 0
nAssociatedSsObjects = 0
associatedSsSources = None
unassociatedSsObjects = None
if len(assocResults.matchedDiaSources) > 0:
toAssociate.append(assocResults.matchedDiaSources)
toAssociate.append(createResults.diaSources)
Expand All @@ -622,7 +632,7 @@ def associateDiaSources(self, diaSourceTable, solarSystemObjectTable, diffIm, di
assocResults.nUnassociatedDiaObjects,
createResults.nNewDiaObjects,
)
return (associatedDiaSources, createResults.newDiaObjects, associatedSsSources)
return (associatedDiaSources, createResults.newDiaObjects, associatedSsSources, unassociatedSsObjects)

@timeMethod
def mergeAssociatedCatalogs(self, preloadedDiaSources, associatedDiaSources, diaObjects, newDiaObjects,
Expand Down
46 changes: 19 additions & 27 deletions python/lsst/ap/association/ssSingleFrameAssociation.py
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,12 @@ class SsSingleFrameAssociationConnections(
storageClass="ArrowAstropy",
dimensions=("instrument", "visit", "detector"),
)
unassociatedSsObjects = connTypes.Output(
doc="Expected locations of an ssObject with no source",
name="ssSingleFrameUnassociatedObjects",
storageClass="ArrowAstropy",
dimensions=("instrument", "visit", "detector"),
)


class SsSingleFrameAssociationConfig(pipeBase.PipelineTaskConfig,
Expand Down Expand Up @@ -122,7 +128,7 @@ def run(self,
----------
exposure : `lsst.afw.image.ExposureF`
Calibrated exposure with wcs and midpoint time.
diaSourceTable : `pandas.DataFrame`
sourceTable : `lsst.afw.table.SourceCatalog`
Newly detected sources.
band : `str`
The band in which the new DiaSources were detected.
Expand All @@ -133,9 +139,16 @@ def run(self,
-------
results : `lsst.pipe.base.Struct`
Results struct with components.
- ``ssoAssocDiaSources`` : DiaSources that were associated with
solar system objects in this visit. (`Astropy.table.Table`)
- ``unAssocDiaSources`` : Set of DiaSources that were not
associated with any solar system object. (`pandas.DataFrame`)
- ``nTotalSsObjects`` : Total number of SolarSystemObjects
contained in the CCD footprint. (`int`)
- ``nAssociatedSsObjects`` : Number of SolarSystemObjects
that were associated with DiaSources.
- ``ssSourceData`` : ssSource table data. (`Astropy.table.Table`)
- ``associatedSsSources`` : Catalog of ssSource records.
(`astropy.table.Table`)
Raises
------
Expand All @@ -144,31 +157,10 @@ def run(self,
"""
if solarSystemObjectTable is None:
raise pipeBase.NoWorkFound("No ephemerides to associate. Skipping ssSingleFrameAssociation.")
else:
# Associate DiaSources with DiaObjects
associatedSsSources = self.associateSources(sourceTable, solarSystemObjectTable, exposure)
return pipeBase.Struct(associatedSsSources=associatedSsSources)

@timeMethod
def associateSources(self, sourceTable, solarSystemObjectTable, exposure):
"""Associate single-image sources with ssObjects.
Parameters
----------
sourceTable : `pandas.DataFrame`
Newly detected sources.
solarSystemObjectTable : `pandas.DataFrame`
Preloaded Solar System objects expected to be visible in the image.
Returns
-------
associatedSsSources : `astropy.table.Table`
Table of new ssSources after association.
"""
# Associate DiaSources with DiaObjects
sourceTable = sourceTable.asAstropy()
sourceTable['ra'] = sourceTable['coord_ra'].to(deg).value
sourceTable['dec'] = sourceTable['coord_dec'].to(deg).value
ssoAssocResult = self.solarSystemAssociator.run(sourceTable.to_pandas(),
solarSystemObjectTable, exposure)
associatedSsSources = ssoAssocResult.ssSourceData
return associatedSsSources
return self.solarSystemAssociator.run(sourceTable.to_pandas(),
solarSystemObjectTable, exposure)
62 changes: 46 additions & 16 deletions python/lsst/ap/association/ssoAssociation.py
Original file line number Diff line number Diff line change
Expand Up @@ -86,18 +86,18 @@ def run(self, diaSourceCatalog, solarSystemObjects, exposure):
resultsStruct : `lsst.pipe.base.Struct`
- ``ssoAssocDiaSources`` : DiaSources that were associated with
solar system objects in this visit. (`pandas.DataFrame`)
solar system objects in this visit. (`Astropy.table.Table`)
- ``unAssocDiaSources`` : Set of DiaSources that were not
associated with any solar system object. (`pandas.DataFrame`)
- ``nTotalSsObjects`` : Total number of SolarSystemObjects
contained in the CCD footprint. (`int`)
- ``nAssociatedSsObjects`` : Number of SolarSystemObjects
that were associated with DiaSources.
that were associated with DiaSources. (`int`)
- ``ssSourceData`` : ssSource table data. (`Astropy.table.Table`)
"""
nSolarSystemObjects = len(solarSystemObjects)
if nSolarSystemObjects <= 0:
return self._return_empty(diaSourceCatalog)
return self._return_empty(diaSourceCatalog, solarSystemObjects)

mjd_midpoint = exposure.visitInfo.date.toAstropy().tai.mjd
ref_time = mjd_midpoint - solarSystemObjects["tmin"].values[0]
Expand Down Expand Up @@ -125,12 +125,11 @@ def run(self, diaSourceCatalog, solarSystemObjects, exposure):
maskedObjects = self._maskToCcdRegion(
solarSystemObjects,
exposure,
solarSystemObjects["Err(arcsec)"].max())
solarSystemObjects["Err(arcsec)"].max()).copy()
nSolarSystemObjects = len(maskedObjects)
if nSolarSystemObjects <= 0:
return self._return_empty(diaSourceCatalog)
return self._return_empty(diaSourceCatalog, maskedObjects)

self.log.info("Attempting to associate %d objects...", nSolarSystemObjects)
maxRadius = np.deg2rad(self.config.maxDistArcSeconds / 3600)

# Transform DIA RADEC coordinates to unit sphere xyz for tree building.
Expand Down Expand Up @@ -163,23 +162,29 @@ def run(self, diaSourceCatalog, solarSystemObjects, exposure):
decs.append(dia_dec)
expected_ras.append(ssObject["ra"])
expected_decs.append(ssObject["dec"])

self.log.info("Successfully associated %d SolarSystemObjects.", nFound)
assocMask = diaSourceCatalog["ssObjectId"] != 0
maskedObjects.loc[index, 'associated'] = True
else:
maskedObjects.loc[index, 'associated'] = False

self.log.info("Successfully associated %d / %d SolarSystemObjects.", nFound, nSolarSystemObjects)
self.metadata['nAssociatedSsObjects'] = nFound
self.metadata['nExpectedSsObjects'] = nSolarSystemObjects
assocSourceMask = diaSourceCatalog["ssObjectId"] != 0
unAssocObjectMask = np.logical_not(maskedObjects['associated'].values)
ssSourceData = pd.DataFrame(ssSourceData, columns=["ssObjectId", "obs_position_x", "obs_position_y",
"obs_position_z", "obj_position_x",
"obj_position_y", "obj_position_z"])
ssSourceData["ra"] = ras
ssSourceData["dec"] = decs
ssSourceData["expected_ra"] = expected_ras
ssSourceData["expected_dec"] = expected_decs

return pipeBase.Struct(
ssoAssocDiaSources=diaSourceCatalog[assocMask].reset_index(drop=True),
unAssocDiaSources=diaSourceCatalog[~assocMask].reset_index(drop=True),
ssoAssocDiaSources=Table.from_pandas(diaSourceCatalog[assocSourceMask].reset_index(drop=True)),
unAssocDiaSources=diaSourceCatalog[~assocSourceMask].reset_index(drop=True),
nTotalSsObjects=nSolarSystemObjects,
nAssociatedSsObjects=nFound,
ssSourceData=Table.from_pandas(ssSourceData))
associatedSsSources=Table.from_pandas(ssSourceData),
unassociatedSsObjects=Table.from_pandas(maskedObjects[unAssocObjectMask]))

def _maskToCcdRegion(self, solarSystemObjects, exposure, marginArcsec):
"""Mask the input SolarSystemObjects to only those in the exposure
Expand Down Expand Up @@ -242,13 +247,38 @@ def _radec_to_xyz(self, ras, decs):

return vectors

def _return_empty(self, diaSourceCatalog):
def _return_empty(self, diaSourceCatalog, emptySolarSystemObjects):
"""Return a struct with all appropriate empty values for no SSO associations.
Parameters
----------
diaSourceCatalog : `pandas.DataFrame`
Used for column names
emptySolarSystemObjects : `pandas.DataFrame`
Used for column names.
Returns
-------
results : `lsst.pipe.base.Struct`
Results struct with components.
- ``ssoAssocDiaSources`` : Empty. (`Astropy.table.Table`)
- ``unAssocDiaSources`` : Input DiaSources. (`Astropy.table.Table`)
- ``nTotalSsObjects`` : Zero. (`int`)
- ``nAssociatedSsObjects`` : Zero.
- ``ssSourceData`` : Empty. (`Astropy.table.Table`)
Raises
------
RuntimeError
Raised if duplicate DiaObjects or duplicate DiaSources are found.
"""
self.log.info("No SolarSystemObjects found in detector bounding box.")
return pipeBase.Struct(
ssoAssocDiaSources=Table(names=diaSourceCatalog.columns),
unAssocDiaSources=diaSourceCatalog,
nTotalSsObjects=0,
nAssociatedSsObjects=0,
ssSourceData=Table(names=["ssObjectId", "ra", "dec", "obs_position", "obj_position",
"residual_ras", "residual_decs"])
associatedSsSources=Table(names=["ssObjectId", "ra", "dec", "obs_position", "obj_position",
"residual_ras", "residual_decs"]),
unassociatedSsObjects=Table(names=emptySolarSystemObjects.columns)
)
3 changes: 2 additions & 1 deletion tests/test_diaPipe.py
Original file line number Diff line number Diff line change
Expand Up @@ -151,7 +151,8 @@ def solarSystemAssociator_run(unAssocDiaSources, solarSystemObjectTable, diffIm)
nAssociatedSsObjects=30,
ssoAssocDiaSources=_makeMockDataFrame(),
unAssocDiaSources=_makeMockDataFrame(),
ssSourceData=_makeMockDataFrame())
associatedSsSources=_makeMockDataFrame(),
unassociatedSsObjects=_makeMockDataFrame())

def associator_run(table, diaObjects):
return lsst.pipe.base.Struct(nUpdatedDiaObjects=2, nUnassociatedDiaObjects=3,
Expand Down

0 comments on commit 7c860aa

Please sign in to comment.