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