Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1# This file is part of daf_butler. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (http://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 <http://www.gnu.org/licenses/>. 

21 

22from __future__ import annotations 

23 

24import os 

25import unittest 

26import tempfile 

27import shutil 

28 

29from typing import TYPE_CHECKING 

30 

31import lsst.utils.tests 

32 

33import lsst.afw.image 

34from lsst.afw.image import LOCAL 

35from lsst.geom import Box2I, Point2I 

36from lsst.daf.base import PropertyList, PropertySet 

37 

38from lsst.daf.butler import Config 

39from lsst.daf.butler import StorageClassFactory 

40from lsst.daf.butler import DatasetType 

41from lsst.daf.butler.tests import DatasetTestHelper, makeTestRepo, addDatasetType, makeTestCollection 

42 

43from lsst.obs.base.exposureAssembler import ExposureAssembler 

44 

45if TYPE_CHECKING: 45 ↛ 46line 45 didn't jump to line 46, because the condition on line 45 was never true

46 from lsst.daf.butler import DatasetRef 

47 

48TESTDIR = os.path.dirname(__file__) 

49 

50BUTLER_CONFIG = """ 

51storageClasses: 

52 ExposureCompositeF: 

53 inheritsFrom: ExposureF 

54datastore: 

55 # Want to check disassembly so can't use InMemory 

56 cls: lsst.daf.butler.datastores.posixDatastore.PosixDatastore 

57 formatters: 

58 ExposureCompositeF: lsst.obs.base.fitsExposureFormatter.FitsExposureFormatter 

59 lossless: 

60 formatter: lsst.obs.base.fitsExposureFormatter.FitsExposureFormatter 

61 parameters: 

62 recipe: lossless 

63 uncompressed: 

64 formatter: lsst.obs.base.fitsExposureFormatter.FitsExposureFormatter 

65 parameters: 

66 recipe: noCompression 

67 lossy: 

68 formatter: lsst.obs.base.fitsExposureFormatter.FitsExposureFormatter 

69 parameters: 

70 recipe: lossyBasic 

71 composites: 

72 disassembled: 

73 ExposureCompositeF: True 

74""" 

75 

76# Components present in the test file 

77COMPONENTS = {"wcs", "image", "mask", "coaddInputs", "psf", "visitInfo", "variance", "metadata", "photoCalib"} 

78 

79 

80class ButlerFitsTests(DatasetTestHelper, lsst.utils.tests.TestCase): 

81 

82 @classmethod 

83 def setUpClass(cls): 

84 """Create a new butler once only.""" 

85 

86 cls.storageClassFactory = StorageClassFactory() 

87 

88 cls.root = tempfile.mkdtemp(dir=TESTDIR) 

89 

90 dataIds = { 

91 "instrument": ["DummyCam"], 

92 "physical_filter": ["d-r"], 

93 "visit": [42, 43, 44], 

94 } 

95 

96 cls.creatorButler = makeTestRepo(cls.root, dataIds, config=Config.fromYaml(BUTLER_CONFIG)) 

97 

98 # Create dataset types used by the tests 

99 for datasetTypeName, storageClassName in (("calexp", "ExposureF"), 

100 ("unknown", "ExposureCompositeF"), 

101 ("testCatalog", "SourceCatalog"), 

102 ("lossless", "ExposureF"), 

103 ("uncompressed", "ExposureF"), 

104 ("lossy", "ExposureF"), 

105 ): 

106 storageClass = cls.storageClassFactory.getStorageClass(storageClassName) 

107 addDatasetType(cls.creatorButler, datasetTypeName, set(dataIds), storageClass) 

108 

109 # And some dataset types that have no dimensions for easy testing 

110 for datasetTypeName, storageClassName in (("ps", "PropertySet"), 

111 ("pl", "PropertyList"), 

112 ): 

113 storageClass = cls.storageClassFactory.getStorageClass(storageClassName) 

114 addDatasetType(cls.creatorButler, datasetTypeName, {}, storageClass) 

115 

116 @classmethod 

117 def tearDownClass(cls): 

118 if cls.root is not None: 

119 shutil.rmtree(cls.root, ignore_errors=True) 

120 

121 def setUp(self): 

122 self.butler = makeTestCollection(self.creatorButler) 

123 

124 def makeExampleCatalog(self) -> lsst.afw.table.SourceCatalog: 

125 catalogPath = os.path.join(TESTDIR, "data", "source_catalog.fits") 

126 return lsst.afw.table.SourceCatalog.readFits(catalogPath) 

127 

128 def assertCatalogEqual(self, inputCatalog: lsst.afw.table.SourceCatalog, 

129 outputCatalog: lsst.afw.table.SourceCatalog) -> None: 

130 self.assertIsInstance(outputCatalog, lsst.afw.table.SourceCatalog) 

131 inputTable = inputCatalog.getTable() 

132 inputRecord = inputCatalog[0] 

133 outputTable = outputCatalog.getTable() 

134 outputRecord = outputCatalog[0] 

135 self.assertEqual(inputRecord.getPsfInstFlux(), outputRecord.getPsfInstFlux()) 

136 self.assertEqual(inputRecord.getPsfFluxFlag(), outputRecord.getPsfFluxFlag()) 

137 self.assertEqual(inputTable.getSchema().getAliasMap().get("slot_Centroid"), 

138 outputTable.getSchema().getAliasMap().get("slot_Centroid")) 

139 self.assertEqual(inputRecord.getCentroid(), outputRecord.getCentroid()) 

140 self.assertFloatsAlmostEqual( 

141 inputRecord.getCentroidErr()[0, 0], 

142 outputRecord.getCentroidErr()[0, 0], rtol=1e-6) 

143 self.assertFloatsAlmostEqual( 

144 inputRecord.getCentroidErr()[1, 1], 

145 outputRecord.getCentroidErr()[1, 1], rtol=1e-6) 

146 self.assertEqual(inputTable.getSchema().getAliasMap().get("slot_Shape"), 

147 outputTable.getSchema().getAliasMap().get("slot_Shape")) 

148 self.assertFloatsAlmostEqual( 

149 inputRecord.getShapeErr()[0, 0], 

150 outputRecord.getShapeErr()[0, 0], rtol=1e-6) 

151 self.assertFloatsAlmostEqual( 

152 inputRecord.getShapeErr()[1, 1], 

153 outputRecord.getShapeErr()[1, 1], rtol=1e-6) 

154 self.assertFloatsAlmostEqual( 

155 inputRecord.getShapeErr()[2, 2], 

156 outputRecord.getShapeErr()[2, 2], rtol=1e-6) 

157 

158 def runFundamentalTypeTest(self, datasetTypeName, entity): 

159 """Put and get the supplied entity and compare.""" 

160 ref = self.butler.put(entity, datasetTypeName) 

161 butler_ps = self.butler.get(ref) 

162 self.assertEqual(butler_ps, entity) 

163 

164 # Break the contact by ensuring that we are writing YAML 

165 uri = self.butler.getURI(ref) 

166 self.assertTrue(uri.path.endswith(".yaml"), f"Check extension of {uri}") 

167 

168 def testFundamentalTypes(self) -> None: 

169 """Ensure that some fundamental stack types round trip.""" 

170 ps = PropertySet() 

171 ps["a.b"] = 5 

172 ps["c.d.e"] = "string" 

173 self.runFundamentalTypeTest("ps", ps) 

174 

175 pl = PropertyList() 

176 pl["A"] = 1 

177 pl.setComment("A", "An int comment") 

178 pl["B"] = "string" 

179 pl.setComment("B", "A string comment") 

180 self.runFundamentalTypeTest("pl", pl) 

181 

182 def testFitsCatalog(self) -> None: 

183 catalog = self.makeExampleCatalog() 

184 dataId = {"visit": 42, "instrument": "DummyCam", "physical_filter": "d-r"} 

185 ref = self.butler.put(catalog, "testCatalog", dataId) 

186 stored = self.butler.get(ref) 

187 self.assertCatalogEqual(catalog, stored) 

188 

189 def testExposureCompositePutGetConcrete(self) -> None: 

190 ref = self.runExposureCompositePutGetTest("calexp") 

191 

192 uri = self.butler.getURI(ref) 

193 self.assertTrue(os.path.exists(uri.path), f"Checking URI {uri} existence") 

194 

195 def testExposureCompositePutGetVirtual(self) -> None: 

196 ref = self.runExposureCompositePutGetTest("unknown") 

197 

198 primary, components = self.butler.getURIs(ref) 

199 self.assertIsNone(primary) 

200 self.assertEqual(set(components), COMPONENTS) 

201 for compName, uri in components.items(): 

202 self.assertTrue(os.path.exists(uri.path), 

203 f"Checking URI {uri} existence for component {compName}") 

204 

205 def runExposureCompositePutGetTest(self, datasetTypeName: str) -> DatasetRef: 

206 example = os.path.join(TESTDIR, "data", "small.fits") 

207 exposure = lsst.afw.image.ExposureF(example) 

208 

209 dataId = {"visit": 42, "instrument": "DummyCam", "physical_filter": "d-r"} 

210 ref = self.butler.put(exposure, datasetTypeName, dataId) 

211 

212 # Get the full thing 

213 composite = self.butler.get(datasetTypeName, dataId) 

214 

215 # There is no assert for Exposure so just look at maskedImage 

216 self.assertMaskedImagesEqual(composite.maskedImage, exposure.maskedImage) 

217 

218 # Helper for extracting components 

219 assembler = ExposureAssembler(ref.datasetType.storageClass) 

220 

221 # Get each component from butler independently 

222 for compName in COMPONENTS: 

223 compTypeName = DatasetType.nameWithComponent(datasetTypeName, compName) 

224 component = self.butler.get(compTypeName, dataId) 

225 

226 reference = assembler.getComponent(exposure, compName) 

227 self.assertIsInstance(component, type(reference), f"Checking type of component {compName}") 

228 

229 if compName in ("image", "variance"): 

230 self.assertImagesEqual(component, reference) 

231 elif compName == "mask": 

232 self.assertMasksEqual(component, reference) 

233 elif compName == "wcs": 

234 self.assertWcsAlmostEqualOverBBox(component, reference, exposure.getBBox()) 

235 elif compName == "coaddInputs": 

236 self.assertEqual(len(component.visits), len(reference.visits), 

237 f"cf visits {component.visits}") 

238 self.assertEqual(len(component.ccds), len(reference.ccds), 

239 f"cf CCDs {component.ccds}") 

240 elif compName == "psf": 

241 # Equality for PSF does not work 

242 pass 

243 elif compName == "visitInfo": 

244 self.assertEqual(component.getExposureId(), reference.getExposureId(), 

245 f"VisitInfo comparison") 

246 elif compName == "metadata": 

247 # The component metadata has extra fields in it so cannot 

248 # compare directly. 

249 for k, v in reference.items(): 

250 self.assertEqual(component[k], v) 

251 elif compName == "photoCalib": 

252 # This example has a 

253 # "spatially constant with mean: inf error: nan" entry 

254 # which does not compare directly. 

255 self.assertEqual(str(component), str(reference)) 

256 self.assertIn("spatially constant with mean: inf", str(component), "Checking photoCalib") 

257 else: 

258 raise RuntimeError(f"Unexpected component '{compName}' encountered in test") 

259 

260 # With parameters 

261 inBBox = Box2I(minimum=Point2I(0, 0), maximum=Point2I(3, 3)) 

262 parameters = dict(bbox=inBBox, origin=LOCAL) 

263 subset = self.butler.get(datasetTypeName, dataId, parameters=parameters) 

264 outBBox = subset.getBBox() 

265 self.assertEqual(inBBox, outBBox) 

266 self.assertImagesEqual(subset.getImage(), exposure.subset(inBBox, origin=LOCAL).getImage()) 

267 

268 return ref 

269 

270 def putFits(self, exposure, datasetTypeName, visit): 

271 """Put different datasetTypes and return information.""" 

272 dataId = {"visit": visit, "instrument": "DummyCam", "physical_filter": "d-r"} 

273 refC = self.butler.put(exposure, datasetTypeName, dataId) 

274 uriC = self.butler.getURI(refC) 

275 stat = os.stat(uriC.path) 

276 size = stat.st_size 

277 meta = self.butler.get(f"{datasetTypeName}.metadata", dataId) 

278 return meta, size 

279 

280 def testCompression(self): 

281 """Test that we can write compressed and uncompressed FITS.""" 

282 example = os.path.join(TESTDIR, "data", "small.fits") 

283 exposure = lsst.afw.image.ExposureF(example) 

284 

285 # Write a lossless compressed 

286 metaC, sizeC = self.putFits(exposure, "lossless", 42) 

287 self.assertEqual(metaC["TTYPE1"], "COMPRESSED_DATA") 

288 self.assertEqual(metaC["ZCMPTYPE"], "GZIP_2") 

289 

290 # Write an uncompressed FITS file 

291 metaN, sizeN = self.putFits(exposure, "uncompressed", 43) 

292 self.assertNotIn("ZCMPTYPE", metaN) 

293 

294 # Write an uncompressed FITS file 

295 metaL, sizeL = self.putFits(exposure, "lossy", 44) 

296 self.assertEqual(metaL["TTYPE1"], "COMPRESSED_DATA") 

297 self.assertEqual(metaL["ZCMPTYPE"], "RICE_1") 

298 

299 self.assertNotEqual(sizeC, sizeN) 

300 # Data file is so small that Lossy and Compressed are dominated 

301 # by the extra compression tables 

302 self.assertEqual(sizeL, sizeC) 

303 

304 

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

306 unittest.main()