Coverage for python / lsst / obs / base / utils.py: 16%
159 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-24 01:14 -0800
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. Protect against abnormally
249 # large number of inputs.
250 prov_dict = provenance.to_flat_dict(
251 ref, prefix="LSST BUTLER", sep=" ", simple_types=True, max_inputs=3_000
252 )
254 # Copy keys into a PropertyList so we have the option of including
255 # comments.
256 extras = PropertyList()
258 for k, v in prov_dict.items():
259 if not allow_long_headers and isinstance(v, str):
260 _store_str_header(extras, k, v, allow_long_headers=allow_long_headers)
261 else:
262 extras.set(k, v)
263 # Only add comments if we think it's safe.
264 if allow_long_headers:
265 comment = ""
266 if re.search(r"\bRUN\b", k):
267 comment = "Run collection"
268 elif re.search(r"\bID\b", k):
269 comment = "Dataset ID"
270 elif re.search(r"\bDATAID\b", k):
271 comment = "Data identifier"
272 elif re.search(r"\bDATASETTYPE\b", k):
273 comment = "Dataset type"
274 if comment:
275 extras.setComment(k, comment)
277 # Purge old headers from metadata (important for data ID and input headers
278 # and to prevent headers accumulating in a PropertyList).
279 strip_provenance_from_fits_header(hdr)
281 # Update the header.
282 hdr.update(extras)
285def strip_provenance_from_fits_header(hdr: MutableMapping | PropertyList) -> None:
286 """Remove all FITS headers relating to butler provenance.
288 Parameters
289 ----------
290 hdr : `lsst.daf.base.PropertyList` or `collections.abc.MutableMapping`
291 The FITS header to modify. Assumes ``HIERARCH`` will be handled
292 implicitly by the writer.
294 Notes
295 -----
296 These headers will have been added by, for example, the butler formatter
297 via a call to `add_provenance_to_fits_header`.
298 """
299 for k in list(hdr):
300 if k.startswith(_PROVENANCE_PREFIX):
301 del hdr[k]
304def iso_date_to_curated_calib_file_root(valid_start: str) -> str:
305 """Parse an ISO time, potentially from a calibration valid start time,
306 and convert to form suitable for use in a curated calibration filename.
308 Parameters
309 ----------
310 valid_start : `str`
311 Validity start time in ISOT format.
313 Returns
314 -------
315 date_str : `str`
316 Form suitable for use in a curated calibration file name in a
317 data package: YYYYMMDDTHHMMSS.
318 """
319 # Parse the ISO string to ensure conformity if there is an edit to
320 # valid_start.
321 valid_date = datetime.datetime.fromisoformat(valid_start)
323 # Drop any fractional seconds.
324 return re.sub(r"\W", "", valid_date.isoformat(timespec="seconds"))
327class TableVStack:
328 """A helper class for stacking astropy tables without having them all in
329 memory at once.
331 Parameters
332 ----------
333 capacity : `int`
334 Full size of the final table.
336 Notes
337 -----
338 Unlike `astropy.table.vstack`, this class requires all tables to have the
339 exact same columns (it's slightly more strict than even the
340 ``join_type="exact"`` argument to `astropy.table.vstack`).
341 """
343 def __init__(self, capacity: int) -> None:
344 self.index: int = 0
345 self.capacity = capacity
346 self.result: astropy.table.Table | None = None
348 @classmethod
349 def set_extra_values(
350 cls,
351 table: astropy.table.Table,
352 key: str,
353 values: Any,
354 capacity: int,
355 slicer: slice | None = None,
356 validate: bool = True,
357 ) -> None:
358 """Set extra column values in a slice of a table.
360 Parameters
361 ----------
362 table : `astropy.table.Table`
363 The table to set values for.
364 key : `str`
365 The column key.
366 values : `Any`
367 The value(s) to set. Can be a scalar.
368 capacity : `int`
369 The size to initialize the column with, if it doesn't exist yet.
370 slicer : `slice` or `None`, optional
371 A slice to select values to update.
372 validate : `bool`, optional
373 If True and the column already exists, will raise if the new values
374 do not match the existing ones.
375 """
376 if key in table.colnames:
377 column = table[key]
378 if validate and not np.all((table[key] if (slicer is None) else table[key][slicer]) == values):
379 raise RuntimeError(
380 f"table already contains {column=} with {key=} but values differ from {values=}"
381 )
382 table[key][slicer] = values
383 else:
384 if slicer is None:
385 table[key] = values
386 else:
387 try:
388 dtype = values.dtype
389 except AttributeError:
390 dtype = np.dtype(type(values))
391 table[key] = np.empty(capacity, dtype=dtype)
392 table[key][slicer] = values
394 @classmethod
395 def from_handles(cls, handles: DeferredHandles) -> Self:
396 """Construct from an iterable of
397 `lsst.daf.butler.DeferredDatasetHandle`.
399 Parameters
400 ----------
401 handles : `~collections.abc.Iterable` [ \
402 `lsst.daf.butler.DeferredDatasetHandle` ]
403 Iterable of handles. Must have a storage class that supports the
404 "rowcount" component, which is all that will be fetched.
406 Returns
407 -------
408 vstack : `TableVStack`
409 An instance of this class, initialized with capacity equal to the
410 sum of the rowcounts of all the given table handles.
411 """
412 capacity = sum(handle.get(component="rowcount") for handle in handles)
413 return cls(capacity=capacity)
415 def extend(self, table: astropy.table.Table, extra_values: dict[str, ArrayLike] | None = None) -> None:
416 """Add a single table to the stack.
418 Parameters
419 ----------
420 table : `astropy.table.Table`
421 An astropy table instance.
422 extra_values : `dict`
423 Dict keyed by column name with an array-like of values to set
424 for this table only.
425 """
426 if extra_values is None:
427 extra_values = {}
428 if self.result is None:
429 self.result = astropy.table.Table()
430 slicer = slice(None, len(table))
431 for name in table.colnames:
432 column = table[name]
433 column_cls = type(column)
434 self.result[name] = column_cls.info.new_like([column], self.capacity, name=name)
435 self.result[name][: len(table)] = column
436 for name, values in extra_values.items():
437 self.set_extra_values(
438 table=self.result,
439 key=name,
440 values=values,
441 capacity=self.capacity,
442 slicer=slicer,
443 )
444 self.index = len(table)
445 self.result.meta = table.meta.copy()
446 else:
447 next_index = self.index + len(table)
448 slicer = slice(self.index, next_index)
449 for name in table.colnames:
450 out_col = self.result[name]
451 in_col = table[name]
452 if out_col.dtype != in_col.dtype:
453 raise TypeError(f"Type mismatch on column {name!r}: {out_col.dtype} != {in_col.dtype}.")
454 self.result[name][slicer] = table[name]
455 for name, values in extra_values.items():
456 self.set_extra_values(
457 table=self.result,
458 key=name,
459 values=values,
460 capacity=self.capacity,
461 slicer=slicer,
462 validate=False,
463 )
464 self.index = next_index
465 # Butler provenance should be stripped on merge. It will be
466 # added by butler on write. No attempt is made here to combine
467 # provenance from multiple input tables.
468 self.result.meta = merge_headers([self.result.meta, table.meta], mode="drop")
469 strip_provenance_from_fits_header(self.result.meta)
471 @classmethod
472 def vstack_handles(
473 cls,
474 handles: DeferredHandles,
475 extra_values: dict[int, dict[str, ArrayLike]] | None = None,
476 kwargs_get: dict[str, Any] | None = None,
477 ) -> astropy.table.Table:
478 """Vertically stack tables represented by deferred dataset handles.
480 Parameters
481 ----------
482 handles : `~collections.abc.Iterable` [ \
483 `lsst.daf.butler.DeferredDatasetHandle` ]
484 Iterable of handles. Must have the "ArrowAstropy" storage class
485 and identical columns.
486 extra_values : `dict`
487 Dictionary keyed by index of handle of additional values
488 to pass to extend.
489 kwargs_get : `dict`[`str`, `Any`]
490 Keyword argument-value pairs to pass to handle.get().
492 Returns
493 -------
494 table : `astropy.table.Table`
495 Concatenated table with the same columns as each input table and
496 the rows of all of them, or an empty table if there are no handles.
497 """
498 if not handles:
499 return astropy.table.Table()
500 if extra_values is None:
501 extra_values = {}
502 if kwargs_get is None:
503 kwargs_get = {}
504 handles = tuple(handles) # guard against single-pass iterators
505 # Ensure that zero length catalogs are not included
506 rowcounts = tuple(handle.get(component="rowcount") for handle in handles)
507 handles = tuple(handle for handle, count in zip(handles, rowcounts) if count > 0)
509 vstack = cls(capacity=np.sum(rowcounts))
510 for idx, handle in enumerate(handles):
511 vstack.extend(handle.get(**kwargs_get), extra_values=extra_values.get(idx))
512 return astropy.table.Table() if vstack.result is None else vstack.result