Coverage for tests/test_qg_builder_dimensions.py: 19%
120 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-04 01:47 -0800
1# This file is part of pipe_base.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (https://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 software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <https://www.gnu.org/licenses/>.
28import logging
29import unittest
31import astropy.table
33import lsst.utils.tests
34from lsst.daf.butler import Butler, DataCoordinate
35from lsst.daf.butler.tests.utils import create_populated_sqlite_registry
36from lsst.pipe.base import PipelineGraph
37from lsst.pipe.base.all_dimensions_quantum_graph_builder import (
38 AllDimensionsQuantumGraphBuilder,
39 DatasetQueryConstraintVariant,
40)
41from lsst.pipe.base.tests.mocks import (
42 DynamicConnectionConfig,
43 DynamicTestPipelineTask,
44 DynamicTestPipelineTaskConfig,
45)
46from lsst.resources import ResourcePath
47from lsst.sphgeom import RangeSet
49_LOG = logging.getLogger(__name__)
52class AllDimensionsQuantumGraphBuilderTestCase(unittest.TestCase):
53 """Tests for AllDimensionsQuantumGraphBuilder with various interesting
54 combinations of dimensions.
55 """
57 @staticmethod
58 def make_butler() -> Butler:
59 DATA_ROOT = ResourcePath("resource://lsst.daf.butler/tests/registry_data", forceDirectory=True)
60 return create_populated_sqlite_registry(
61 *[DATA_ROOT.join(filename) for filename in ("base.yaml", "spatial.yaml")]
62 )
64 def setUp(self):
65 self.butler = self.make_butler()
67 def tearDown(self):
68 del self.butler
69 return super().tearDown()
71 def test_one_to_one(self) -> None:
72 """Test building a QG with a single task whose inputs and outputs and
73 quanta all have the same dimensions.
74 """
75 config = DynamicTestPipelineTaskConfig()
76 config.dimensions = ["visit", "detector"]
77 config.inputs["i1"] = DynamicConnectionConfig(
78 dataset_type_name="d1", dimensions=["visit", "detector"]
79 )
80 config.outputs["o2"] = DynamicConnectionConfig(
81 dataset_type_name="d2", dimensions=["visit", "detector"]
82 )
83 pipeline_graph = PipelineGraph(universe=self.butler.dimensions)
84 pipeline_graph.add_task("t1", DynamicTestPipelineTask, config=config)
85 pipeline_graph.resolve(self.butler.registry)
86 self._insert_overall_inputs(pipeline_graph)
87 builder = AllDimensionsQuantumGraphBuilder(
88 pipeline_graph,
89 self.butler,
90 input_collections=["c1"],
91 output_run="c2",
92 )
93 qg = builder.finish(attach_datastore_records=False).assemble()
94 quanta = qg.quanta_by_task["t1"]
95 self.assertEqual(len(quanta), 8) # 2 visits x 4 detectors
97 def test_patch_to_hpx_to_global(self) -> None:
98 """Test building a QG with patch inputs and a hierarchy of healpix
99 outputs and a final global step.
100 """
101 pipeline_graph = self._make_hpx_pipeline_graph(
102 patch_to_hpx8_and_hpx11=True,
103 hpx8_to_hpx5=True,
104 hpx5_to_global=True,
105 )
106 self._insert_overall_inputs(pipeline_graph)
107 builder = AllDimensionsQuantumGraphBuilder(
108 pipeline_graph,
109 self.butler,
110 input_collections=["c1"],
111 output_run="c2",
112 )
113 qg = builder.finish(attach_datastore_records=False).assemble()
114 self.assertEqual(len(qg.quanta_by_task["patch_to_hpx8_and_hpx11"]), 22)
115 self.assertEqual(len(qg.quanta_by_task["hpx8_to_hpx5"]), 2)
116 self.assertEqual(len(qg.quanta_by_task["hpx5_to_global"]), 1)
118 def test_patch_to_hpx_to_instrument(self) -> None:
119 """Test building a QG with patch inputs and a hierarchy of healpix
120 outputs and a final per-instrument step.
121 """
122 pipeline_graph = self._make_hpx_pipeline_graph(
123 patch_to_hpx8_and_hpx11=True,
124 hpx8_to_hpx5=True,
125 hpx5_to_instrument=True,
126 )
127 self._insert_overall_inputs(pipeline_graph)
128 builder = AllDimensionsQuantumGraphBuilder(
129 pipeline_graph,
130 self.butler,
131 input_collections=["c1"],
132 output_run="c2",
133 )
134 qg = builder.finish(attach_datastore_records=False).assemble()
135 self.assertEqual(len(qg.quanta_by_task["patch_to_hpx8_and_hpx11"]), 22)
136 self.assertEqual(len(qg.quanta_by_task["hpx8_to_hpx5"]), 2)
137 self.assertEqual(len(qg.quanta_by_task["hpx5_to_instrument"]), 1)
139 def test_hpx_to_global_dataset_constraint(self) -> None:
140 """Test building a QG with healpix inputs and a hierarchy of healpix
141 outputs and a final global step, with the healpix input used as a
142 dataset constraint (i.e. the default behavior).
143 """
144 pipeline_graph = self._make_hpx_pipeline_graph(
145 hpx8_to_hpx5=True,
146 hpx5_to_global=True,
147 )
148 self._insert_overall_inputs(pipeline_graph)
149 builder = AllDimensionsQuantumGraphBuilder(
150 pipeline_graph,
151 self.butler,
152 input_collections=["c1"],
153 output_run="c2",
154 )
155 qg = builder.finish(attach_datastore_records=False).assemble()
156 self.assertEqual(len(qg.quanta_by_task["hpx8_to_hpx5"]), 2)
157 self.assertEqual(len(qg.quanta_by_task["hpx5_to_global"]), 1)
159 @unittest.expectedFailure
160 def test_hpx_to_global_where_constraint(self) -> None:
161 """Test building a QG with healpix inputs and a hierarchy of healpix
162 outputs and a final global step, with a 'where' constraint instead of
163 a dataset constraint.
165 It would be nice for this to work, but we do not currently expect it to
166 given the limitations of the butler query system.
167 """
168 pipeline_graph = self._make_hpx_pipeline_graph(
169 hpx8_to_hpx5=True,
170 hpx5_to_global=True,
171 )
172 self._insert_overall_inputs(pipeline_graph)
173 builder = AllDimensionsQuantumGraphBuilder(
174 pipeline_graph,
175 self.butler,
176 input_collections=["c1"],
177 output_run="c2",
178 dataset_query_constraint=DatasetQueryConstraintVariant.OFF,
179 where="healpix5 IN (4864, 4522)",
180 )
181 qg = builder.finish(attach_datastore_records=False).assemble()
182 self.assertEqual(len(qg.quanta_by_task["hpx8_to_hpx5"]), 2)
183 self.assertEqual(len(qg.quanta_by_task["hpx5_to_global"]), 1)
185 def test_hpx_to_global_data_id_table(self) -> None:
186 """Test building a QG with healpix inputs and a hierarchy of healpix
187 outputs and a final global step, with healpix IDs provided via a data
188 ID table.
189 """
190 pipeline_graph = self._make_hpx_pipeline_graph(
191 hpx8_to_hpx5=True,
192 hpx5_to_global=True,
193 )
194 self._insert_overall_inputs(pipeline_graph)
195 tbl = astropy.table.Table(rows=[(4864,), (4522,)], names=["healpix5"])
196 builder = AllDimensionsQuantumGraphBuilder(
197 pipeline_graph,
198 self.butler,
199 input_collections=["c1"],
200 output_run="c2",
201 dataset_query_constraint=DatasetQueryConstraintVariant.OFF,
202 data_id_tables=[tbl],
203 )
204 qg = builder.finish(attach_datastore_records=False).assemble()
205 self.assertEqual(len(qg.quanta_by_task["hpx8_to_hpx5"]), 2)
206 self.assertEqual(len(qg.quanta_by_task["hpx5_to_global"]), 1)
208 def _make_hpx_pipeline_graph(
209 self,
210 *,
211 patch_to_hpx8_and_hpx11: bool = False,
212 hpx8_to_hpx5: bool = False,
213 hpx5_to_global=False,
214 hpx5_to_instrument: bool = False,
215 ) -> PipelineGraph:
216 """Generate a pipeline graph with tasks that use healpix dimensions
217 in various ways.
218 """
219 pipeline_graph = PipelineGraph(universe=self.butler.dimensions)
220 if patch_to_hpx8_and_hpx11:
221 config = DynamicTestPipelineTaskConfig()
222 config.dimensions = {"healpix8"}
223 config.inputs["i1"] = DynamicConnectionConfig(
224 dataset_type_name="patch_dataset",
225 dimensions={"patch"},
226 multiple=True,
227 )
228 config.outputs["o1"] = DynamicConnectionConfig(
229 dataset_type_name="hpx11_dataset",
230 dimensions={"healpix11"},
231 multiple=True,
232 )
233 config.outputs["o2"] = DynamicConnectionConfig(
234 dataset_type_name="hpx8_dataset",
235 dimensions={"healpix8"},
236 )
237 pipeline_graph.add_task("patch_to_hpx8_and_hpx11", DynamicTestPipelineTask, config=config)
238 if hpx8_to_hpx5:
239 config = DynamicTestPipelineTaskConfig()
240 config.dimensions = {"healpix5"}
241 config.inputs["i1"] = DynamicConnectionConfig(
242 dataset_type_name="hpx8_dataset",
243 dimensions={"healpix8"},
244 multiple=True,
245 )
246 config.outputs["o1"] = DynamicConnectionConfig(
247 dataset_type_name="hpx5_dataset",
248 dimensions={"healpix5"},
249 )
250 pipeline_graph.add_task("hpx8_to_hpx5", DynamicTestPipelineTask, config=config)
251 if hpx5_to_global:
252 config = DynamicTestPipelineTaskConfig()
253 config.dimensions = {}
254 config.inputs["i1"] = DynamicConnectionConfig(
255 dataset_type_name="hpx5_dataset",
256 dimensions={"healpix5"},
257 multiple=True,
258 )
259 config.outputs["o1"] = DynamicConnectionConfig(
260 dataset_type_name="global_dataset",
261 dimensions={},
262 )
263 pipeline_graph.add_task("hpx5_to_global", DynamicTestPipelineTask, config=config)
264 if hpx5_to_instrument:
265 config = DynamicTestPipelineTaskConfig()
266 config.dimensions = {"instrument"}
267 config.inputs["i1"] = DynamicConnectionConfig(
268 dataset_type_name="hpx5_dataset",
269 dimensions={"healpix5"},
270 multiple=True,
271 )
272 config.outputs["o1"] = DynamicConnectionConfig(
273 dataset_type_name="instrument_dataset",
274 dimensions={"instrument"},
275 )
276 pipeline_graph.add_task("hpx5_to_instrument", DynamicTestPipelineTask, config=config)
277 pipeline_graph.resolve(self.butler.registry)
278 return pipeline_graph
280 def _insert_overall_inputs(self, pipeline_graph: PipelineGraph) -> None:
281 """Insert overall-input datasets for a pipeline graph."""
282 self.butler.collections.register("c1")
283 for _, node in pipeline_graph.iter_overall_inputs():
284 self.butler.registry.registerDatasetType(node.dataset_type)
285 if node.dimensions.skypix:
286 if len(node.dimensions) == 1:
287 (skypix_name,) = node.dimensions.skypix
288 pixelization = node.dimensions.universe.skypix_dimensions[skypix_name].pixelization
289 ranges = RangeSet()
290 for patch_record in self.butler.query_dimension_records("patch"):
291 ranges |= pixelization.envelope(patch_record.region)
292 data_ids = []
293 for begin, end in ranges:
294 for index in range(begin, end):
295 data_ids.append(DataCoordinate.from_required_values(node.dimensions, (index,)))
296 else:
297 raise NotImplementedError(
298 "Can only generate data IDs for queryable dimensions and isolated skypix."
299 )
300 else:
301 data_ids = self.butler.query_data_ids(node.dimensions, explain=False)
302 self.butler.registry.insertDatasets(node.dataset_type, data_ids, run="c1")
305if __name__ == "__main__":
306 lsst.utils.tests.init()
307 unittest.main()