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 shutil 

26import tempfile 

27from typing import Any 

28import unittest 

29 

30import astropy.time 

31 

32from lsst.daf.butler import ( 

33 Butler, 

34 ButlerConfig, 

35 CollectionType, 

36 Registry, 

37 Timespan, 

38) 

39from lsst.daf.butler.registry import RegistryConfig 

40from lsst.daf.butler.tests import DatastoreMock 

41 

42 

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

44 

45 

46class SimpleButlerTestCase(unittest.TestCase): 

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

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

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

50 """ 

51 

52 def setUp(self): 

53 self.root = tempfile.mkdtemp() 

54 

55 def tearDown(self): 

56 if self.root is not None and os.path.exists(self.root): 

57 shutil.rmtree(self.root, ignore_errors=True) 

58 

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

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

61 """ 

62 config = ButlerConfig() 

63 

64 # make separate temporary directory for registry of this instance 

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

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

67 config["root"] = self.root 

68 

69 # have to make a registry first 

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

71 Registry.createFromConfig(registryConfig) 

72 

73 butler = Butler(config, **kwargs) 

74 DatastoreMock.apply(butler) 

75 return butler 

76 

77 def testReadBackwardsCompatibility(self): 

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

79 and commit to the daf_butler git repo. 

80 

81 Notes 

82 ----- 

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

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

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

86 the export format required for CALIBRATION collections to land. 

87 """ 

88 butler = self.makeButler(writeable=True) 

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

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

91 # the above does not raise. 

92 self.assertGreaterEqual( 

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

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

95 ) 

96 self.assertGreaterEqual( 

97 { 

98 (record.id, record.physical_filter) 

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

100 }, 

101 { 

102 (27136, 'HSC-Z'), 

103 (11694, 'HSC-G'), 

104 (23910, 'HSC-R'), 

105 (11720, 'HSC-Y'), 

106 (23900, 'HSC-R'), 

107 (22646, 'HSC-Y'), 

108 (1248, 'HSC-I'), 

109 (19680, 'HSC-I'), 

110 (1240, 'HSC-I'), 

111 (424, 'HSC-Y'), 

112 (19658, 'HSC-I'), 

113 (344, 'HSC-Y'), 

114 (1218, 'HSC-R'), 

115 (1190, 'HSC-Z'), 

116 (23718, 'HSC-R'), 

117 (11700, 'HSC-G'), 

118 (26036, 'HSC-G'), 

119 (23872, 'HSC-R'), 

120 (1170, 'HSC-Z'), 

121 (1876, 'HSC-Y'), 

122 } 

123 ) 

124 

125 def testDatasetTransfers(self): 

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

127 back in again. 

128 """ 

129 # Import data to play with. 

130 butler1 = self.makeButler(writeable=True) 

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

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

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

134 # Export all datasets. 

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

136 exporter.saveDatasets( 

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

138 ) 

139 # Import it all again. 

140 butler2 = self.makeButler(writeable=True) 

141 butler2.import_(filename=file.name) 

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

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

144 # rewritten. 

145 self.assertCountEqual( 

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

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

148 ) 

149 

150 def testCollectionTransfers(self): 

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

152 """ 

153 # Populate a registry with some datasets. 

154 butler1 = self.makeButler(writeable=True) 

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

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

157 registry1 = butler1.registry 

158 # Add some more collections. 

159 registry1.registerRun("run1") 

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

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

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

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

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

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

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

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

168 registry1.associate("tag1", flats1) 

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

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

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

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

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

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

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

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

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

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

179 

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

181 # Export all collections, and some datasets. 

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

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

184 # intentionally not topological order. 

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

186 exporter.saveCollection(collection) 

187 exporter.saveDatasets(flats1) 

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

189 # Import them into a new registry. 

190 butler2 = self.makeButler(writeable=True) 

191 butler2.import_(filename=file.name) 

192 registry2 = butler2.registry 

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

194 # themselves. 

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

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

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

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

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

200 self.assertEqual( 

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

202 ["tag1", "run1", "chain2"], 

203 ) 

204 self.assertEqual( 

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

206 ["calibration1", "run1"], 

207 ) 

208 # Check that tag collection contents are the same. 

209 self.maxDiff = None 

210 self.assertCountEqual( 

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

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

213 ) 

214 # Check that calibration collection contents are the same. 

215 self.assertCountEqual( 

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

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

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

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

220 ) 

221 

222 def testGetCalibration(self): 

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

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

225 extra dimensions with temporal information. 

226 """ 

227 # Import data to play with. 

228 butler = self.makeButler(writeable=True) 

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

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

231 # Certify some biases into a CALIBRATION collection. 

232 registry = butler.registry 

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

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

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

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

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

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

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

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

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

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

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

244 # Insert some exposure dimension data. 

245 registry.insertDimensionData( 

246 "exposure", 

247 { 

248 "instrument": "Cam1", 

249 "id": 3, 

250 "obs_id": "three", 

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

252 "physical_filter": "Cam1-G", 

253 }, 

254 { 

255 "instrument": "Cam1", 

256 "id": 4, 

257 "obs_id": "four", 

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

259 "physical_filter": "Cam1-G", 

260 }, 

261 ) 

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

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

264 collections="calibs") 

265 self.assertEqual(bias2a_id, bias2a.id) 

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

267 collections="calibs") 

268 self.assertEqual(bias3b_id, bias3b.id) 

269 

270 # Get using the kwarg form 

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

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

273 collections="calibs") 

274 self.assertEqual(bias3b_id, bias3b.id) 

275 

276 # Do it again but using the record information 

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

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

279 collections="calibs") 

280 self.assertEqual(bias2a_id, bias2a.id) 

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

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

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

284 self.assertEqual(bias3b_id, bias3b.id) 

285 

286 

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

288 unittest.main()