Coverage for python/lsst/obs/base/utils.py: 16%
159 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-10-14 01:46 -0700
1# This file is part of obs_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 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/>.
22__all__ = (
23 "InitialSkyWcsError",
24 "TableVStack",
25 "add_provenance_to_fits_header",
26 "bboxFromIraf",
27 "createInitialSkyWcs",
28 "createInitialSkyWcsFromBoresight",
29 "iso_date_to_curated_calib_file_root",
30 "strip_provenance_from_fits_header",
31)
33import datetime
34import re
35from collections.abc import Iterable, MutableMapping
36from typing import Any, Self
38import astropy.table
39import numpy as np
40from astro_metadata_translator.headers import merge_headers
41from numpy.typing import ArrayLike
43import lsst.geom as geom
44import lsst.pex.exceptions
45from lsst.afw.cameraGeom import FIELD_ANGLE, PIXELS
46from lsst.afw.geom.skyWcs import makeSkyWcs
47from lsst.afw.image import RotType
48from lsst.daf.base import PropertyList
49from lsst.daf.butler import DatasetProvenance, DatasetRef
51_PROVENANCE_PREFIX = "LSST BUTLER"
52"""The prefix used in FITS headers for butler provenance."""
54DeferredHandles = Iterable[lsst.daf.butler.DeferredDatasetHandle]
57class InitialSkyWcsError(Exception):
58 """For handling failures when creating a SkyWcs from a camera geometry and
59 boresight.
61 Typically used as a chained exception from a lower level exception.
62 """
64 pass
67def createInitialSkyWcs(visitInfo, detector, flipX=None):
68 """Create a SkyWcs from the visit information and detector geometry.
70 A typical use case for this is to create the initial WCS for a newly-read
71 raw exposure.
74 Parameters
75 ----------
76 visitInfo : `lsst.afw.image.VisitInfo`
77 Where to get the telescope boresight and rotator angle from.
78 detector : `lsst.afw.cameraGeom.Detector`
79 Where to get the camera geometry from.
80 flipX : `bool` or `None`, optional
81 If `False`, +X is along W, if `True` +X is along E. If `None`, the
82 focal plane parity in the camera geometry is assumed to be correct,
83 which may not be true for old camera datasets.
85 Returns
86 -------
87 skyWcs : `lsst.afw.geom.SkyWcs`
88 The new composed WCS.
90 Raises
91 ------
92 InitialSkyWcsError
93 Raised if there is an error generating the SkyWcs, chained from the
94 lower-level exception if available.
95 """
96 if visitInfo.getRotType() != RotType.SKY:
97 msg = (
98 "Cannot create SkyWcs from camera geometry: rotator angle defined using "
99 f"RotType={visitInfo.getRotType()} instead of SKY."
100 )
101 raise InitialSkyWcsError(msg)
102 orientation = visitInfo.getBoresightRotAngle()
103 boresight = visitInfo.getBoresightRaDec()
104 return createInitialSkyWcsFromBoresight(boresight, orientation, detector, flipX)
107def createInitialSkyWcsFromBoresight(boresight, orientation, detector, flipX=None):
108 """Create a SkyWcs from the telescope boresight and detector geometry.
110 A typical usecase for this is to create the initial WCS for a newly-read
111 raw exposure.
113 Parameters
114 ----------
115 boresight : `lsst.geom.SpherePoint`
116 The ICRS boresight RA/Dec
117 orientation : `lsst.geom.Angle`
118 The rotation angle of the focal plane on the sky.
119 detector : `lsst.afw.cameraGeom.Detector`
120 Where to get the camera geometry from.
121 flipX : `bool` or `None`, optional
122 If `False`, +X is along W, if `True` +X is along E. If `None`, the
123 focal plane parity in the camera geometry is assumed to be correct,
124 which may not be true for old camera datasets.
126 Returns
127 -------
128 skyWcs : `lsst.afw.geom.SkyWcs`
129 The new composed WCS.
131 Raises
132 ------
133 InitialSkyWcsError
134 Raised if there is an error generating the SkyWcs, chained from the
135 lower-level exception if available.
136 """
137 camera_parity = detector.getTransformMap().getFocalPlaneParity()
138 actual_flip_x = False
139 if flipX and not camera_parity:
140 # This is probably an old camera definition from before those could
141 # hold parity flips. We need to work around this since data
142 # repositories aren't necessarily updated in sync with code.
143 actual_flip_x = True
144 elif flipX is not None and not flipX and camera_parity:
145 raise InitialSkyWcsError(
146 "Camera geometry reports a parity flip between FOCAL_PLANE and FIELD_ANGLE, but this "
147 "was not expected by caller for this instrument."
148 )
149 try:
150 pixelsToFieldAngle = detector.getTransform(
151 detector.makeCameraSys(PIXELS), detector.makeCameraSys(FIELD_ANGLE)
152 )
153 except lsst.pex.exceptions.InvalidParameterError as e:
154 raise InitialSkyWcsError("Cannot compute PIXELS to FIELD_ANGLE Transform.") from e
155 return makeSkyWcs(pixelsToFieldAngle, orientation, actual_flip_x, boresight)
158def bboxFromIraf(irafBBoxStr):
159 """Return a Box2I corresponding to an IRAF-style BBOX.
161 [x0:x1,y0:y1] where x0 and x1 are the one-indexed start and end columns,
162 and correspondingly y0 and y1 are the start and end rows.
163 """
164 mat = re.search(r"^\[([-\d]+):([-\d]+),([-\d]+):([-\d]+)\]$", irafBBoxStr)
165 if not mat:
166 raise RuntimeError(f'Unable to parse IRAF-style bbox "{irafBBoxStr}"')
167 x0, x1, y0, y1 = (int(_) for _ in mat.groups())
169 return geom.BoxI(geom.PointI(x0 - 1, y0 - 1), geom.PointI(x1 - 1, y1 - 1))
172def _store_str_header(
173 hdr: PropertyList, key: str, value: str, comment: str | None = None, allow_long_headers: bool = True
174) -> None:
175 """Examine string header and return value that can be used.
177 Parameters
178 ----------
179 key : `str`
180 The key to use in the FITS header (without HIERARCH).
181 value : `str`
182 The value that is to be stored in the header.
183 allow_long_headers : `bool`, optional
184 If `True` the value will be used unchanged. If `False` the value
185 could have some content elided and a modified version used that will
186 fit in a FITS header. If `False` the maximum length of ``key`` is
187 58 characters to prevent inconsistency where a very short string will
188 fit in 80 characters with ``HIERARCH`` but a very long string will
189 not.
190 """
191 if not allow_long_headers:
192 # Declare that we do not allow keys longer than a fixed number of
193 # characters in this mode. This is enough to allow
194 # HIERARCH {key} = 'abc '
195 # to fit into 80 characters.
196 max_card_length = 80
197 if len(key) > 58:
198 raise ValueError(
199 f"Given keyword {key} is too long when requiring {max_card_length}-character header cards"
200 )
202 # FITS pads strings to 8 characters.
203 min_len = 8
204 formatted = f"HIERARCH {key} = '{value:{min_len}s}'"
205 if (n_over := len(formatted) - max_card_length) > 0:
206 # Elide some content.
207 # Remove the over run + 3 for the ... and 1 for rounding.
208 n_remove = n_over + 1 + 3
209 middle = len(value) // 2
210 half_remove = n_remove // 2
211 value = f"{value[: middle - half_remove]}...{value[middle + half_remove :]}"
213 # Do not forward comment if we have elided.
214 comment = None
216 hdr.set(key, value, comment)
219def add_provenance_to_fits_header(
220 hdr: PropertyList | MutableMapping | None,
221 ref: DatasetRef,
222 provenance: DatasetProvenance | None = None,
223 allow_long_headers: bool = True,
224) -> None:
225 """Modify the given header to include provenance headers.
227 Parameters
228 ----------
229 hdr : `lsst.daf.base.PropertyList` or `collections.abc.MutableMapping`
230 The FITS header to modify. Assumes ``HIERARCH`` will be handled
231 implicitly by the writer.
232 ref : `lsst.daf.butler.DatasetRef`
233 The butler dataset associated with this FITS file.
234 provenance : `lsst.daf.butler.DatasetProvenance` or `None`, optional
235 Provenance for this dataset.
236 allow_long_headers : `bool`, optional
237 If `True` it is assumed that there is no limit to the length of the
238 values being stored. If `False`, assumes that FITS header cards must be
239 kept within 80 characters including quoting, ``HIERARCH``, and `=`.
240 """
241 # Some datasets do not define a header.
242 if hdr is None:
243 return
245 if provenance is None:
246 provenance = DatasetProvenance()
248 # Get the flat dict in form suitable for FITS.
249 prov_dict = provenance.to_flat_dict(ref, prefix="LSST BUTLER", sep=" ", simple_types=True)
251 # Copy keys into a PropertyList so we have the option of including
252 # comments.
253 extras = PropertyList()
255 for k, v in prov_dict.items():
256 if not allow_long_headers and isinstance(v, str):
257 _store_str_header(extras, k, v, allow_long_headers=allow_long_headers)
258 else:
259 extras.set(k, v)
260 # Only add comments if we think it's safe.
261 if allow_long_headers:
262 comment = ""
263 if re.search(r"\bRUN\b", k):
264 comment = "Run collection"
265 elif re.search(r"\bID\b", k):
266 comment = "Dataset ID"
267 elif re.search(r"\bDATAID\b", k):
268 comment = "Data identifier"
269 elif re.search(r"\bDATASETTYPE\b", k):
270 comment = "Dataset type"
271 if comment:
272 extras.setComment(k, comment)
274 # Purge old headers from metadata (important for data ID and input headers
275 # and to prevent headers accumulating in a PropertyList).
276 strip_provenance_from_fits_header(hdr)
278 # Update the header.
279 hdr.update(extras)
282def strip_provenance_from_fits_header(hdr: MutableMapping | PropertyList) -> None:
283 """Remove all FITS headers relating to butler provenance.
285 Parameters
286 ----------
287 hdr : `lsst.daf.base.PropertyList` or `collections.abc.MutableMapping`
288 The FITS header to modify. Assumes ``HIERARCH`` will be handled
289 implicitly by the writer.
291 Notes
292 -----
293 These headers will have been added by, for example, the butler formatter
294 via a call to `add_provenance_to_fits_header`.
295 """
296 for k in list(hdr):
297 if k.startswith(_PROVENANCE_PREFIX):
298 del hdr[k]
301def iso_date_to_curated_calib_file_root(valid_start: str) -> str:
302 """Parse an ISO time, potentially from a calibration valid start time,
303 and convert to form suitable for use in a curated calibration filename.
305 Parameters
306 ----------
307 valid_start : `str`
308 Validity start time in ISOT format.
310 Returns
311 -------
312 date_str : `str`
313 Form suitable for use in a curated calibration file name in a
314 data package: YYYYMMDDTHHMMSS.
315 """
316 # Parse the ISO string to ensure conformity if there is an edit to
317 # valid_start.
318 valid_date = datetime.datetime.fromisoformat(valid_start)
320 # Drop any fractional seconds.
321 return re.sub(r"\W", "", valid_date.isoformat(timespec="seconds"))
324class TableVStack:
325 """A helper class for stacking astropy tables without having them all in
326 memory at once.
328 Parameters
329 ----------
330 capacity : `int`
331 Full size of the final table.
333 Notes
334 -----
335 Unlike `astropy.table.vstack`, this class requires all tables to have the
336 exact same columns (it's slightly more strict than even the
337 ``join_type="exact"`` argument to `astropy.table.vstack`).
338 """
340 def __init__(self, capacity: int) -> None:
341 self.index: int = 0
342 self.capacity = capacity
343 self.result: astropy.table.Table | None = None
345 @classmethod
346 def set_extra_values(
347 cls,
348 table: astropy.table.Table,
349 key: str,
350 values: Any,
351 capacity: int,
352 slicer: slice | None = None,
353 validate: bool = True,
354 ) -> None:
355 """Set extra column values in a slice of a table.
357 Parameters
358 ----------
359 table : `astropy.table.Table`
360 The table to set values for.
361 key : `str`
362 The column key.
363 values : `Any`
364 The value(s) to set. Can be a scalar.
365 capacity : `int`
366 The size to initialize the column with, if it doesn't exist yet.
367 slicer : `slice` or `None`, optional
368 A slice to select values to update.
369 validate : `bool`, optional
370 If True and the column already exists, will raise if the new values
371 do not match the existing ones.
372 """
373 if key in table.colnames:
374 column = table[key]
375 if validate and not np.all((table[key] if (slicer is None) else table[key][slicer]) == values):
376 raise RuntimeError(
377 f"table already contains {column=} with {key=} but values differ from {values=}"
378 )
379 table[key][slicer] = values
380 else:
381 if slicer is None:
382 table[key] = values
383 else:
384 try:
385 dtype = values.dtype
386 except AttributeError:
387 dtype = np.dtype(type(values))
388 table[key] = np.empty(capacity, dtype=dtype)
389 table[key][slicer] = values
391 @classmethod
392 def from_handles(cls, handles: DeferredHandles) -> Self:
393 """Construct from an iterable of
394 `lsst.daf.butler.DeferredDatasetHandle`.
396 Parameters
397 ----------
398 handles : `~collections.abc.Iterable` [ \
399 `lsst.daf.butler.DeferredDatasetHandle` ]
400 Iterable of handles. Must have a storage class that supports the
401 "rowcount" component, which is all that will be fetched.
403 Returns
404 -------
405 vstack : `TableVStack`
406 An instance of this class, initialized with capacity equal to the
407 sum of the rowcounts of all the given table handles.
408 """
409 capacity = sum(handle.get(component="rowcount") for handle in handles)
410 return cls(capacity=capacity)
412 def extend(self, table: astropy.table.Table, extra_values: dict[str, ArrayLike] | None = None) -> None:
413 """Add a single table to the stack.
415 Parameters
416 ----------
417 table : `astropy.table.Table`
418 An astropy table instance.
419 extra_values : `dict`
420 Dict keyed by column name with an array-like of values to set
421 for this table only.
422 """
423 if extra_values is None:
424 extra_values = {}
425 if self.result is None:
426 self.result = astropy.table.Table()
427 slicer = slice(None, len(table))
428 for name in table.colnames:
429 column = table[name]
430 column_cls = type(column)
431 self.result[name] = column_cls.info.new_like([column], self.capacity, name=name)
432 self.result[name][: len(table)] = column
433 for name, values in extra_values.items():
434 self.set_extra_values(
435 table=self.result,
436 key=name,
437 values=values,
438 capacity=self.capacity,
439 slicer=slicer,
440 )
441 self.index = len(table)
442 self.result.meta = table.meta.copy()
443 else:
444 next_index = self.index + len(table)
445 slicer = slice(self.index, next_index)
446 for name in table.colnames:
447 out_col = self.result[name]
448 in_col = table[name]
449 if out_col.dtype != in_col.dtype:
450 raise TypeError(f"Type mismatch on column {name!r}: {out_col.dtype} != {in_col.dtype}.")
451 self.result[name][slicer] = table[name]
452 for name, values in extra_values.items():
453 self.set_extra_values(
454 table=self.result,
455 key=name,
456 values=values,
457 capacity=self.capacity,
458 slicer=slicer,
459 validate=False,
460 )
461 self.index = next_index
462 # Butler provenance should be stripped on merge. It will be
463 # added by butler on write. No attempt is made here to combine
464 # provenance from multiple input tables.
465 self.result.meta = merge_headers([self.result.meta, table.meta], mode="drop")
466 strip_provenance_from_fits_header(self.result.meta)
468 @classmethod
469 def vstack_handles(
470 cls,
471 handles: DeferredHandles,
472 extra_values: dict[int, dict[str, ArrayLike]] | None = None,
473 kwargs_get: dict[str, Any] | None = None,
474 ) -> astropy.table.Table:
475 """Vertically stack tables represented by deferred dataset handles.
477 Parameters
478 ----------
479 handles : `~collections.abc.Iterable` [ \
480 `lsst.daf.butler.DeferredDatasetHandle` ]
481 Iterable of handles. Must have the "ArrowAstropy" storage class
482 and identical columns.
483 extra_values : `dict`
484 Dictionary keyed by index of handle of additional values
485 to pass to extend.
486 kwargs_get : `dict`[`str`, `Any`]
487 Keyword argument-value pairs to pass to handle.get().
489 Returns
490 -------
491 table : `astropy.table.Table`
492 Concatenated table with the same columns as each input table and
493 the rows of all of them, or an empty table if there are no handles.
494 """
495 if not handles:
496 return astropy.table.Table()
497 if extra_values is None:
498 extra_values = {}
499 if kwargs_get is None:
500 kwargs_get = {}
501 handles = tuple(handles) # guard against single-pass iterators
502 # Ensure that zero length catalogs are not included
503 rowcounts = tuple(handle.get(component="rowcount") for handle in handles)
504 handles = tuple(handle for handle, count in zip(handles, rowcounts) if count > 0)
506 vstack = cls(capacity=np.sum(rowcounts))
507 for idx, handle in enumerate(handles):
508 vstack.extend(handle.get(**kwargs_get), extra_values=extra_values.get(idx))
509 return astropy.table.Table() if vstack.result is None else vstack.result