Merge pull request #33 from lsst/tickets/DM-40563

DM-40563: Persist ObservationIdentifiers

.pre-commit-config.yaml

@@ -1,18 +1,18 @@
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.5.0
+    rev: v4.6.0
     hooks:
       - id: check-yaml
       - id: end-of-file-fixer
       - id: trailing-whitespace
   - repo: https://github.com/psf/black
-    rev: 24.1.1
+    rev: 24.4.2
     hooks:
       - id: black
         language_version: python3.11
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.2.0
+    rev: v0.4.3
     hooks:
       - id: ruff
   - repo: https://github.com/pycqa/isort

doc/changes/DM-40563.api.rst

@@ -0,0 +1 @@
+Replace "BAND" in the header with a more standard "FILTER" keyword.

doc/changes/DM-40563.feature.rst

@@ -0,0 +1 @@
+Persist information about the input visits (+detectors) to each cell in separate HDUs.

python/lsst/cell_coadds/_fits.py

@@ -85,7 +85,8 @@
 
 import logging
 import os
-from collections.abc import Mapping
+from collections.abc import Iterable, Mapping
+from dataclasses import dataclass
 from typing import Any
 
 import lsst.afw.geom as afwGeom
@@ -97,14 +98,16 @@
 from lsst.geom import Box2I, Extent2I, Point2I
 from lsst.obs.base.formatters.fitsGeneric import FitsGenericFormatter
 from lsst.skymap import Index2D
+from packaging import version
 
 from ._common_components import CoaddUnits, CommonComponents
-from ._identifiers import CellIdentifiers, PatchIdentifiers
+from ._grid_container import GridContainer
+from ._identifiers import CellIdentifiers, ObservationIdentifiers, PatchIdentifiers
 from ._image_planes import OwnedImagePlanes
 from ._multiple_cell_coadd import MultipleCellCoadd, SingleCellCoadd
 from ._uniform_grid import UniformGrid
 
-FILE_FORMAT_VERSION = "0.2"
+FILE_FORMAT_VERSION = "0.3"
 """Version number for the file format as persisted, presented as a string of
 the form M.m, where M is the major version, m is the minor version.
 """
@@ -118,6 +121,18 @@ class IncompatibleVersionError(RuntimeError):
     """
 
 
+@dataclass
+class VisitRecord:
+    """A dataclass to hold relevant info about a visit.
+
+    This is intended for use with this module.
+    """
+
+    visit: int
+    day_obs: int
+    physical_filter: str
+
+
 class CellCoaddFitsFormatter(FitsGenericFormatter):
     """Interface for writing and reading cell coadds to/from FITS files.
 
@@ -172,12 +187,13 @@ class can read a file, based on the VERSION in its header.
         -----
         This accepts the other version as a positional argument only.
         """
+        written_version_object = version.parse(written_version)
         for min_version, max_version in zip(
             cls.MINIMUM_FILE_FORMAT_VERSIONS,
             cls.MAXIMUM_FILE_FORMAT_VERSIONS,
             strict=True,
         ):
-            if min_version <= written_version < max_version:
+            if version.parse(min_version) <= written_version_object < version.parse(max_version):
                 return True
 
         return False
@@ -207,6 +223,12 @@ def readAsMultipleCellCoadd(self) -> MultipleCellCoadd:
                     FILE_FORMAT_VERSION,
                 )
 
+            written_version = version.parse(written_version)
+
+            # TODO: Remove this when FILE_FORMAT_VERSION is bumped to 1.0
+            if written_version < version.parse("0.3"):
+                header.rename_keyword("BAND", "FILTER")
+
             data = hdu_list[1].data
 
             # Read in WCS
@@ -218,12 +240,12 @@ def readAsMultipleCellCoadd(self) -> MultipleCellCoadd:
             common = CommonComponents(
                 units=CoaddUnits(1),  # TODO: read from FITS TUNIT1 (DM-40562)
                 wcs=wcs,
-                band=header["BAND"],
+                band=header["FILTER"],
                 identifiers=PatchIdentifiers(
                     skymap=header["SKYMAP"],
                     tract=header["TRACT"],
                     patch=Index2D(x=header["PATCH_X"], y=header["PATCH_Y"]),
-                    band=header["BAND"],
+                    band=header["FILTER"],
                 ),
             )
 
@@ -241,12 +263,46 @@ def readAsMultipleCellCoadd(self) -> MultipleCellCoadd:
             outer_cell_size = Extent2I(header["OCELL1"], header["OCELL2"])
             psf_image_size = Extent2I(header["PSFSIZE1"], header["PSFSIZE2"])
 
+            # Attempt to get inputs for each cell.
+            inputs = GridContainer[list[ObservationIdentifiers]](shape=grid.shape)
+            if written_version >= version.parse("0.3"):
+                visit_dict = {
+                    row["visit"]: VisitRecord(
+                        visit=row["visit"],
+                        physical_filter=row["physical_filter"],
+                        day_obs=row["day_obs"],
+                    )
+                    for row in hdu_list[hdu_list.index_of("VISIT")].data
+                }
+                link_table = hdu_list[hdu_list.index_of("CELL")].data
+                for link_row in link_table:
+                    cell_id = Index2D(link_row["cell_x"], link_row["cell_y"])
+                    visit = link_row["visit"]
+                    obs_id = ObservationIdentifiers(
+                        instrument=header["INSTRUME"],
+                        visit=visit,
+                        detector=link_row["detector"],
+                        day_obs=visit_dict[visit].day_obs,
+                        physical_filter=visit_dict[visit].physical_filter,
+                    )
+                    if cell_id in inputs:
+                        inputs[cell_id] += [obs_id]
+                    else:
+                        inputs[cell_id] = [obs_id]
+            else:
+                logger.info(
+                    "Cell inputs are available for VERSION=0.3 or later. The file provided has ",
+                    "VERSION = %s",
+                    written_version,
+                )
+
             coadd = MultipleCellCoadd(
                 (
                     self._readSingleCellCoadd(
                         data=row,
                         header=header,
                         common=common,
+                        inputs=inputs[Index2D(row["cell_id"][0], row["cell_id"][1])],
                         outer_cell_size=outer_cell_size,
                         psf_image_size=psf_image_size,
                         inner_cell_size=grid_cell_size,
@@ -268,6 +324,7 @@ def _readSingleCellCoadd(
         common: CommonComponents,
         header: Mapping[str, Any],
         *,
+        inputs: Iterable[ObservationIdentifiers],
         outer_cell_size: Extent2I,
         inner_cell_size: Extent2I,
         psf_image_size: Extent2I,
@@ -281,6 +338,11 @@ def _readSingleCellCoadd(
             table representation.
         common : `CommonComponents`
             The common components of the coadd.
+        header : `Mapping`
+            The header of the FITS file as a dictionary.
+        inputs : `Iterable` [`ObservationIdentifiers`]
+            Any iterable of ObservationIdentifiers instances that contributed
+            to this cell.
         outer_cell_size : `Extent2I`
             The size of the outer cell.
         psf_image_size : `Extent2I`
@@ -335,8 +397,7 @@ def _readSingleCellCoadd(
             ),
             common=common,
             identifiers=identifiers,
-            # TODO: Pass a sensible value here in DM-40563.
-            inputs=None,  # type: ignore[arg-type]
+            inputs=inputs,
         )
 
     def readWcs(self) -> afwGeom.SkyWcs:
@@ -360,7 +421,7 @@ def writeMultipleCellCoaddAsFits(
     filename: str,
     overwrite: bool = False,
     metadata: PropertySet | None = None,
-) -> None:
+) -> fits.HDUList:
     """Write a MultipleCellCoadd object to a FITS file.
 
     Parameters
@@ -374,11 +435,56 @@ def writeMultipleCellCoaddAsFits(
     metadata : `~lsst.daf.base.PropertySet`, optional
         Additional metadata to write to the FITS file.
 
+    Returns
+    -------
+    hdu_list : `~astropy.io.fits.HDUList`
+        The FITS file as an HDUList.
+
     Notes
     -----
     Changes to this function that modify the way the file is written to disk
     must be accompanied with a change to FILE_FORMAT_VERSION.
     """
+    # Create metadata tables:
+    # 1. Visit table containing information about the visits.
+    # 2. Cell table containing info about the visit+detector for each cell.
+    visit_records: list[Any] = []
+    cell_records: list[Any] = []
+    instrument_set = set()
+    for cell_id, single_cell_coadd in multiple_cell_coadd.cells.items():
+        for observation_id in single_cell_coadd.inputs:
+            visit_records.append(
+                (observation_id.visit, observation_id.physical_filter, observation_id.day_obs)
+            )
+            cell_records.append((cell_id.x, cell_id.y, observation_id.visit, observation_id.detector))
+            instrument_set.add(observation_id.instrument)
+
+    assert len(instrument_set) == 1, "All cells must have the same instrument."
+    instrument = instrument_set.pop()
+
+    visit_recarray = np.rec.fromrecords(
+        recList=sorted(set(visit_records), key=lambda x: x[0]),  # Sort by visit.
+        formats=None,  # formats has specified to please mypy. See numpy#26376.
+        names=(
+            "visit",
+            "physical_filter",
+            "day_obs",
+        ),
+    )
+    cell_recarray = np.rec.fromrecords(
+        recList=cell_records,
+        formats=None,  # formats has specified to please mypy. See numpy#26376.
+        names=(
+            "cell_x",
+            "cell_y",
+            "visit",
+            "detector",
+        ),
+    )
+
+    visit_hdu = fits.BinTableHDU.from_columns(visit_recarray, name="VISIT")
+    cell_hdu = fits.BinTableHDU.from_columns(cell_recarray, name="CELL")
+
     cell_id = fits.Column(
         name="cell_id",
         format="2I",
@@ -464,7 +570,8 @@ def writeMultipleCellCoaddAsFits(
     hdu.header["TUNIT1"] = multiple_cell_coadd.common.units.name
     # This assumed to be the same as multiple_cell_coadd.common.identifers.band
     # See DM-38843.
-    hdu.header["BAND"] = multiple_cell_coadd.common.band
+    hdu.header["INSTRUME"] = instrument
+    hdu.header["FILTER"] = multiple_cell_coadd.common.band
     hdu.header["SKYMAP"] = multiple_cell_coadd.common.identifiers.skymap
     hdu.header["TRACT"] = multiple_cell_coadd.common.identifiers.tract
     hdu.header["PATCH_X"] = multiple_cell_coadd.common.identifiers.patch.x
@@ -473,5 +580,7 @@ def writeMultipleCellCoaddAsFits(
     if metadata is not None:
         hdu.header.extend(metadata.toDict())
 
-    hdu_list = fits.HDUList([primary_hdu, hdu])
+    hdu_list = fits.HDUList([primary_hdu, hdu, cell_hdu, visit_hdu])
     hdu_list.writeto(filename, overwrite=overwrite)
+
+    return hdu_list

python/lsst/cell_coadds/_single_cell_coadd.py

@@ -85,7 +85,7 @@ def __init__(
         self._common = common
         # Remove any duplicate elements in the input, sorted them and make
         # them an immutable sequence.
-        # TODO: Remove the conditioning in DM-40563.
+        # TODO: Remove support for inputs as None when bumping to v1.0 .
         self._inputs = tuple(sorted(set(inputs))) if inputs else ()
         self._identifiers = identifiers
 
@@ -107,7 +107,7 @@ def psf_image(self) -> ImageF:
         return self._psf
 
     @property
-    # TODO: Remove the conditioning in DM-40563.
+    # TODO: Remove the option of returning empty tuple in v1.0.
     def inputs(self) -> tuple[ObservationIdentifiers, ...] | tuple[()]:
         """Identifiers for the input images that contributed to this cell,
         sorted by their `visit` attribute first, and then by `detector`.

tests/test_coadds.py

@@ -268,11 +268,13 @@ def assertMultipleCellCoaddsEqual(self, mcc1: MultipleCellCoadd, mcc2: MultipleC
             self.assertMasksEqual(mcc1.cells[idx].outer.mask, mcc2.cells[idx].outer.mask)
             self.assertImagesEqual(mcc1.cells[idx].outer.variance, mcc2.cells[idx].outer.variance)
             self.assertImagesEqual(mcc1.cells[idx].psf_image, mcc2.cells[idx].psf_image)
+            self.assertEqual(mcc1.cells[idx].inputs, mcc2.cells[idx].inputs)
 
             self.assertEqual(mcc1.cells[idx].band, mcc1.band)
             self.assertEqual(mcc1.cells[idx].common, mcc1.common)
             self.assertEqual(mcc1.cells[idx].units, mcc2.units)
             self.assertEqual(mcc1.cells[idx].wcs, mcc1.wcs)
+
             # Identifiers differ because of the ``cell`` component.
             # Check the other attributes within the identifiers.
             for attr in ("skymap", "tract", "patch", "band"):