Coverage for python/lsst/dax/images/cutout/projection_finders.py: 28%
102 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-06-25 08:42 +0000
1# This file is part of dax_images_cutout.
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 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 <https://www.gnu.org/licenses/>.
22from __future__ import annotations
24__all__ = (
25 "Chain",
26 "MatchDatasetTypeName",
27 "ProjectionFinder",
28 "ReadComponents",
29 "ReadComponentsAstropyFits",
30 "TryComponentParents",
31 "UseSkyMap",
32)
34import logging
35import re
36from abc import ABC, abstractmethod
37from collections.abc import Iterable
38from typing import TYPE_CHECKING, cast
40import astropy.io.fits
41import astropy.units as u
43from lsst.daf.butler import Butler, DatasetRef
44from lsst.images import Box, GeneralFrame, SkyProjection
45from lsst.utils.timer import time_this
47from ._fits_projection import projection_and_bbox_from_fits_header
49if TYPE_CHECKING:
50 # lsst.skymap ships only with the Science Pipelines, not on PyPI; the
51 # mypy config ignores it when missing, and it is never imported at runtime.
52 from lsst.skymap import BaseSkyMap
54# Pixel coordinate frame for projections built from afw SkyWcs objects.
55_PIXEL_FRAME = GeneralFrame(unit=u.pix)
57_LOG = logging.getLogger(__name__)
58_TIMER_LOG_LEVEL = logging.INFO
61class ProjectionFinder(ABC):
62 """An interface for objects that can find the WCS and bounding box of a
63 butler dataset.
65 Notes
66 -----
67 Concrete `ProjectionFinder` implementations are intended to be composed to
68 define rules for how to find this projection information for a particular
69 dataset type; a finder that cannot handle a particular `DatasetRef` should
70 implement `find_projection` to return `None`, and implementations that
71 compose (e.g. `Chain`) can use this to decide when to try another nested
72 finder.
73 """
75 @abstractmethod
76 def find_projection(
77 self, ref: DatasetRef, butler: Butler, logger: logging.Logger | None = None
78 ) -> tuple[SkyProjection, Box] | None:
79 """Run the finder on the given dataset with the given butler.
81 Parameters
82 ----------
83 ref : `DatasetRef`
84 Fully-resolved reference to the dataset.
85 butler : `Butler`
86 Butler client to use for reads. Need not support writes, and any
87 default search collections will be ignored.
88 logger : `logging.Logger`, optional
89 Logger to use for timing messages. If `None`, a default logger
90 will be used.
92 Returns
93 -------
94 projection : `lsst.images.SkyProjection` (only if result is not `None`)
95 Mapping from sky to pixel coordinates for this dataset.
96 bbox : `lsst.images.Box` (only if result is not `None`)
97 Bounding box of the image dataset (or an image closely associated
98 with the dataset) in pixel coordinates.
99 """
100 raise NotImplementedError()
102 def __call__(
103 self, ref: DatasetRef, butler: Butler, logger: logging.Logger | None = None
104 ) -> tuple[SkyProjection, Box]:
105 """Call `find_projection` but raise `LookupError` when no projection
106 information is found.
108 Parameters
109 ----------
110 ref : `DatasetRef`
111 Fully-resolved reference to the dataset.
112 butler : `Butler`
113 Butler client to use for reads. Need not support writes, and any
114 default search collections will be ignored.
115 logger : `logging.Logger`, optional
116 Logger to use for timing messages. If `None`, a default logger
117 will be used.
119 Returns
120 -------
121 projection : `lsst.images.SkyProjection`
122 Mapping from sky to pixel coordinates for this dataset.
123 bbox : `lsst.images.Box`
124 Bounding box of the image dataset (or an image closely associated
125 with the dataset) in pixel coordinates.
126 """
127 result = self.find_projection(ref, butler, logger=logger)
128 if result is None:
129 raise LookupError(f"No way to obtain WCS and bounding box information for ref {ref}.")
130 return result
132 @staticmethod
133 def make_default() -> ProjectionFinder:
134 """Return a concrete finder appropriate for most pipelines.
136 Returns
137 -------
138 finder : `ProjectionFinder`
139 A finder that prefers to read, use, and cache a skymap when the
140 data ID includes tract or patch, and falls back to reading the WCS
141 and bbox from the dataset itself (or its parent, if the dataset is
142 a component).
143 """
144 return TryComponentParents(
145 Chain(
146 UseSkyMap(),
147 ReadComponentsAstropyFits(),
148 ReadComponents(),
149 )
150 )
153class ReadComponents(ProjectionFinder):
154 """A `ProjectionFinder` implementation that reads ``wcs`` and ``bbox``
155 from datasets that have them (e.g. ``Exposure``).
157 Notes
158 -----
159 This should usually be the final finder attempted in any chain; it's the
160 one most likely to work, but in many cases will not be the most efficient
161 or yield the most accurate WCS.
162 """
164 def find_projection(
165 self, ref: DatasetRef, butler: Butler, logger: logging.Logger | None = None
166 ) -> tuple[SkyProjection, Box] | None:
167 # Docstring inherited.
168 if {"wcs", "bbox"}.issubset(ref.datasetType.storageClass.allComponents().keys()):
169 logger = logger if logger is not None else _LOG
170 with time_this(_LOG, msg="Read projection info from butler components", level=_TIMER_LOG_LEVEL):
171 wcs = butler.get(ref.makeComponentRef("wcs"))
172 bbox = butler.get(ref.makeComponentRef("bbox"))
173 return SkyProjection.from_legacy(wcs, _PIXEL_FRAME), Box.from_legacy(bbox)
174 return None
177class ReadComponentsAstropyFits(ProjectionFinder):
178 """A `ProjectionFinder` implementation that reads ``wcs`` and ``bbox``
179 from datasets that have them (e.g. ``Exposure``) and assumes there is
180 a WCS associated with an IMAGE HDU in a FITS file.
182 Notes
183 -----
184 This might be more efficient for remote datasets where a full file download
185 is needed for butler to work using AFW.
186 """
188 def find_projection(
189 self, ref: DatasetRef, butler: Butler, logger: logging.Logger | None = None
190 ) -> tuple[SkyProjection, Box] | None:
191 # Docstring inherited.
192 logger = logger if logger is not None else _LOG
193 with time_this(logger, msg="Read projection info using Astropy", level=_TIMER_LOG_LEVEL):
194 if {"wcs", "bbox"}.issubset(ref.datasetType.storageClass.allComponents().keys()):
195 try:
196 fs, fspath = butler.getURI(ref).to_fsspec()
197 with (
198 fs.open(fspath) as f,
199 astropy.io.fits.open(f) as fits_obj,
200 ):
201 # Look for first pixel HDU.
202 pixel_components = {"mask", "image", "variance"}
203 for i, hdu in enumerate(fits_obj):
204 if i == 0:
205 # Assumes WCS is in the IMAGE extension
206 # and not stored in the primary.
207 continue
208 hdr = hdu.header
209 extname = hdr.get("EXTNAME")
210 if extname and extname.lower() in pixel_components:
211 return projection_and_bbox_from_fits_header(hdr, hdu.shape)
212 except Exception:
213 # Any failure and we will try the next option.
214 pass
215 return None
218class TryComponentParents(ProjectionFinder):
219 """A composite `ProjectionFinder` that walks from component dataset to its
220 parent composite until its nested finder succeeds.
222 Parameters
223 ----------
224 nested : `ProjectionFinder`
225 Nested finder to delegate to.
227 Notes
228 -----
229 This is a good choice for the outermost composite finder, so that the same
230 sequence of nested rules are applied to each level of the
231 component-composite tree.
232 """
234 def __init__(self, nested: ProjectionFinder):
235 self._nested = nested
237 def find_projection(
238 self, ref: DatasetRef, butler: Butler, logger: logging.Logger | None = None
239 ) -> tuple[SkyProjection, Box] | None:
240 # Docstring inherited.
241 while True:
242 if (result := self._nested.find_projection(ref, butler, logger=logger)) is not None:
243 return result
244 if ref.isComponent():
245 ref = ref.makeCompositeRef()
246 else:
247 return None
250class UseSkyMap(ProjectionFinder):
251 """A `ProjectionFinder` implementation that reads and caches
252 `lsst.skymap.BaseSkyMap` instances, allowing projections for coadds to be
253 found without requiring I/O for each one.
255 Parameters
256 ----------
257 dataset_type_name : `str`, optional
258 Name of the dataset type used to load `BaseSkyMap` instances.
259 collections : `Iterable` [ `str` ]
260 Collection search path for skymap datasets.
262 Notes
263 -----
264 This finder assumes any dataset with ``patch`` or ``tract`` dimensions
265 should get its WCS and bounding box from the skymap, and that datasets
266 without these dimensions never should (i.e. `find_projection` will return
267 `None` for these).
269 `BaseSkyMap` instances are never removed from the cache after being loaded;
270 we expect the number of distinct skymaps to be very small.
271 """
273 def __init__(self, dataset_type_name: str = "skyMap", collections: Iterable[str] = ("skymaps",)):
274 self._dataset_type_name = dataset_type_name
275 self._collections = tuple(collections)
276 self._cache: dict[str, BaseSkyMap] = {}
278 def find_projection(
279 self, ref: DatasetRef, butler: Butler, logger: logging.Logger | None = None
280 ) -> tuple[SkyProjection, Box] | None:
281 # Docstring inherited.
282 if "tract" in ref.dataId.dimensions:
283 assert "skymap" in ref.dataId.dimensions, "Guaranteed by expected dimension schema."
284 if (skymap := self._cache.get(cast(str, ref.dataId["skymap"]))) is None:
285 skymap = butler.get(
286 self._dataset_type_name, skymap=ref.dataId["skymap"], collections=self._collections
287 )
288 tractInfo = skymap[ref.dataId["tract"]]
289 if "patch" in ref.dataId.dimensions:
290 patchInfo = tractInfo[ref.dataId["patch"]]
291 return (
292 SkyProjection.from_legacy(patchInfo.wcs, _PIXEL_FRAME),
293 Box.from_legacy(patchInfo.outer_bbox),
294 )
295 else:
296 return (
297 SkyProjection.from_legacy(tractInfo.wcs, _PIXEL_FRAME),
298 Box.from_legacy(tractInfo.bbox),
299 )
300 return None
303class Chain(ProjectionFinder):
304 """A composite `ProjectionFinder` that attempts each finder in a sequence
305 until one succeeds.
307 Parameters
308 ----------
309 *nested : `ProjectionFinder`
310 Nested finders to delegate to, in order.
311 """
313 def __init__(self, *nested: ProjectionFinder):
314 self._nested = tuple(nested)
316 def find_projection(
317 self, ref: DatasetRef, butler: Butler, logger: logging.Logger | None = None
318 ) -> tuple[SkyProjection, Box] | None:
319 # Docstring inherited.
320 for f in self._nested:
321 if (result := f.find_projection(ref, butler, logger=logger)) is not None:
322 return result
323 return None
326class MatchDatasetTypeName(ProjectionFinder):
327 """A composite `ProjectionFinder` that delegates to different nested
328 finders based on whether the dataset type name matches a regular
329 expression.
331 Parameters
332 ----------
333 regex : `str`
334 Regular expression the dataset type name must match (in full).
335 on_match : `ProjectionFinder`, optional
336 Finder to try when the match succeeds, or `None` to return `None`.
337 otherwise : `ProjectionFinder`, optional
338 Finder to try when the match does not succeed, or `None` to return
339 `None`.
340 """
342 def __init__(
343 self,
344 regex: str,
345 on_match: ProjectionFinder | None = None,
346 otherwise: ProjectionFinder | None = None,
347 ):
348 self._regex = re.compile(regex)
349 self._on_match = on_match
350 self._otherwise = otherwise
352 def find_projection(
353 self, ref: DatasetRef, butler: Butler, logger: logging.Logger | None = None
354 ) -> tuple[SkyProjection, Box] | None:
355 # Docstring inherited.
356 if self._regex.match(ref.datasetType.name):
357 if self._on_match is not None:
358 return self._on_match.find_projection(ref, butler, logger=logger)
359 else:
360 if self._otherwise is not None:
361 return self._otherwise.find_projection(ref, butler, logger=logger)
362 return None