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