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 

80 

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

82 

83 @classmethod 

84 def setUpClass(cls): 

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

86 

87 cls.storageClassFactory = StorageClassFactory() 

88 

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

90 

91 dataIds = { 

92 "instrument": ["DummyCam"], 

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

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

95 } 

96 

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

98 

99 # Create dataset types used by the tests 

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

101 ("unknown", "ExposureCompositeF"), 

102 ("testCatalog", "SourceCatalog"), 

103 ("lossless", "ExposureF"), 

104 ("uncompressed", "ExposureF"), 

105 ("lossy", "ExposureF"), 

106 ): 

107 storageClass = cls.storageClassFactory.getStorageClass(storageClassName) 

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

109 

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

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

112 ("pl", "PropertyList"), 

113 ("pkg", "Packages") 

114 ): 

115 storageClass = cls.storageClassFactory.getStorageClass(storageClassName) 

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

117 

118 @classmethod 

119 def tearDownClass(cls): 

120 if cls.root is not None: 

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

122 

123 def setUp(self): 

124 self.butler = makeTestCollection(self.creatorButler) 

125 

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

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

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

129 

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

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

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

133 inputTable = inputCatalog.getTable() 

134 inputRecord = inputCatalog[0] 

135 outputTable = outputCatalog.getTable() 

136 outputRecord = outputCatalog[0] 

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

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

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

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

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

142 self.assertFloatsAlmostEqual( 

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

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

145 self.assertFloatsAlmostEqual( 

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

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

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

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

150 self.assertFloatsAlmostEqual( 

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

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

153 self.assertFloatsAlmostEqual( 

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

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

156 self.assertFloatsAlmostEqual( 

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

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

159 

160 def runFundamentalTypeTest(self, datasetTypeName, entity): 

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

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

163 butler_ps = self.butler.get(ref) 

164 self.assertEqual(butler_ps, entity) 

165 

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

167 uri = self.butler.getURI(ref) 

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

169 

170 def testFundamentalTypes(self) -> None: 

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

172 ps = PropertySet() 

173 ps["a.b"] = 5 

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

175 self.runFundamentalTypeTest("ps", ps) 

176 

177 pl = PropertyList() 

178 pl["A"] = 1 

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

180 pl["B"] = "string" 

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

182 self.runFundamentalTypeTest("pl", pl) 

183 

184 pkg = Packages.fromSystem() 

185 self.runFundamentalTypeTest("pkg", pkg) 

186 

187 def testFitsCatalog(self) -> None: 

188 catalog = self.makeExampleCatalog() 

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

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

191 stored = self.butler.get(ref) 

192 self.assertCatalogEqual(catalog, stored) 

193 

194 def testExposureCompositePutGetConcrete(self) -> None: 

195 ref = self.runExposureCompositePutGetTest("calexp") 

196 

197 uri = self.butler.getURI(ref) 

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

199 

200 def testExposureCompositePutGetVirtual(self) -> None: 

201 ref = self.runExposureCompositePutGetTest("unknown") 

202 

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

204 self.assertIsNone(primary) 

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

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

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

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

209 

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

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

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

213 

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

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

216 

217 # Get the full thing 

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

219 

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

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

222 

223 # Helper for extracting components 

224 assembler = ExposureAssembler(ref.datasetType.storageClass) 

225 

226 # Get each component from butler independently 

227 for compName in COMPONENTS: 

228 compTypeName = DatasetType.nameWithComponent(datasetTypeName, compName) 

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

230 

231 reference = assembler.getComponent(exposure, compName) 

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

233 

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

235 self.assertImagesEqual(component, reference) 

236 elif compName == "mask": 

237 self.assertMasksEqual(component, reference) 

238 elif compName == "wcs": 

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

240 elif compName == "coaddInputs": 

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

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

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

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

245 elif compName == "psf": 

246 # Equality for PSF does not work 

247 pass 

248 elif compName == "visitInfo": 

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

250 f"VisitInfo comparison") 

251 elif compName == "metadata": 

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

253 # compare directly. 

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

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

256 elif compName == "photoCalib": 

257 # This example has a 

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

259 # which does not compare directly. 

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

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

262 else: 

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

264 

265 # With parameters 

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

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

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

269 outBBox = subset.getBBox() 

270 self.assertEqual(inBBox, outBBox) 

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

272 

273 return ref 

274 

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

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

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

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

279 uriC = self.butler.getURI(refC) 

280 stat = os.stat(uriC.path) 

281 size = stat.st_size 

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

283 return meta, size 

284 

285 def testCompression(self): 

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

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

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

289 

290 # Write a lossless compressed 

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

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

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

294 

295 # Write an uncompressed FITS file 

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

297 self.assertNotIn("ZCMPTYPE", metaN) 

298 

299 # Write an uncompressed FITS file 

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

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

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

303 

304 self.assertNotEqual(sizeC, sizeN) 

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

306 # by the extra compression tables 

307 self.assertEqual(sizeL, sizeC) 

308 

309 

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

311 unittest.main()