Coverage for python/lsst/images/cells/_provenance.py: 44%
164 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-18 10:45 -0700
1# This file is part of lsst-images.
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# Use of this source code is governed by a 3-clause BSD-style
10# license that can be found in the LICENSE file.
12from __future__ import annotations
14__all__ = ("CoaddProvenance", "CoaddProvenanceSerializationModel")
16from collections.abc import Iterable
17from typing import TYPE_CHECKING, Any, ClassVar
19import astropy.table
20import astropy.units as u
21import numpy as np
22import pydantic
24from .._cell_grid import CellIJ
25from .._polygon import Polygon
26from ..serialization import ArchiveTree, InputArchive, InvalidParameterError, OutputArchive, TableModel
28if TYPE_CHECKING:
29 try:
30 from lsst.afw.geom import Polygon as LegacyPolygon
31 from lsst.cell_coadds import CoaddInputs as LegacyCellCoaddInputs
32 from lsst.cell_coadds import MultipleCellCoadd as LegacyMultipleCellCoadd
33 from lsst.cell_coadds import ObservationIdentifiers as LegacyObservationIdentifiers
34 from lsst.skymap import Index2D as LegacyIndex2D
35 except ImportError:
36 type LegacyIndex2D = Any # type: ignore[no-redef]
37 type LegacyCellCoaddInputs = Any # type: ignore[no-redef]
38 type LegacyPolygon = Any # type: ignore[no-redef]
39 type LegacyMultipleCellCoadd = Any # type: ignore[no-redef]
40 type LegacyObservationIdentifiers = Any # type: ignore[no-redef]
43class CoaddProvenance:
44 """A pair of tables that record the inputs to a cell-based coadd.
46 Parameters
47 ----------
48 inputs
49 A table of {visit, detector} combinations that contribute to any cell
50 in the coadd.
51 contributions
52 A table of {visit, detector, cell} combinations that describe how an
53 observation contributed to a cell.
55 Notes
56 -----
57 This object can represent the provenance of a whole patch, a single cell,
58 or anything in between. In the single-cell case, the ``inputs`` and
59 ``contributions`` tables have the same number of rows (but may not be
60 ordered the same way!).
61 """
63 def __init__(self, inputs: astropy.table.Table, contributions: astropy.table.Table):
64 self._inputs = inputs
65 self._contributions = contributions
67 _INPUT_TABLE_COLUMNS: ClassVar[list[tuple[str, type, str]]] = [
68 ("instrument", np.object_, "Name of the instrument."),
69 ("visit", np.uint64, "ID of the visit."),
70 ("detector", np.uint16, "ID of the detector."),
71 ("physical_filter", np.object_, "Full name of the bandpass filter."),
72 ("day_obs", np.uint32, "Observation night as a YYYYMMDD integer."),
73 (
74 "polygon",
75 np.object_,
76 (
77 "Polygon that approximates the overlap of the observation and the coadd patch, "
78 "in coadd coordinates."
79 ),
80 ),
81 ]
83 _CONTRIBUTION_TABLE_COLUMNS: ClassVar[list[tuple[str, type, str, u.UnitBase | None]]] = [
84 ("cell_i", np.uint16, "Y-axis index of the cell within the patch.", None),
85 ("cell_j", np.uint16, "X-axis index of the cell within the patch.", None),
86 ("instrument", np.object_, "Name of the instrument.", None),
87 ("visit", np.uint64, "ID of the visit.", None),
88 ("detector", np.uint16, "ID of the detector.", None),
89 ("overlaps_center", np.bool_, "Whether a this observation overlaps the center of the cell.", None),
90 ("overlap_fraction", np.float64, "Fraction of the cell that is covered by the overlap region.", None),
91 ("unmasked_fraction", np.float64, "Fraction of the cell propagated to the coadd.", None),
92 ("weight", np.float64, "Weight to be used for this input in this cell.", None),
93 ("psf_shape_xx", np.float64, "Second order moments of the PSF.", u.pix**2),
94 ("psf_shape_yy", np.float64, "Second order moments of the PSF.", u.pix**2),
95 ("psf_shape_xy", np.float64, "Second order moments of the PSF.", u.pix**2),
96 (
97 "psf_shape_flag",
98 np.bool_,
99 "Flag indicating whether the PSF shape measurement was successful.",
100 None,
101 ),
102 ]
104 @classmethod
105 def make_empty_input_table(cls, n_rows: int) -> astropy.table.Table:
106 """Make an empty `inputs` table with a set number of rows."""
107 return astropy.table.Table(
108 [
109 astropy.table.Column(name=name, length=n_rows, dtype=dtype, description=description)
110 for name, dtype, description in cls._INPUT_TABLE_COLUMNS
111 ]
112 )
114 @classmethod
115 def make_empty_contribution_table(cls, n_rows: int) -> astropy.table.Table:
116 """Make an empty `contributions` table with a set number of rows."""
117 return astropy.table.Table(
118 [
119 astropy.table.Column(
120 name=name, length=n_rows, dtype=dtype, description=description, unit=unit
121 )
122 for name, dtype, description, unit in cls._CONTRIBUTION_TABLE_COLUMNS
123 ]
124 )
126 @property
127 def inputs(self) -> astropy.table.Table:
128 """A table of {visit, detector} combinations that contribute to any
129 cell in the coadd.
130 """
131 return self._inputs
133 @property
134 def contributions(self) -> astropy.table.Table:
135 """A table of {visit, detector, cell} combinations that describe how an
136 observation contributed to a cell.
137 """
138 return self._contributions
140 def __getitem__(self, cell: CellIJ) -> CoaddProvenance:
141 return self.subset([cell])
143 def subset(self, cells: Iterable[CellIJ]) -> CoaddProvenance:
144 """Return a new provenance object with just the given cells."""
145 cells_to_keep = astropy.table.Table(
146 rows=[(index.i, index.j) for index in cells],
147 names=["cell_i", "cell_j"],
148 dtype=[np.uint16, np.uint16],
149 )
150 contributions = astropy.table.join(self._contributions, cells_to_keep)
151 assert contributions.columns.keys() == {name for name, _, _, _ in self._CONTRIBUTION_TABLE_COLUMNS}
152 inputs = astropy.table.join(contributions["instrument", "visit", "detector"], self._inputs)
153 assert inputs.columns.keys() == {name for name, _, _ in self._INPUT_TABLE_COLUMNS}
154 return CoaddProvenance(inputs=inputs, contributions=contributions)
156 def serialize(self, archive: OutputArchive[Any]) -> CoaddProvenanceSerializationModel:
157 """Serialize the provenance to an output archive.
159 Parameters
160 ----------
161 archive
162 Archive to write to.
163 """
164 inputs = self._inputs.copy(copy_data=False)
165 contributions = self._contributions.copy(copy_data=False)
166 instrument = CoaddProvenanceSerializationModel._fix_str_for_serialization(
167 "instrument", inputs, contributions
168 )
169 physical_filter = CoaddProvenanceSerializationModel._fix_str_for_serialization(
170 "physical_filter", inputs
171 )
172 CoaddProvenanceSerializationModel._fix_polygon_for_serialization(inputs)
173 inputs_model = archive.add_table(inputs, name="inputs")
174 contributions_model = archive.add_table(contributions, name="contributions")
175 return CoaddProvenanceSerializationModel(
176 instrument=instrument,
177 physical_filter=physical_filter,
178 inputs=inputs_model,
179 contributions=contributions_model,
180 )
182 @staticmethod
183 def from_legacy(legacy_cell_coadd: LegacyMultipleCellCoadd) -> CoaddProvenance:
184 """Extract provenance from a legacy
185 `lsst.cell_coadds.MultipleCellCoadd` object.
186 """
187 inputs = CoaddProvenance.make_empty_input_table(len(legacy_cell_coadd.common.visit_polygons))
188 for n, (legacy_identifiers, legacy_polygon) in enumerate(
189 legacy_cell_coadd.common.visit_polygons.items()
190 ):
191 inputs["instrument"][n] = legacy_identifiers.instrument
192 inputs["visit"][n] = legacy_identifiers.visit
193 inputs["detector"][n] = legacy_identifiers.detector
194 inputs["physical_filter"][n] = legacy_identifiers.physical_filter
195 inputs["day_obs"][n] = legacy_identifiers.day_obs
196 inputs["polygon"][n] = Polygon.from_legacy(legacy_polygon)
197 n_contributions = 0
198 for legacy_cell in legacy_cell_coadd.cells.values():
199 n_contributions += len(legacy_cell.inputs)
200 contributions = CoaddProvenance.make_empty_contribution_table(n_contributions)
201 n = 0
202 for legacy_cell in legacy_cell_coadd.cells.values():
203 for legacy_identifiers, legacy_inputs in legacy_cell.inputs.items():
204 contributions["cell_i"][n] = legacy_cell.identifiers.cell.y
205 contributions["cell_j"][n] = legacy_cell.identifiers.cell.x
206 contributions["instrument"][n] = legacy_identifiers.instrument
207 contributions["visit"][n] = legacy_identifiers.visit
208 contributions["detector"][n] = legacy_identifiers.detector
209 contributions["overlaps_center"][n] = legacy_inputs.overlaps_center
210 contributions["overlap_fraction"][n] = legacy_inputs.overlap_fraction
211 contributions["unmasked_fraction"][n] = legacy_inputs.unmasked_overlap_fraction
212 contributions["weight"][n] = legacy_inputs.weight
213 contributions["psf_shape_xx"][n] = legacy_inputs.psf_shape.getIxx()
214 contributions["psf_shape_yy"][n] = legacy_inputs.psf_shape.getIyy()
215 contributions["psf_shape_xy"][n] = legacy_inputs.psf_shape.getIxy()
216 contributions["psf_shape_flag"][n] = legacy_inputs.psf_shape_flag
217 n += 1
218 return CoaddProvenance(inputs=inputs, contributions=contributions)
220 def to_legacy_polygon_map(self) -> dict[LegacyObservationIdentifiers, LegacyPolygon]:
221 """Construct a legacy mapping from
222 `lsst.cell_coadds.ObservationIdentifiers` to `lsst.afw.geom.Polygon`
223 from the `inputs` table.
224 """
225 from lsst.cell_coadds import ObservationIdentifiers as LegacyObservationIdentifiers
227 return {
228 LegacyObservationIdentifiers(
229 instrument=str(row["instrument"]),
230 physical_filter=str(row["physical_filter"]),
231 visit=int(row["visit"]),
232 day_obs=int(row["day_obs"]),
233 detector=int(row["detector"]),
234 ): row["polygon"].to_legacy()
235 for row in self.inputs
236 }
238 def to_legacy_cell_coadd_inputs(
239 self, observations: Iterable[LegacyObservationIdentifiers] | None
240 ) -> dict[LegacyIndex2D, dict[LegacyObservationIdentifiers, LegacyCellCoaddInputs]]:
241 """Construct a mapping from legacy cell index to the list of legacy
242 input structs for that cell.
243 """
244 from lsst.afw.geom.ellipses import Quadrupole
245 from lsst.cell_coadds import CoaddInputs as LegacyCoaddInputs
246 from lsst.skymap import Index2D as LegacyIndex2D
248 if observations is None:
249 observations = self.to_legacy_polygon_map().keys()
250 observations_by_key: dict[tuple[str, int, int], LegacyObservationIdentifiers] = {
251 (obs.instrument, obs.visit, obs.detector): obs for obs in observations
252 }
253 result: dict[LegacyIndex2D, dict[LegacyObservationIdentifiers, LegacyCoaddInputs]] = {}
254 for row in self.contributions:
255 obs_key = (str(row["instrument"]), int(row["visit"]), int(row["detector"]))
256 obs = observations_by_key[obs_key]
257 cell_inputs = result.setdefault(LegacyIndex2D(x=int(row["cell_j"]), y=int(row["cell_i"])), {})
258 cell_inputs[obs] = LegacyCoaddInputs(
259 overlaps_center=bool(row["overlaps_center"]),
260 overlap_fraction=float(row["overlap_fraction"]),
261 unmasked_overlap_fraction=float(row["unmasked_fraction"]),
262 weight=float(row["weight"]),
263 psf_shape=Quadrupole(row["psf_shape_xx"], row["psf_shape_yy"], row["psf_shape_xy"]),
264 psf_shape_flag=bool(row["psf_shape_flag"]),
265 )
266 return result
269class CoaddProvenanceSerializationModel(ArchiveTree):
270 """A Pydantic model used to represent a serialized `CoaddProvenance`.
272 Notes
273 -----
274 We can't rewrite the Astropy tables directly into the archive (e.g. as
275 FITS binary tables for a FITS archive), because:
277 - `str` columns are a huge pain in both Numpy and FITS;
278 - the polygon columns need to be rewritten as array-valued columns.
280 To deal with the string columns (``instrument`` and ``physical_filter``)
281 we do dictionary compression: we map each distinct value of those columns
282 to an integer, and then we save that mapping to the model while saving
283 an integer version of that column in the table. But if there is actually
284 only one value in that column (the most common case by far) we just drop
285 the column and store that value directly in the model.
286 """
288 SCHEMA_NAME: ClassVar[str] = "coadd_provenance"
289 SCHEMA_VERSION: ClassVar[str] = "1.0.0"
290 MIN_READ_VERSION: ClassVar[int] = 1
291 PUBLIC_TYPE: ClassVar[type] = CoaddProvenance
293 instrument: str | dict[str, int] = pydantic.Field(
294 description=(
295 "Instrument name for all inputs to this coadd, or a mapping from "
296 "instrument name to the integer used in its place in the tables."
297 )
298 )
299 physical_filter: str | dict[str, int] = pydantic.Field(
300 description="Physical filter name for all inputs to this coadd."
301 )
302 inputs: TableModel = pydantic.Field(description="Table of all inputs to the coadd.")
303 contributions: TableModel = pydantic.Field(description="Table of per-cell contributions to the coadd.")
305 def deserialize(self, archive: InputArchive[Any], **kwargs: Any) -> CoaddProvenance:
306 """Deserialize a provenance from an input archive.
308 Parameters
309 ----------
310 archive
311 Archive to read from.
313 Notes
314 -----
315 While `CoaddProvenance.subset` can be used to filter provenance
316 information down to just certain cells, there is no advantage to be
317 had from doing this during deserialization (the table data is not
318 ordered by cell, and hence there's read-slicing we can do).
319 """
320 if kwargs: 320 ↛ 321line 320 didn't jump to line 321 because the condition on line 320 was never true
321 raise InvalidParameterError(f"Unrecognized parameters for CoaddProvenance: {set(kwargs.keys())}.")
322 inputs = archive.get_table(self.inputs)
323 contributions = archive.get_table(self.contributions)
324 CoaddProvenanceSerializationModel._fix_str_for_deserialization(
325 "instrument", self.instrument, inputs, contributions
326 )
327 CoaddProvenanceSerializationModel._fix_str_for_deserialization(
328 "physical_filter", self.physical_filter, inputs
329 )
330 CoaddProvenanceSerializationModel._fix_polygon_for_deserialization(inputs)
331 for name, _, description in CoaddProvenance._INPUT_TABLE_COLUMNS:
332 inputs.columns[name].description = description
333 for name, _, description, unit in CoaddProvenance._CONTRIBUTION_TABLE_COLUMNS:
334 contributions.columns[name].description = description
335 contributions.columns[name].unit = unit
336 return CoaddProvenance(inputs=inputs, contributions=contributions)
338 @staticmethod
339 def _fix_str_for_serialization(column: str, *tables: astropy.table.Table) -> str | dict[str, int]:
340 """Rewrite a string column as an integer column or drop it.
342 Parameters
343 ----------
344 column
345 Name of the column to rewrite.
346 *tables
347 One or more astropy tables to rewrite. The first table is assumed
348 to have all values for this column that might appear in any other
349 tables.
351 Returns
352 -------
353 `str` | `dict` [`str`, `int`]
354 If there is only one unique value for this column in the first
355 table, that value (and the column will have been dropped from
356 all givne tables). If the tables are empty, the column is
357 dropped and an empty `dict` is returned. In all other cases the
358 given column is replaced with an integer column in all given
359 tables and the mapping from strings to integers is returned.
360 """
361 result: str | dict[str, int] = {name: n for n, name in enumerate(sorted(set(tables[0][column])))}
362 match len(result):
363 case 0:
364 pass
365 case 1:
366 (result,) = result.keys() # type: ignore[union-attr]
367 case _:
368 for table in tables:
369 table.columns[column] = astropy.table.Column(
370 data=[result[k] for k in table.columns[column]],
371 name=column,
372 dtype=np.uint8,
373 description=f"Integer mapped to {column} name.",
374 )
375 return result
376 # If we didn't remap to an integer (case 0 and 1 above), delete the
377 # column.
378 for table in tables:
379 del table.columns[column]
380 return result
382 @staticmethod
383 def _fix_str_for_deserialization(
384 column: str, value: str | dict[str, int], *tables: astropy.table.Table
385 ) -> None:
386 """Rewrite an integer column back to a string one.
388 Parameters
389 ----------
390 column
391 Name of the column to rewrite.
392 value
393 Value or mapping of values returned by
394 `_fix_str_for_serialization`.
395 tables
396 Tables to rewrite this column in.
397 """
398 match value:
399 case str(): 399 ↛ 402line 399 didn't jump to line 402 because the pattern on line 399 always matched
400 for table in tables:
401 table.columns[column] = astropy.table.Column([value] * len(table), dtype=object)
402 case dict():
403 mapping = {v: k for k, v in value.items()}
404 for table in tables:
405 table.columns[column] = astropy.table.Column(
406 [mapping[k] for k in table[column]], dtype=object
407 )
409 @staticmethod
410 def _fix_polygon_for_serialization(inputs: astropy.table.Table) -> None:
411 """Rewrite a polygon `object` column as a pair of array-valued columns
412 and an array-size column.
414 Parameters
415 ----------
416 inputs
417 A copy of the in-memory coadd inputs table to modify in-place into
418 its serialization form.
419 """
420 max_n_vertices = max(p.n_vertices for p in inputs["polygon"])
421 inputs["n_vertices"] = astropy.table.Column(
422 [p.n_vertices for p in inputs["polygon"]],
423 name="n_vertices",
424 dtype=np.uint8,
425 description="Number of polygon vertices.",
426 )
427 inputs["x_vertices"] = astropy.table.Column(
428 name="x_vertices",
429 dtype=np.float64,
430 length=len(inputs),
431 shape=(max_n_vertices,),
432 description="X coordinates of polygon vertices, in tract coordinates.",
433 )
434 inputs["x_vertices"][:, :] = np.nan
435 inputs["y_vertices"] = astropy.table.Column(
436 name="y_vertices",
437 dtype=np.float64,
438 length=len(inputs),
439 shape=(max_n_vertices,),
440 description="Y coordinates of polygon vertices, in tract coordinates.",
441 )
442 inputs["y_vertices"][:, :] = np.nan
443 for i, polygon in enumerate(inputs["polygon"]):
444 inputs["n_vertices"][i] = polygon.n_vertices
445 inputs["x_vertices"][i][: polygon.n_vertices] = polygon.x_vertices
446 inputs["y_vertices"][i][: polygon.n_vertices] = polygon.y_vertices
447 del inputs["polygon"]
449 @staticmethod
450 def _fix_polygon_for_deserialization(inputs: astropy.table.Table) -> None:
451 """Rewrite a a pair of array-valued columns and an array-size column
452 into a polygon `object` column.
454 Parameters
455 ----------
456 inputs
457 The serialized version of the coadd inputs table, to be modified
458 in-place into its in-memory form.
459 """
460 polygons = [
461 Polygon(x_vertices=x_vertices[:n_vertices], y_vertices=y_vertices[:n_vertices])
462 for n_vertices, x_vertices, y_vertices in zip(
463 inputs["n_vertices"], inputs["x_vertices"], inputs["y_vertices"]
464 )
465 ]
466 del inputs["n_vertices"]
467 del inputs["x_vertices"]
468 del inputs["y_vertices"]
469 inputs["polygon"] = astropy.table.Column(polygons, name="polygon", dtype=np.object_)