Coverage for tests/test_drp-v2.py: 19%
145 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-05 01:17 -0800
1# This file is part of drp_pipe.
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/>.
22from __future__ import annotations
24import logging
25import os
26import re
27import tempfile
28import unittest
29from collections.abc import Set
30from typing import Any
32from lsst.daf.butler import Butler, Config, DatasetType
33from lsst.daf.butler.tests.utils import makeTestTempDir, removeTestTempDir
34from lsst.pipe.base import Pipeline, PipelineGraph
36from lsst.drp.pipe.tests.correspondence import Correspondence
38_LOG = logging.getLogger(__name__)
41PIPELINES_DIR = os.path.join(os.path.dirname(__file__), "..", "pipelines")
42TEST_DIR = os.path.abspath(os.path.dirname(__file__))
45COMCAM_REFCAT = "the_monster_20250219"
46COMCAM_INPUTS = {
47 COMCAM_REFCAT,
48 "raw",
49 "flat",
50 "dark",
51 "bias",
52 "bfk",
53 "linearizer",
54 "cti",
55 "ptc",
56 "crosstalk",
57 "camera",
58 "defects",
59 "skyMap",
60 "fgcmLookUpTable",
61 "pretrainedModelPackage",
62 "preloaded_DRP_SsObjects",
63 "illuminationCorrection",
64}
67class DrpV2TestCase(unittest.TestCase):
68 """Self-consistency and overall-input checks for the DRP-v2 pipeline
69 variants.
70 """
72 def setUp(self):
73 self.root = makeTestTempDir(TEST_DIR)
74 self.maxDiff = None
76 def tearDown(self):
77 removeTestTempDir(self.root)
79 def make_butler(self, **kwargs: Any) -> Butler:
80 """Construct a butler repository and return a `Butler` client."""
81 config = Config()
83 # make separate temporary directory for registry of this instance
84 tmpdir = tempfile.mkdtemp(dir=self.root)
85 config["registry", "db"] = f"sqlite:///{tmpdir}/gen3.sqlite3"
86 config = Butler.makeRepo(self.root, config)
87 butler = Butler.from_config(config, **kwargs)
88 return butler
90 def register_refcat(self, butler: Butler, name: str) -> None:
91 butler.registry.registerDatasetType(
92 DatasetType(
93 name,
94 {"htm7"},
95 "SimpleCatalog",
96 isCalibration=False,
97 universe=butler.dimensions,
98 )
99 )
101 def register_dataset(self, butler: Butler, name: str, dimensions: set(str), storageClass: str) -> None:
102 butler.registry.registerDatasetType(
103 DatasetType(
104 name,
105 dimensions,
106 storageClass,
107 isCalibration=False,
108 universe=butler.dimensions,
109 )
110 )
112 def test_comcam_full_load(self) -> None:
113 """Test the LSSTComCam/DRP-v2 pipeline after reading the full pipeline
114 at once.
115 """
116 butler = self.make_butler(writeable=True)
117 self.register_refcat(butler, COMCAM_REFCAT)
118 pipeline = Pipeline.from_uri(
119 os.path.join(PIPELINES_DIR, "LSSTComCam/DRP-v2.yaml")
120 )
121 # Just constructing and resolving the pipeline graph does a lot of
122 # validation, since it checks that the step flow and dimensions are
123 # valid (step flow is what the PipelineStepTester checks for the v1
124 # pipelines).
125 pipeline_graph = pipeline.to_graph(registry=butler.registry)
126 self.check_stage(
127 pipeline_graph, pipeline_graph.task_subsets["stage1-single-visit"], "step1"
128 )
129 self.check_stage(
130 pipeline_graph, pipeline_graph.task_subsets["stage2-recalibrate"], "step2"
131 )
132 self.check_stage(
133 pipeline_graph, pipeline_graph.task_subsets["stage3-coadd"], "step3"
134 )
135 self.check_stage(
136 pipeline_graph,
137 pipeline_graph.task_subsets["stage4-measure-variability"],
138 "step4",
139 )
140 # Check that the overall inputs are only the ones we expect.
141 overall_inputs = {name for name, _ in pipeline_graph.iter_overall_inputs()}
142 self.assertEqual(overall_inputs, COMCAM_INPUTS)
144 def test_comcam_stage_load(self) -> None:
145 """Test the LSSTComCam/DRP-v2 pipeline after reading the pipeline
146 one stage at a time.
147 """
148 butler = self.make_butler(writeable=True)
149 self.register_refcat(butler, COMCAM_REFCAT)
150 # Pre-define visit_table.
151 # It is stored as ArrowAstropy but read as a DataFrame.
152 self.register_dataset(butler, "visit_table", {"instrument"}, "ArrowAstropy")
153 pipeline_graph_1 = Pipeline.from_uri(
154 os.path.join(PIPELINES_DIR, "LSSTComCam/DRP-v2.yaml#stage1-single-visit")
155 ).to_graph(registry=butler.registry)
156 pipeline_graph_1.register_dataset_types(butler)
157 self.check_stage(pipeline_graph_1, pipeline_graph_1.tasks.keys(), "")
158 pipeline_graph_2 = Pipeline.from_uri(
159 os.path.join(PIPELINES_DIR, "LSSTComCam/DRP-v2.yaml#stage2-recalibrate")
160 ).to_graph(registry=butler.registry)
161 pipeline_graph_2.register_dataset_types(butler)
162 self.check_stage(
163 pipeline_graph_2, pipeline_graph_2.task_subsets["stage2-recalibrate"], ""
164 )
165 pipeline_graph_3 = Pipeline.from_uri(
166 os.path.join(PIPELINES_DIR, "LSSTComCam/DRP-v2.yaml#stage3-coadd")
167 ).to_graph(registry=butler.registry)
168 pipeline_graph_3.register_dataset_types(butler)
169 self.check_stage(
170 pipeline_graph_3, pipeline_graph_3.task_subsets["stage3-coadd"], ""
171 )
172 pipeline_graph_4 = Pipeline.from_uri(
173 os.path.join(
174 PIPELINES_DIR, "LSSTComCam/DRP-v2.yaml#stage4-measure-variability"
175 )
176 ).to_graph(registry=butler.registry)
177 pipeline_graph_4.register_dataset_types(butler)
178 self.check_stage(
179 pipeline_graph_4,
180 pipeline_graph_4.task_subsets["stage4-measure-variability"],
181 "",
182 )
183 # Spot-check a few prominent outputs for each stage.
184 self.assertIsNotNone(pipeline_graph_1.producer_of("single_visit_star"))
185 self.assertIsNotNone(pipeline_graph_1.producer_of("preliminary_visit_image"))
186 self.assertIsNotNone(pipeline_graph_1.producer_of("preliminary_visit_summary"))
187 self.assertIsNotNone(pipeline_graph_1.producer_of("preliminary_visit_table"))
188 self.assertIsNotNone(
189 pipeline_graph_1.producer_of("preliminary_visit_detector_table")
190 )
191 self.assertIsNotNone(
192 pipeline_graph_1.producer_of("single_visit_star_association_metrics")
193 )
194 self.assertIsNotNone(pipeline_graph_2.producer_of("recalibrated_star"))
195 self.assertIsNotNone(
196 pipeline_graph_2.producer_of("recalibrated_star_association_metrics")
197 )
198 self.assertIsNotNone(pipeline_graph_2.producer_of("visit_summary"))
199 self.assertIsNotNone(pipeline_graph_2.producer_of("visit_table"))
200 self.assertIsNotNone(pipeline_graph_2.producer_of("visit_detector_table"))
201 self.assertIsNotNone(pipeline_graph_3.producer_of("deep_coadd"))
202 self.assertIsNotNone(pipeline_graph_3.producer_of("template_coadd"))
203 self.assertIsNotNone(pipeline_graph_3.producer_of("object"))
204 self.assertIsNotNone(pipeline_graph_4.producer_of("dia_source"))
205 self.assertIsNotNone(pipeline_graph_4.producer_of("dia_object"))
206 self.assertIsNotNone(pipeline_graph_4.producer_of("source"))
207 self.assertIsNotNone(
208 pipeline_graph_4.producer_of("analysis_source_association_metrics")
209 )
210 self.assertIsNotNone(pipeline_graph_4.producer_of("visit_image"))
211 self.assertIsNotNone(pipeline_graph_4.producer_of("object_forced_source"))
212 self.assertIsNotNone(pipeline_graph_4.producer_of("dia_object_forced_source"))
214 def check_stage(
215 self, pipeline_graph: PipelineGraph, stage_members: Set[str], step_prefix: str
216 ) -> None:
217 """Check that the formal steps that should subdivide a stage subset
218 actually do.
220 Parameters
221 ----------
222 pipeline_graph : `lsst.pipe.base.pipeline_graph.PipelineGraph`
223 A pipeline graph that has step definitions.
224 stage_members : `Set` [ `str` ]
225 The task labels that belong to the stage being tested.
226 step_prefix : `str`
227 String prefix for all steps that are part of this stage. May be
228 an empty string to match all steps in the pipeline and check that
229 only the desired steps are in the pipeline.
230 """
231 step_labels = [
232 label for label in pipeline_graph.steps if label.startswith(step_prefix)
233 ]
234 # Steps should be sorted alphabetically.
235 self.assertEqual(sorted(step_labels), step_labels)
236 # Test that the checkpoint/stage subset is the unions of its steps.
237 step_member_union = set()
238 for step_label in step_labels:
239 step_member_union.update(pipeline_graph.task_subsets[step_label])
240 self.assertSetEqual(step_member_union, set(stage_members))
242 def test_comcam_correspondence(self) -> None:
243 butler = self.make_butler(writeable=True)
244 self.register_refcat(butler, COMCAM_REFCAT)
245 correspondence_filename = os.path.join(PIPELINES_DIR, "LSSTComCam/DRP-v2.yaml")
246 new_pipeline_graph = Pipeline.from_uri(correspondence_filename).to_graph(
247 registry=butler.registry
248 )
249 old_steps = [
250 "step1",
251 "step2a",
252 "step2b",
253 "step2c",
254 "step2d",
255 "step2e",
256 "step3a",
257 "step3b",
258 "step4",
259 "step5",
260 "step6",
261 "step7",
262 ]
263 old_pipeline_graph = Pipeline.from_uri(
264 os.path.join(PIPELINES_DIR, f"LSSTComCam/DRP.yaml#{','.join(old_steps)}")
265 ).to_graph(registry=butler.registry)
266 ignore_edges = {
267 # Use recalibrated_star instead of single_visit_star for calib
268 # flag propagation.
269 ("measureObjectUnforced", "sourceTableHandles"),
270 # Use deep_coadd instead of deep_coadd_preliminary for WCS.
271 ("measureObjectForced", "refWcs"),
272 }
273 correspondence_filename = os.path.join(
274 TEST_DIR, "migration_data", "LSSTComCam", "DRP-v2.json"
275 )
276 correspondence = Correspondence.read(correspondence_filename)
277 if correspondence.check(
278 new_pipeline_graph,
279 old_pipeline_graph,
280 "LSSTComCam/DRP-v2",
281 "LSSTComCam/DRP",
282 ignore_edges=ignore_edges,
283 ):
284 new_correspondence = correspondence.find_matches(
285 new_pipeline_graph, old_pipeline_graph
286 ).sorted(new_pipeline_graph, old_pipeline_graph)
287 new_correspondence_filename = (
288 f"{os.path.splitext(correspondence_filename)[0]}.new.json"
289 )
290 new_correspondence.write(new_correspondence_filename)
291 print(
292 "An attempt to fix the correspondence file has "
293 f"been written to {new_correspondence_filename!r}."
294 )
295 if messages := new_correspondence.check(
296 new_pipeline_graph,
297 old_pipeline_graph,
298 "LSSTComCam/DRP-v2",
299 "LSSTComCam/DRP",
300 ignore_edges=ignore_edges,
301 ):
302 print("Some problems remain that require manual edits:")
303 for message in messages:
304 print(message)
305 else:
306 print("This file should be reviewed and moved to replace the original.")
307 raise AssertionError(
308 f"Pipeline correspondence file {correspondence_filename!r} is out of date."
309 )
310 for new_label, old_label in correspondence.tasks_new_to_old.items():
311 new_node = new_pipeline_graph.tasks[new_label]
312 old_node = old_pipeline_graph.tasks[old_label]
313 task_config_diff = correspondence.diff_task_configs(new_node, old_node)
314 if task_config_diff:
315 print(task_config_diff)
316 raise AssertionError(
317 f"Differences detected in configs for task {new_label!r} (previously {old_label!r})."
318 )
320 def test_comcam_compat(self) -> None:
321 """Test the LSSTComCam/DRP-v2-compat pipeline, which should be
322 identical to DRP-v2 aside from replacing 'source' -> 'source2'
323 throughout.
324 """
325 butler = self.make_butler(writeable=True)
326 self.register_refcat(butler, COMCAM_REFCAT)
327 pipeline = Pipeline.from_uri(
328 os.path.join(PIPELINES_DIR, "LSSTComCam/DRP-v2.yaml")
329 )
330 pipeline_graph = pipeline.to_graph(registry=butler.registry)
331 pipeline_compat = Pipeline.from_uri(
332 os.path.join(PIPELINES_DIR, "LSSTComCam/DRP-v2-compat.yaml")
333 )
334 pipeline_graph_compat = pipeline_compat.to_graph(registry=butler.registry)
335 self.assertFalse(
336 Correspondence(
337 dataset_types_new_to_old={"source2": "source"},
338 ).find_matches(pipeline_graph_compat, pipeline_graph).check(
339 pipeline_graph_compat, pipeline_graph, "DRP-v2-compat", "DRP-v2"
340 )
341 )
343 def test_comcam_release_id_parameter(self) -> None:
344 """Test that changing the release_id parameter affects all appropriate
345 config options in the pipeline.
346 """
347 pipeline = Pipeline.fromString(
348 """
349 description: test pipeline
350 imports:
351 - $DRP_PIPE_DIR/pipelines/LSSTComCam/DRP-v2.yaml
352 parameters:
353 release_id: 52
354 """
355 )
356 pipeline_graph = pipeline.to_graph()
357 regex = re.compile(r"config\.(.+)\.release_id\=")
358 release_id_options = []
359 for task_node in pipeline_graph.tasks.values():
360 for match_string in regex.findall(task_node.get_config_str()):
361 attribute_path: list[str] = match_string.split(".")
362 if any(not term.isidentifier() for term in attribute_path):
363 _LOG.warning(
364 f"Not checking config option {task_node.label}:{match_string}.release_id "
365 "because the test code is not sophisticated enough. Please improve it."
366 )
367 continue
368 config_attribute = task_node.config
369 for term in attribute_path:
370 config_attribute = getattr(config_attribute, term)
371 self.assertEqual(config_attribute.release_id, 52)
372 release_id_options.append((task_node.label, attribute_path))
373 # Spot check a few expected entries to make sure a logic bug or bad
374 # regex isn't preventing this test from doing anything.
375 self.assertIn(("calibrateImage", ["id_generator"]), release_id_options)
376 self.assertIn(("detectCoaddPeaks", ["idGenerator"]), release_id_options)
379if __name__ == "__main__": 379 ↛ 380line 379 didn't jump to line 380 because the condition on line 379 was never true
380 unittest.main()