Coverage for python/lsst/daf/butler/script/exportCalibs.py: 11%

49 statements  

« prev     ^ index     » next       coverage.py v6.4, created at 2022-05-24 02:27 -0700

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 

22import logging 

23import os 

24 

25from astropy.table import Table 

26 

27from .._butler import Butler 

28from ..registry import CollectionType 

29 

30log = logging.getLogger(__name__) 

31 

32 

33def parseCalibrationCollection(registry, collection, datasetTypes): 

34 """Search a calibration collection for calibration datasets. 

35 

36 Parameters 

37 ---------- 

38 registry : `lsst.daf.butler.Registry` 

39 Butler registry to use. 

40 collection : `str` 

41 Collection to search. This should be a CALIBRATION 

42 collection. 

43 datasetTypes : `list` [`lsst.daf.Butler.DatasetType`] 

44 List of calibration dataset types. 

45 

46 Returns 

47 ------- 

48 exportCollections : `list` [`str`] 

49 List of collections to save on export. 

50 exportDatasets : `list` [`lsst.daf.butler.queries.DatasetQueryResults`] 

51 Datasets to save on export. 

52 

53 Raises 

54 ------ 

55 RuntimeError 

56 Raised if the collection to search is not a CALIBRATION collection. 

57 """ 

58 if registry.getCollectionType(collection) != CollectionType.CALIBRATION: 

59 raise RuntimeError(f"Collection {collection} is not a CALIBRATION collection.") 

60 

61 exportCollections = [] 

62 exportDatasets = [] 

63 for calibType in datasetTypes: 

64 associations = registry.queryDatasetAssociations( 

65 calibType, collections=collection, collectionTypes=[CollectionType.CALIBRATION] 

66 ) 

67 for result in associations: 

68 exportDatasets.append(result.ref) 

69 exportCollections.append(result.ref.run) 

70 return exportCollections, exportDatasets 

71 

72 

73def exportCalibs(repo, directory, collections): 

74 """Certify a set of calibrations with a validity range. 

75 

76 Parameters 

77 ---------- 

78 repo : `str` 

79 URI to the location of the repo or URI to a config file 

80 describing the repo and its location. 

81 directory : `str` 

82 URI string of the directory to write the exported 

83 calibrations. 

84 collections : `list` [`str`] 

85 Data collections to pull calibrations from. Must be an 

86 existing `~CollectionType.CHAINED` or 

87 `~CollectionType.CALIBRATION` collection. 

88 

89 Returns 

90 ------- 

91 datasetTable : `astropy.table.Table` 

92 A table containing relevant information about the calibrations 

93 exported. 

94 

95 Raises 

96 ------ 

97 RuntimeError : 

98 Raised if the output directory already exists. 

99 """ 

100 butler = Butler(repo, writeable=False) 

101 

102 calibTypes = [ 

103 datasetType for datasetType in butler.registry.queryDatasetTypes(...) if datasetType.isCalibration() 

104 ] 

105 

106 collectionsToExport = [] 

107 datasetsToExport = [] 

108 

109 for collection in butler.registry.queryCollections( 

110 collections, 

111 flattenChains=True, 

112 includeChains=True, 

113 collectionTypes={CollectionType.CALIBRATION, CollectionType.CHAINED}, 

114 ): 

115 log.info("Checking collection: %s", collection) 

116 

117 # Get collection information. 

118 collectionsToExport.append(collection) 

119 collectionType = butler.registry.getCollectionType(collection) 

120 if collectionType == CollectionType.CALIBRATION: 

121 exportCollections, exportDatasets = parseCalibrationCollection( 

122 butler.registry, collection, calibTypes 

123 ) 

124 collectionsToExport.extend(exportCollections) 

125 datasetsToExport.extend(exportDatasets) 

126 

127 if os.path.exists(directory): 

128 raise RuntimeError(f"Export directory exists: {directory}") 

129 os.makedirs(directory) 

130 with butler.export(directory=directory, format="yaml", transfer="auto") as export: 

131 collectionsToExport = list(set(collectionsToExport)) 

132 datasetsToExport = list(set(datasetsToExport)) 

133 

134 for exportable in collectionsToExport: 

135 try: 

136 export.saveCollection(exportable) 

137 except Exception as e: 

138 log.warning("Did not save collection %s due to %s.", exportable, e) 

139 

140 log.info("Saving %d dataset(s)", len(datasetsToExport)) 

141 export.saveDatasets(datasetsToExport) 

142 

143 sortedDatasets = sorted(datasetsToExport, key=lambda x: x.datasetType.name) 

144 

145 requiredDimensions = set() 

146 for ref in sortedDatasets: 

147 requiredDimensions.update(ref.dimensions.names) 

148 dimensionColumns = { 

149 dimensionName: [ref.dataId.get(dimensionName, "") for ref in sortedDatasets] 

150 for dimensionName in requiredDimensions 

151 } 

152 

153 return Table( 

154 { 

155 "calibrationType": [ref.datasetType.name for ref in sortedDatasets], 

156 "run": [ref.run for ref in sortedDatasets], 

157 **dimensionColumns, 

158 } 

159 )