Coverage for python/lsst/ap/association/transformDiaSourceCatalog.py: 87%
183 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-11-08 01:50 -0800
1# This file is part of ap_association
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (https://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This program is free software: you can redistribute it and/or modify
10# it under the terms of the GNU General Public License as published by
11# the Free Software Foundation, either version 3 of the License, or
12# (at your option) any later version.
13#
14# This program is distributed in the hope that it will be useful,
15# but WITHOUT ANY WARRANTY; without even the implied warranty of
16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
17# GNU General Public License for more details.
18#
19# You should have received a copy of the GNU General Public License
20# along with this program. If not, see <https://www.gnu.org/licenses/>.
22__all__ = ("TransformDiaSourceCatalogConnections",
23 "TransformDiaSourceCatalogConfig",
24 "TransformDiaSourceCatalogTask",
25 "UnpackApdbFlags")
27import dataclasses
28import os
29import yaml
31import numpy as np
33from lsst.resources import ResourcePath
34from lsst.daf.base import DateTime
35import lsst.pex.config as pexConfig
36import lsst.pipe.base as pipeBase
37import lsst.pipe.base.connectionTypes as connTypes
38from lsst.pipe.tasks.postprocess import TransformCatalogBaseTask, TransformCatalogBaseConfig
39from lsst.pipe.tasks.functors import Column
40from lsst.utils.timer import timeMethod
42from lsst.pipe.tasks.schemaUtils import convertDataFrameToSdmSchema, readSdmSchemaFile
44_UNKNOWN_INSTRUMENT = "UNKNOWN"
45"""Name of the template instrument when it is not known, as recorded by
46`lsst.images.DifferenceImageTemplateInfo`.
47"""
50class TransformDiaSourceCatalogConnections(pipeBase.PipelineTaskConnections,
51 dimensions=("instrument", "visit", "detector"),
52 defaultTemplates={"coaddName": "deep", "fakesType": ""}):
53 diaSourceSchema = connTypes.InitInput(
54 doc="Schema for DIASource catalog output by ImageDifference.",
55 storageClass="SourceCatalog",
56 name="{fakesType}{coaddName}Diff_diaSrc_schema",
57 )
58 diaSourceCat = connTypes.Input(
59 doc="Catalog of DiaSources produced during image differencing.",
60 name="{fakesType}{coaddName}Diff_candidateDiaSrc",
61 storageClass="SourceCatalog",
62 dimensions=("instrument", "visit", "detector"),
63 )
64 diffIm = connTypes.Input(
65 doc="Difference image on which the DiaSources were detected.",
66 name="{fakesType}{coaddName}Diff_differenceExp",
67 storageClass="ExposureF",
68 dimensions=("instrument", "visit", "detector"),
69 )
70 diaSourceTable = connTypes.Output(
71 doc=".",
72 name="{fakesType}{coaddName}Diff_diaSrcTable",
73 storageClass="ArrowAstropy",
74 dimensions=("instrument", "visit", "detector"),
75 )
77 def __init__(self, *, config=None):
78 super().__init__(config=config)
79 if config.input_image_type == "future":
80 self.diffIm = dataclasses.replace(self.diffIm, storageClass="DifferenceImage")
83class TransformDiaSourceCatalogConfig(TransformCatalogBaseConfig,
84 pipelineConnections=TransformDiaSourceCatalogConnections):
85 input_image_type = pexConfig.ChoiceField[str](
86 doc="Type of the input difference image. run() takes an"
87 " lsst.afw.image.Exposure either way.",
88 allowed={
89 "legacy": "Read an `lsst.afw.image.Exposure`.",
90 "future": "Read an `lsst.images.DifferenceImage`.",
91 },
92 optional=False,
93 default="legacy",
94 )
95 flagMap = pexConfig.Field(
96 dtype=str,
97 doc="Yaml file specifying SciencePipelines flag fields to bit packs.",
98 default=os.path.join("${AP_ASSOCIATION_DIR}",
99 "data",
100 "association-flag-map.yaml"),
101 )
102 flagRenameMap = pexConfig.Field(
103 dtype=str,
104 doc="Yaml file specifying specifying rules to rename flag names",
105 default=os.path.join("${AP_ASSOCIATION_DIR}",
106 "data",
107 "flag-rename-rules.yaml"),
108 )
109 doRemoveSkySources = pexConfig.Field(
110 dtype=bool,
111 default=False,
112 doc="Input DiaSource catalog contains SkySources that should be "
113 "removed before storing the output DiaSource catalog.",
114 )
115 # TODO: remove on DM-41532
116 doPackFlags = pexConfig.Field(
117 dtype=bool,
118 default=False,
119 doc="Do pack the flags into one integer column named 'flags'."
120 "If False, instead produce one boolean column per flag.",
121 deprecated="This field is no longer used. Will be removed after v28."
122 )
123 doUseApdbSchema = pexConfig.Field(
124 dtype=bool,
125 default=False,
126 doc="Use the APDB schema to coerce the data types of the output columns.",
127 deprecated="This field has been renamed to doUseSchema, and will be "
128 "removed after v30."
129 )
130 doUseSchema = pexConfig.Field(
131 dtype=bool,
132 default=False,
133 doc="Use an existing schema to coerce the data types of the output columns."
134 )
135 schemaDir = pexConfig.Field(
136 dtype=str,
137 doc="Path to the directory containing schema definitions.",
138 default=os.path.join("${SDM_SCHEMAS_DIR}",
139 "yml"),
140 )
141 schemaFile = pexConfig.Field(
142 dtype=str,
143 doc="Yaml file specifying the schema of the output catalog.",
144 default="apdb.yaml",
145 )
146 schemaName = pexConfig.Field(
147 dtype=str,
148 doc="Name of the table in the schema file to read.",
149 default="ApdbSchema",
150 deprecated="This config is no longer used, and will be removed after v30"
151 )
153 def setDefaults(self):
154 super().setDefaults()
155 self.functorFile = os.path.join("${AP_ASSOCIATION_DIR}",
156 "data",
157 "DiaSource.yaml")
160class TransformDiaSourceCatalogTask(TransformCatalogBaseTask):
161 """Transform a DiaSource catalog by calibrating and renaming columns to
162 produce a table ready to insert into the Apdb.
164 Parameters
165 ----------
166 initInputs : `dict`
167 Must contain ``diaSourceSchema`` as the schema for the input catalog.
168 """
169 ConfigClass = TransformDiaSourceCatalogConfig
170 _DefaultName = "transformDiaSourceCatalog"
171 # Needed to create a valid TransformCatalogBaseTask, but unused
172 inputDataset = "deepDiff_diaSrc"
173 outputDataset = "deepDiff_diaSrcTable"
175 def __init__(self, initInputs, **kwargs):
176 super().__init__(**kwargs)
177 self.funcs = self.getFunctors()
178 self.inputSchema = initInputs['diaSourceSchema'].schema
179 self._create_bit_pack_mappings()
180 if self.config.doUseSchema:
181 schemaFile = os.path.join(self.config.schemaDir, self.config.schemaFile)
182 self.schema = readSdmSchemaFile(schemaFile)
183 else:
184 self.schema = None
186 if not self.config.doPackFlags: 186 ↛ exitline 186 didn't return from function '__init__' because the condition on line 186 was always true
187 # get the flag rename rules
188 with open(os.path.expandvars(self.config.flagRenameMap)) as yaml_stream:
189 self.rename_rules = list(yaml.safe_load_all(yaml_stream))
191 def _create_bit_pack_mappings(self):
192 """Setup all flag bit packings.
193 """
194 self.bit_pack_columns = []
195 flag_map_file = os.path.expandvars(self.config.flagMap)
196 with open(flag_map_file) as yaml_stream:
197 table_list = list(yaml.safe_load_all(yaml_stream))
198 for table in table_list: 198 ↛ 206line 198 didn't jump to line 206
199 if table['tableName'] == 'DiaSource': 199 ↛ 198line 199 didn't jump to line 198 because the condition on line 199 was always true
200 self.bit_pack_columns = table['columns']
201 break
203 # Test that all flags requested are present in the input schemas.
204 # Output schemas are flexible, however if names are not specified in
205 # the Apdb schema, flag columns will not be persisted.
206 for outputFlag in self.bit_pack_columns:
207 bitList = outputFlag['bitList']
208 for bit in bitList:
209 try:
210 self.inputSchema.find(bit['name'])
211 except KeyError:
212 raise KeyError(
213 "Requested column %s not found in input DiaSource "
214 "schema. Please check that the requested input "
215 "column exists." % bit['name'])
217 def runQuantum(self, butlerQC, inputRefs, outputRefs):
218 inputs = butlerQC.get(inputRefs)
219 inputs["band"] = butlerQC.quantum.dataId["band"]
220 inputs["instrument"] = butlerQC.quantum.dataId["instrument"]
221 if self.config.input_image_type == "future": 221 ↛ 228line 221 didn't jump to line 228 because the condition on line 221 was always true
222 # Templates are required to all use the same instrument.
223 templates = getattr(inputs["diffIm"], "templates", None)
224 if templates:
225 inputs["templateInstrument"] = templates[0].instrument
226 inputs["diffIm"] = inputs["diffIm"].to_legacy()
228 outputs = self.run(**inputs)
230 butlerQC.put(outputs, outputRefs)
232 @timeMethod
233 def run(self,
234 diaSourceCat,
235 diffIm,
236 band,
237 reliability=None,
238 instrument=None,
239 templateInstrument=None):
240 """Convert input catalog to ParquetTable/Pandas and run functors.
242 Additionally, add new columns for stripping information from the
243 exposure and into the DiaSource catalog.
245 Parameters
246 ----------
247 diaSourceCat : `lsst.afw.table.SourceCatalog`
248 Catalog of sources measured on the difference image.
249 diffIm : `lsst.afw.image.Exposure`
250 Result of subtracting template and science images.
251 band : `str`
252 Filter band of the science image.
253 reliability : `lsst.afw.table.SourceCatalog`
254 Reliability (e.g. real/bogus) scores, row-matched to
255 ``diaSourceCat``.
256 instrument : `str`, optional
257 Name of the instrument that took the science image.
258 templateInstrument : `str`, optional
259 Name of the instrument recorded for the template, or `None` if it
260 is not known.
262 Returns
263 -------
264 results : `lsst.pipe.base.Struct`
265 Results struct with components.
267 - ``diaSourceTable`` : Catalog of DiaSources with calibrated values
268 and renamed columns.
269 (`lsst.pipe.tasks.ParquetTable` or `pandas.DataFrame`)
270 """
271 self.log.info(
272 "Transforming/standardizing the DiaSource table for visit,detector: %i, %i",
273 diffIm.visitInfo.id, diffIm.detector.getId())
275 diaSourceDf = diaSourceCat.asAstropy().to_pandas()
276 if self.config.doRemoveSkySources: 276 ↛ 277line 276 didn't jump to line 277 because the condition on line 276 was never true
277 diaSourceDf = diaSourceDf[~diaSourceDf["sky_source"]]
278 diaSourceCat = diaSourceCat[~diaSourceCat["sky_source"]]
280 # Need UTC time but without a timezone because pandas requires a
281 # naive datetime.
282 diaSourceDf["timeProcessedMjdTai"] = DateTime.now().get(system=DateTime.MJD, scale=DateTime.TAI)
283 diaSourceDf["snr"] = getSignificance(diaSourceCat)
284 diaSourceDf["bboxSize"] = self.computeBBoxSizes(diaSourceCat)
285 diaSourceDf["visit"] = diffIm.visitInfo.id
286 # int16 instead of uint8 because databases don't like unsigned bytes.
287 diaSourceDf["detector"] = np.int16(diffIm.detector.getId())
288 diaSourceDf["band"] = band
289 diaSourceDf["midpointMjdTai"] = diffIm.visitInfo.date.get(system=DateTime.MJD)
290 diaSourceDf["exposureTime"] = diffIm.visitInfo.exposureTime
291 diaSourceDf["diaObjectId"] = 0
292 diaSourceDf["ssObjectId"] = 0
293 templateInstrument = templateInstrument or _UNKNOWN_INSTRUMENT
294 mixedInstruments = (instrument is not None
295 and templateInstrument not in (_UNKNOWN_INSTRUMENT, instrument))
296 diaSourceDf["templateInstrument"] = templateInstrument
297 diaSourceDf["mixedInstruments"] = mixedInstruments
298 # Only templates from another instrument have a nonzero version for now.
299 diaSourceDf["templateVersion"] = 1 if mixedInstruments else 0
301 # TODO: this has been formally deprecated and should be removed too
302 if self.config.doPackFlags: 302 ↛ 304line 302 didn't jump to line 304 because the condition on line 302 was never true
303 # either bitpack the flags
304 self.bitPackFlags(diaSourceDf)
305 else:
306 # or add the individual flag functors
307 self.addUnpackedFlagFunctors()
308 # and remove the packed flag functor
309 if 'flags' in self.funcs.funcDict:
310 del self.funcs.funcDict['flags']
312 df = self.transform(band,
313 diaSourceDf,
314 self.funcs,
315 dataId=None).df
316 if self.config.doUseSchema:
317 df = convertDataFrameToSdmSchema(self.schema, df, tableName="DiaSource")
319 return pipeBase.Struct(
320 diaSourceTable=df,
321 )
323 def addUnpackedFlagFunctors(self):
324 """Add Column functor for each of the flags to the internal functor
325 dictionary.
326 """
327 for flag in self.bit_pack_columns[0]['bitList']:
328 flagName = flag['name']
329 targetName = self.funcs.renameCol(flagName, self.rename_rules[0]['flag_rename_rules'])
330 self.funcs.update({targetName: Column(flagName)})
332 def computeBBoxSizes(self, inputCatalog):
333 """Compute the size of a square bbox that fully contains the detection
334 footprint.
336 Parameters
337 ----------
338 inputCatalog : `lsst.afw.table.SourceCatalog`
339 Catalog containing detected footprints.
341 Returns
342 -------
343 outputBBoxSizes : `np.ndarray`, (N,)
344 Array of bbox sizes.
345 """
346 # Schema validation requires that this field is int.
347 outputBBoxSizes = np.empty(len(inputCatalog), dtype=int)
348 for i, record in enumerate(inputCatalog):
349 footprintBBox = record.getFootprint().getBBox()
350 # Compute twice the size of the largest dimension of the footprint
351 # bounding box. This is the largest footprint we should need to cover
352 # the complete DiaSource assuming the centroid is within the bounding
353 # box.
354 maxSize = 2 * np.max([footprintBBox.getWidth(),
355 footprintBBox.getHeight()])
356 recX = record.getCentroid().x
357 recY = record.getCentroid().y
358 bboxSize = int(
359 np.ceil(2 * np.max(np.fabs([footprintBBox.maxX - recX,
360 footprintBBox.minX - recX,
361 footprintBBox.maxY - recY,
362 footprintBBox.minY - recY]))))
363 if bboxSize > maxSize: 363 ↛ 364line 363 didn't jump to line 364 because the condition on line 363 was never true
364 bboxSize = maxSize
365 outputBBoxSizes[i] = bboxSize
367 return outputBBoxSizes
369 def bitPackFlags(self, df):
370 """Pack requested flag columns in inputRecord into single columns in
371 outputRecord.
373 Parameters
374 ----------
375 df : `pandas.DataFrame`
376 DataFrame to read bits from and pack them into.
377 """
378 for outputFlag in self.bit_pack_columns:
379 bitList = outputFlag['bitList']
380 value = np.zeros(len(df), dtype=np.uint64)
381 for bit in bitList:
382 # Hard type the bit arrays.
383 value += (df[bit['name']]*2**bit['bit']).to_numpy().astype(np.uint64)
384 df[outputFlag['columnName']] = value
387class UnpackApdbFlags:
388 """Class for unpacking bits from integer flag fields stored in the Apdb.
390 Attributes
391 ----------
392 flag_map_file : `lsst.resources.ResourcePathExpression`
393 Absolute or relative URI to a yaml file specifiying mappings of flags
394 to integer bits.
395 table_name : `str`
396 Name of the Apdb table the integer bit data are coming from.
397 """
399 def __init__(self, flag_map_file, table_name):
400 self.bit_pack_columns = []
401 flag_map_file = os.path.expandvars(flag_map_file)
402 with ResourcePath(flag_map_file, forceDirectory=False).open("r") as yaml_stream:
403 table_list = list(yaml.safe_load_all(yaml_stream))
404 for table in table_list: 404 ↛ 409line 404 didn't jump to line 409
405 if table['tableName'] == table_name: 405 ↛ 404line 405 didn't jump to line 404 because the condition on line 405 was always true
406 self.bit_pack_columns = table['columns']
407 break
409 self.output_flag_columns = {}
411 for column in self.bit_pack_columns:
412 names = {}
413 for bit in column["bitList"]:
414 names[bit["name"]] = bit["bit"]
415 self.output_flag_columns[column["columnName"]] = names
417 def unpack(self, input_flag_values, flag_name):
418 """Determine individual boolean flags from an input array of unsigned
419 ints.
421 Parameters
422 ----------
423 input_flag_values : array-like of type uint
424 Array of integer packed bit flags to unpack.
425 flag_name : `str`
426 Apdb column name from the loaded file, e.g. "flags".
428 Returns
429 -------
430 output_flags : `numpy.ndarray`
431 Numpy structured array of booleans, one column per flag in the
432 loaded file.
433 """
434 output_flags = np.zeros(len(input_flag_values),
435 dtype=[(name, bool) for name in self.output_flag_columns[flag_name]])
437 for name in self.output_flag_columns[flag_name]:
438 masked_bits = np.bitwise_and(input_flag_values,
439 2**self.output_flag_columns[flag_name][name])
440 output_flags[name] = masked_bits
442 return output_flags
444 def flagExists(self, flagName, columnName='flags'):
445 """Check if named flag is in the bitpacked flag set.
447 Parameters:
448 ----------
449 flagName : `str`
450 Flag name to search for.
451 columnName : `str`, optional
452 Name of bitpacked flag column to search in.
454 Returns
455 -------
456 flagExists : `bool`
457 `True` if `flagName` is present in `columnName`.
459 Raises
460 ------
461 ValueError
462 Raised if `columnName` is not defined.
463 """
464 if columnName not in self.output_flag_columns:
465 raise ValueError(f'column {columnName} not in flag map: {self.output_flag_columns}')
467 return flagName in [c for c in self.output_flag_columns[columnName]]
469 def makeFlagBitMask(self, flagNames, columnName='flags'):
470 """Return a bitmask corresponding to the supplied flag names.
472 Parameters:
473 ----------
474 flagNames : `list` [`str`]
475 Flag names to include in the bitmask.
476 columnName : `str`, optional
477 Name of bitpacked flag column.
479 Returns
480 -------
481 bitmask : `np.unit64`
482 Bitmask corresponding to the supplied flag names given the loaded configuration.
484 Raises
485 ------
486 ValueError
487 Raised if a flag in `flagName` is not included in `columnName`.
488 """
489 bitmask = np.uint64(0)
491 for flag in flagNames:
492 if not self.flagExists(flag, columnName=columnName):
493 raise ValueError(f"flag '{flag}' not included in '{columnName}' flag column")
495 for outputFlag in self.bit_pack_columns:
496 if outputFlag['columnName'] == columnName: 496 ↛ 495line 496 didn't jump to line 495 because the condition on line 496 was always true
497 bitList = outputFlag['bitList']
498 for bit in bitList:
499 if bit['name'] in flagNames:
500 bitmask += np.uint64(2**bit['bit'])
502 return bitmask
505def getSignificance(catalog):
506 """Return the significance value of the first peak in each source
507 footprint, or NaN for peaks without a significance field.
509 Parameters
510 ----------
511 catalog : `lsst.afw.table.SourceCatalog`
512 Catalog to process.
514 Returns
515 -------
516 significance : `np.ndarray`, (N,)
517 Signficance of the first peak in each source footprint.
518 """
519 result = np.full(len(catalog), np.nan)
520 for i, record in enumerate(catalog):
521 peaks = record.getFootprint().peaks
522 if "significance" in peaks.schema:
523 result[i] = peaks[0]["significance"]
524 return result