Coverage for python/lsst/images/_cell_grid.py: 80%
132 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-16 15:24 -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# 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, assert_type, overload
32import numpy as np
33import numpy.typing as npt
34import pydantic
36from ._geom import XY, YX, Bounds, Box
38if TYPE_CHECKING:
39 try:
40 from lsst.cell_coadds import UniformGrid as LegacyUniformGrid
41 from lsst.skymap import Index2D as LegacyIndex2D
42 except ImportError:
43 type LegacyUniformGrid = Any # type: ignore[no-redef]
44 type LegacyIndex2D = Any # type: ignore[no-redef]
47@dataclasses.dataclass(frozen=True, order=True)
48class CellIJ:
49 """An index in a grid of cells.
51 Notes
52 -----
53 This is deliberately not a `tuple` or other `~collections.abc.Sequence` in
54 order to make it typing-incompatible with sequence-based pixel coordinate
55 pairs (e.g. `.YX`). This also allows it to have addition and subtraction
56 operators.
57 """
59 i: int
60 """The y / row object."""
62 j: int
63 """The x / column object."""
65 def __add__(self, other: CellIJ) -> CellIJ:
66 return CellIJ(i=self.i + other.i, j=self.j + other.j)
68 def __sub__(self, other: CellIJ) -> CellIJ:
69 return CellIJ(i=self.i - other.i, j=self.j - other.j)
71 @staticmethod
72 def from_legacy(legacy_index: LegacyIndex2D) -> CellIJ:
73 """Convert from a legacy `lsst.skymap.Index2D` instance.
75 Parameters
76 ----------
77 legacy_index
78 Legacy `lsst.skymap.Index2D` to convert.
80 Notes
81 -----
82 `lsst.skymap.Index2D` is ordered ``(x, y)``, i.e. ``(j, i)``.
83 """
84 return CellIJ(i=legacy_index.y, j=legacy_index.x)
86 def to_legacy(self) -> LegacyIndex2D:
87 """Convert to a legacy `lsst.skymap.Index2D` instance.
89 Notes
90 -----
91 `lsst.skymap.Index2D` is ordered ``(x, y)``, i.e. ``(j, i)``.
92 """
93 from lsst.skymap import Index2D as LegacyIndex2D
95 return LegacyIndex2D(x=self.j, y=self.i)
97 def as_tuple(self) -> tuple[int, int]:
98 """Convert to an (i, j) `tuple`."""
99 return (self.i, self.j)
102class CellGrid(pydantic.BaseModel, frozen=True):
103 """A grid of rectangular cells with no overlaps or space between cells.
105 Notes
106 -----
107 A cell grid usually corresponds to a full patch, but we do not explicitly
108 encode this in the type to permit full-tract grids, which would have to
109 drop the cells in patch overlap regions and re-label all cells.
111 Subsets of grids are usually represented via `CellGridBounds`.
112 """
114 bbox: Box = pydantic.Field(
115 description=(
116 "Bounding box of the grid of cells (snapped to cell boundaries. "
117 "The cell with index (i=0, j=0) always has a corner at ``(y=bbox.y.min, x=bbox.x.min)`` "
118 "but there is no expectation that ``(y=bbox.y.min, x=bbox.x.min)`` be ``(y=0, x=0)``."
119 )
120 )
121 cell_shape: YX[int] = pydantic.Field(description="Shape of each cell in pixels.")
123 @property
124 def grid_size(self) -> CellIJ:
125 """The number of cells in each dimension (`CellIJ`)."""
126 return CellIJ(i=self.bbox.y.size // self.cell_shape.y, j=self.bbox.x.size // self.cell_shape.x)
128 def index_of(self, *, y: int, x: int) -> CellIJ:
129 """Return the 2-d index of the cell that contains the given pixel.
131 Parameters
132 ----------
133 y
134 Y cell index.
135 x
136 X cell index.
137 """
138 return CellIJ(
139 i=(y - self.bbox.y.start) // self.cell_shape.y,
140 j=(x - self.bbox.x.start) // self.cell_shape.x,
141 )
143 def bbox_of(self, cell: CellIJ) -> Box:
144 """Return the bounding box of the given cell.
146 Parameters
147 ----------
148 cell
149 Index of the cell whose bounding box is returned.
150 """
151 return Box.from_shape(
152 self.cell_shape,
153 start=YX(
154 y=cell.i * self.cell_shape.y + self.bbox.y.start,
155 x=cell.j * self.cell_shape.x + self.bbox.x.start,
156 ),
157 )
159 @staticmethod
160 def from_legacy(legacy: LegacyUniformGrid) -> CellGrid:
161 """Construct from a legacy `lsst.cell_coadds.UniformGrid` object.
163 Parameters
164 ----------
165 legacy
166 Legacy grid to convert.
167 """
168 if legacy.padding:
169 raise ValueError("Only cell grids with no padding are supported.")
170 bbox = Box.from_legacy(legacy.bbox)
171 cell_shape = YX(y=legacy.cell_size.y, x=legacy.cell_size.x)
172 return CellGrid(bbox=bbox, cell_shape=cell_shape)
174 def to_legacy(self) -> LegacyUniformGrid:
175 """Convert to a legacy `lsst.cell_coadds.UniformGrid` object."""
176 from lsst.cell_coadds import UniformGrid as LegacyUniformGrid
178 return LegacyUniformGrid(
179 self.cell_shape.to_legacy_int_extent(),
180 self.grid_size.to_legacy(),
181 min=self.bbox.min.to_legacy_int_point(),
182 )
185class CellGridBounds(pydantic.BaseModel, frozen=True):
186 """A region of pixels defined by a set of cells within a grid.
188 Notes
189 -----
190 This data structure is optimized for the case where a continguous
191 rectangular region of the grid (the `bbox` attribute) is populated with
192 only a few exceptions (the `missing` set).
194 Slicing a `CellGridBounds` with a `.Box` returns a new `CellGridBounds`
195 with just the cells that overlap that box. As always,
196 `CellGridBounds.bbox` will be snapped to the outer boundaries of those
197 cells, so it will contain (and not generally equal) the given box.
198 """
200 grid: CellGrid = pydantic.Field(description="Definition of the grid that defines the cells.")
201 bbox: Box = pydantic.Field(description="Pixel bounding box of the region (snapped to cell boundaries).")
202 missing: frozenset[CellIJ] = pydantic.Field(
203 default=frozenset(),
204 description=(
205 "Indices of cells that are missing, where (i=0, j=0) is the cell that starts at grid.bbox.start."
206 ),
207 )
209 @cached_property
210 def subgrid_start(self) -> CellIJ:
211 """The index of the first cell in this bounds' bounding box within
212 its grid.
213 """
214 return self.grid.index_of(y=self.bbox.y.start, x=self.bbox.x.start)
216 @cached_property
217 def subgrid_stop(self) -> CellIJ:
218 """One-past-the-last indices for the cells in these bounds, within
219 its grid.
220 """
221 return self.grid.index_of(y=self.bbox.y.stop, x=self.bbox.x.stop)
223 @cached_property
224 def subgrid_size(self) -> CellIJ:
225 """Number of cells within these bounds in both dimensions, not
226 accounting for `missing`.
227 """
228 return self.subgrid_stop - self.subgrid_start
230 @overload
231 def contains(self, point: XY[int | float] | YX[int | float], /) -> bool: ... 231 ↛ exitline 231 didn't return from function 'contains' because
233 @overload
234 def contains(self, point: XY[npt.ArrayLike] | YX[npt.ArrayLike], /) -> np.ndarray: ... 234 ↛ exitline 234 didn't return from function 'contains' because
236 @overload
237 def contains(self, *, x: int | float, y: int | float) -> bool: ... 237 ↛ exitline 237 didn't return from function 'contains' because
239 @overload
240 def contains(self, *, x: npt.ArrayLike, y: npt.ArrayLike) -> np.ndarray: ... 240 ↛ exitline 240 didn't return from function 'contains' because
242 def contains(self, point: XY[Any] | YX[Any] | None = None, /, *, x: Any = None, y: Any = None) -> Any: # type: ignore[misc]
243 """Test whether these bounds contain one or more points.
245 Parameters
246 ----------
247 point
248 An `.XY` or `.YX` coordinate pair to test for containment.
249 Mutually exclusive with ``x`` and ``y``.
250 x
251 One or more X coordinates to test for containment, as a scalar or
252 any array-like. Results are broadcast against ``y``.
253 Mutually exclusive with ``point``.
254 y
255 One or more Y coordinates to test for containment, as a scalar or
256 any array-like. Results are broadcast against ``x``.
257 Mutually exclusive with ``point``.
259 Returns
260 -------
261 `bool` | `numpy.ndarray`
262 If ``x`` and ``y`` are both scalars, a single `bool` value. If
263 ``x`` and ``y`` are array-like, a boolean array with their
264 broadcasted shape.
265 """
266 match point:
267 case None:
268 if x is None or y is None: 268 ↛ 269line 268 didn't jump to line 269 because the condition on line 268 was never true
269 raise TypeError("Pass either a point or both x= and y= to 'CellGridBounds.contains'.")
270 case XY() | YX(): 270 ↛ 276line 270 didn't jump to line 276 because the pattern on line 270 always matched
271 if x is not None or y is not None: 271 ↛ 272line 271 didn't jump to line 272 because the condition on line 271 was never true
272 raise TypeError(
273 "'CellGridBounds.contains' point argument is mutually exclusive with x= and y=."
274 )
275 x, y = point.x, point.y
276 case _:
277 raise TypeError(f"Unexpected positional argument type: {type(point)!r}.")
278 result = self.bbox.contains(x=x, y=y)
279 if not self.missing: 279 ↛ 280line 279 didn't jump to line 280 because the condition on line 279 was never true
280 return result
281 match result:
282 case False:
283 return False
284 case True:
285 return self.grid.index_of(x=x, y=y) not in self.missing
286 case np.ndarray(): 286 ↛ 289line 286 didn't jump to line 289 because the pattern on line 286 always matched
287 for box in self.missing_boxes():
288 result = np.logical_and(result, np.logical_not(box.contains(x=x, y=y)))
289 return result
291 def intersection(self, other: Bounds) -> Bounds:
292 """Compute the intersection of this bounds object with another.
294 Parameters
295 ----------
296 other
297 Bounds to intersect with this one.
298 """
299 from ._concrete_bounds import _intersect_cgb
301 return _intersect_cgb(self, other)
303 def contains_cell(self, index: CellIJ) -> bool:
304 """Test whether the given cell is in the bounds.
306 Parameters
307 ----------
308 index
309 Index of the cell to test.
310 """
311 return (
312 (index.i >= self.subgrid_start.i and index.i < self.subgrid_stop.i)
313 and (index.j >= self.subgrid_start.j and index.j < self.subgrid_stop.j)
314 and index not in self.missing
315 )
317 def missing_boxes(self) -> Iterator[Box]:
318 """Iterate over the bounding boxes of the missing cells."""
319 for index in sorted(self.missing):
320 yield self.grid.bbox_of(index)
322 def cell_indices(self) -> Iterator[CellIJ]:
323 """Iterate over the indices of the cells in these bounds."""
324 for i in range(self.subgrid_start.i, self.subgrid_stop.i):
325 for j in range(self.subgrid_start.j, self.subgrid_stop.j):
326 index = CellIJ(i=i, j=j)
327 if index not in self.missing:
328 yield index
330 def __getitem__(self, bbox: Box) -> CellGridBounds:
331 if not self.bbox.contains(bbox): 331 ↛ 332line 331 didn't jump to line 332 because the condition on line 331 was never true
332 raise ValueError(
333 f"Original grid bounding box {self.bbox} does not contain the subset bounding box {bbox}."
334 )
335 c = self.grid.cell_shape
336 s = self.grid.bbox.start
337 i1 = (bbox.y.start - s.y) // c.y
338 j1 = (bbox.x.start - s.x) // c.x
339 i2 = math.ceil((bbox.y.stop - s.y) / c.y)
340 j2 = math.ceil((bbox.x.stop - s.x) / c.x)
341 subset_bbox = Box.factory[i1 * c.y + s.y : i2 * c.y + s.y, j1 * c.x + s.x : j2 * c.x + s.x]
342 grid_as_box = Box.factory[i1:i2, j1:j2]
343 subset_missing = {index for index in self.missing if grid_as_box.contains(y=index.i, x=index.j)}
344 return CellGridBounds(grid=self.grid, bbox=subset_bbox, missing=frozenset(subset_missing))
346 def serialize(self) -> CellGridBounds:
347 """Convert a bounds instance into a serializable object."""
348 return self
350 def deserialize(self) -> CellGridBounds:
351 """Deserialize a bounds object on the assumption it is a
352 `CellGridBounds`.
354 This method just returns the `CellGridBounds` itself, since that
355 already provides Pydantic serialization hooks. It exists for
356 compatibility with the `.Bounds` protocol.
357 """
358 return self
361class PatchDefinition(pydantic.BaseModel, frozen=True):
362 """Identifiers and geometry for a full patch."""
364 id: int = pydantic.Field(description="ID for the patch.")
365 index: YX[int] = pydantic.Field(description="2-d index of this patch within the tract.")
366 inner_bbox: Box = pydantic.Field(description="Inner bounding box of this patch.")
367 cells: CellGrid = pydantic.Field(description="Cell grid for the full patch.")
369 @property
370 def outer_bbox(self) -> Box:
371 """The outer bounding box of this patch (`.Box`)."""
372 return self.cells.bbox
375if TYPE_CHECKING:
377 def _test_types() -> None:
378 arr = np.zeros(3)
379 bbox = Box.from_shape((100, 200))
380 grid = CellGrid(bbox=bbox, cell_shape=YX(10, 20))
381 cgb = CellGridBounds(grid=grid, bbox=bbox)
383 # CellGridBounds satisfies the Bounds Protocol.
384 bounds: Bounds = cgb
386 # CellGridBounds.contains: XY/YX, scalar, array-like
387 assert_type(cgb.contains(x=1, y=2), bool)
388 assert_type(cgb.contains(x=1.0, y=2.0), bool)
389 assert_type(cgb.contains(x=arr, y=arr), np.ndarray)
390 assert_type(cgb.contains(XY(1, 2)), bool)
391 assert_type(cgb.contains(YX(2, 1)), bool)
392 assert_type(cgb.contains(XY(arr, arr)), np.ndarray)
393 assert_type(cgb.contains(YX(arr, arr)), np.ndarray)
395 # Via the Bounds Protocol view, same signatures hold.
396 assert_type(bounds.contains(x=1, y=1), bool)
397 assert_type(bounds.contains(x=1.0, y=1.0), bool)
398 assert_type(bounds.contains(x=arr, y=arr), np.ndarray)
399 assert_type(bounds.contains(XY(1, 1)), bool)
400 assert_type(bounds.contains(YX(1, 1)), bool)
401 assert_type(bounds.contains(XY(arr, arr)), np.ndarray)
402 assert_type(bounds.contains(YX(arr, arr)), np.ndarray)