Coverage for tests/test_utils.py: 29%

92 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-03-23 05:18 -0700

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 expanded_pipeline = merged_pipeline.toExpandedPipeline() 

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 taskDef in expanded_pipeline: 

108 conns = taskDef.connections 

109 if taskDef.label == "isr": 

110 self.assertEqual(conns.outputExposure.name, "postISRCCD") 

111 elif taskDef.label == "inject_exposure": 

112 self.assertEqual(conns.input_exposure.name, "postISRCCD") 

113 self.assertEqual(conns.output_exposure.name, "injected_postISRCCD") 

114 self.assertEqual(conns.output_catalog.name, "injected_postISRCCD_catalog") 

115 elif taskDef.label == "characterizeImage": 

116 self.assertEqual(conns.exposure.name, "injected_postISRCCD") 

117 self.assertEqual(conns.characterized.name, "injected_icExp") 

118 self.assertEqual(conns.backgroundModel.name, "injected_icExpBackground") 

119 self.assertEqual(conns.sourceCat.name, "injected_icSrc") 

120 

121 def test_ingest_injection_catalog(self): 

122 input_dataset_refs = ingest_injection_catalog( 

123 writeable_butler=self.writeable_butler, 

124 table=self.injection_catalog, 

125 band="g", 

126 output_collection="test_collection", 

127 dataset_type_name="injection_catalog", 

128 log_level=logging.DEBUG, 

129 ) 

130 output_dataset_refs = self.writeable_butler.registry.queryDatasets( 

131 "injection_catalog", 

132 collections="test_collection", 

133 ) 

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

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

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

137 self.assertEqual(input_ids, output_ids) 

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

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

140 

141 def test_consolidate_injected_catalogs(self): 

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

143 output_catalog = consolidate_injected_deepCoadd_catalogs( 

144 catalog_dict=catalog_dict, 

145 skymap=self.skyMap, 

146 tract=9, 

147 pixel_match_radius=0.1, 

148 get_catalogs_from_butler=False, 

149 col_ra="ra", 

150 col_dec="dec", 

151 col_mag="mag", 

152 isPatchInnerKey="injected_isPatchInner", 

153 isTractInnerKey="injected_isTractInner", 

154 isPrimaryKey="injected_isPrimary", 

155 injectionKey="injection_flag", 

156 ) 

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

158 expected_columns = { 

159 "injection_id", 

160 "injected_id", 

161 "ra", 

162 "dec", 

163 "source_type", 

164 "g_mag", 

165 "r_mag", 

166 "injection_draw_size", 

167 "injection_flag", 

168 "injected_isPatchInner", 

169 "injected_isTractInner", 

170 "injected_isPrimary", 

171 } 

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

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

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

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

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

177 

178 

179class MemoryTestCase(MemoryTestCase): 

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

181 

182 pass 

183 

184 

185def setup_module(module): 

186 """Configure pytest.""" 

187 lsst.utils.tests.init() 

188 

189 

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

191 lsst.utils.tests.init() 

192 unittest.main()