Coverage for python / lsst / images / _cell_grid.py: 44%
111 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-05-07 08:34 +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, Self, cast, overload
32import numpy as np
33import pydantic
35from ._geom import YX, Bounds, Box
37if TYPE_CHECKING:
38 from ._concrete_bounds import SerializableBounds
40 try:
41 from lsst.cell_coadds import UniformGrid
42 from lsst.skymap import Index2D
43 except ImportError:
44 type UniformGrid = Any # type: ignore[no-redef]
45 type Index2D = Any # type: ignore[no-redef]
48@dataclasses.dataclass(frozen=True, order=True)
49class CellIJ:
50 """An index in a grid of cells.
52 Notes
53 -----
54 This is deliberately not a `tuple` or other `~collections.abc.Sequence` in
55 order to make it typing-incompatible with sequence-based pixel coordinate
56 pairs (e.g. `.YX`). This also allows it to have addition and subtraction
57 operators.
58 """
60 i: int
61 """The y / row object."""
63 j: int
64 """The x / column object."""
66 def __add__(self, other: CellIJ) -> CellIJ:
67 return CellIJ(i=self.i + other.i, j=self.j + other.j)
69 def __sub__(self, other: CellIJ) -> CellIJ:
70 return CellIJ(i=self.i - other.i, j=self.j - other.j)
72 @staticmethod
73 def from_legacy(legacy_index: Index2D) -> CellIJ:
74 """Convert from a legacy `lsst.skymap.Index2D` instance.
76 Notes
77 -----
78 `lsst.skymap.Index2D` is ordered ``(x, y)``, i.e. ``(j, i)``.
79 """
80 return CellIJ(i=legacy_index.y, j=legacy_index.x)
82 def to_legacy(self) -> Index2D:
83 """Convert to a legacy `lsst.skymap.Index2D` instance.
85 Notes
86 -----
87 `lsst.skymap.Index2D` is ordered ``(x, y)``, i.e. ``(j, i)``.
88 """
89 from lsst.skymap import Index2D
91 return Index2D(x=self.j, y=self.i)
94class CellGrid(pydantic.BaseModel, frozen=True):
95 """A grid of rectangular cells with no overlaps or space between cells.
97 Notes
98 -----
99 A cell grid usually corresponds to a full patch, but we do not explicitly
100 encode this in the type to permit full-tract grids, which would have to
101 drop the cells in patch overlap regions and re-label all cells.
103 Subsets of grids are usually represented via `CellGridBounds`.
104 """
106 bbox: Box = pydantic.Field(
107 description=(
108 "Bounding box of the grid of cells (snapped to cell boundaries. "
109 "The cell with index (i=0, j=0) always has a corner at ``(y=bbox.y.min, x=bbox.x.min)`` "
110 "but there is no expectation that ``(y=bbox.y.min, x=bbox.x.min)`` be ``(y=0, x=0)``."
111 )
112 )
113 cell_shape: YX[int] = pydantic.Field(description="Shape of each cell in pixels.")
115 @property
116 def grid_shape(self) -> CellIJ:
117 """The number of cells in each dimension (`CellIJ`)."""
118 return CellIJ(i=self.bbox.y.size // self.cell_shape.y, j=self.bbox.x.size // self.cell_shape.x)
120 def index_of(self, *, y: int, x: int) -> CellIJ:
121 """Return the 2-d index of the cell that contains the given pixel.
123 Parameters
124 ----------
125 y
126 Y cell index.
127 x
128 X cell index.
129 """
130 return CellIJ(
131 i=(y - self.bbox.y.start) // self.cell_shape.y,
132 j=(x - self.bbox.x.start) // self.cell_shape.x,
133 )
135 def bbox_of(self, cell: CellIJ) -> Box:
136 """Return the bounding box of the given cell."""
137 return Box.from_shape(
138 self.cell_shape,
139 start=YX(
140 y=cell.i * self.cell_shape.y + self.bbox.y.start,
141 x=cell.j * self.cell_shape.x + self.bbox.x.start,
142 ),
143 )
145 @staticmethod
146 def from_legacy(legacy: UniformGrid) -> CellGrid:
147 """Construct from a legacy grid object.
149 Parameters
150 ----------
151 legacy
152 Legacy grid to convert.
153 """
154 if legacy.padding:
155 raise ValueError("Only cell grids with no padding are supported.")
156 bbox = Box.from_legacy(legacy.bbox)
157 cell_shape = YX(y=legacy.cell_size.y, x=legacy.cell_size.x)
158 return CellGrid(bbox=bbox, cell_shape=cell_shape)
161class CellGridBounds(pydantic.BaseModel, frozen=True):
162 """A region of pixels defined by a set of cells within a grid.
164 Notes
165 -----
166 This data structure is optimized for the case where a continguous
167 rectangular region of the grid (the `bbox` attribute) is populated with
168 only a few exceptions (the `missing` set).
170 Slicing a `CellGridBounds` with a `.Box` returns a new `CellGridBounds`
171 with just the cells that overlap that box. As always,
172 `CellGridBounds.bbox` will be snapped to the outer boundaries of those
173 cells, so it will contain (and not generally equal) the given box.
174 """
176 grid: CellGrid = pydantic.Field(description="Definition of the grid that defines the cells.")
177 bbox: Box = pydantic.Field(description="Pixel bounding box of the region (snapped to cell boundaries).")
178 missing: frozenset[CellIJ] = pydantic.Field(
179 default=frozenset(),
180 description=(
181 "Indices of cells that are missing, where (i=0, j=0) is the cell that starts at grid.bbox.start."
182 ),
183 )
185 @cached_property
186 def grid_start(self) -> CellIJ:
187 """The index of the first cell in this bounds' bounding box within
188 its grid.
189 """
190 return self.grid.index_of(y=self.bbox.y.start, x=self.bbox.x.start)
192 @cached_property
193 def grid_stop(self) -> CellIJ:
194 """One-past-the-last indices for the cells in these bounds, within
195 its grid.
196 """
197 return self.grid.index_of(y=self.bbox.y.stop, x=self.bbox.x.stop)
199 @overload
200 def contains(self, *, x: int, y: int) -> bool: ... 200 ↛ exitline 200 didn't return from function 'contains' because
202 @overload
203 def contains(self, *, x: np.ndarray, y: np.ndarray) -> np.ndarray: ... 203 ↛ exitline 203 didn't return from function 'contains' because
205 def contains(self, *, x: Any, y: Any) -> Any:
206 """Test whether these bounds contain one or more points.
208 Parameters
209 ----------
210 x
211 One or more integer X coordinates to test for containment.
212 If an array, an array of results will be returned.
213 y
214 One or more integer Y coordinates to test for containment.
215 If an array, an array of results will be returned.
217 Returns
218 -------
219 `bool` | `numpy.ndarray`
220 If ``x`` and ``y`` are both scalars, a single `bool` value. If
221 ``x`` and ``y`` are arrays, a boolean array with their broadcasted
222 shape.
223 """
224 result = self.bbox.contains(x=x, y=y)
225 if not self.missing:
226 return result
227 match result:
228 case False:
229 return False
230 case True:
231 return self.grid.index_of(x=x, y=y) not in self.missing
232 case np.ndarray():
233 for box in self.missing_boxes():
234 result = np.logical_and(result, np.logical_not(box.contains(x=x, y=y)))
235 return result
237 def intersection(self, other: Bounds) -> Bounds:
238 """Compute the intersection of this bounds object with another."""
239 from ._concrete_bounds import _intersect_cgb
241 return _intersect_cgb(self, other)
243 def contains_cell(self, index: CellIJ) -> bool:
244 """Test whether the given cell is in the bounds."""
245 return (
246 (index.i >= self.grid_start.i and index.i < self.grid_stop.i)
247 and (index.j >= self.grid_start.j and index.j < self.grid_stop.j)
248 and index not in self.missing
249 )
251 def missing_boxes(self) -> Iterator[Box]:
252 """Iterate over the bounding boxes of the missing cells."""
253 for index in sorted(self.missing):
254 yield self.grid.bbox_of(index)
256 def cell_indices(self) -> Iterator[CellIJ]:
257 """Iterate over the indices of the cells in these bounds."""
258 for i in range(self.grid_start.i, self.grid_stop.i):
259 for j in range(self.grid_start.j, self.grid_stop.j):
260 index = CellIJ(i=i, j=j)
261 if index not in self.missing:
262 yield index
264 def __getitem__(self, bbox: Box) -> CellGridBounds:
265 if not self.bbox.contains(bbox):
266 raise ValueError(
267 f"Original grid bounding box {self.bbox} does not contain the subset bounding box {bbox}."
268 )
269 c = self.grid.cell_shape
270 s = self.grid.bbox.start
271 i1 = (bbox.y.start - s.y) // c.y
272 j1 = (bbox.x.start - s.x) // c.x
273 i2 = math.ceil((bbox.y.stop - s.y) / c.y)
274 j2 = math.ceil((bbox.x.stop - s.x) / c.x)
275 subset_bbox = Box.factory[i1 * c.y + s.y : i2 * c.y + s.y, j1 * c.x + s.x : j2 * c.x + s.x]
276 grid_as_box = Box.factory[i1:i2, j1:j2]
277 subset_missing = {index for index in self.missing if grid_as_box.contains(y=index.i, x=index.j)}
278 return CellGridBounds(grid=self.grid, bbox=subset_bbox, missing=frozenset(subset_missing))
280 def serialize(self) -> SerializableBounds:
281 """Convert a bounds instance into a serializable object."""
282 return self
284 @classmethod
285 def deserialize(cls, serialized: SerializableBounds) -> Self:
286 """Convert a serialized bounds object into its in-memory form."""
287 from ._concrete_bounds import deserialize_bounds
289 return cast(Self, deserialize_bounds(serialized))
292class PatchDefinition(pydantic.BaseModel, frozen=True):
293 """Identifiers and geometry for a full patch."""
295 id: int = pydantic.Field(description="ID for the patch.")
296 index: YX[int] = pydantic.Field(description="2-d index of this patch within the tract.")
297 inner_bbox: Box = pydantic.Field(description="Inner bounding box of this patch.")
298 cells: CellGrid = pydantic.Field(description="Cell grid for the full patch.")
300 @property
301 def outer_bbox(self) -> Box:
302 """The outer bounding box of this patch (`.Box`)."""
303 return self.cells.bbox