Coverage for python/lsst/ap/association/transformDiaSourceCatalog.py: 21%

158 statements  

« prev     ^ index     » next       coverage.py v7.10.6, created at 2025-10-14 02:23 -0700

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/>. 

21 

22__all__ = ("TransformDiaSourceCatalogConnections", 

23 "TransformDiaSourceCatalogConfig", 

24 "TransformDiaSourceCatalogTask", 

25 "UnpackApdbFlags") 

26 

27import os 

28import yaml 

29 

30import numpy as np 

31 

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 

39 

40from .utils import convertDataFrameToSdmSchema, readSdmSchemaFile 

41 

42 

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 ) 

69 

70 

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 deprecated="This field is no longer used. Will be removed after v29." 

93 ) 

94 # TODO: remove on DM-41532 

95 doPackFlags = pexConfig.Field( 

96 dtype=bool, 

97 default=False, 

98 doc="Do pack the flags into one integer column named 'flags'." 

99 "If False, instead produce one boolean column per flag.", 

100 deprecated="This field is no longer used. Will be removed after v28." 

101 ) 

102 doIncludeReliability = pexConfig.Field( 

103 dtype=bool, 

104 default=False, 

105 doc="Include the reliability (e.g. real/bogus) classifications in the output.", 

106 deprecated="This field is no longer used. Will be removed after v29." 

107 ) 

108 doUseApdbSchema = pexConfig.Field( 

109 dtype=bool, 

110 default=False, 

111 doc="Use the APDB schema to coerce the data types of the output columns." 

112 ) 

113 schemaFile = pexConfig.Field( 

114 dtype=str, 

115 doc="Yaml file specifying the schema of the output catalog.", 

116 default=os.path.join("${SDM_SCHEMAS_DIR}", 

117 "yml", 

118 "apdb.yaml"), 

119 ) 

120 schemaName = pexConfig.Field( 

121 dtype=str, 

122 doc="Name of the table in the schema file to read.", 

123 default="ApdbSchema", 

124 ) 

125 

126 def setDefaults(self): 

127 super().setDefaults() 

128 self.functorFile = os.path.join("${AP_ASSOCIATION_DIR}", 

129 "data", 

130 "DiaSource.yaml") 

131 

132 

133class TransformDiaSourceCatalogTask(TransformCatalogBaseTask): 

134 """Transform a DiaSource catalog by calibrating and renaming columns to 

135 produce a table ready to insert into the Apdb. 

136 

137 Parameters 

138 ---------- 

139 initInputs : `dict` 

140 Must contain ``diaSourceSchema`` as the schema for the input catalog. 

141 """ 

142 ConfigClass = TransformDiaSourceCatalogConfig 

143 _DefaultName = "transformDiaSourceCatalog" 

144 # Needed to create a valid TransformCatalogBaseTask, but unused 

145 inputDataset = "deepDiff_diaSrc" 

146 outputDataset = "deepDiff_diaSrcTable" 

147 

148 def __init__(self, initInputs, **kwargs): 

149 super().__init__(**kwargs) 

150 self.funcs = self.getFunctors() 

151 self.inputSchema = initInputs['diaSourceSchema'].schema 

152 self._create_bit_pack_mappings() 

153 self.apdbSchema = readSdmSchemaFile(self.config.schemaFile, self.config.schemaName) 

154 

155 if not self.config.doPackFlags: 

156 # get the flag rename rules 

157 with open(os.path.expandvars(self.config.flagRenameMap)) as yaml_stream: 

158 self.rename_rules = list(yaml.safe_load_all(yaml_stream)) 

159 

160 def _create_bit_pack_mappings(self): 

161 """Setup all flag bit packings. 

162 """ 

163 self.bit_pack_columns = [] 

164 flag_map_file = os.path.expandvars(self.config.flagMap) 

165 with open(flag_map_file) as yaml_stream: 

166 table_list = list(yaml.safe_load_all(yaml_stream)) 

167 for table in table_list: 

168 if table['tableName'] == 'DiaSource': 

169 self.bit_pack_columns = table['columns'] 

170 break 

171 

172 # Test that all flags requested are present in the input schemas. 

173 # Output schemas are flexible, however if names are not specified in 

174 # the Apdb schema, flag columns will not be persisted. 

175 for outputFlag in self.bit_pack_columns: 

176 bitList = outputFlag['bitList'] 

177 for bit in bitList: 

178 try: 

179 self.inputSchema.find(bit['name']) 

180 except KeyError: 

181 raise KeyError( 

182 "Requested column %s not found in input DiaSource " 

183 "schema. Please check that the requested input " 

184 "column exists." % bit['name']) 

185 

186 def runQuantum(self, butlerQC, inputRefs, outputRefs): 

187 inputs = butlerQC.get(inputRefs) 

188 inputs["band"] = butlerQC.quantum.dataId["band"] 

189 

190 outputs = self.run(**inputs) 

191 

192 butlerQC.put(outputs, outputRefs) 

193 

194 @timeMethod 

195 def run(self, 

196 diaSourceCat, 

197 diffIm, 

198 band, 

199 reliability=None): 

200 """Convert input catalog to ParquetTable/Pandas and run functors. 

201 

202 Additionally, add new columns for stripping information from the 

203 exposure and into the DiaSource catalog. 

204 

205 Parameters 

206 ---------- 

207 diaSourceCat : `lsst.afw.table.SourceCatalog` 

208 Catalog of sources measured on the difference image. 

209 diffIm : `lsst.afw.image.Exposure` 

210 Result of subtracting template and science images. 

211 band : `str` 

212 Filter band of the science image. 

213 reliability : `lsst.afw.table.SourceCatalog` 

214 Reliability (e.g. real/bogus) scores, row-matched to 

215 ``diaSourceCat``. 

216 

217 Returns 

218 ------- 

219 results : `lsst.pipe.base.Struct` 

220 Results struct with components. 

221 

222 - ``diaSourceTable`` : Catalog of DiaSources with calibrated values 

223 and renamed columns. 

224 (`lsst.pipe.tasks.ParquetTable` or `pandas.DataFrame`) 

225 """ 

226 self.log.info( 

227 "Transforming/standardizing the DiaSource table for visit,detector: %i, %i", 

228 diffIm.visitInfo.id, diffIm.detector.getId()) 

229 

230 diaSourceDf = diaSourceCat.asAstropy().to_pandas() 

231 if self.config.doRemoveSkySources: 

232 diaSourceDf = diaSourceDf[~diaSourceDf["sky_source"]] 

233 diaSourceCat = diaSourceCat[~diaSourceCat["sky_source"]] 

234 

235 # Need UTC time but without a timezone because pandas requires a 

236 # naive datetime. 

237 diaSourceDf["timeProcessedMjdTai"] = DateTime.now().get(system=DateTime.MJD, scale=DateTime.TAI) 

238 diaSourceDf["snr"] = getSignificance(diaSourceCat) 

239 diaSourceDf["bboxSize"] = self.computeBBoxSizes(diaSourceCat) 

240 diaSourceDf["visit"] = diffIm.visitInfo.id 

241 # int16 instead of uint8 because databases don't like unsigned bytes. 

242 diaSourceDf["detector"] = np.int16(diffIm.detector.getId()) 

243 diaSourceDf["band"] = band 

244 diaSourceDf["midpointMjdTai"] = diffIm.visitInfo.date.get(system=DateTime.MJD) 

245 diaSourceDf["diaObjectId"] = 0 

246 diaSourceDf["ssObjectId"] = 0 

247 

248 # TODO: this has been formally deprecated and should be removed too 

249 if self.config.doPackFlags: 

250 # either bitpack the flags 

251 self.bitPackFlags(diaSourceDf) 

252 else: 

253 # or add the individual flag functors 

254 self.addUnpackedFlagFunctors() 

255 # and remove the packed flag functor 

256 if 'flags' in self.funcs.funcDict: 

257 del self.funcs.funcDict['flags'] 

258 

259 df = self.transform(band, 

260 diaSourceDf, 

261 self.funcs, 

262 dataId=None).df 

263 if self.config.doUseApdbSchema: 

264 df = convertDataFrameToSdmSchema(self.apdbSchema, df, tableName="DiaSource") 

265 

266 return pipeBase.Struct( 

267 diaSourceTable=df, 

268 ) 

269 

270 def addUnpackedFlagFunctors(self): 

271 """Add Column functor for each of the flags to the internal functor 

272 dictionary. 

273 """ 

274 for flag in self.bit_pack_columns[0]['bitList']: 

275 flagName = flag['name'] 

276 targetName = self.funcs.renameCol(flagName, self.rename_rules[0]['flag_rename_rules']) 

277 self.funcs.update({targetName: Column(flagName)}) 

278 

279 def computeBBoxSizes(self, inputCatalog): 

280 """Compute the size of a square bbox that fully contains the detection 

281 footprint. 

282 

283 Parameters 

284 ---------- 

285 inputCatalog : `lsst.afw.table.SourceCatalog` 

286 Catalog containing detected footprints. 

287 

288 Returns 

289 ------- 

290 outputBBoxSizes : `np.ndarray`, (N,) 

291 Array of bbox sizes. 

292 """ 

293 # Schema validation requires that this field is int. 

294 outputBBoxSizes = np.empty(len(inputCatalog), dtype=int) 

295 for i, record in enumerate(inputCatalog): 

296 footprintBBox = record.getFootprint().getBBox() 

297 # Compute twice the size of the largest dimension of the footprint 

298 # bounding box. This is the largest footprint we should need to cover 

299 # the complete DiaSource assuming the centroid is within the bounding 

300 # box. 

301 maxSize = 2 * np.max([footprintBBox.getWidth(), 

302 footprintBBox.getHeight()]) 

303 recX = record.getCentroid().x 

304 recY = record.getCentroid().y 

305 bboxSize = int( 

306 np.ceil(2 * np.max(np.fabs([footprintBBox.maxX - recX, 

307 footprintBBox.minX - recX, 

308 footprintBBox.maxY - recY, 

309 footprintBBox.minY - recY])))) 

310 if bboxSize > maxSize: 

311 bboxSize = maxSize 

312 outputBBoxSizes[i] = bboxSize 

313 

314 return outputBBoxSizes 

315 

316 def bitPackFlags(self, df): 

317 """Pack requested flag columns in inputRecord into single columns in 

318 outputRecord. 

319 

320 Parameters 

321 ---------- 

322 df : `pandas.DataFrame` 

323 DataFrame to read bits from and pack them into. 

324 """ 

325 for outputFlag in self.bit_pack_columns: 

326 bitList = outputFlag['bitList'] 

327 value = np.zeros(len(df), dtype=np.uint64) 

328 for bit in bitList: 

329 # Hard type the bit arrays. 

330 value += (df[bit['name']]*2**bit['bit']).to_numpy().astype(np.uint64) 

331 df[outputFlag['columnName']] = value 

332 

333 

334class UnpackApdbFlags: 

335 """Class for unpacking bits from integer flag fields stored in the Apdb. 

336 

337 Attributes 

338 ---------- 

339 flag_map_file : `str` 

340 Absolute or relative path to a yaml file specifiying mappings of flags 

341 to integer bits. 

342 table_name : `str` 

343 Name of the Apdb table the integer bit data are coming from. 

344 """ 

345 

346 def __init__(self, flag_map_file, table_name): 

347 self.bit_pack_columns = [] 

348 flag_map_file = os.path.expandvars(flag_map_file) 

349 with open(flag_map_file) as yaml_stream: 

350 table_list = list(yaml.safe_load_all(yaml_stream)) 

351 for table in table_list: 

352 if table['tableName'] == table_name: 

353 self.bit_pack_columns = table['columns'] 

354 break 

355 

356 self.output_flag_columns = {} 

357 

358 for column in self.bit_pack_columns: 

359 names = {} 

360 for bit in column["bitList"]: 

361 names[bit["name"]] = bit["bit"] 

362 self.output_flag_columns[column["columnName"]] = names 

363 

364 def unpack(self, input_flag_values, flag_name): 

365 """Determine individual boolean flags from an input array of unsigned 

366 ints. 

367 

368 Parameters 

369 ---------- 

370 input_flag_values : array-like of type uint 

371 Array of integer packed bit flags to unpack. 

372 flag_name : `str` 

373 Apdb column name from the loaded file, e.g. "flags". 

374 

375 Returns 

376 ------- 

377 output_flags : `numpy.ndarray` 

378 Numpy structured array of booleans, one column per flag in the 

379 loaded file. 

380 """ 

381 output_flags = np.zeros(len(input_flag_values), 

382 dtype=[(name, bool) for name in self.output_flag_columns[flag_name]]) 

383 

384 for name in self.output_flag_columns[flag_name]: 

385 masked_bits = np.bitwise_and(input_flag_values, 

386 2**self.output_flag_columns[flag_name][name]) 

387 output_flags[name] = masked_bits 

388 

389 return output_flags 

390 

391 def flagExists(self, flagName, columnName='flags'): 

392 """Check if named flag is in the bitpacked flag set. 

393 

394 Parameters: 

395 ---------- 

396 flagName : `str` 

397 Flag name to search for. 

398 columnName : `str`, optional 

399 Name of bitpacked flag column to search in. 

400 

401 Returns 

402 ------- 

403 flagExists : `bool` 

404 `True` if `flagName` is present in `columnName`. 

405 

406 Raises 

407 ------ 

408 ValueError 

409 Raised if `columnName` is not defined. 

410 """ 

411 if columnName not in self.output_flag_columns: 

412 raise ValueError(f'column {columnName} not in flag map: {self.output_flag_columns}') 

413 

414 return flagName in [c for c in self.output_flag_columns[columnName]] 

415 

416 def makeFlagBitMask(self, flagNames, columnName='flags'): 

417 """Return a bitmask corresponding to the supplied flag names. 

418 

419 Parameters: 

420 ---------- 

421 flagNames : `list` [`str`] 

422 Flag names to include in the bitmask. 

423 columnName : `str`, optional 

424 Name of bitpacked flag column. 

425 

426 Returns 

427 ------- 

428 bitmask : `np.unit64` 

429 Bitmask corresponding to the supplied flag names given the loaded configuration. 

430 

431 Raises 

432 ------ 

433 ValueError 

434 Raised if a flag in `flagName` is not included in `columnName`. 

435 """ 

436 bitmask = np.uint64(0) 

437 

438 for flag in flagNames: 

439 if not self.flagExists(flag, columnName=columnName): 

440 raise ValueError(f"flag '{flag}' not included in '{columnName}' flag column") 

441 

442 for outputFlag in self.bit_pack_columns: 

443 if outputFlag['columnName'] == columnName: 

444 bitList = outputFlag['bitList'] 

445 for bit in bitList: 

446 if bit['name'] in flagNames: 

447 bitmask += np.uint64(2**bit['bit']) 

448 

449 return bitmask 

450 

451 

452def getSignificance(catalog): 

453 """Return the significance value of the first peak in each source 

454 footprint, or NaN for peaks without a significance field. 

455 

456 Parameters 

457 ---------- 

458 catalog : `lsst.afw.table.SourceCatalog` 

459 Catalog to process. 

460 

461 Returns 

462 ------- 

463 significance : `np.ndarray`, (N,) 

464 Signficance of the first peak in each source footprint. 

465 """ 

466 result = np.full(len(catalog), np.nan) 

467 for i, record in enumerate(catalog): 

468 peaks = record.getFootprint().peaks 

469 if "significance" in peaks.schema: 

470 result[i] = peaks[0]["significance"] 

471 return result