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.base import Packages 

37from lsst.daf.base import PropertyList, PropertySet 

38 

39from lsst.daf.butler import Config 

40from lsst.daf.butler import StorageClassFactory 

41from lsst.daf.butler import DatasetType 

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

43 

44from lsst.obs.base.exposureAssembler import ExposureAssembler 

45 

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

47 from lsst.daf.butler import DatasetRef 

48 

49TESTDIR = os.path.dirname(__file__) 

50 

51BUTLER_CONFIG = """ 

52storageClasses: 

53 ExposureCompositeF: 

54 inheritsFrom: ExposureF 

55datastore: 

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

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

58 formatters: 

59 ExposureCompositeF: lsst.obs.base.formatters.fitsExposure.FitsExposureFormatter 

60 lossless: 

61 formatter: lsst.obs.base.formatters.fitsExposure.FitsExposureFormatter 

62 parameters: 

63 recipe: lossless 

64 uncompressed: 

65 formatter: lsst.obs.base.formatters.fitsExposure.FitsExposureFormatter 

66 parameters: 

67 recipe: noCompression 

68 lossy: 

69 formatter: lsst.obs.base.formatters.fitsExposure.FitsExposureFormatter 

70 parameters: 

71 recipe: lossyBasic 

72 composites: 

73 disassembled: 

74 ExposureCompositeF: True 

75""" 

76 

77# Components present in the test file 

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

79 "filter"} 

80 

81 

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

83 

84 @classmethod 

85 def setUpClass(cls): 

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

87 

88 cls.storageClassFactory = StorageClassFactory() 

89 

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

91 

92 dataIds = { 

93 "instrument": ["DummyCam"], 

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

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

96 } 

97 

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

99 

100 # Create dataset types used by the tests 

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

102 ("unknown", "ExposureCompositeF"), 

103 ("testCatalog", "SourceCatalog"), 

104 ("lossless", "ExposureF"), 

105 ("uncompressed", "ExposureF"), 

106 ("lossy", "ExposureF"), 

107 ): 

108 storageClass = cls.storageClassFactory.getStorageClass(storageClassName) 

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

110 

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

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

113 ("pl", "PropertyList"), 

114 ("pkg", "Packages") 

115 ): 

116 storageClass = cls.storageClassFactory.getStorageClass(storageClassName) 

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

118 

119 @classmethod 

120 def tearDownClass(cls): 

121 if cls.root is not None: 

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

123 

124 def setUp(self): 

125 self.butler = makeTestCollection(self.creatorButler) 

126 

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

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

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

130 

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

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

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

134 inputTable = inputCatalog.getTable() 

135 inputRecord = inputCatalog[0] 

136 outputTable = outputCatalog.getTable() 

137 outputRecord = outputCatalog[0] 

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

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

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

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

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

143 self.assertFloatsAlmostEqual( 

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

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

146 self.assertFloatsAlmostEqual( 

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

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

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

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

151 self.assertFloatsAlmostEqual( 

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

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

154 self.assertFloatsAlmostEqual( 

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

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

157 self.assertFloatsAlmostEqual( 

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

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

160 

161 def runFundamentalTypeTest(self, datasetTypeName, entity): 

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

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

164 butler_ps = self.butler.get(ref) 

165 self.assertEqual(butler_ps, entity) 

166 

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

168 uri = self.butler.getURI(ref) 

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

170 

171 def testFundamentalTypes(self) -> None: 

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

173 ps = PropertySet() 

174 ps["a.b"] = 5 

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

176 self.runFundamentalTypeTest("ps", ps) 

177 

178 pl = PropertyList() 

179 pl["A"] = 1 

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

181 pl["B"] = "string" 

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

183 self.runFundamentalTypeTest("pl", pl) 

184 

185 pkg = Packages.fromSystem() 

186 self.runFundamentalTypeTest("pkg", pkg) 

187 

188 def testFitsCatalog(self) -> None: 

189 catalog = self.makeExampleCatalog() 

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

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

192 stored = self.butler.get(ref) 

193 self.assertCatalogEqual(catalog, stored) 

194 

195 def testExposureCompositePutGetConcrete(self) -> None: 

196 ref = self.runExposureCompositePutGetTest("calexp") 

197 

198 uri = self.butler.getURI(ref) 

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

200 

201 def testExposureCompositePutGetVirtual(self) -> None: 

202 ref = self.runExposureCompositePutGetTest("unknown") 

203 

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

205 self.assertIsNone(primary) 

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

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

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

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

210 

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

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

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

214 

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

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

217 

218 # Get the full thing 

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

220 

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

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

223 

224 # Helper for extracting components 

225 assembler = ExposureAssembler(ref.datasetType.storageClass) 

226 

227 # Get each component from butler independently 

228 for compName in COMPONENTS: 

229 compTypeName = DatasetType.nameWithComponent(datasetTypeName, compName) 

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

231 

232 reference = assembler.getComponent(exposure, compName) 

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

234 

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

236 self.assertImagesEqual(component, reference) 

237 elif compName == "mask": 

238 self.assertMasksEqual(component, reference) 

239 elif compName == "wcs": 

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

241 elif compName == "coaddInputs": 

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

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

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

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

246 elif compName == "psf": 

247 # Equality for PSF does not work 

248 pass 

249 elif compName == "filter": 

250 self.assertEqual(component.getCanonicalName(), reference.getCanonicalName()) 

251 elif compName == "visitInfo": 

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

253 f"VisitInfo comparison") 

254 elif compName == "metadata": 

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

256 # compare directly. 

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

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

259 elif compName == "photoCalib": 

260 # This example has a 

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

262 # which does not compare directly. 

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

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

265 else: 

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

267 

268 # With parameters 

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

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

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

272 outBBox = subset.getBBox() 

273 self.assertEqual(inBBox, outBBox) 

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

275 

276 return ref 

277 

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

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

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

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

282 uriC = self.butler.getURI(refC) 

283 stat = os.stat(uriC.path) 

284 size = stat.st_size 

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

286 return meta, size 

287 

288 def testCompression(self): 

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

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

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

292 

293 # Write a lossless compressed 

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

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

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

297 

298 # Write an uncompressed FITS file 

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

300 self.assertNotIn("ZCMPTYPE", metaN) 

301 

302 # Write an uncompressed FITS file 

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

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

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

306 

307 self.assertNotEqual(sizeC, sizeN) 

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

309 # by the extra compression tables 

310 self.assertEqual(sizeL, sizeC) 

311 

312 

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

314 unittest.main()