Coverage for tests/test_butlerFits.py: 16%

250 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-06 04:13 -0700

1# This file is part of obs_base. 

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 shutil 

26import tempfile 

27import unittest 

28from typing import TYPE_CHECKING 

29 

30import astropy.table 

31import lsst.afw.cameraGeom.testUtils # for test asserts injected into TestCase 

32import lsst.afw.image 

33import lsst.pex.config 

34import lsst.utils.tests 

35from lsst.afw.fits import readMetadata 

36from lsst.afw.image import LOCAL, ExposureFitsReader, MaskedImageFitsReader 

37from lsst.afw.math import flipImage 

38from lsst.daf.base import PropertyList, PropertySet 

39from lsst.daf.butler import Config, DatasetType, StorageClassFactory 

40from lsst.daf.butler.tests import addDatasetType, makeTestCollection, makeTestRepo 

41from lsst.geom import Box2I, Extent2I, Point2I 

42from lsst.obs.base.exposureAssembler import ExposureAssembler 

43from lsst.obs.base.tests import make_ramp_exposure_trimmed, make_ramp_exposure_untrimmed 

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.fileDatastore.FileDatastore 

57 formatters: 

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

59 lossless: 

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

61 parameters: 

62 recipe: lossless 

63 uncompressed: 

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

65 parameters: 

66 recipe: noCompression 

67 lossy: 

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

69 parameters: 

70 recipe: lossyBasic 

71 composites: 

72 disassembled: 

73 ExposureCompositeF: True 

74""" 

75 

76# Components present in the test file 

77COMPONENTS = { 

78 "wcs", 

79 "image", 

80 "mask", 

81 "coaddInputs", 

82 "psf", 

83 "visitInfo", 

84 "variance", 

85 "metadata", 

86 "photoCalib", 

87 "filter", 

88 "validPolygon", 

89 "transmissionCurve", 

90 "detector", 

91 "apCorrMap", 

92 "summaryStats", 

93 "id", 

94} 

95READ_COMPONENTS = { 

96 "bbox", 

97 "xy0", 

98 "dimensions", 

99} 

100 

101 

102class ButlerFitsTests(lsst.utils.tests.TestCase): 

103 """Tests for butler interaction with FITS files.""" 

104 

105 @classmethod 

106 def setUpClass(cls): 

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

108 cls.storageClassFactory = StorageClassFactory() 

109 

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

111 

112 dataIds = { 

113 "instrument": ["DummyCam"], 

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

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

116 } 

117 

118 # Ensure that we test in a directory that will include some 

119 # metacharacters 

120 subdir = "sub?#dir" 

121 butlerRoot = os.path.join(cls.root, subdir) 

122 

123 cls.creatorButler = makeTestRepo(butlerRoot, dataIds, config=Config.fromYaml(BUTLER_CONFIG)) 

124 

125 # Create dataset types used by the tests 

126 for datasetTypeName, storageClassName in ( 

127 ("calexp", "ExposureF"), 

128 ("noise", "MaskedImageF"), 

129 ("unknown", "ExposureCompositeF"), 

130 ("testCatalog", "SourceCatalog"), 

131 ("lossless", "ExposureF"), 

132 ("uncompressed", "ExposureF"), 

133 ("lossy", "ExposureF"), 

134 ): 

135 storageClass = cls.storageClassFactory.getStorageClass(storageClassName) 

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

137 

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

139 for datasetTypeName, storageClassName in ( 

140 ("ps", "PropertySet"), 

141 ("pl", "PropertyList"), 

142 ("int_exp_trimmed", "ExposureI"), 

143 ("int_exp_untrimmed", "ExposureI"), 

144 ): 

145 storageClass = cls.storageClassFactory.getStorageClass(storageClassName) 

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

147 

148 @classmethod 

149 def tearDownClass(cls): 

150 if cls.root is not None: 

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

152 

153 def setUp(self): 

154 self.butler = makeTestCollection(self.creatorButler, uniqueId=self.id()) 

155 

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

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

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

159 

160 def assertCatalogEqual( 

161 self, inputCatalog: lsst.afw.table.SourceCatalog, outputCatalog: lsst.afw.table.SourceCatalog 

162 ) -> None: 

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

164 inputTable = inputCatalog.getTable() 

165 inputRecord = inputCatalog[0] 

166 outputTable = outputCatalog.getTable() 

167 outputRecord = outputCatalog[0] 

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

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

170 self.assertEqual( 

171 inputTable.getSchema().getAliasMap().get("slot_Centroid"), 

172 outputTable.getSchema().getAliasMap().get("slot_Centroid"), 

173 ) 

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

175 self.assertFloatsAlmostEqual( 

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

177 ) 

178 self.assertFloatsAlmostEqual( 

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

180 ) 

181 self.assertEqual( 

182 inputTable.getSchema().getAliasMap().get("slot_Shape"), 

183 outputTable.getSchema().getAliasMap().get("slot_Shape"), 

184 ) 

185 self.assertFloatsAlmostEqual( 

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

187 ) 

188 self.assertFloatsAlmostEqual( 

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

190 ) 

191 self.assertFloatsAlmostEqual( 

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

193 ) 

194 

195 def runFundamentalTypeTest(self, datasetTypeName, entity): 

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

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

198 butler_ps = self.butler.get(ref) 

199 self.assertEqual(butler_ps, entity) 

200 

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

202 uri = self.butler.getURI(ref) 

203 self.assertEqual(uri.getExtension(), ".yaml", f"Check extension of {uri}") 

204 

205 def testFundamentalTypes(self) -> None: 

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

207 ps = PropertySet() 

208 ps["a.b"] = 5 

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

210 self.runFundamentalTypeTest("ps", ps) 

211 

212 pl = PropertyList() 

213 pl["A"] = 1 

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

215 pl["B"] = "string" 

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

217 self.runFundamentalTypeTest("pl", pl) 

218 

219 def testFitsCatalog(self) -> None: 

220 """Test reading of a FITS catalog.""" 

221 catalog = self.makeExampleCatalog() 

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

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

224 stored = self.butler.get(ref) 

225 self.assertCatalogEqual(catalog, stored) 

226 

227 # Override the storage class. 

228 astropy_table = self.butler.get(ref, storageClass="AstropyTable") 

229 self.assertIsInstance(astropy_table, astropy.table.Table) 

230 self.assertEqual(len(astropy_table), len(stored)) 

231 

232 def testExposureCompositePutGetConcrete(self) -> None: 

233 """Test composite with no disassembly.""" 

234 ref = self.runExposureCompositePutGetTest("calexp") 

235 

236 uri = self.butler.getURI(ref) 

237 self.assertTrue(uri.exists(), f"Checking URI {uri} existence") 

238 

239 def testExposureCompositePutGetVirtual(self) -> None: 

240 """Testing composite disassembly.""" 

241 ref = self.runExposureCompositePutGetTest("unknown") 

242 

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

244 self.assertIsNone(primary) 

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

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

247 self.assertTrue(uri.exists(), f"Checking URI {uri} existence for component {compName}") 

248 

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

250 example = os.path.join(TESTDIR, "data", "calexp.fits") 

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

252 

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

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

255 

256 # Get the full thing 

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

258 

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

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

261 

262 # Helper for extracting components 

263 assembler = ExposureAssembler(ref.datasetType.storageClass) 

264 

265 # Check all possible components that can be read 

266 allComponents = set() 

267 allComponents.update(COMPONENTS, READ_COMPONENTS) 

268 

269 # Get each component from butler independently 

270 for compName in allComponents: 

271 compTypeName = DatasetType.nameWithComponent(datasetTypeName, compName) 

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

273 

274 reference = assembler.getComponent(exposure, compName) 

275 

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

277 

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

279 self.assertImagesEqual(component, reference) 

280 elif compName == "mask": 

281 self.assertMasksEqual(component, reference) 

282 elif compName == "wcs": 

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

284 elif compName == "coaddInputs": 

285 self.assertEqual( 

286 len(component.visits), len(reference.visits), f"cf visits {component.visits}" 

287 ) 

288 self.assertEqual(len(component.ccds), len(reference.ccds), f"cf CCDs {component.ccds}") 

289 elif compName == "psf": 

290 # Equality for PSF does not work 

291 pass 

292 elif compName == "filter": 

293 self.assertEqual(component, reference) 

294 elif compName == "id": 

295 self.assertEqual(component, reference) 

296 elif compName == "visitInfo": 

297 self.assertEqual(component, reference, "VisitInfo comparison") 

298 elif compName == "metadata": 

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

300 # compare directly. 

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

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

303 elif compName == "photoCalib": 

304 # This example has a 

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

306 # which does not compare directly. 

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

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

309 elif compName in ("bbox", "xy0", "dimensions", "validPolygon"): 

310 self.assertEqual(component, reference) 

311 elif compName == "apCorrMap": 

312 self.assertEqual(set(component.keys()), set(reference.keys())) 

313 elif compName == "transmissionCurve": 

314 self.assertEqual(component.getThroughputAtBounds(), reference.getThroughputAtBounds()) 

315 elif compName == "detector": 

316 c_amps = {a.getName() for a in component.getAmplifiers()} 

317 r_amps = {a.getName() for a in reference.getAmplifiers()} 

318 self.assertEqual(c_amps, r_amps) 

319 elif compName == "summaryStats": 

320 self.assertEqual(component.psfSigma, reference.psfSigma) 

321 else: 

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

323 

324 # Full Exposure with parameters 

325 inBBox = Box2I(minimum=Point2I(3, 3), maximum=Point2I(21, 16)) 

326 parameters = {"bbox": inBBox, "origin": LOCAL} 

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

328 outBBox = subset.getBBox() 

329 self.assertEqual(inBBox, outBBox) 

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

331 

332 return ref 

333 

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

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

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

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

338 uriC = self.butler.getURI(refC) 

339 stat = os.stat(uriC.ospath) 

340 size = stat.st_size 

341 # We can't use butler's Exposure storage class metadata component, 

342 # because that intentionally strips keywords that science code 

343 # shouldn't ever read (to at least weaken our assumptions that we write 

344 # to FITS). Instead use a lower-level method on the URI for this test. 

345 meta = readMetadata(uriC.ospath) 

346 return meta, size 

347 

348 def testCompression(self): 

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

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

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

352 

353 # Write a lossless compressed 

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

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

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

357 

358 # Write an uncompressed FITS file 

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

360 self.assertNotIn("ZCMPTYPE", metaN) 

361 

362 # Write an uncompressed FITS file 

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

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

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

366 

367 self.assertNotEqual(sizeC, sizeN) 

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

369 # by the extra compression tables 

370 self.assertEqual(sizeL, sizeC) 

371 

372 def testExposureFormatterAmpParameter(self): 

373 """Test the FitsExposureFormatter implementation of the Exposure 

374 StorageClass's 'amp' and 'detector' parameters. 

375 """ 

376 # Our example exposure file has a realistic detector (looks like HSC), 

377 # but the image itself doesn't match it. So we just load the detector, 

378 # and use it to make our own more useful images and put them. 

379 detector = ExposureFitsReader(os.path.join(TESTDIR, "data", "calexp.fits")).readDetector() 

380 trimmed_full = make_ramp_exposure_trimmed(detector) 

381 untrimmed_full = make_ramp_exposure_untrimmed(detector) 

382 trimmed_ref = self.butler.put(trimmed_full, "int_exp_trimmed") 

383 untrimmed_ref = self.butler.put(untrimmed_full, "int_exp_untrimmed") 

384 for n, amp in enumerate(detector): 

385 # Try to read each amp as it is on disk, with a variety of 

386 # parameters that should all do the same thing. 

387 for amp_parameter in [amp, amp.getName(), n]: 

388 for parameters in [{"amp": amp_parameter}, {"amp": amp_parameter, "detector": detector}]: 

389 with self.subTest(parameters=parameters): 

390 test_trimmed = self.butler.get(trimmed_ref, parameters=parameters) 

391 test_untrimmed = self.butler.get(untrimmed_ref, parameters=parameters) 

392 self.assertImagesEqual(test_trimmed.image, trimmed_full[amp.getBBox()].image) 

393 self.assertImagesEqual(test_untrimmed.image, untrimmed_full[amp.getRawBBox()].image) 

394 self.assertEqual(len(test_trimmed.getDetector()), 1) 

395 self.assertEqual(len(test_untrimmed.getDetector()), 1) 

396 self.assertAmplifiersEqual(test_trimmed.getDetector()[0], amp) 

397 self.assertAmplifiersEqual(test_untrimmed.getDetector()[0], amp) 

398 # Try to read various transformed versions of the original amp, 

399 # to make sure flips and offsets are applied correctly. 

400 # First flip X only. 

401 amp_t1 = amp.rebuild().transform(outFlipX=True).finish() 

402 test_t1_trimmed = self.butler.get(trimmed_ref, parameters={"amp": amp_t1}) 

403 self.assertImagesEqual( 

404 test_t1_trimmed.image, flipImage(trimmed_full[amp.getBBox()].image, flipLR=True, flipTB=False) 

405 ) 

406 self.assertAmplifiersEqual(test_t1_trimmed.getDetector()[0], amp_t1) 

407 test_t1_untrimmed = self.butler.get(untrimmed_ref, parameters={"amp": amp_t1}) 

408 self.assertImagesEqual( 

409 test_t1_untrimmed.image, 

410 flipImage(untrimmed_full[amp.getRawBBox()].image, flipLR=True, flipTB=False), 

411 ) 

412 self.assertAmplifiersEqual(test_t1_trimmed.getDetector()[0], amp_t1) 

413 # Flip Y only. 

414 amp_t2 = amp.rebuild().transform(outFlipY=True).finish() 

415 test_t2_trimmed = self.butler.get(trimmed_ref, parameters={"amp": amp_t2}) 

416 self.assertImagesEqual( 

417 test_t2_trimmed.image, flipImage(trimmed_full[amp.getBBox()].image, flipLR=False, flipTB=True) 

418 ) 

419 self.assertAmplifiersEqual(test_t2_trimmed.getDetector()[0], amp_t2) 

420 test_t2_untrimmed = self.butler.get(untrimmed_ref, parameters={"amp": amp_t2}) 

421 self.assertImagesEqual( 

422 test_t2_untrimmed.image, 

423 flipImage(untrimmed_full[amp.getRawBBox()].image, flipLR=False, flipTB=True), 

424 ) 

425 self.assertAmplifiersEqual(test_t2_trimmed.getDetector()[0], amp_t2) 

426 # Add an XY offset only. 

427 amp_t3 = amp.rebuild().transform(outOffset=Extent2I(5, 4)).finish() 

428 test_t3_trimmed = self.butler.get(trimmed_ref, parameters={"amp": amp_t3}) 

429 self.assertImagesEqual(test_t3_trimmed.image, trimmed_full[amp.getBBox()].image) 

430 self.assertAmplifiersEqual(test_t3_trimmed.getDetector()[0], amp_t3) 

431 test_t3_untrimmed = self.butler.get(untrimmed_ref, parameters={"amp": amp_t3}) 

432 self.assertImagesEqual(test_t3_untrimmed.image, untrimmed_full[amp.getRawBBox()].image) 

433 self.assertAmplifiersEqual(test_t3_trimmed.getDetector()[0], amp_t3) 

434 

435 def testMaskedImageFormatter(self): 

436 """Test that a MaskedImage can be persisted and read from a Butler.""" 

437 # Read in an Exposure as MaskedImage using MaskedImageFitsReader. 

438 reader = MaskedImageFitsReader(os.path.join(TESTDIR, "data", "calexp.fits")) 

439 mi = reader.read() 

440 

441 # Put the MaskedImage into the Butler and get a reference to it. 

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

443 ref = self.butler.put(mi, "noise", dataId) 

444 

445 # Check that the MaskedImage can be retrieved from the butler. 

446 maskedImage = self.butler.get(ref) 

447 self.assertImagesEqual(maskedImage.image, mi.image) 

448 self.assertImagesEqual(maskedImage.mask, mi.mask) 

449 self.assertImagesEqual(maskedImage.variance, mi.variance) 

450 

451 # Get a DeferredDatasetHandle to load parts of the MaskedImage. 

452 handle = self.butler.getDeferred(ref) 

453 

454 for parameters in ( 

455 {}, 

456 {"bbox": reader.readBBox()}, 

457 {"bbox": Box2I(minimum=Point2I(3, 3), maximum=Point2I(21, 16))}, 

458 {"bbox": Box2I(minimum=Point2I(3, 3), maximum=Point2I(21, 16)), "origin": LOCAL}, 

459 ): 

460 bbox = parameters.get("bbox", reader.readBBox()) 

461 with self.subTest(parameters=parameters): 

462 # Check that the reader supports reading sub-regions. 

463 subMaskedImage = reader.read(**parameters) 

464 self.assertImagesEqual(subMaskedImage.image, mi.image[bbox]) 

465 self.assertImagesEqual(subMaskedImage.mask, mi.mask[bbox]) 

466 self.assertImagesEqual(subMaskedImage.variance, mi.variance[bbox]) 

467 

468 # Get a maskedImage within a bounding box from the butler. 

469 subMaskedImage = handle.get(parameters=parameters) 

470 self.assertImagesEqual(subMaskedImage.image, mi.image[bbox]) 

471 self.assertImagesEqual(subMaskedImage.mask, mi.mask[bbox]) 

472 self.assertImagesEqual(subMaskedImage.variance, mi.variance[bbox]) 

473 

474 # Get one component at a time from the butler. 

475 subImage = handle.get(parameters=parameters, component="image") 

476 subMask = handle.get(parameters=parameters, component="mask") 

477 subVariance = handle.get(parameters=parameters, component="variance") 

478 self.assertImagesEqual(subImage, mi.image[bbox]) 

479 self.assertImagesEqual(subMask, mi.mask[bbox]) 

480 self.assertImagesEqual(subVariance, mi.variance[bbox]) 

481 

482 

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

484 unittest.main()