Coverage for python/lsst/images/_cell_grid.py: 41%
151 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-03 01:36 +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, assert_type, overload
32import numpy as np
33import numpy.typing as npt
34import pydantic
36from ._geom import XY, YX, Bounds, Box, SerializableYX
38if TYPE_CHECKING:
39 try:
40 from lsst.cell_coadds import UniformGrid as LegacyUniformGrid
41 from lsst.skymap import Index2D as LegacyIndex2D
42 from lsst.skymap import PatchInfo as LegacyPatchInfo
43 except ImportError:
44 type LegacyUniformGrid = Any # type: ignore[no-redef]
45 type LegacyIndex2D = Any # type: ignore[no-redef]
46 type LegacyPatchInfo = Any # type: ignore[no-redef]
49@dataclasses.dataclass(frozen=True, order=True)
50class CellIJ:
51 """An index in a grid of cells.
53 Notes
54 -----
55 This is deliberately not a `tuple` or other `~collections.abc.Sequence` in
56 order to make it typing-incompatible with sequence-based pixel coordinate
57 pairs (e.g. `.YX`). This also allows it to have addition and subtraction
58 operators.
59 """
61 i: int
62 """The y / row object."""
64 j: int
65 """The x / column object."""
67 def __add__(self, other: CellIJ) -> CellIJ:
68 return CellIJ(i=self.i + other.i, j=self.j + other.j)
70 def __sub__(self, other: CellIJ) -> CellIJ:
71 return CellIJ(i=self.i - other.i, j=self.j - other.j)
73 @staticmethod
74 def from_legacy(legacy_index: LegacyIndex2D) -> CellIJ:
75 """Convert from a legacy `lsst.skymap.Index2D` instance.
77 Parameters
78 ----------
79 legacy_index
80 Legacy `lsst.skymap.Index2D` to convert.
82 Notes
83 -----
84 `lsst.skymap.Index2D` is ordered ``(x, y)``, i.e. ``(j, i)``.
85 """
86 return CellIJ(i=legacy_index.y, j=legacy_index.x)
88 def to_legacy(self) -> LegacyIndex2D:
89 """Convert to a legacy `lsst.skymap.Index2D` instance.
91 Notes
92 -----
93 `lsst.skymap.Index2D` is ordered ``(x, y)``, i.e. ``(j, i)``.
94 """
95 from lsst.skymap import Index2D as LegacyIndex2D
97 return LegacyIndex2D(x=self.j, y=self.i)
99 def as_tuple(self) -> tuple[int, int]:
100 """Convert to an (i, j) `tuple`."""
101 return (self.i, self.j)
104class CellGrid(pydantic.BaseModel, frozen=True):
105 """A grid of rectangular cells with no overlaps or space between cells.
107 Notes
108 -----
109 A cell grid usually corresponds to a full patch, but we do not explicitly
110 encode this in the type to permit full-tract grids, which would have to
111 drop the cells in patch overlap regions and re-label all cells.
113 Subsets of grids are usually represented via `CellGridBounds`.
114 """
116 bbox: Box = pydantic.Field(
117 description=(
118 "Bounding box of the grid of cells (snapped to cell boundaries. "
119 "The cell with index (i=0, j=0) always has a corner at ``(y=bbox.y.min, x=bbox.x.min)`` "
120 "but there is no expectation that ``(y=bbox.y.min, x=bbox.x.min)`` be ``(y=0, x=0)``."
121 )
122 )
123 cell_shape: SerializableYX[int] = pydantic.Field(description="Shape of each cell in pixels.")
125 def __str__(self) -> str:
126 return f"{self.bbox}, cell_shape=(y={self.cell_shape.y}, x={self.cell_shape.x})"
128 @property
129 def grid_size(self) -> CellIJ:
130 """The number of cells in each dimension (`CellIJ`)."""
131 return CellIJ(i=self.bbox.y.size // self.cell_shape.y, j=self.bbox.x.size // self.cell_shape.x)
133 def index_of(self, *, y: int, x: int) -> CellIJ:
134 """Return the 2-d index of the cell that contains the given pixel.
136 Parameters
137 ----------
138 y
139 Y cell index.
140 x
141 X cell index.
142 """
143 return CellIJ(
144 i=(y - self.bbox.y.start) // self.cell_shape.y,
145 j=(x - self.bbox.x.start) // self.cell_shape.x,
146 )
148 def bbox_of(self, cell: CellIJ) -> Box:
149 """Return the bounding box of the given cell.
151 Parameters
152 ----------
153 cell
154 Index of the cell whose bounding box is returned.
155 """
156 return Box.from_shape(
157 self.cell_shape,
158 start=YX(
159 y=cell.i * self.cell_shape.y + self.bbox.y.start,
160 x=cell.j * self.cell_shape.x + self.bbox.x.start,
161 ),
162 )
164 @staticmethod
165 def from_legacy(legacy: LegacyUniformGrid) -> CellGrid:
166 """Construct from a legacy `lsst.cell_coadds.UniformGrid` object.
168 Parameters
169 ----------
170 legacy
171 Legacy grid to convert.
172 """
173 if legacy.padding:
174 raise ValueError("Only cell grids with no padding are supported.")
175 bbox = Box.from_legacy(legacy.bbox)
176 cell_shape = YX(y=legacy.cell_size.y, x=legacy.cell_size.x)
177 return CellGrid(bbox=bbox, cell_shape=cell_shape)
179 def to_legacy(self) -> LegacyUniformGrid:
180 """Convert to a legacy `lsst.cell_coadds.UniformGrid` object."""
181 from lsst.cell_coadds import UniformGrid as LegacyUniformGrid
183 return LegacyUniformGrid(
184 self.cell_shape.to_legacy_int_extent(),
185 self.grid_size.to_legacy(),
186 min=self.bbox.min.to_legacy_int_point(),
187 )
190class CellGridBounds(pydantic.BaseModel, frozen=True):
191 """A region of pixels defined by a set of cells within a grid.
193 Notes
194 -----
195 This data structure is optimized for the case where a continguous
196 rectangular region of the grid (the `bbox` attribute) is populated with
197 only a few exceptions (the `missing` set).
199 Slicing a `CellGridBounds` with a `.Box` returns a new `CellGridBounds`
200 with just the cells that overlap that box. As always,
201 `CellGridBounds.bbox` will be snapped to the outer boundaries of those
202 cells, so it will contain (and not generally equal) the given box.
203 """
205 grid: CellGrid = pydantic.Field(description="Definition of the grid that defines the cells.")
206 bbox: Box = pydantic.Field(description="Pixel bounding box of the region (snapped to cell boundaries).")
207 missing: frozenset[CellIJ] = pydantic.Field(
208 default=frozenset(),
209 description=(
210 "Indices of cells that are missing, where (i=0, j=0) is the cell that starts at grid.bbox.start."
211 ),
212 )
214 def __str__(self) -> str:
215 text = f"{self.bbox} in grid ({self.grid})"
216 if self.missing:
217 missing = ", ".join(f"(i={c.i}, j={c.j})" for c in sorted(self.missing))
218 text = f"{text}, missing={{{missing}}}"
219 return text
221 @cached_property
222 def subgrid_start(self) -> CellIJ:
223 """The index of the first cell in this bounds' bounding box within
224 its grid.
225 """
226 return self.grid.index_of(y=self.bbox.y.start, x=self.bbox.x.start)
228 @cached_property
229 def subgrid_stop(self) -> CellIJ:
230 """One-past-the-last indices for the cells in these bounds, within
231 its grid.
232 """
233 return self.grid.index_of(y=self.bbox.y.stop, x=self.bbox.x.stop)
235 @cached_property
236 def subgrid_size(self) -> CellIJ:
237 """Number of cells within these bounds in both dimensions, not
238 accounting for `missing`.
239 """
240 return self.subgrid_stop - self.subgrid_start
242 @property
243 def n_cells(self) -> int:
244 """Return the number of cells within this bounds, accounting for
245 `missing`.
246 """
247 return self.subgrid_size.i * self.subgrid_size.j - len(self.missing)
249 @overload
250 def contains(self, point: XY[int | float] | YX[int | float], /) -> bool: ...
252 @overload
253 def contains(self, point: XY[npt.ArrayLike] | YX[npt.ArrayLike], /) -> np.ndarray: ...
255 @overload
256 def contains(self, *, x: int | float, y: int | float) -> bool: ...
258 @overload
259 def contains(self, *, x: npt.ArrayLike, y: npt.ArrayLike) -> np.ndarray: ...
261 def contains(self, point: XY[Any] | YX[Any] | None = None, /, *, x: Any = None, y: Any = None) -> Any: # type: ignore[misc]
262 """Test whether these bounds contain one or more points.
264 Parameters
265 ----------
266 point
267 An `.XY` or `.YX` coordinate pair to test for containment.
268 Mutually exclusive with ``x`` and ``y``.
269 x
270 One or more X coordinates to test for containment, as a scalar or
271 any array-like. Results are broadcast against ``y``.
272 Mutually exclusive with ``point``.
273 y
274 One or more Y coordinates to test for containment, as a scalar or
275 any array-like. Results are broadcast against ``x``.
276 Mutually exclusive with ``point``.
278 Returns
279 -------
280 `bool` | `numpy.ndarray`
281 If ``x`` and ``y`` are both scalars, a single `bool` value. If
282 ``x`` and ``y`` are array-like, a boolean array with their
283 broadcasted shape.
284 """
285 match point:
286 case None:
287 if x is None or y is None: 287 ↛ 288line 287 didn't jump to line 288 because the condition on line 287 was never true
288 raise TypeError("Pass either a point or both x= and y= to 'CellGridBounds.contains'.")
289 case XY() | YX(): 289 ↛ 295line 289 didn't jump to line 295 because the pattern on line 289 always matched
290 if x is not None or y is not None: 290 ↛ 291line 290 didn't jump to line 291 because the condition on line 290 was never true
291 raise TypeError(
292 "'CellGridBounds.contains' point argument is mutually exclusive with x= and y=."
293 )
294 x, y = point.x, point.y
295 case _:
296 raise TypeError(f"Unexpected positional argument type: {type(point)!r}.")
297 result = self.bbox.contains(x=x, y=y)
298 if not self.missing: 298 ↛ 299line 298 didn't jump to line 299 because the condition on line 298 was never true
299 return result
300 match result:
301 case False:
302 return False
303 case True:
304 return self.grid.index_of(x=x, y=y) not in self.missing
305 case np.ndarray(): 305 ↛ 308line 305 didn't jump to line 308 because the pattern on line 305 always matched
306 for box in self.missing_boxes():
307 result = np.logical_and(result, np.logical_not(box.contains(x=x, y=y)))
308 return result
310 def intersection(self, other: Bounds) -> Bounds:
311 """Compute the intersection of this bounds object with another.
313 Parameters
314 ----------
315 other
316 Bounds to intersect with this one.
317 """
318 from ._concrete_bounds import _intersect_cgb
320 return _intersect_cgb(self, other)
322 def contains_cell(self, index: CellIJ) -> bool:
323 """Test whether the given cell is in the bounds.
325 Parameters
326 ----------
327 index
328 Index of the cell to test.
329 """
330 return (
331 (index.i >= self.subgrid_start.i and index.i < self.subgrid_stop.i)
332 and (index.j >= self.subgrid_start.j and index.j < self.subgrid_stop.j)
333 and index not in self.missing
334 )
336 def missing_boxes(self) -> Iterator[Box]:
337 """Iterate over the bounding boxes of the missing cells."""
338 for index in sorted(self.missing):
339 yield self.grid.bbox_of(index)
341 def cell_indices(self) -> Iterator[CellIJ]:
342 """Iterate over the indices of the cells in these bounds."""
343 for i in range(self.subgrid_start.i, self.subgrid_stop.i):
344 for j in range(self.subgrid_start.j, self.subgrid_stop.j):
345 index = CellIJ(i=i, j=j)
346 if index not in self.missing:
347 yield index
349 def __getitem__(self, bbox: Box) -> CellGridBounds:
350 if not self.bbox.contains(bbox): 350 ↛ 351line 350 didn't jump to line 351 because the condition on line 350 was never true
351 raise ValueError(
352 f"Original grid bounding box {self.bbox} does not contain the subset bounding box {bbox}."
353 )
354 c = self.grid.cell_shape
355 s = self.grid.bbox.start
356 i1 = (bbox.y.start - s.y) // c.y
357 j1 = (bbox.x.start - s.x) // c.x
358 i2 = math.ceil((bbox.y.stop - s.y) / c.y)
359 j2 = math.ceil((bbox.x.stop - s.x) / c.x)
360 subset_bbox = Box.factory[i1 * c.y + s.y : i2 * c.y + s.y, j1 * c.x + s.x : j2 * c.x + s.x]
361 grid_as_box = Box.factory[i1:i2, j1:j2]
362 subset_missing = {index for index in self.missing if grid_as_box.contains(y=index.i, x=index.j)}
363 return CellGridBounds(grid=self.grid, bbox=subset_bbox, missing=frozenset(subset_missing))
365 def serialize(self) -> CellGridBounds:
366 """Convert a bounds instance into a serializable object."""
367 return self
369 def deserialize(self) -> CellGridBounds:
370 """Deserialize a bounds object on the assumption it is a
371 `CellGridBounds`.
373 This method just returns the `CellGridBounds` itself, since that
374 already provides Pydantic serialization hooks. It exists for
375 compatibility with the `.Bounds` protocol.
376 """
377 return self
380class PatchDefinition(pydantic.BaseModel, frozen=True):
381 """Identifiers and geometry for a full patch."""
383 id: int = pydantic.Field(description="ID for the patch.")
384 index: SerializableYX[int] = pydantic.Field(description="2-d index of this patch within the tract.")
385 inner_bbox: Box = pydantic.Field(description="Inner bounding box of this patch.")
386 cells: CellGrid = pydantic.Field(description="Cell grid for the full patch.")
388 def __str__(self) -> str:
389 return (
390 f"id={self.id}, index=(y={self.index.y}, x={self.index.x}), "
391 f"inner_bbox={self.inner_bbox}, cells=({self.cells})"
392 )
394 @property
395 def outer_bbox(self) -> Box:
396 """The outer bounding box of this patch (`.Box`)."""
397 return self.cells.bbox
399 @staticmethod
400 def from_legacy_patch_info(patch_info: LegacyPatchInfo) -> PatchDefinition:
401 """Convert from a legacy `lsst.skymap.PatchInfo` object.
403 Parameters
404 ----------
405 patch_info
406 Legacy patch definition to convert.
407 """
408 index = patch_info.getIndex()
409 cell_extent = patch_info.getCellInnerDimensions()
410 cell_grid = CellGrid(
411 bbox=Box.from_legacy(patch_info.getOuterBBox()),
412 cell_shape=YX(y=cell_extent.y, x=cell_extent.x),
413 )
414 return PatchDefinition(
415 id=patch_info.getSequentialIndex(),
416 index=YX(y=index.y, x=index.x),
417 inner_bbox=Box.from_legacy(patch_info.getInnerBBox()),
418 cells=cell_grid,
419 )
422if TYPE_CHECKING:
424 def _test_types() -> None:
425 arr = np.zeros(3)
426 bbox = Box.from_shape((100, 200))
427 grid = CellGrid(bbox=bbox, cell_shape=YX(10, 20))
428 cgb = CellGridBounds(grid=grid, bbox=bbox)
430 # CellGridBounds satisfies the Bounds Protocol.
431 bounds: Bounds = cgb
433 # CellGridBounds.contains: XY/YX, scalar, array-like
434 assert_type(cgb.contains(x=1, y=2), bool)
435 assert_type(cgb.contains(x=1.0, y=2.0), bool)
436 assert_type(cgb.contains(x=arr, y=arr), np.ndarray)
437 assert_type(cgb.contains(XY(1, 2)), bool)
438 assert_type(cgb.contains(YX(2, 1)), bool)
439 assert_type(cgb.contains(XY(arr, arr)), np.ndarray)
440 assert_type(cgb.contains(YX(arr, arr)), np.ndarray)
442 # Via the Bounds Protocol view, same signatures hold.
443 assert_type(bounds.contains(x=1, y=1), bool)
444 assert_type(bounds.contains(x=1.0, y=1.0), bool)
445 assert_type(bounds.contains(x=arr, y=arr), np.ndarray)
446 assert_type(bounds.contains(XY(1, 1)), bool)
447 assert_type(bounds.contains(YX(1, 1)), bool)
448 assert_type(bounds.contains(XY(arr, arr)), np.ndarray)
449 assert_type(bounds.contains(YX(arr, arr)), np.ndarray)