Coverage for tests/test_utils.py: 30%

91 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-19 13:56 +0000

1# This file is part of source_injection. 

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 

22import logging 

23import os 

24import unittest 

25 

26import lsst.utils.tests 

27from lsst.daf.butler.tests import makeTestCollection, makeTestRepo 

28from lsst.daf.butler.tests.utils import makeTestTempDir, removeTestTempDir 

29from lsst.obs.base.instrument_tests import DummyCam 

30from lsst.pipe.base import Pipeline 

31from lsst.skymap.ringsSkyMap import RingsSkyMap, RingsSkyMapConfig 

32from lsst.source.injection import ( 

33 ExposureInjectTask, 

34 consolidate_injected_deepCoadd_catalogs, 

35 ingest_injection_catalog, 

36 make_injection_pipeline, 

37) 

38from lsst.source.injection.utils.test_utils import ( 

39 make_test_exposure, 

40 make_test_injection_catalog, 

41 make_test_reference_pipeline, 

42) 

43from lsst.utils.tests import MemoryTestCase, TestCase 

44 

45TEST_DIR = os.path.abspath(os.path.dirname(__file__)) 

46 

47 

48class SourceInjectionUtilsTestCase(TestCase): 

49 """Test the utility functions in the source_injection package.""" 

50 

51 @classmethod 

52 def setUpClass(cls): 

53 cls.root = makeTestTempDir(TEST_DIR) 

54 cls.creator_butler = makeTestRepo(cls.root) 

55 cls.writeable_butler = makeTestCollection(cls.creator_butler) 

56 # Register an instrument so we can get some bands. 

57 DummyCam().register(cls.writeable_butler.registry) 

58 skyMapConfig = RingsSkyMapConfig() 

59 skyMapConfig.numRings = 3 

60 cls.skyMap = RingsSkyMap(config=skyMapConfig) 

61 

62 @classmethod 

63 def tearDownClass(cls): 

64 del cls.writeable_butler 

65 del cls.creator_butler 

66 del cls.skyMap 

67 removeTestTempDir(cls.root) 

68 

69 def setUp(self): 

70 self.exposure = make_test_exposure() 

71 self.injection_catalog = make_test_injection_catalog( 

72 self.exposure.getWcs(), 

73 self.exposure.getBBox(), 

74 ) 

75 self.reference_pipeline = make_test_reference_pipeline() 

76 self.injected_catalog = self.injection_catalog.copy() 

77 self.injected_catalog.add_columns(cols=[0, 0], names=["injection_draw_size", "injection_flag"]) 

78 self.injected_catalog["injection_flag"][:5] = 1 

79 

80 def tearDown(self): 

81 del self.exposure 

82 del self.injection_catalog 

83 del self.reference_pipeline 

84 del self.injected_catalog 

85 

86 def test_generate_injection_catalog(self): 

87 self.assertEqual(len(self.injection_catalog), 30) 

88 expected_columns = {"injection_id", "ra", "dec", "source_type", "mag"} 

89 self.assertEqual(set(self.injection_catalog.columns), expected_columns) 

90 

91 def test_make_injection_pipeline(self): 

92 injection_pipeline = Pipeline("reference_pipeline") 

93 injection_pipeline.addTask(ExposureInjectTask, "inject_exposure") 

94 merged_pipeline = make_injection_pipeline( 

95 dataset_type_name="postISRCCD", 

96 reference_pipeline=self.reference_pipeline, 

97 injection_pipeline=injection_pipeline, 

98 exclude_subsets=False, 

99 prefix="injected_", 

100 instrument="lsst.obs.subaru.HyperSuprimeCam", 

101 log_level=logging.DEBUG, 

102 ) 

103 pipeline_graph = merged_pipeline.to_graph() 

104 expected_subset_tasks = ["isr", "inject_exposure", "characterizeImage"] 

105 merged_task_subsets = [merged_pipeline.findSubsetsWithLabel(x) for x in expected_subset_tasks] 

106 self.assertEqual(len(merged_task_subsets), len(expected_subset_tasks)) 

107 for task_node in pipeline_graph.tasks.values(): 

108 if task_node.label == "isr": 

109 self.assertEqual(task_node.outputs["outputExposure"].dataset_type_name, "postISRCCD") 

110 elif task_node.label == "inject_exposure": 

111 self.assertEqual(task_node.inputs["input_exposure"].dataset_type_name, "postISRCCD") 

112 self.assertEqual( 

113 task_node.outputs["output_exposure"].dataset_type_name, "injected_postISRCCD" 

114 ) 

115 self.assertEqual( 

116 task_node.outputs["output_catalog"].dataset_type_name, "injected_postISRCCD_catalog" 

117 ) 

118 elif task_node.label == "characterizeImage": 

119 self.assertEqual(task_node.inputs["exposure"].dataset_type_name, "injected_postISRCCD") 

120 self.assertEqual(task_node.outputs["characterized"].dataset_type_name, "injected_icExp") 

121 self.assertEqual( 

122 task_node.outputs["backgroundModel"].dataset_type_name, "injected_icExpBackground" 

123 ) 

124 self.assertEqual(task_node.outputs["sourceCat"].dataset_type_name, "injected_icSrc") 

125 

126 def test_ingest_injection_catalog(self): 

127 input_dataset_refs = ingest_injection_catalog( 

128 writeable_butler=self.writeable_butler, 

129 table=self.injection_catalog, 

130 band="g", 

131 output_collection="test_collection", 

132 dataset_type_name="injection_catalog", 

133 log_level=logging.DEBUG, 

134 ) 

135 output_dataset_refs = self.writeable_butler.registry.queryDatasets( 

136 "injection_catalog", 

137 collections="test_collection", 

138 ) 

139 self.assertEqual(len(input_dataset_refs), output_dataset_refs.count()) 

140 input_ids = {x.id for x in input_dataset_refs} 

141 output_ids = {x.id for x in output_dataset_refs} 

142 self.assertEqual(input_ids, output_ids) 

143 injected_catalog = self.writeable_butler.get(input_dataset_refs[0]) 

144 self.assertTrue(all(self.injection_catalog == injected_catalog)) 

145 

146 def test_consolidate_injected_catalogs(self): 

147 catalog_dict = {"g": self.injected_catalog, "r": self.injected_catalog} 

148 output_catalog = consolidate_injected_deepCoadd_catalogs( 

149 catalog_dict=catalog_dict, 

150 skymap=self.skyMap, 

151 tract=9, 

152 pixel_match_radius=0.1, 

153 get_catalogs_from_butler=False, 

154 col_ra="ra", 

155 col_dec="dec", 

156 col_mag="mag", 

157 isPatchInnerKey="injected_isPatchInner", 

158 isTractInnerKey="injected_isTractInner", 

159 isPrimaryKey="injected_isPrimary", 

160 injectionKey="injection_flag", 

161 ) 

162 self.assertEqual(len(output_catalog), 30) 

163 expected_columns = { 

164 "injection_id", 

165 "injected_id", 

166 "ra", 

167 "dec", 

168 "source_type", 

169 "g_mag", 

170 "r_mag", 

171 "injection_draw_size", 

172 "injection_flag", 

173 "injected_isPatchInner", 

174 "injected_isTractInner", 

175 "injected_isPrimary", 

176 } 

177 self.assertEqual(set(output_catalog.columns), expected_columns) 

178 self.assertEqual(sum(output_catalog["injection_flag"]), 5) 

179 self.assertEqual(sum(output_catalog["injected_isPatchInner"]), 30) 

180 self.assertEqual(sum(output_catalog["injected_isTractInner"]), 30) 

181 self.assertEqual(sum(output_catalog["injected_isPrimary"]), 25) 

182 

183 

184class MemoryTestCase(MemoryTestCase): 

185 """Test memory usage of functions in this script.""" 

186 

187 pass 

188 

189 

190def setup_module(module): 

191 """Configure pytest.""" 

192 lsst.utils.tests.init() 

193 

194 

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

196 lsst.utils.tests.init() 

197 unittest.main()