Rewrite SkyMapDimensionPacker to be more configurable.

python/lsst/skymap/packers.py

@@ -19,9 +19,50 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
+from __future__ import annotations
+
 __all__ = ("SkyMapDimensionPacker",)
 
+from collections.abc import Mapping
+
+from lsst.pex.config import Config, Field, DictField, ConfigurableField
 from lsst.daf.butler import DimensionPacker, DimensionGraph, DataCoordinate
+from deprecated.sphinx import deprecated
+
+
+class SkyMapDimensionPackerConfig(Config):
+    bands = DictField(
+        "Mapping from band name to integer to use in the packed ID. "
+        "The default (None) is to use a hard-coded list of common bands; "
+        "pipelines that can enumerate the set of bands they are likely to see "
+        "should override this.",
+        keytype=str,
+        itemtype=int,
+        default=None,
+        optional=True,
+    )
+    n_bands = Field(
+        "Number of bands to reserve space for. "
+        "If zero, bands are not included in the packed integer at all. "
+        "If `None`, the size of 'bands' is used.",
+        dtype=int,
+        optional=True,
+        default=0,
+    )
+    n_tracts = Field(
+        "Number of tracts, or, more precisely, one greater than the maximum tract ID."
+        "Default (None) obtains this value from the skymap dimension record.",
+        dtype=int,
+        optional=True,
+        default=None,
+    )
+    n_patches = Field(
+        "Number of patches per tract, or, more precisely, one greater than the maximum patch ID."
+        "Default (None) obtains this value from the skymap dimension record.",
+        dtype=int,
+        optional=True,
+        default=None,
+    )
 
 
 class SkyMapDimensionPacker(DimensionPacker):
@@ -31,10 +72,50 @@ class SkyMapDimensionPacker(DimensionPacker):
     Parameters
     ----------
     fixed : `lsst.daf.butler.DataCoordinate`
-        Expanded data ID that must include at least the skymap dimension.
-    dimensions : `lsst.daf.butler.DimensionGraph`
+        Data ID that identifies just the ``skymap`` dimension.  Must have
+        dimension records attached unless ``n_tracts`` and ``n_patches`` are
+        not `None`.
+    dimensions : `lsst.daf.butler.DimensionGraph`, optional
         The dimensions of data IDs packed by this instance.  Must include
-        skymap, tract, and patch, and may include band.
+        ``{skymap, tract, patch}``, and may include ``band``.  If not provided,
+        this will be set to include ``band`` if ``n_bands != 0``.
+    bands : `~collections.abc.Mapping` [ `str`, `int` ] or `None`, optional
+        Mapping from band name to integer to use in the packed ID.  `None` uses
+        a fixed set of bands defined in this class.  When calling code can
+        enumerate the bands it is likely to see, providing an explicit mapping
+        is preferable.
+    n_bands : `int` or `None`, optional
+        The number of bands to leave room for in the packed ID.  If `None`,
+        this will be set to ``len(bands)``.  If ``0``, the band will not be
+        included in the dimensions at all.  If ``1``, the band will be included
+        in the dimensions but will not occupy any extra bits in the packed ID.
+        This may be larger or smaller than ``len(bands)``, to reserve extra
+        space for the future or align to byte boundaries, or support a subset
+        of a larger mapping, respectively.
+    n_tracts : `int` or `None`, optional
+        The number of tracts to leave room for in the packed ID.  If `None`,
+        this will be set via the ``skymap`` dimension record in ``fixed``.
+    n_patches : `int` or `None`, optional
+        The number of patches (per tract) to leave room for in the packed ID.
+        If `None`, this will be set via the ``skymap`` dimension record in
+        ``fixed``.
+
+    Notes
+    -----
+    The standard pattern for constructing instances of this class is to use
+    `make_config_field`::
+
+        class SomeConfig(lsst.pex.config.Config):
+            packer = ObservationDimensionPacker.make_config_field()
+
+        class SomeTask(lsst.pipe.base.Task):
+            ConfigClass = SomeConfig
+
+            def run(self, ..., data_id: DataCoordinate):
+                packer = self.config.packer.apply(data_id)
+                packed_id = packer.pack(data_id)
+                ...
+
     """
 
     SUPPORTED_FILTERS = (
@@ -43,13 +124,22 @@ class SkyMapDimensionPacker(DimensionPacker):
         + [f"N{d}" for d in (387, 515, 656, 816, 921, 1010)]  # HSC narrow-bands
         + [f"N{d}" for d in (419, 540, 708, 964)]  # DECam narrow-bands
     )
-    """band names supported by this packer.
+    """Sequence of supported bands used to construct a mapping from band name
+    to integer when the 'bands' config option is `None` or no config is
+    provided.
 
-    New filters should be added to the end of the list to maximize
-    compatibility with existing IDs.
+    This variable should no longer be modified to add new filters; pass
+    ``bands`` at construction or use `from_config` instead.
     """
 
+    ConfigClass = SkyMapDimensionPackerConfig
+
     @classmethod
+    @deprecated(
+        reason="This classmethod cannot reflect all __init__ args and will be removed after v27.",
+        version="v26.0",
+        category=FutureWarning,
+    )
     def getIntFromFilter(cls, name):
         """Return an integer that represents the band with the given
         name.
@@ -60,47 +150,162 @@ def getIntFromFilter(cls, name):
             raise NotImplementedError(f"band '{name}' not supported by this ID packer.")
 
     @classmethod
+    @deprecated(
+        reason="This classmethod cannot reflect all __init__ args and will be removed after v27.",
+        version="v26.0",
+        category=FutureWarning,
+    )
     def getFilterNameFromInt(cls, num):
-        """Return an band name from its integer representation.
-        """
+        """Return an band name from its integer representation."""
         return cls.SUPPORTED_FILTERS[num]
 
     @classmethod
+    @deprecated(
+        reason="This classmethod cannot reflect all __init__ args and will be removed after v27.",
+        version="v26.0",
+        category=FutureWarning,
+    )
     def getMaxIntForFilters(cls):
         return len(cls.SUPPORTED_FILTERS)
 
-    def __init__(self, fixed: DataCoordinate, dimensions: DimensionGraph):
-        super().__init__(fixed, dimensions)
-        record = fixed.records["skymap"]
-        self._skyMapName = record.name
-        self._patchMax = record.patch_nx_max * record.patch_ny_max
-        self._tractPatchMax = self._patchMax*record.tract_max
-        if "band" in dimensions:
-            self._filterMax = self.getMaxIntForFilters()
+    def __init__(
+        self,
+        fixed: DataCoordinate,
+        dimensions: DimensionGraph | None = None,
+        bands: Mapping[str, int] | None = None,
+        n_bands: int | None = None,
+        n_tracts: int | None = None,
+        n_patches: int | None = None,
+    ):
+        if bands is None:
+            bands = {b: i for i, b in enumerate(self.SUPPORTED_FILTERS)}
+        if dimensions is None:
+            if n_bands is None:
+                n_bands = len(bands)
+            dimension_names = ["tract", "patch"]
+            if n_bands != 0:
+                dimension_names.append("band")
+            dimensions = fixed.universe.extract(dimension_names)
         else:
-            self._filterMax = None
+            if "band" not in dimensions.names:
+                n_bands = 0
+                if dimensions.names != {"tract", "patch", "skymap"}:
+                    raise ValueError(
+                        f"Invalid dimensions for skymap dimension packer with n_bands=0: {dimensions}."
+                    )
+            else:
+                if dimensions.names != {"tract", "patch", "skymap", "band"}:
+                    raise ValueError(
+                        f"Invalid dimensions for skymap dimension packer with n_bands>0: {dimensions}."
+                    )
+                if n_bands is None:
+                    n_bands = len(bands)
+        if n_tracts is None:
+            n_tracts = fixed.records["skymap"].tract_max
+        if n_patches is None:
+            n_patches = (
+                fixed.records["skymap"].patch_nx_max
+                * fixed.records["skymap"].patch_ny_max
+            )
+        super().__init__(fixed, dimensions)
+        self._bands = bands
+        self._n_bands = n_bands
+        self._n_tracts = n_tracts
+        self._n_patches = n_patches
+        self._bands_list = None
+
+    @classmethod
+    def make_config_field(
+        cls,
+        doc: str = "How to pack tract, patch, and possibly band into an integer."
+    ) -> ConfigurableField:
+        """Make a config field to control how skymap data IDs are packed.
+
+        Parameters
+        ----------
+        doc : `str`, optional
+            Documentation for the config field.
+
+        Returns
+        -------
+        field : `lsst.pex.config.ConfigurableField`
+            A config field whose instance values are [wrapper proxies to]
+            `SkyMapDimensionPackerConfig` instances.
+        """
+        return ConfigurableField(doc, target=cls.from_config, ConfigClass=cls.ConfigClass)
+
+    @classmethod
+    def from_config(
+        cls, data_id: DataCoordinate, config: SkyMapDimensionPackerConfig
+    ) -> SkyMapDimensionPacker:
+        """Construct a dimension packer from a config object and a data ID.
+
+        Parameters
+        ----------
+        data_id : `lsst.daf.butler.DataCoordinate`
+            Data ID that identifies at least the ``skymap`` dimension.  Must
+            have dimension records attached unless ``config.n_tracts`` and
+            ``config.n_patches`` are both not `None`.
+        config : `SkyMapDimensionPackerConfig`
+            Configuration object.
+
+        Returns
+        -------
+        packer : `SkyMapDimensionPackerConfig`
+            New dimension packer.
+
+        Notes
+        -----
+        This interface is provided for consistency with the `lsst.pex.config`
+        "Configurable" concept, and is invoked when ``apply(data_id)`` is
+        called on a config instance attribute that corresponds to a field
+        created by `make_config_field`.  The constructor signature cannot play
+        this role easily for backwards compatibility reasons.
+        """
+        return cls(
+            data_id.subset(data_id.universe.extract(["skymap"])),
+            n_bands=config.n_bands,
+            bands=config.bands,
+            n_tracts=config.n_tracts,
+            n_patches=config.n_patches,
+        )
 
     @property
     def maxBits(self) -> int:
         # Docstring inherited from DataIdPacker.maxBits
-        packedMax = self._tractPatchMax
-        if self._filterMax is not None:
-            packedMax *= self._filterMax
-        return packedMax.bit_length()
+        packedMax = self._n_tracts * self._n_patches
+        if self._n_bands:
+            packedMax *= self._n_bands
+        return (packedMax - 1).bit_length()
 
     def _pack(self, dataId: DataCoordinate) -> int:
         # Docstring inherited from DataIdPacker.pack
-        packed = dataId["patch"] + self._patchMax*dataId["tract"]
-        if self._filterMax is not None:
-            packed += self.getIntFromFilter(dataId["band"])*self._tractPatchMax
+        if dataId["patch"] >= self._n_patches:
+            raise ValueError(f"Patch ID {dataId['patch']} is out of bounds; expected <{self._n_patches}.")
+        if dataId["tract"] >= self._n_tracts:
+            raise ValueError(f"Tract ID {dataId['tract']} is out of bounds; expected <{self._n_tracts}.")
+        packed = dataId["patch"] + self._n_patches * dataId["tract"]
+        if self._n_bands:
+            if (band_index := self._bands.get(dataId["band"])) is None:
+                raise ValueError(
+                    f"Band {dataId['band']!r} is not supported by SkyMapDimensionPacker "
+                    f"configuration; expected one of {list(self._bands)}."
+                )
+            if band_index >= self._n_bands:
+                raise ValueError(
+                    f"Band index {band_index} for {dataId['band']!r} is out of bounds; "
+                    f"expected <{self._n_bands}."
+                )
+            packed += self._bands[dataId["band"]] * self._n_patches * self._n_tracts
         return packed
 
     def unpack(self, packedId: int) -> DataCoordinate:
         # Docstring inherited from DataIdPacker.unpack
-        d = {"skymap": self._skyMapName}
-        if self._filterMax is not None:
-            d["band"] = self.getFilterNameFromInt(packedId // self._tractPatchMax)
-            packedId %= self._tractPatchMax
-        d["tract"] = packedId // self._patchMax
-        d["patch"] = packedId % self._patchMax
+        d = {"skymap": self.fixed["skymap"]}
+        if self._n_bands:
+            index, packedId = divmod(packedId, (self._n_tracts * self._n_patches))
+            if self._bands_list is None:
+                self._bands_list = list(self._bands)
+            d["band"] = self._bands_list[index]
+        d["tract"], d["patch"] = divmod(packedId, self._n_patches)
         return DataCoordinate.standardize(d, graph=self.dimensions)

tests/test_packers.py

@@ -78,6 +78,68 @@ def testWithFilter(self):
         self.assertLessEqual(packedId.bit_length(), packer.maxBits)
         self.assertEqual(packer.unpack(packedId), dataId)
 
+    def test_bad_dimensions(self):
+        with self.assertRaises(ValueError):
+            SkyMapDimensionPacker(
+                self.fixed,
+                DimensionGraph(universe=self.universe, names=["tract", "patch", "visit"]),
+            )
+        with self.assertRaises(ValueError):
+            SkyMapDimensionPacker(
+                self.fixed,
+                DimensionGraph(universe=self.universe, names=["tract", "patch", "detector"]),
+            )
+
+    def test_from_config(self):
+        data_id = DataCoordinate.standardize(
+            skymap=self.fixed["skymap"],
+            tract=2,
+            patch=6,
+            band="g",
+            universe=self.universe
+        )
+        config = SkyMapDimensionPacker.ConfigClass()
+        config.n_tracts = 5
+        config.n_patches = 9
+        config.n_bands = 3
+        config.bands = {"r": 0, "g": 1}
+        packer = SkyMapDimensionPacker.from_config(data_id, config=config)
+        packed_id = packer.pack(data_id)
+        self.assertLessEqual(packed_id.bit_length(), packer.maxBits)
+        self.assertEqual(packer.unpack(packed_id), data_id)
+
+    def test_from_config_no_bands(self):
+        data_id = DataCoordinate.standardize(
+            skymap=self.fixed["skymap"],
+            tract=2,
+            patch=6,
+            universe=self.universe
+        )
+        config = SkyMapDimensionPacker.ConfigClass()
+        config.n_tracts = 5
+        config.n_patches = 9
+        packer = SkyMapDimensionPacker.from_config(data_id, config=config)
+        packed_id = packer.pack(data_id)
+        self.assertLessEqual(packed_id.bit_length(), packer.maxBits)
+        self.assertEqual(packer.unpack(packed_id), data_id)
+
+    def test_from_config_default_bands(self):
+        data_id = DataCoordinate.standardize(
+            skymap=self.fixed["skymap"],
+            tract=2,
+            patch=6,
+            band="g",
+            universe=self.universe
+        )
+        config = SkyMapDimensionPacker.ConfigClass()
+        config.n_tracts = 5
+        config.n_patches = 9
+        config.n_bands = None
+        packer = SkyMapDimensionPacker.from_config(data_id, config=config)
+        packed_id = packer.pack(data_id)
+        self.assertLessEqual(packed_id.bit_length(), packer.maxBits)
+        self.assertEqual(packer.unpack(packed_id), data_id)
+
 
 class MemoryTester(lsst.utils.tests.MemoryTestCase):
     pass