Coverage for tests/test_produce_and_standardize_prompt_source.py: 93%

92 statements  

« prev     ^ index     » next       coverage.py v7.16.0, created at 2026-09-19 03:10 -0700

1# This file is part of drp_tasks. 

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"""End-to-end test of the prompt_source production chain. 

23 

24This exercises the sequence of tasks that ApPipe.yaml runs to turn a 

25calibrated visit image into the ``prompt_source`` table: 

26 

271. ``SingleFrameDetectAndMeasureTask`` detects and measures 

28 sources on a synthetic exposure, producing the 

29 ``single_visit_star_reprocessed_unstandardized`` catalog. 

302. ``TransformSourceTableTask`` applies the 

31 ``prompt_source.yaml`` functors from ``pipe_tasks``. 

323. ``ConsolidateSourceTableTask`` concatenates the per-detector tables. 

334. ``SplitPrimaryTask`` keeps the primary rows and drops sky sources. 

34""" 

35 

36import os 

37import unittest 

38 

39import numpy as np 

40 

41import lsst.afw.geom 

42import lsst.afw.image 

43import lsst.geom 

44import lsst.meas.algorithms 

45import lsst.meas.base.tests 

46import lsst.utils.tests 

47from lsst.drp.tasks.single_frame_detect_and_measure import ( 

48 SingleFrameDetectAndMeasureConfig, 

49 SingleFrameDetectAndMeasureTask, 

50) 

51from lsst.pipe.base import InMemoryDatasetHandle 

52from lsst.pipe.tasks.postprocess import ( 

53 TableVStack, 

54 TransformSourceTableConfig, 

55 TransformSourceTableTask, 

56) 

57from lsst.pipe.tasks.schemaUtils import ( 

58 checkSdmSchemaColumns, 

59 column_dtype, 

60 readSdmSchemaFile, 

61) 

62from lsst.pipe.tasks.split_primary import SplitPrimaryTask 

63from lsst.utils import getPackageDir 

64 

65FUNCTOR_FILE = os.path.join(getPackageDir("pipe_tasks"), "schemas", "prompt_source.yaml") 

66SCHEMA_FILE = os.path.join("${SDM_SCHEMAS_DIR}", "yml", "ap_extra.yaml") 

67TABLE_NAME = "PromptSource" 

68SPLIT_DISCARD_PRIMARY_COLUMNS = ["sky_source"] 

69 

70VISIT = 98765 

71DETECTOR = 42 

72 

73 

74def make_exposure_and_background(): 

75 """Build a synthetic calibrated exposure and its background model. 

76 

77 Modeled on ``test_reprocess_visit_image.py``: a `TestDataset` provides the 

78 PSF, WCS and PhotoCalib, a handful of point sources are added, and a 

79 background is fit to the realized image. 

80 

81 Returns 

82 ------- 

83 exposure : `lsst.afw.image.ExposureF` 

84 Calibrated exposure ready for detection and measurement. 

85 background : `lsst.afw.math.BackgroundList` 

86 Background model for the exposure. 

87 """ 

88 bbox = lsst.geom.Box2I(lsst.geom.Point2I(5, 4), lsst.geom.Point2I(205, 184)) 

89 dataset = lsst.meas.base.tests.TestDataset( 

90 bbox, 

91 crval=lsst.geom.SpherePoint(245.0, -45.0, lsst.geom.degrees), 

92 calibration=12.3, 

93 detector=DETECTOR, 

94 visitId=VISIT, 

95 ) 

96 psf_scale = np.sqrt(4 * np.pi * (dataset.psfShape.getDeterminantRadius()) ** 2) 

97 noise = 10.0 

98 for flux, centroid in [ 

99 (45 * noise * psf_scale, (40, 70)), 

100 (150 * noise * psf_scale, (50, 120)), 

101 (400 * noise * psf_scale, (92, 35)), 

102 (1000 * noise * psf_scale, (175, 154)), 

103 ]: 

104 dataset.addSource(instFlux=flux, centroid=lsst.geom.Point2D(*centroid)) 

105 truth_exposure, _ = dataset.realize(noise=noise, schema=dataset.makeMinimalSchema()) 

106 

107 # Build an input exposure carrying the calibrations the task reads. 

108 exposure = lsst.afw.image.ExposureF(truth_exposure.maskedImage.clone()) 

109 exposure.mask.clearMaskPlane(exposure.mask.getMaskPlane("DETECTED")) 

110 exposure.setPsf(truth_exposure.psf) 

111 exposure.setWcs(truth_exposure.wcs) 

112 exposure.setPhotoCalib(truth_exposure.photoCalib) 

113 exposure.info.setApCorrMap(lsst.afw.image.ApCorrMap()) 

114 exposure.info.setVisitInfo(truth_exposure.visitInfo) 

115 exposure.info.setDetector(truth_exposure.getDetector()) 

116 

117 bkgConfig = lsst.meas.algorithms.SubtractBackgroundTask.ConfigClass() 

118 # Small test image; fit a simple background model. 

119 bkgConfig.approxOrderX = 1 

120 bkgTask = lsst.meas.algorithms.SubtractBackgroundTask(config=bkgConfig) 

121 background = bkgTask.run(truth_exposure).background 

122 

123 return exposure, background 

124 

125 

126class PromptSourceEndToEndTestCase(lsst.utils.tests.TestCase): 

127 """Run SingleFrameDetectAndMeasure output through the prompt_source steps 

128 and validate the final schema. 

129 """ 

130 

131 @classmethod 

132 def setUpClass(cls): 

133 super().setUpClass() 

134 

135 exposure, background = make_exposure_and_background() 

136 

137 # 1. Detect and measure direct sources. 

138 sfdmConfig = SingleFrameDetectAndMeasureConfig() 

139 sfdmConfig.detection.background.approxOrderX = 1 

140 sfdmConfig.sky_sources.nSources = 2 

141 sfdmTask = SingleFrameDetectAndMeasureTask(config=sfdmConfig) 

142 cls.sources = sfdmTask.run(exposure=exposure, input_background=background).sources 

143 

144 schemaFile = os.path.expandvars(SCHEMA_FILE) 

145 cls.schema = readSdmSchemaFile(schemaFile) 

146 

147 # 2. Run standardization 

148 transformConfig = TransformSourceTableConfig() 

149 transformConfig.functorFile = FUNCTOR_FILE 

150 transformTask = TransformSourceTableTask(config=transformConfig) 

151 detectorTable = transformTask.run( 

152 handle=InMemoryDatasetHandle(cls.sources.to_pandas(), storageClass="DataFrame"), 

153 funcs=transformTask.funcs, 

154 dataId={"visit": VISIT, "detector": DETECTOR, "band": "r", "physical_filter": "r_03"}, 

155 ).outputCatalog 

156 

157 # 3. Mimic the stacking done by the consolidate task 

158 consolidated = TableVStack.vstack_handles( 

159 [InMemoryDatasetHandle(detectorTable, storageClass="ArrowAstropy")] 

160 ) 

161 

162 # 4. Drop non-primary rows from blends, and sky sources 

163 splitConfig = SplitPrimaryTask.ConfigClass() 

164 splitConfig.discard_primary_columns = SPLIT_DISCARD_PRIMARY_COLUMNS 

165 splitTask = SplitPrimaryTask(config=splitConfig) 

166 cls.promptSource = splitTask.run(full=consolidated).primary 

167 

168 def testProducesPrimarySources(self): 

169 """prompt_source holds the primary rows minus the dropped columns.""" 

170 nPrimary = int(np.sum(self.sources["detect_isPrimary"])) 

171 self.assertGreater(nPrimary, 0) 

172 self.assertEqual(len(self.promptSource), nPrimary) 

173 self.assertNotIn("detect_isPrimary", self.promptSource.colnames) 

174 self.assertNotIn("sky_source", self.promptSource.colnames) 

175 

176 def testCalibratedFluxesComputed(self): 

177 """The calibrated flux functors ran and produced finite values.""" 

178 for column in ("psfFlux", "calibFlux"): 

179 self.assertIn(column, self.promptSource.colnames) 

180 self.assertTrue(np.isfinite(self.promptSource[column]).any()) 

181 

182 def testSchema(self): 

183 """Check that prompt_source column names and dtypes match the 

184 PromptSource schema. 

185 """ 

186 dataframe = self.promptSource.to_pandas() 

187 # There should be no columns beyond those defined for PromptSource. 

188 extra = checkSdmSchemaColumns(self.schema, list(dataframe.columns), TABLE_NAME) 

189 self.assertEqual(extra, [], f"prompt_source has columns absent from {TABLE_NAME}: {extra}") 

190 

191 # Every schema column is present with the expected dtype. 

192 mismatches = {} 

193 for columnDef in self.schema[TABLE_NAME].columns: 

194 # Extract the expected datatype for each column from the schema 

195 expected = column_dtype(columnDef.datatype) 

196 if columnDef.name not in dataframe.columns: 196 ↛ 197line 196 didn't jump to line 197 because the condition on line 196 was never true

197 mismatches[columnDef.name] = f"missing (expected {expected})" 

198 elif str(dataframe[columnDef.name].dtype) != expected: 198 ↛ 199line 198 didn't jump to line 199 because the condition on line 198 was never true

199 mismatches[columnDef.name] = f"{dataframe[columnDef.name].dtype} != {expected}" 

200 self.assertEqual(mismatches, {}, f"prompt_source dtypes do not match {TABLE_NAME}: {mismatches}") 

201 

202 

203class MemoryTester(lsst.utils.tests.MemoryTestCase): 

204 pass 

205 

206 

207def setup_module(module): 

208 lsst.utils.tests.init() 

209 

210 

211if __name__ == "__main__": 211 ↛ 212line 211 didn't jump to line 212 because the condition on line 211 was never true

212 lsst.utils.tests.init() 

213 unittest.main()