Coverage for python/lsst/ap/association/transformDiaSourceCatalog.py: 86%
172 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-21 09:23 +0000
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
45class TransformDiaSourceCatalogConnections(pipeBase.PipelineTaskConnections,
46 dimensions=("instrument", "visit", "detector"),
47 defaultTemplates={"coaddName": "deep", "fakesType": ""}):
48 diaSourceSchema = connTypes.InitInput(
49 doc="Schema for DIASource catalog output by ImageDifference.",
50 storageClass="SourceCatalog",
51 name="{fakesType}{coaddName}Diff_diaSrc_schema",
52 )
53 diaSourceCat = connTypes.Input(
54 doc="Catalog of DiaSources produced during image differencing.",
55 name="{fakesType}{coaddName}Diff_candidateDiaSrc",
56 storageClass="SourceCatalog",
57 dimensions=("instrument", "visit", "detector"),
58 )
59 diffIm = connTypes.Input(
60 doc="Difference image on which the DiaSources were detected.",
61 name="{fakesType}{coaddName}Diff_differenceExp",
62 storageClass="ExposureF",
63 dimensions=("instrument", "visit", "detector"),
64 )
65 diaSourceTable = connTypes.Output(
66 doc=".",
67 name="{fakesType}{coaddName}Diff_diaSrcTable",
68 storageClass="ArrowAstropy",
69 dimensions=("instrument", "visit", "detector"),
70 )
72 def __init__(self, *, config=None):
73 super().__init__(config=config)
74 if config.input_image_type == "future":
75 self.diffIm = dataclasses.replace(self.diffIm, storageClass="DifferenceImage")
78class TransformDiaSourceCatalogConfig(TransformCatalogBaseConfig,
79 pipelineConnections=TransformDiaSourceCatalogConnections):
80 input_image_type = pexConfig.ChoiceField[str](
81 doc="Type of the input difference image. run() takes an"
82 " lsst.afw.image.Exposure either way.",
83 allowed={
84 "legacy": "Read an `lsst.afw.image.Exposure`.",
85 "future": "Read an `lsst.images.DifferenceImage`.",
86 },
87 optional=False,
88 default="legacy",
89 )
90 flagMap = pexConfig.Field(
91 dtype=str,
92 doc="Yaml file specifying SciencePipelines flag fields to bit packs.",
93 default=os.path.join("${AP_ASSOCIATION_DIR}",
94 "data",
95 "association-flag-map.yaml"),
96 )
97 flagRenameMap = pexConfig.Field(
98 dtype=str,
99 doc="Yaml file specifying specifying rules to rename flag names",
100 default=os.path.join("${AP_ASSOCIATION_DIR}",
101 "data",
102 "flag-rename-rules.yaml"),
103 )
104 doRemoveSkySources = pexConfig.Field(
105 dtype=bool,
106 default=False,
107 doc="Input DiaSource catalog contains SkySources that should be "
108 "removed before storing the output DiaSource catalog.",
109 )
110 # TODO: remove on DM-41532
111 doPackFlags = pexConfig.Field(
112 dtype=bool,
113 default=False,
114 doc="Do pack the flags into one integer column named 'flags'."
115 "If False, instead produce one boolean column per flag.",
116 deprecated="This field is no longer used. Will be removed after v28."
117 )
118 doUseApdbSchema = pexConfig.Field(
119 dtype=bool,
120 default=False,
121 doc="Use the APDB schema to coerce the data types of the output columns.",
122 deprecated="This field has been renamed to doUseSchema, and will be "
123 "removed after v30."
124 )
125 doUseSchema = pexConfig.Field(
126 dtype=bool,
127 default=False,
128 doc="Use an existing schema to coerce the data types of the output columns."
129 )
130 schemaDir = pexConfig.Field(
131 dtype=str,
132 doc="Path to the directory containing schema definitions.",
133 default=os.path.join("${SDM_SCHEMAS_DIR}",
134 "yml"),
135 )
136 schemaFile = pexConfig.Field(
137 dtype=str,
138 doc="Yaml file specifying the schema of the output catalog.",
139 default="apdb.yaml",
140 )
141 schemaName = pexConfig.Field(
142 dtype=str,
143 doc="Name of the table in the schema file to read.",
144 default="ApdbSchema",
145 deprecated="This config is no longer used, and will be removed after v30"
146 )
148 def setDefaults(self):
149 super().setDefaults()
150 self.functorFile = os.path.join("${AP_ASSOCIATION_DIR}",
151 "data",
152 "DiaSource.yaml")
155class TransformDiaSourceCatalogTask(TransformCatalogBaseTask):
156 """Transform a DiaSource catalog by calibrating and renaming columns to
157 produce a table ready to insert into the Apdb.
159 Parameters
160 ----------
161 initInputs : `dict`
162 Must contain ``diaSourceSchema`` as the schema for the input catalog.
163 """
164 ConfigClass = TransformDiaSourceCatalogConfig
165 _DefaultName = "transformDiaSourceCatalog"
166 # Needed to create a valid TransformCatalogBaseTask, but unused
167 inputDataset = "deepDiff_diaSrc"
168 outputDataset = "deepDiff_diaSrcTable"
170 def __init__(self, initInputs, **kwargs):
171 super().__init__(**kwargs)
172 self.funcs = self.getFunctors()
173 self.inputSchema = initInputs['diaSourceSchema'].schema
174 self._create_bit_pack_mappings()
175 if self.config.doUseSchema:
176 schemaFile = os.path.join(self.config.schemaDir, self.config.schemaFile)
177 self.schema = readSdmSchemaFile(schemaFile)
178 else:
179 self.schema = None
181 if not self.config.doPackFlags: 181 ↛ exitline 181 didn't return from function '__init__' because the condition on line 181 was always true
182 # get the flag rename rules
183 with open(os.path.expandvars(self.config.flagRenameMap)) as yaml_stream:
184 self.rename_rules = list(yaml.safe_load_all(yaml_stream))
186 def _create_bit_pack_mappings(self):
187 """Setup all flag bit packings.
188 """
189 self.bit_pack_columns = []
190 flag_map_file = os.path.expandvars(self.config.flagMap)
191 with open(flag_map_file) as yaml_stream:
192 table_list = list(yaml.safe_load_all(yaml_stream))
193 for table in table_list: 193 ↛ 201line 193 didn't jump to line 201
194 if table['tableName'] == 'DiaSource': 194 ↛ 193line 194 didn't jump to line 193 because the condition on line 194 was always true
195 self.bit_pack_columns = table['columns']
196 break
198 # Test that all flags requested are present in the input schemas.
199 # Output schemas are flexible, however if names are not specified in
200 # the Apdb schema, flag columns will not be persisted.
201 for outputFlag in self.bit_pack_columns:
202 bitList = outputFlag['bitList']
203 for bit in bitList:
204 try:
205 self.inputSchema.find(bit['name'])
206 except KeyError:
207 raise KeyError(
208 "Requested column %s not found in input DiaSource "
209 "schema. Please check that the requested input "
210 "column exists." % bit['name'])
212 def runQuantum(self, butlerQC, inputRefs, outputRefs):
213 inputs = butlerQC.get(inputRefs)
214 inputs["band"] = butlerQC.quantum.dataId["band"]
215 if self.config.input_image_type == "future": 215 ↛ 218line 215 didn't jump to line 218 because the condition on line 215 was always true
216 inputs["diffIm"] = inputs["diffIm"].to_legacy()
218 outputs = self.run(**inputs)
220 butlerQC.put(outputs, outputRefs)
222 @timeMethod
223 def run(self,
224 diaSourceCat,
225 diffIm,
226 band,
227 reliability=None):
228 """Convert input catalog to ParquetTable/Pandas and run functors.
230 Additionally, add new columns for stripping information from the
231 exposure and into the DiaSource catalog.
233 Parameters
234 ----------
235 diaSourceCat : `lsst.afw.table.SourceCatalog`
236 Catalog of sources measured on the difference image.
237 diffIm : `lsst.afw.image.Exposure`
238 Result of subtracting template and science images.
239 band : `str`
240 Filter band of the science image.
241 reliability : `lsst.afw.table.SourceCatalog`
242 Reliability (e.g. real/bogus) scores, row-matched to
243 ``diaSourceCat``.
245 Returns
246 -------
247 results : `lsst.pipe.base.Struct`
248 Results struct with components.
250 - ``diaSourceTable`` : Catalog of DiaSources with calibrated values
251 and renamed columns.
252 (`lsst.pipe.tasks.ParquetTable` or `pandas.DataFrame`)
253 """
254 self.log.info(
255 "Transforming/standardizing the DiaSource table for visit,detector: %i, %i",
256 diffIm.visitInfo.id, diffIm.detector.getId())
258 diaSourceDf = diaSourceCat.asAstropy().to_pandas()
259 if self.config.doRemoveSkySources: 259 ↛ 260line 259 didn't jump to line 260 because the condition on line 259 was never true
260 diaSourceDf = diaSourceDf[~diaSourceDf["sky_source"]]
261 diaSourceCat = diaSourceCat[~diaSourceCat["sky_source"]]
263 # Need UTC time but without a timezone because pandas requires a
264 # naive datetime.
265 diaSourceDf["timeProcessedMjdTai"] = DateTime.now().get(system=DateTime.MJD, scale=DateTime.TAI)
266 diaSourceDf["snr"] = getSignificance(diaSourceCat)
267 diaSourceDf["bboxSize"] = self.computeBBoxSizes(diaSourceCat)
268 diaSourceDf["visit"] = diffIm.visitInfo.id
269 # int16 instead of uint8 because databases don't like unsigned bytes.
270 diaSourceDf["detector"] = np.int16(diffIm.detector.getId())
271 diaSourceDf["band"] = band
272 diaSourceDf["midpointMjdTai"] = diffIm.visitInfo.date.get(system=DateTime.MJD)
273 diaSourceDf["exposureTime"] = diffIm.visitInfo.exposureTime
274 diaSourceDf["diaObjectId"] = 0
275 diaSourceDf["ssObjectId"] = 0
277 # TODO: this has been formally deprecated and should be removed too
278 if self.config.doPackFlags: 278 ↛ 280line 278 didn't jump to line 280 because the condition on line 278 was never true
279 # either bitpack the flags
280 self.bitPackFlags(diaSourceDf)
281 else:
282 # or add the individual flag functors
283 self.addUnpackedFlagFunctors()
284 # and remove the packed flag functor
285 if 'flags' in self.funcs.funcDict: 285 ↛ 288line 285 didn't jump to line 288 because the condition on line 285 was always true
286 del self.funcs.funcDict['flags']
288 df = self.transform(band,
289 diaSourceDf,
290 self.funcs,
291 dataId=None).df
292 if self.config.doUseSchema:
293 df = convertDataFrameToSdmSchema(self.schema, df, tableName="DiaSource")
295 return pipeBase.Struct(
296 diaSourceTable=df,
297 )
299 def addUnpackedFlagFunctors(self):
300 """Add Column functor for each of the flags to the internal functor
301 dictionary.
302 """
303 for flag in self.bit_pack_columns[0]['bitList']:
304 flagName = flag['name']
305 targetName = self.funcs.renameCol(flagName, self.rename_rules[0]['flag_rename_rules'])
306 self.funcs.update({targetName: Column(flagName)})
308 def computeBBoxSizes(self, inputCatalog):
309 """Compute the size of a square bbox that fully contains the detection
310 footprint.
312 Parameters
313 ----------
314 inputCatalog : `lsst.afw.table.SourceCatalog`
315 Catalog containing detected footprints.
317 Returns
318 -------
319 outputBBoxSizes : `np.ndarray`, (N,)
320 Array of bbox sizes.
321 """
322 # Schema validation requires that this field is int.
323 outputBBoxSizes = np.empty(len(inputCatalog), dtype=int)
324 for i, record in enumerate(inputCatalog):
325 footprintBBox = record.getFootprint().getBBox()
326 # Compute twice the size of the largest dimension of the footprint
327 # bounding box. This is the largest footprint we should need to cover
328 # the complete DiaSource assuming the centroid is within the bounding
329 # box.
330 maxSize = 2 * np.max([footprintBBox.getWidth(),
331 footprintBBox.getHeight()])
332 recX = record.getCentroid().x
333 recY = record.getCentroid().y
334 bboxSize = int(
335 np.ceil(2 * np.max(np.fabs([footprintBBox.maxX - recX,
336 footprintBBox.minX - recX,
337 footprintBBox.maxY - recY,
338 footprintBBox.minY - recY]))))
339 if bboxSize > maxSize: 339 ↛ 340line 339 didn't jump to line 340 because the condition on line 339 was never true
340 bboxSize = maxSize
341 outputBBoxSizes[i] = bboxSize
343 return outputBBoxSizes
345 def bitPackFlags(self, df):
346 """Pack requested flag columns in inputRecord into single columns in
347 outputRecord.
349 Parameters
350 ----------
351 df : `pandas.DataFrame`
352 DataFrame to read bits from and pack them into.
353 """
354 for outputFlag in self.bit_pack_columns:
355 bitList = outputFlag['bitList']
356 value = np.zeros(len(df), dtype=np.uint64)
357 for bit in bitList:
358 # Hard type the bit arrays.
359 value += (df[bit['name']]*2**bit['bit']).to_numpy().astype(np.uint64)
360 df[outputFlag['columnName']] = value
363class UnpackApdbFlags:
364 """Class for unpacking bits from integer flag fields stored in the Apdb.
366 Attributes
367 ----------
368 flag_map_file : `lsst.resources.ResourcePathExpression`
369 Absolute or relative URI to a yaml file specifiying mappings of flags
370 to integer bits.
371 table_name : `str`
372 Name of the Apdb table the integer bit data are coming from.
373 """
375 def __init__(self, flag_map_file, table_name):
376 self.bit_pack_columns = []
377 flag_map_file = os.path.expandvars(flag_map_file)
378 with ResourcePath(flag_map_file, forceDirectory=False).open("r") as yaml_stream:
379 table_list = list(yaml.safe_load_all(yaml_stream))
380 for table in table_list: 380 ↛ 385line 380 didn't jump to line 385
381 if table['tableName'] == table_name: 381 ↛ 380line 381 didn't jump to line 380 because the condition on line 381 was always true
382 self.bit_pack_columns = table['columns']
383 break
385 self.output_flag_columns = {}
387 for column in self.bit_pack_columns:
388 names = {}
389 for bit in column["bitList"]:
390 names[bit["name"]] = bit["bit"]
391 self.output_flag_columns[column["columnName"]] = names
393 def unpack(self, input_flag_values, flag_name):
394 """Determine individual boolean flags from an input array of unsigned
395 ints.
397 Parameters
398 ----------
399 input_flag_values : array-like of type uint
400 Array of integer packed bit flags to unpack.
401 flag_name : `str`
402 Apdb column name from the loaded file, e.g. "flags".
404 Returns
405 -------
406 output_flags : `numpy.ndarray`
407 Numpy structured array of booleans, one column per flag in the
408 loaded file.
409 """
410 output_flags = np.zeros(len(input_flag_values),
411 dtype=[(name, bool) for name in self.output_flag_columns[flag_name]])
413 for name in self.output_flag_columns[flag_name]:
414 masked_bits = np.bitwise_and(input_flag_values,
415 2**self.output_flag_columns[flag_name][name])
416 output_flags[name] = masked_bits
418 return output_flags
420 def flagExists(self, flagName, columnName='flags'):
421 """Check if named flag is in the bitpacked flag set.
423 Parameters:
424 ----------
425 flagName : `str`
426 Flag name to search for.
427 columnName : `str`, optional
428 Name of bitpacked flag column to search in.
430 Returns
431 -------
432 flagExists : `bool`
433 `True` if `flagName` is present in `columnName`.
435 Raises
436 ------
437 ValueError
438 Raised if `columnName` is not defined.
439 """
440 if columnName not in self.output_flag_columns:
441 raise ValueError(f'column {columnName} not in flag map: {self.output_flag_columns}')
443 return flagName in [c for c in self.output_flag_columns[columnName]]
445 def makeFlagBitMask(self, flagNames, columnName='flags'):
446 """Return a bitmask corresponding to the supplied flag names.
448 Parameters:
449 ----------
450 flagNames : `list` [`str`]
451 Flag names to include in the bitmask.
452 columnName : `str`, optional
453 Name of bitpacked flag column.
455 Returns
456 -------
457 bitmask : `np.unit64`
458 Bitmask corresponding to the supplied flag names given the loaded configuration.
460 Raises
461 ------
462 ValueError
463 Raised if a flag in `flagName` is not included in `columnName`.
464 """
465 bitmask = np.uint64(0)
467 for flag in flagNames:
468 if not self.flagExists(flag, columnName=columnName):
469 raise ValueError(f"flag '{flag}' not included in '{columnName}' flag column")
471 for outputFlag in self.bit_pack_columns:
472 if outputFlag['columnName'] == columnName: 472 ↛ 471line 472 didn't jump to line 471 because the condition on line 472 was always true
473 bitList = outputFlag['bitList']
474 for bit in bitList:
475 if bit['name'] in flagNames:
476 bitmask += np.uint64(2**bit['bit'])
478 return bitmask
481def getSignificance(catalog):
482 """Return the significance value of the first peak in each source
483 footprint, or NaN for peaks without a significance field.
485 Parameters
486 ----------
487 catalog : `lsst.afw.table.SourceCatalog`
488 Catalog to process.
490 Returns
491 -------
492 significance : `np.ndarray`, (N,)
493 Signficance of the first peak in each source footprint.
494 """
495 result = np.full(len(catalog), np.nan)
496 for i, record in enumerate(catalog):
497 peaks = record.getFootprint().peaks
498 if "significance" in peaks.schema:
499 result[i] = peaks[0]["significance"]
500 return result