Coverage for python/lsst/daf/butler/script/exportCalibs.py: 11%
Shortcuts 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
Shortcuts 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/>.
22import logging
23import os
25from astropy.table import Table
27from .._butler import Butler
28from ..registry import CollectionType
30log = logging.getLogger(__name__)
33def parseCalibrationCollection(registry, collection, datasetTypes):
34 """Search a calibration collection for calibration datasets.
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.
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.
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.")
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
73def exportCalibs(repo, directory, collections):
74 """Certify a set of calibrations with a validity range.
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.
89 Returns
90 -------
91 datasetTable : `astropy.table.Table`
92 A table containing relevant information about the calibrations
93 exported.
95 Raises
96 ------
97 RuntimeError :
98 Raised if the output directory already exists.
99 """
100 butler = Butler(repo, writeable=False)
102 calibTypes = [
103 datasetType for datasetType in butler.registry.queryDatasetTypes(...) if datasetType.isCalibration()
104 ]
106 collectionsToExport = []
107 datasetsToExport = []
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)
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)
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))
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)
140 log.info("Saving %d dataset(s)", len(datasetsToExport))
141 export.saveDatasets(datasetsToExport)
143 sortedDatasets = sorted(datasetsToExport, key=lambda x: x.datasetType.name)
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 }
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 )