Coverage for python/lsst/images/_cell_grid.py: 65%
117 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-04 17:17 +0000
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# This module is conceptually part of the 'cells' subpackage, but we don't
15# want the stuff in '_concrete_bounds' to depend on all of that. So the
16# basic CellGrid and CellGridBounds objects are defined here, used in both
17# places, and exported from 'cells'.
19__all__ = (
20 "CellGrid",
21 "CellGridBounds",
22 "CellIJ",
23 "PatchDefinition",
24)
26import dataclasses
27import math
28from collections.abc import Iterator
29from functools import cached_property
30from typing import TYPE_CHECKING, Any, overload
32import numpy as np
33import pydantic
35from ._geom import YX, Bounds, Box
37if TYPE_CHECKING:
38 try:
39 from lsst.cell_coadds import UniformGrid as LegacyUniformGrid
40 from lsst.skymap import Index2D as LegacyIndex2D
41 except ImportError:
42 type LegacyUniformGrid = Any # type: ignore[no-redef]
43 type LegacyIndex2D = Any # type: ignore[no-redef]
46@dataclasses.dataclass(frozen=True, order=True)
47class CellIJ:
48 """An index in a grid of cells.
50 Notes
51 -----
52 This is deliberately not a `tuple` or other `~collections.abc.Sequence` in
53 order to make it typing-incompatible with sequence-based pixel coordinate
54 pairs (e.g. `.YX`). This also allows it to have addition and subtraction
55 operators.
56 """
58 i: int
59 """The y / row object."""
61 j: int
62 """The x / column object."""
64 def __add__(self, other: CellIJ) -> CellIJ:
65 return CellIJ(i=self.i + other.i, j=self.j + other.j)
67 def __sub__(self, other: CellIJ) -> CellIJ:
68 return CellIJ(i=self.i - other.i, j=self.j - other.j)
70 @staticmethod
71 def from_legacy(legacy_index: LegacyIndex2D) -> CellIJ:
72 """Convert from a legacy `lsst.skymap.Index2D` instance.
74 Parameters
75 ----------
76 legacy_index
77 Legacy `lsst.skymap.Index2D` to convert.
79 Notes
80 -----
81 `lsst.skymap.Index2D` is ordered ``(x, y)``, i.e. ``(j, i)``.
82 """
83 return CellIJ(i=legacy_index.y, j=legacy_index.x)
85 def to_legacy(self) -> LegacyIndex2D:
86 """Convert to a legacy `lsst.skymap.Index2D` instance.
88 Notes
89 -----
90 `lsst.skymap.Index2D` is ordered ``(x, y)``, i.e. ``(j, i)``.
91 """
92 from lsst.skymap import Index2D as LegacyIndex2D
94 return LegacyIndex2D(x=self.j, y=self.i)
96 def as_tuple(self) -> tuple[int, int]:
97 """Convert to an (i, j) `tuple`."""
98 return (self.i, self.j)
101class CellGrid(pydantic.BaseModel, frozen=True):
102 """A grid of rectangular cells with no overlaps or space between cells.
104 Notes
105 -----
106 A cell grid usually corresponds to a full patch, but we do not explicitly
107 encode this in the type to permit full-tract grids, which would have to
108 drop the cells in patch overlap regions and re-label all cells.
110 Subsets of grids are usually represented via `CellGridBounds`.
111 """
113 bbox: Box = pydantic.Field(
114 description=(
115 "Bounding box of the grid of cells (snapped to cell boundaries. "
116 "The cell with index (i=0, j=0) always has a corner at ``(y=bbox.y.min, x=bbox.x.min)`` "
117 "but there is no expectation that ``(y=bbox.y.min, x=bbox.x.min)`` be ``(y=0, x=0)``."
118 )
119 )
120 cell_shape: YX[int] = pydantic.Field(description="Shape of each cell in pixels.")
122 @property
123 def grid_size(self) -> CellIJ:
124 """The number of cells in each dimension (`CellIJ`)."""
125 return CellIJ(i=self.bbox.y.size // self.cell_shape.y, j=self.bbox.x.size // self.cell_shape.x)
127 def index_of(self, *, y: int, x: int) -> CellIJ:
128 """Return the 2-d index of the cell that contains the given pixel.
130 Parameters
131 ----------
132 y
133 Y cell index.
134 x
135 X cell index.
136 """
137 return CellIJ(
138 i=(y - self.bbox.y.start) // self.cell_shape.y,
139 j=(x - self.bbox.x.start) // self.cell_shape.x,
140 )
142 def bbox_of(self, cell: CellIJ) -> Box:
143 """Return the bounding box of the given cell.
145 Parameters
146 ----------
147 cell
148 Index of the cell whose bounding box is returned.
149 """
150 return Box.from_shape(
151 self.cell_shape,
152 start=YX(
153 y=cell.i * self.cell_shape.y + self.bbox.y.start,
154 x=cell.j * self.cell_shape.x + self.bbox.x.start,
155 ),
156 )
158 @staticmethod
159 def from_legacy(legacy: LegacyUniformGrid) -> CellGrid:
160 """Construct from a legacy `lsst.cell_coadds.UniformGrid` object.
162 Parameters
163 ----------
164 legacy
165 Legacy grid to convert.
166 """
167 if legacy.padding:
168 raise ValueError("Only cell grids with no padding are supported.")
169 bbox = Box.from_legacy(legacy.bbox)
170 cell_shape = YX(y=legacy.cell_size.y, x=legacy.cell_size.x)
171 return CellGrid(bbox=bbox, cell_shape=cell_shape)
173 def to_legacy(self) -> LegacyUniformGrid:
174 """Convert to a legacy `lsst.cell_coadds.UniformGrid` object."""
175 from lsst.cell_coadds import UniformGrid as LegacyUniformGrid
177 return LegacyUniformGrid(
178 self.cell_shape.to_legacy_int_extent(),
179 self.grid_size.to_legacy(),
180 min=self.bbox.min.to_legacy_int_point(),
181 )
184class CellGridBounds(pydantic.BaseModel, frozen=True):
185 """A region of pixels defined by a set of cells within a grid.
187 Notes
188 -----
189 This data structure is optimized for the case where a continguous
190 rectangular region of the grid (the `bbox` attribute) is populated with
191 only a few exceptions (the `missing` set).
193 Slicing a `CellGridBounds` with a `.Box` returns a new `CellGridBounds`
194 with just the cells that overlap that box. As always,
195 `CellGridBounds.bbox` will be snapped to the outer boundaries of those
196 cells, so it will contain (and not generally equal) the given box.
197 """
199 grid: CellGrid = pydantic.Field(description="Definition of the grid that defines the cells.")
200 bbox: Box = pydantic.Field(description="Pixel bounding box of the region (snapped to cell boundaries).")
201 missing: frozenset[CellIJ] = pydantic.Field(
202 default=frozenset(),
203 description=(
204 "Indices of cells that are missing, where (i=0, j=0) is the cell that starts at grid.bbox.start."
205 ),
206 )
208 @cached_property
209 def subgrid_start(self) -> CellIJ:
210 """The index of the first cell in this bounds' bounding box within
211 its grid.
212 """
213 return self.grid.index_of(y=self.bbox.y.start, x=self.bbox.x.start)
215 @cached_property
216 def subgrid_stop(self) -> CellIJ:
217 """One-past-the-last indices for the cells in these bounds, within
218 its grid.
219 """
220 return self.grid.index_of(y=self.bbox.y.stop, x=self.bbox.x.stop)
222 @cached_property
223 def subgrid_size(self) -> CellIJ:
224 """Number of cells within these bounds in both dimensions, not
225 accounting for `missing`.
226 """
227 return self.subgrid_stop - self.subgrid_start
229 @overload
230 def contains(self, *, x: int, y: int) -> bool: ... 230 ↛ exitline 230 didn't return from function 'contains' because
232 @overload
233 def contains(self, *, x: np.ndarray, y: np.ndarray) -> np.ndarray: ... 233 ↛ exitline 233 didn't return from function 'contains' because
235 def contains(self, *, x: Any, y: Any) -> Any:
236 """Test whether these bounds contain one or more points.
238 Parameters
239 ----------
240 x
241 One or more integer X coordinates to test for containment.
242 If an array, an array of results will be returned.
243 y
244 One or more integer Y coordinates to test for containment.
245 If an array, an array of results will be returned.
247 Returns
248 -------
249 `bool` | `numpy.ndarray`
250 If ``x`` and ``y`` are both scalars, a single `bool` value. If
251 ``x`` and ``y`` are arrays, a boolean array with their broadcasted
252 shape.
253 """
254 result = self.bbox.contains(x=x, y=y)
255 if not self.missing:
256 return result
257 match result:
258 case False:
259 return False
260 case True:
261 return self.grid.index_of(x=x, y=y) not in self.missing
262 case np.ndarray():
263 for box in self.missing_boxes():
264 result = np.logical_and(result, np.logical_not(box.contains(x=x, y=y)))
265 return result
267 def intersection(self, other: Bounds) -> Bounds:
268 """Compute the intersection of this bounds object with another.
270 Parameters
271 ----------
272 other
273 Bounds to intersect with this one.
274 """
275 from ._concrete_bounds import _intersect_cgb
277 return _intersect_cgb(self, other)
279 def contains_cell(self, index: CellIJ) -> bool:
280 """Test whether the given cell is in the bounds.
282 Parameters
283 ----------
284 index
285 Index of the cell to test.
286 """
287 return (
288 (index.i >= self.subgrid_start.i and index.i < self.subgrid_stop.i)
289 and (index.j >= self.subgrid_start.j and index.j < self.subgrid_stop.j)
290 and index not in self.missing
291 )
293 def missing_boxes(self) -> Iterator[Box]:
294 """Iterate over the bounding boxes of the missing cells."""
295 for index in sorted(self.missing):
296 yield self.grid.bbox_of(index)
298 def cell_indices(self) -> Iterator[CellIJ]:
299 """Iterate over the indices of the cells in these bounds."""
300 for i in range(self.subgrid_start.i, self.subgrid_stop.i):
301 for j in range(self.subgrid_start.j, self.subgrid_stop.j):
302 index = CellIJ(i=i, j=j)
303 if index not in self.missing:
304 yield index
306 def __getitem__(self, bbox: Box) -> CellGridBounds:
307 if not self.bbox.contains(bbox): 307 ↛ 308line 307 didn't jump to line 308 because the condition on line 307 was never true
308 raise ValueError(
309 f"Original grid bounding box {self.bbox} does not contain the subset bounding box {bbox}."
310 )
311 c = self.grid.cell_shape
312 s = self.grid.bbox.start
313 i1 = (bbox.y.start - s.y) // c.y
314 j1 = (bbox.x.start - s.x) // c.x
315 i2 = math.ceil((bbox.y.stop - s.y) / c.y)
316 j2 = math.ceil((bbox.x.stop - s.x) / c.x)
317 subset_bbox = Box.factory[i1 * c.y + s.y : i2 * c.y + s.y, j1 * c.x + s.x : j2 * c.x + s.x]
318 grid_as_box = Box.factory[i1:i2, j1:j2]
319 subset_missing = {index for index in self.missing if grid_as_box.contains(y=index.i, x=index.j)}
320 return CellGridBounds(grid=self.grid, bbox=subset_bbox, missing=frozenset(subset_missing))
322 def serialize(self) -> CellGridBounds:
323 """Convert a bounds instance into a serializable object."""
324 return self
326 def deserialize(self) -> CellGridBounds:
327 """Deserialize a bounds object on the assumption it is a
328 `CellGridBounds`.
330 This method just returns the `CellGridBounds` itself, since that
331 already provides Pydantic serialization hooks. It exists for
332 compatibility with the `.Bounds` protocol.
333 """
334 return self
337class PatchDefinition(pydantic.BaseModel, frozen=True):
338 """Identifiers and geometry for a full patch."""
340 id: int = pydantic.Field(description="ID for the patch.")
341 index: YX[int] = pydantic.Field(description="2-d index of this patch within the tract.")
342 inner_bbox: Box = pydantic.Field(description="Inner bounding box of this patch.")
343 cells: CellGrid = pydantic.Field(description="Cell grid for the full patch.")
345 @property
346 def outer_bbox(self) -> Box:
347 """The outer bounding box of this patch (`.Box`)."""
348 return self.cells.bbox