Coverage for python/lsst/images/cells/_provenance.py: 44%

164 statements  

« prev     ^ index     » next       coverage.py v7.14.3, created at 2026-07-09 02:27 -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. 

11 

12from __future__ import annotations 

13 

14__all__ = ("CoaddProvenance", "CoaddProvenanceSerializationModel") 

15 

16from collections.abc import Iterable 

17from typing import TYPE_CHECKING, Any, ClassVar 

18 

19import astropy.table 

20import astropy.units as u 

21import numpy as np 

22import pydantic 

23 

24from .._cell_grid import CellIJ 

25from .._polygon import Polygon 

26from ..serialization import ArchiveTree, InputArchive, InvalidParameterError, OutputArchive, TableModel 

27 

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] 

41 

42 

43class CoaddProvenance: 

44 """A pair of tables that record the inputs to a cell-based coadd. 

45 

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. 

54 

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 """ 

62 

63 def __init__(self, inputs: astropy.table.Table, contributions: astropy.table.Table) -> None: 

64 self._inputs = inputs 

65 self._contributions = contributions 

66 

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 ] 

82 

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 ] 

103 

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 

108 Parameters 

109 ---------- 

110 n_rows 

111 Number of rows in the new table. 

112 """ 

113 return astropy.table.Table( 

114 [ 

115 astropy.table.Column(name=name, length=n_rows, dtype=dtype, description=description) 

116 for name, dtype, description in cls._INPUT_TABLE_COLUMNS 

117 ] 

118 ) 

119 

120 @classmethod 

121 def make_empty_contribution_table(cls, n_rows: int) -> astropy.table.Table: 

122 """Make an empty `contributions` table with a set number of rows. 

123 

124 Parameters 

125 ---------- 

126 n_rows 

127 Number of rows in the new table. 

128 """ 

129 return astropy.table.Table( 

130 [ 

131 astropy.table.Column( 

132 name=name, length=n_rows, dtype=dtype, description=description, unit=unit 

133 ) 

134 for name, dtype, description, unit in cls._CONTRIBUTION_TABLE_COLUMNS 

135 ] 

136 ) 

137 

138 @property 

139 def inputs(self) -> astropy.table.Table: 

140 """A table of {visit, detector} combinations that contribute to any 

141 cell in the coadd. 

142 """ 

143 return self._inputs 

144 

145 @property 

146 def contributions(self) -> astropy.table.Table: 

147 """A table of {visit, detector, cell} combinations that describe how an 

148 observation contributed to a cell. 

149 """ 

150 return self._contributions 

151 

152 def __getitem__(self, cell: CellIJ) -> CoaddProvenance: 

153 return self.subset([cell]) 

154 

155 def subset(self, cells: Iterable[CellIJ]) -> CoaddProvenance: 

156 """Return a new provenance object with just the given cells. 

157 

158 Parameters 

159 ---------- 

160 cells 

161 Cells to keep in the returned provenance. 

162 """ 

163 cells_to_keep = astropy.table.Table( 

164 rows=[(index.i, index.j) for index in cells], 

165 names=["cell_i", "cell_j"], 

166 dtype=[np.uint16, np.uint16], 

167 ) 

168 contributions = astropy.table.join(self._contributions, cells_to_keep) 

169 assert contributions.columns.keys() == {name for name, _, _, _ in self._CONTRIBUTION_TABLE_COLUMNS} 

170 inputs = astropy.table.join(contributions["instrument", "visit", "detector"], self._inputs) 

171 assert inputs.columns.keys() == {name for name, _, _ in self._INPUT_TABLE_COLUMNS} 

172 return CoaddProvenance(inputs=inputs, contributions=contributions) 

173 

174 def serialize(self, archive: OutputArchive[Any]) -> CoaddProvenanceSerializationModel: 

175 """Serialize the provenance to an output archive. 

176 

177 Parameters 

178 ---------- 

179 archive 

180 Archive to write to. 

181 """ 

182 inputs = self._inputs.copy(copy_data=False) 

183 contributions = self._contributions.copy(copy_data=False) 

184 instrument = CoaddProvenanceSerializationModel._fix_str_for_serialization( 

185 "instrument", inputs, contributions 

186 ) 

187 physical_filter = CoaddProvenanceSerializationModel._fix_str_for_serialization( 

188 "physical_filter", inputs 

189 ) 

190 CoaddProvenanceSerializationModel._fix_polygon_for_serialization(inputs) 

191 inputs_model = archive.add_table(inputs, name="inputs") 

192 contributions_model = archive.add_table(contributions, name="contributions") 

193 return CoaddProvenanceSerializationModel( 

194 instrument=instrument, 

195 physical_filter=physical_filter, 

196 inputs=inputs_model, 

197 contributions=contributions_model, 

198 ) 

199 

200 @staticmethod 

201 def from_legacy(legacy_cell_coadd: LegacyMultipleCellCoadd) -> CoaddProvenance: 

202 """Extract provenance from a legacy 

203 `lsst.cell_coadds.MultipleCellCoadd` object. 

204 

205 Parameters 

206 ---------- 

207 legacy_cell_coadd 

208 Legacy cell coadd to extract provenance from. 

209 """ 

210 inputs = CoaddProvenance.make_empty_input_table(len(legacy_cell_coadd.common.visit_polygons)) 

211 for n, (legacy_identifiers, legacy_polygon) in enumerate( 

212 legacy_cell_coadd.common.visit_polygons.items() 

213 ): 

214 inputs["instrument"][n] = legacy_identifiers.instrument 

215 inputs["visit"][n] = legacy_identifiers.visit 

216 inputs["detector"][n] = legacy_identifiers.detector 

217 inputs["physical_filter"][n] = legacy_identifiers.physical_filter 

218 inputs["day_obs"][n] = legacy_identifiers.day_obs 

219 inputs["polygon"][n] = Polygon.from_legacy(legacy_polygon) 

220 n_contributions = 0 

221 for legacy_cell in legacy_cell_coadd.cells.values(): 

222 n_contributions += len(legacy_cell.inputs) 

223 contributions = CoaddProvenance.make_empty_contribution_table(n_contributions) 

224 n = 0 

225 for legacy_cell in legacy_cell_coadd.cells.values(): 

226 for legacy_identifiers, legacy_inputs in legacy_cell.inputs.items(): 

227 contributions["cell_i"][n] = legacy_cell.identifiers.cell.y 

228 contributions["cell_j"][n] = legacy_cell.identifiers.cell.x 

229 contributions["instrument"][n] = legacy_identifiers.instrument 

230 contributions["visit"][n] = legacy_identifiers.visit 

231 contributions["detector"][n] = legacy_identifiers.detector 

232 contributions["overlaps_center"][n] = legacy_inputs.overlaps_center 

233 contributions["overlap_fraction"][n] = legacy_inputs.overlap_fraction 

234 contributions["unmasked_fraction"][n] = legacy_inputs.unmasked_overlap_fraction 

235 contributions["weight"][n] = legacy_inputs.weight 

236 contributions["psf_shape_xx"][n] = legacy_inputs.psf_shape.getIxx() 

237 contributions["psf_shape_yy"][n] = legacy_inputs.psf_shape.getIyy() 

238 contributions["psf_shape_xy"][n] = legacy_inputs.psf_shape.getIxy() 

239 contributions["psf_shape_flag"][n] = legacy_inputs.psf_shape_flag 

240 n += 1 

241 return CoaddProvenance(inputs=inputs, contributions=contributions) 

242 

243 def to_legacy_polygon_map(self) -> dict[LegacyObservationIdentifiers, LegacyPolygon]: 

244 """Construct a legacy mapping from 

245 `lsst.cell_coadds.ObservationIdentifiers` to `lsst.afw.geom.Polygon` 

246 from the `inputs` table. 

247 """ 

248 from lsst.cell_coadds import ObservationIdentifiers as LegacyObservationIdentifiers 

249 

250 return { 

251 LegacyObservationIdentifiers( 

252 instrument=str(row["instrument"]), 

253 physical_filter=str(row["physical_filter"]), 

254 visit=int(row["visit"]), 

255 day_obs=int(row["day_obs"]), 

256 detector=int(row["detector"]), 

257 ): row["polygon"].to_legacy() 

258 for row in self.inputs 

259 } 

260 

261 def to_legacy_cell_coadd_inputs( 

262 self, observations: Iterable[LegacyObservationIdentifiers] | None 

263 ) -> dict[LegacyIndex2D, dict[LegacyObservationIdentifiers, LegacyCellCoaddInputs]]: 

264 """Construct a mapping from legacy cell index to the list of legacy 

265 input structs for that cell. 

266 

267 Parameters 

268 ---------- 

269 observations 

270 Observations to include, or `None` to include all observations 

271 in the `inputs` table. 

272 """ 

273 from lsst.afw.geom.ellipses import Quadrupole 

274 from lsst.cell_coadds import CoaddInputs as LegacyCoaddInputs 

275 from lsst.skymap import Index2D as LegacyIndex2D 

276 

277 if observations is None: 

278 observations = self.to_legacy_polygon_map().keys() 

279 observations_by_key: dict[tuple[str, int, int], LegacyObservationIdentifiers] = { 

280 (obs.instrument, obs.visit, obs.detector): obs for obs in observations 

281 } 

282 result: dict[LegacyIndex2D, dict[LegacyObservationIdentifiers, LegacyCoaddInputs]] = {} 

283 for row in self.contributions: 

284 obs_key = (str(row["instrument"]), int(row["visit"]), int(row["detector"])) 

285 obs = observations_by_key[obs_key] 

286 cell_inputs = result.setdefault(LegacyIndex2D(x=int(row["cell_j"]), y=int(row["cell_i"])), {}) 

287 cell_inputs[obs] = LegacyCoaddInputs( 

288 overlaps_center=bool(row["overlaps_center"]), 

289 overlap_fraction=float(row["overlap_fraction"]), 

290 unmasked_overlap_fraction=float(row["unmasked_fraction"]), 

291 weight=float(row["weight"]), 

292 psf_shape=Quadrupole(row["psf_shape_xx"], row["psf_shape_yy"], row["psf_shape_xy"]), 

293 psf_shape_flag=bool(row["psf_shape_flag"]), 

294 ) 

295 return result 

296 

297 

298class CoaddProvenanceSerializationModel(ArchiveTree): 

299 """A Pydantic model used to represent a serialized `CoaddProvenance`. 

300 

301 Notes 

302 ----- 

303 We can't rewrite the Astropy tables directly into the archive (e.g. as 

304 FITS binary tables for a FITS archive), because: 

305 

306 - `str` columns are a huge pain in both Numpy and FITS; 

307 - the polygon columns need to be rewritten as array-valued columns. 

308 

309 To deal with the string columns (``instrument`` and ``physical_filter``) 

310 we do dictionary compression: we map each distinct value of those columns 

311 to an integer, and then we save that mapping to the model while saving 

312 an integer version of that column in the table. But if there is actually 

313 only one value in that column (the most common case by far) we just drop 

314 the column and store that value directly in the model. 

315 """ 

316 

317 SCHEMA_NAME: ClassVar[str] = "coadd_provenance" 

318 SCHEMA_VERSION: ClassVar[str] = "1.0.0" 

319 MIN_READ_VERSION: ClassVar[int] = 1 

320 PUBLIC_TYPE: ClassVar[type] = CoaddProvenance 

321 

322 instrument: str | dict[str, int] = pydantic.Field( 

323 description=( 

324 "Instrument name for all inputs to this coadd, or a mapping from " 

325 "instrument name to the integer used in its place in the tables." 

326 ) 

327 ) 

328 physical_filter: str | dict[str, int] = pydantic.Field( 

329 description="Physical filter name for all inputs to this coadd." 

330 ) 

331 inputs: TableModel = pydantic.Field(description="Table of all inputs to the coadd.") 

332 contributions: TableModel = pydantic.Field(description="Table of per-cell contributions to the coadd.") 

333 

334 def deserialize(self, archive: InputArchive[Any], **kwargs: Any) -> CoaddProvenance: 

335 """Deserialize a provenance from an input archive. 

336 

337 Parameters 

338 ---------- 

339 archive 

340 Archive to read from. 

341 **kwargs 

342 Unsupported keyword arguments are accepted only to provide 

343 better error messages (raising 

344 `.serialization.InvalidParameterError`). 

345 

346 Notes 

347 ----- 

348 While `CoaddProvenance.subset` can be used to filter provenance 

349 information down to just certain cells, there is no advantage to be 

350 had from doing this during deserialization (the table data is not 

351 ordered by cell, and hence there's read-slicing we can do). 

352 """ 

353 if kwargs: 353 ↛ 354line 353 didn't jump to line 354 because the condition on line 353 was never true

354 raise InvalidParameterError(f"Unrecognized parameters for CoaddProvenance: {set(kwargs.keys())}.") 

355 inputs = archive.get_table(self.inputs) 

356 contributions = archive.get_table(self.contributions) 

357 CoaddProvenanceSerializationModel._fix_str_for_deserialization( 

358 "instrument", self.instrument, inputs, contributions 

359 ) 

360 CoaddProvenanceSerializationModel._fix_str_for_deserialization( 

361 "physical_filter", self.physical_filter, inputs 

362 ) 

363 CoaddProvenanceSerializationModel._fix_polygon_for_deserialization(inputs) 

364 for name, _, description in CoaddProvenance._INPUT_TABLE_COLUMNS: 

365 inputs.columns[name].description = description 

366 for name, _, description, unit in CoaddProvenance._CONTRIBUTION_TABLE_COLUMNS: 

367 contributions.columns[name].description = description 

368 contributions.columns[name].unit = unit 

369 return CoaddProvenance(inputs=inputs, contributions=contributions) 

370 

371 @staticmethod 

372 def _fix_str_for_serialization(column: str, *tables: astropy.table.Table) -> str | dict[str, int]: 

373 """Rewrite a string column as an integer column or drop it. 

374 

375 Parameters 

376 ---------- 

377 column 

378 Name of the column to rewrite. 

379 *tables 

380 One or more astropy tables to rewrite. The first table is assumed 

381 to have all values for this column that might appear in any other 

382 tables. 

383 

384 Returns 

385 ------- 

386 `str` | `dict` [`str`, `int`] 

387 If there is only one unique value for this column in the first 

388 table, that value (and the column will have been dropped from 

389 all givne tables). If the tables are empty, the column is 

390 dropped and an empty `dict` is returned. In all other cases the 

391 given column is replaced with an integer column in all given 

392 tables and the mapping from strings to integers is returned. 

393 """ 

394 result: str | dict[str, int] = {name: n for n, name in enumerate(sorted(set(tables[0][column])))} 

395 match len(result): 

396 case 0: 

397 pass 

398 case 1: 

399 (result,) = result.keys() # type: ignore[union-attr] 

400 case _: 

401 for table in tables: 

402 table.columns[column] = astropy.table.Column( 

403 data=[result[k] for k in table.columns[column]], 

404 name=column, 

405 dtype=np.uint8, 

406 description=f"Integer mapped to {column} name.", 

407 ) 

408 return result 

409 # If we didn't remap to an integer (case 0 and 1 above), delete the 

410 # column. 

411 for table in tables: 

412 del table.columns[column] 

413 return result 

414 

415 @staticmethod 

416 def _fix_str_for_deserialization( 

417 column: str, value: str | dict[str, int], *tables: astropy.table.Table 

418 ) -> None: 

419 """Rewrite an integer column back to a string one. 

420 

421 Parameters 

422 ---------- 

423 column 

424 Name of the column to rewrite. 

425 value 

426 Value or mapping of values returned by 

427 `_fix_str_for_serialization`. 

428 tables 

429 Tables to rewrite this column in. 

430 """ 

431 match value: 

432 case str(): 432 ↛ 435line 432 didn't jump to line 435 because the pattern on line 432 always matched

433 for table in tables: 

434 table.columns[column] = astropy.table.Column([value] * len(table), dtype=object) 

435 case dict(): 

436 mapping = {v: k for k, v in value.items()} 

437 for table in tables: 

438 table.columns[column] = astropy.table.Column( 

439 [mapping[k] for k in table[column]], dtype=object 

440 ) 

441 

442 @staticmethod 

443 def _fix_polygon_for_serialization(inputs: astropy.table.Table) -> None: 

444 """Rewrite a polygon `object` column as a pair of array-valued columns 

445 and an array-size column. 

446 

447 Parameters 

448 ---------- 

449 inputs 

450 A copy of the in-memory coadd inputs table to modify in-place into 

451 its serialization form. 

452 """ 

453 max_n_vertices = max(p.n_vertices for p in inputs["polygon"]) 

454 inputs["n_vertices"] = astropy.table.Column( 

455 [p.n_vertices for p in inputs["polygon"]], 

456 name="n_vertices", 

457 dtype=np.uint8, 

458 description="Number of polygon vertices.", 

459 ) 

460 inputs["x_vertices"] = astropy.table.Column( 

461 name="x_vertices", 

462 dtype=np.float64, 

463 length=len(inputs), 

464 shape=(max_n_vertices,), 

465 description="X coordinates of polygon vertices, in tract coordinates.", 

466 ) 

467 inputs["x_vertices"][:, :] = np.nan 

468 inputs["y_vertices"] = astropy.table.Column( 

469 name="y_vertices", 

470 dtype=np.float64, 

471 length=len(inputs), 

472 shape=(max_n_vertices,), 

473 description="Y coordinates of polygon vertices, in tract coordinates.", 

474 ) 

475 inputs["y_vertices"][:, :] = np.nan 

476 for i, polygon in enumerate(inputs["polygon"]): 

477 inputs["n_vertices"][i] = polygon.n_vertices 

478 inputs["x_vertices"][i][: polygon.n_vertices] = polygon.x_vertices 

479 inputs["y_vertices"][i][: polygon.n_vertices] = polygon.y_vertices 

480 del inputs["polygon"] 

481 

482 @staticmethod 

483 def _fix_polygon_for_deserialization(inputs: astropy.table.Table) -> None: 

484 """Rewrite a a pair of array-valued columns and an array-size column 

485 into a polygon `object` column. 

486 

487 Parameters 

488 ---------- 

489 inputs 

490 The serialized version of the coadd inputs table, to be modified 

491 in-place into its in-memory form. 

492 """ 

493 polygons = [ 

494 Polygon(x_vertices=x_vertices[:n_vertices], y_vertices=y_vertices[:n_vertices]) 

495 for n_vertices, x_vertices, y_vertices in zip( 

496 inputs["n_vertices"], inputs["x_vertices"], inputs["y_vertices"] 

497 ) 

498 ] 

499 del inputs["n_vertices"] 

500 del inputs["x_vertices"] 

501 del inputs["y_vertices"] 

502 inputs["polygon"] = astropy.table.Column(polygons, name="polygon", dtype=np.object_)