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 tempfile 

26from typing import Any 

27import unittest 

28 

29try: 

30 import numpy as np 

31except ImportError: 

32 np = None 

33 

34import astropy.time 

35 

36from lsst.daf.butler import ( 

37 Butler, 

38 ButlerConfig, 

39 CollectionType, 

40 Registry, 

41 Timespan, 

42) 

43from lsst.daf.butler.registry import RegistryConfig 

44from lsst.daf.butler.tests import DatastoreMock 

45from lsst.daf.butler.tests.utils import makeTestTempDir, removeTestTempDir 

46 

47 

48TESTDIR = os.path.abspath(os.path.dirname(__file__)) 

49 

50 

51class SimpleButlerTestCase(unittest.TestCase): 

52 """Tests for butler (including import/export functionality) that should not 

53 depend on the Registry Database backend or Datastore implementation, and 

54 can instead utilize an in-memory SQLite Registry and a mocked Datastore. 

55 """ 

56 

57 def setUp(self): 

58 self.root = makeTestTempDir(TESTDIR) 

59 

60 def tearDown(self): 

61 removeTestTempDir(self.root) 

62 

63 def makeButler(self, **kwargs: Any) -> Butler: 

64 """Return new Butler instance on each call. 

65 """ 

66 config = ButlerConfig() 

67 

68 # make separate temporary directory for registry of this instance 

69 tmpdir = tempfile.mkdtemp(dir=self.root) 

70 config["registry", "db"] = f"sqlite:///{tmpdir}/gen3.sqlite3" 

71 config["root"] = self.root 

72 

73 # have to make a registry first 

74 registryConfig = RegistryConfig(config.get("registry")) 

75 Registry.createFromConfig(registryConfig) 

76 

77 butler = Butler(config, **kwargs) 

78 DatastoreMock.apply(butler) 

79 return butler 

80 

81 def testReadBackwardsCompatibility(self): 

82 """Test that we can read an export file written by a previous version 

83 and commit to the daf_butler git repo. 

84 

85 Notes 

86 ----- 

87 At present this export file includes only dimension data, not datasets, 

88 which greatly limits the usefulness of this test. We should address 

89 this at some point, but I think it's best to wait for the changes to 

90 the export format required for CALIBRATION collections to land. 

91 """ 

92 butler = self.makeButler(writeable=True) 

93 butler.import_(filename=os.path.join(TESTDIR, "data", "registry", "hsc-rc2-subset.yaml")) 

94 # Spot-check a few things, but the most important test is just that 

95 # the above does not raise. 

96 self.assertGreaterEqual( 

97 set(record.id for record in butler.registry.queryDimensionRecords("detector", instrument="HSC")), 

98 set(range(104)), # should have all science CCDs; may have some focus ones. 

99 ) 

100 self.assertGreaterEqual( 

101 { 

102 (record.id, record.physical_filter) 

103 for record in butler.registry.queryDimensionRecords("visit", instrument="HSC") 

104 }, 

105 { 

106 (27136, 'HSC-Z'), 

107 (11694, 'HSC-G'), 

108 (23910, 'HSC-R'), 

109 (11720, 'HSC-Y'), 

110 (23900, 'HSC-R'), 

111 (22646, 'HSC-Y'), 

112 (1248, 'HSC-I'), 

113 (19680, 'HSC-I'), 

114 (1240, 'HSC-I'), 

115 (424, 'HSC-Y'), 

116 (19658, 'HSC-I'), 

117 (344, 'HSC-Y'), 

118 (1218, 'HSC-R'), 

119 (1190, 'HSC-Z'), 

120 (23718, 'HSC-R'), 

121 (11700, 'HSC-G'), 

122 (26036, 'HSC-G'), 

123 (23872, 'HSC-R'), 

124 (1170, 'HSC-Z'), 

125 (1876, 'HSC-Y'), 

126 } 

127 ) 

128 

129 def testDatasetTransfers(self): 

130 """Test exporting all datasets from a repo and then importing them all 

131 back in again. 

132 """ 

133 # Import data to play with. 

134 butler1 = self.makeButler(writeable=True) 

135 butler1.import_(filename=os.path.join(TESTDIR, "data", "registry", "base.yaml")) 

136 butler1.import_(filename=os.path.join(TESTDIR, "data", "registry", "datasets.yaml")) 

137 with tempfile.NamedTemporaryFile(mode='w', suffix=".yaml") as file: 

138 # Export all datasets. 

139 with butler1.export(filename=file.name) as exporter: 

140 exporter.saveDatasets( 

141 butler1.registry.queryDatasets(..., collections=...) 

142 ) 

143 # Import it all again. 

144 butler2 = self.makeButler(writeable=True) 

145 butler2.import_(filename=file.name) 

146 # Check that it all round-tripped. Use unresolved() to make 

147 # comparison not care about dataset_id values, which may be 

148 # rewritten. 

149 self.assertCountEqual( 

150 [ref.unresolved() for ref in butler1.registry.queryDatasets(..., collections=...)], 

151 [ref.unresolved() for ref in butler2.registry.queryDatasets(..., collections=...)], 

152 ) 

153 

154 def testCollectionTransfers(self): 

155 """Test exporting and then importing collections of various types. 

156 """ 

157 # Populate a registry with some datasets. 

158 butler1 = self.makeButler(writeable=True) 

159 butler1.import_(filename=os.path.join(TESTDIR, "data", "registry", "base.yaml")) 

160 butler1.import_(filename=os.path.join(TESTDIR, "data", "registry", "datasets.yaml")) 

161 registry1 = butler1.registry 

162 # Add some more collections. 

163 registry1.registerRun("run1") 

164 registry1.registerCollection("tag1", CollectionType.TAGGED) 

165 registry1.registerCollection("calibration1", CollectionType.CALIBRATION) 

166 registry1.registerCollection("chain1", CollectionType.CHAINED) 

167 registry1.registerCollection("chain2", CollectionType.CHAINED) 

168 registry1.setCollectionChain("chain1", ["tag1", "run1", "chain2"]) 

169 registry1.setCollectionChain("chain2", ["calibration1", "run1"]) 

170 # Associate some datasets into the TAGGED and CALIBRATION collections. 

171 flats1 = list(registry1.queryDatasets("flat", collections=...)) 

172 registry1.associate("tag1", flats1) 

173 t1 = astropy.time.Time('2020-01-01T01:00:00', format="isot", scale="tai") 

174 t2 = astropy.time.Time('2020-01-01T02:00:00', format="isot", scale="tai") 

175 t3 = astropy.time.Time('2020-01-01T03:00:00', format="isot", scale="tai") 

176 bias2a = registry1.findDataset("bias", instrument="Cam1", detector=2, collections="imported_g") 

177 bias3a = registry1.findDataset("bias", instrument="Cam1", detector=3, collections="imported_g") 

178 bias2b = registry1.findDataset("bias", instrument="Cam1", detector=2, collections="imported_r") 

179 bias3b = registry1.findDataset("bias", instrument="Cam1", detector=3, collections="imported_r") 

180 registry1.certify("calibration1", [bias2a, bias3a], Timespan(t1, t2)) 

181 registry1.certify("calibration1", [bias2b], Timespan(t2, None)) 

182 registry1.certify("calibration1", [bias3b], Timespan(t2, t3)) 

183 

184 with tempfile.NamedTemporaryFile(mode='w', suffix=".yaml") as file: 

185 # Export all collections, and some datasets. 

186 with butler1.export(filename=file.name) as exporter: 

187 # Sort results to put chain1 before chain2, which is 

188 # intentionally not topological order. 

189 for collection in sorted(registry1.queryCollections()): 

190 exporter.saveCollection(collection) 

191 exporter.saveDatasets(flats1) 

192 exporter.saveDatasets([bias2a, bias2b, bias3a, bias3b]) 

193 # Import them into a new registry. 

194 butler2 = self.makeButler(writeable=True) 

195 butler2.import_(filename=file.name) 

196 registry2 = butler2.registry 

197 # Check that it all round-tripped, starting with the collections 

198 # themselves. 

199 self.assertIs(registry2.getCollectionType("run1"), CollectionType.RUN) 

200 self.assertIs(registry2.getCollectionType("tag1"), CollectionType.TAGGED) 

201 self.assertIs(registry2.getCollectionType("calibration1"), CollectionType.CALIBRATION) 

202 self.assertIs(registry2.getCollectionType("chain1"), CollectionType.CHAINED) 

203 self.assertIs(registry2.getCollectionType("chain2"), CollectionType.CHAINED) 

204 self.assertEqual( 

205 list(registry2.getCollectionChain("chain1")), 

206 ["tag1", "run1", "chain2"], 

207 ) 

208 self.assertEqual( 

209 list(registry2.getCollectionChain("chain2")), 

210 ["calibration1", "run1"], 

211 ) 

212 # Check that tag collection contents are the same. 

213 self.maxDiff = None 

214 self.assertCountEqual( 

215 [ref.unresolved() for ref in registry1.queryDatasets(..., collections="tag1")], 

216 [ref.unresolved() for ref in registry2.queryDatasets(..., collections="tag1")], 

217 ) 

218 # Check that calibration collection contents are the same. 

219 self.assertCountEqual( 

220 [(assoc.ref.unresolved(), assoc.timespan) 

221 for assoc in registry1.queryDatasetAssociations("bias", collections="calibration1")], 

222 [(assoc.ref.unresolved(), assoc.timespan) 

223 for assoc in registry2.queryDatasetAssociations("bias", collections="calibration1")], 

224 ) 

225 

226 def testButlerGet(self): 

227 """Test that butler.get can work with different variants.""" 

228 

229 # Import data to play with. 

230 butler = self.makeButler(writeable=True) 

231 butler.import_(filename=os.path.join(TESTDIR, "data", "registry", "base.yaml")) 

232 butler.import_(filename=os.path.join(TESTDIR, "data", "registry", "datasets.yaml")) 

233 

234 # Find the DatasetRef for a flat 

235 coll = "imported_g" 

236 flat2g = butler.registry.findDataset("flat", instrument="Cam1", detector=2, physical_filter="Cam1-G", 

237 collections=coll) 

238 

239 # Create a numpy integer to check that works fine 

240 detector_np = np.int64(2) if np else 2 

241 print(type(detector_np)) 

242 

243 # Try to get it using different variations of dataId + keyword 

244 # arguments 

245 # Note that instrument.class_name does not work 

246 variants = ( 

247 (None, {"instrument": "Cam1", "detector": 2, "physical_filter": "Cam1-G"}), 

248 (None, {"instrument": "Cam1", "detector": detector_np, "physical_filter": "Cam1-G"}), 

249 ({"instrument": "Cam1", "detector": 2, "physical_filter": "Cam1-G"}, {}), 

250 ({"instrument": "Cam1", "detector": detector_np, "physical_filter": "Cam1-G"}, {}), 

251 ({"instrument": "Cam1", "detector": 2}, {"physical_filter": "Cam1-G"}), 

252 ({"detector.full_name": "Ab"}, {"instrument": "Cam1", "physical_filter": "Cam1-G"}), 

253 ({"full_name": "Ab"}, {"instrument": "Cam1", "physical_filter": "Cam1-G"}), 

254 (None, {"full_name": "Ab", "instrument": "Cam1", "physical_filter": "Cam1-G"}), 

255 ({"name_in_raft": "b", "raft": "A"}, {"instrument": "Cam1", "physical_filter": "Cam1-G"}), 

256 ({"name_in_raft": "b"}, {"raft": "A", "instrument": "Cam1", "physical_filter": "Cam1-G"}), 

257 (None, {"name_in_raft": "b", "raft": "A", "instrument": "Cam1", "physical_filter": "Cam1-G"}), 

258 ({"detector.name_in_raft": "b", "detector.raft": "A"}, 

259 {"instrument": "Cam1", "physical_filter": "Cam1-G"}), 

260 ({"detector.name_in_raft": "b", "detector.raft": "A", 

261 "instrument": "Cam1", "physical_filter": "Cam1-G"}, {}), 

262 ) 

263 

264 for dataId, kwds in variants: 

265 try: 

266 flat_id, _ = butler.get("flat", dataId=dataId, collections=coll, **kwds) 

267 except Exception as e: 

268 raise type(e)(f"{str(e)}: dataId={dataId}, kwds={kwds}") from e 

269 self.assertEqual(flat_id, flat2g.id, msg=f"DataId: {dataId}, kwds: {kwds}") 

270 

271 def testGetCalibration(self): 

272 """Test that `Butler.get` can be used to fetch from 

273 `~CollectionType.CALIBRATION` collections if the data ID includes 

274 extra dimensions with temporal information. 

275 """ 

276 # Import data to play with. 

277 butler = self.makeButler(writeable=True) 

278 butler.import_(filename=os.path.join(TESTDIR, "data", "registry", "base.yaml")) 

279 butler.import_(filename=os.path.join(TESTDIR, "data", "registry", "datasets.yaml")) 

280 # Certify some biases into a CALIBRATION collection. 

281 registry = butler.registry 

282 registry.registerCollection("calibs", CollectionType.CALIBRATION) 

283 t1 = astropy.time.Time('2020-01-01T01:00:00', format="isot", scale="tai") 

284 t2 = astropy.time.Time('2020-01-01T02:00:00', format="isot", scale="tai") 

285 t3 = astropy.time.Time('2020-01-01T03:00:00', format="isot", scale="tai") 

286 bias2a = registry.findDataset("bias", instrument="Cam1", detector=2, collections="imported_g") 

287 bias3a = registry.findDataset("bias", instrument="Cam1", detector=3, collections="imported_g") 

288 bias2b = registry.findDataset("bias", instrument="Cam1", detector=2, collections="imported_r") 

289 bias3b = registry.findDataset("bias", instrument="Cam1", detector=3, collections="imported_r") 

290 registry.certify("calibs", [bias2a, bias3a], Timespan(t1, t2)) 

291 registry.certify("calibs", [bias2b], Timespan(t2, None)) 

292 registry.certify("calibs", [bias3b], Timespan(t2, t3)) 

293 # Insert some exposure dimension data. 

294 registry.insertDimensionData( 

295 "exposure", 

296 { 

297 "instrument": "Cam1", 

298 "id": 3, 

299 "obs_id": "three", 

300 "timespan": Timespan(t1, t2), 

301 "physical_filter": "Cam1-G", 

302 "day_obs": 20201114, 

303 "seq_num": 55, 

304 }, 

305 { 

306 "instrument": "Cam1", 

307 "id": 4, 

308 "obs_id": "four", 

309 "timespan": Timespan(t2, t3), 

310 "physical_filter": "Cam1-G", 

311 "day_obs": 20211114, 

312 "seq_num": 42, 

313 }, 

314 ) 

315 # Get some biases from raw-like data IDs. 

316 bias2a_id, _ = butler.get("bias", {"instrument": "Cam1", "exposure": 3, "detector": 2}, 

317 collections="calibs") 

318 self.assertEqual(bias2a_id, bias2a.id) 

319 bias3b_id, _ = butler.get("bias", {"instrument": "Cam1", "exposure": 4, "detector": 3}, 

320 collections="calibs") 

321 self.assertEqual(bias3b_id, bias3b.id) 

322 

323 # Get using the kwarg form 

324 bias3b_id, _ = butler.get("bias", 

325 instrument="Cam1", exposure=4, detector=3, 

326 collections="calibs") 

327 self.assertEqual(bias3b_id, bias3b.id) 

328 

329 # Do it again but using the record information 

330 bias2a_id, _ = butler.get("bias", {"instrument": "Cam1", "exposure.obs_id": "three", 

331 "detector.full_name": "Ab"}, 

332 collections="calibs") 

333 self.assertEqual(bias2a_id, bias2a.id) 

334 bias3b_id, _ = butler.get("bias", {"exposure.obs_id": "four", 

335 "detector.full_name": "Ba"}, 

336 collections="calibs", instrument="Cam1") 

337 self.assertEqual(bias3b_id, bias3b.id) 

338 

339 # And again but this time using the alternate value rather than 

340 # the primary. 

341 bias3b_id, _ = butler.get("bias", {"exposure": "four", 

342 "detector": "Ba"}, 

343 collections="calibs", instrument="Cam1") 

344 self.assertEqual(bias3b_id, bias3b.id) 

345 

346 # And again but this time using the alternate value rather than 

347 # the primary and do it in the keyword arguments. 

348 bias3b_id, _ = butler.get("bias", 

349 exposure="four", detector="Ba", 

350 collections="calibs", instrument="Cam1") 

351 self.assertEqual(bias3b_id, bias3b.id) 

352 

353 # Now with implied record columns 

354 bias3b_id, _ = butler.get("bias", day_obs=20211114, seq_num=42, 

355 raft="B", name_in_raft="a", 

356 collections="calibs", instrument="Cam1") 

357 self.assertEqual(bias3b_id, bias3b.id) 

358 

359 

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

361 unittest.main()