Coverage for python/lsst/images/_cell

1# This file is part of lsst-images.

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.

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 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)

100

101

102class CellGrid(pydantic.BaseModel, frozen=True):

103 """A grid of rectangular cells with no overlaps or space between cells.

104

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.

110

111 Subsets of grids are usually represented via `CellGridBounds`.

112 """

113

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: SerializableYX[int] = pydantic.Field(description="Shape of each cell in pixels.")

122

123 def __str__(self) -> str:

124 return f"{self.bbox}, cell_shape=(y={self.cell_shape.y}, x={self.cell_shape.x})"

125

126 @property

127 def grid_size(self) -> CellIJ:

128 """The number of cells in each dimension (`CellIJ`)."""

129 return CellIJ(i=self.bbox.y.size // self.cell_shape.y, j=self.bbox.x.size // self.cell_shape.x)

130

131 def index_of(self, *, y: int, x: int) -> CellIJ:

132 """Return the 2-d index of the cell that contains the given pixel.

133

134 Parameters

135 ----------

136 y

137 Y cell index.

138 x

139 X cell index.

140 """

141 return CellIJ(

142 i=(y - self.bbox.y.start) // self.cell_shape.y,

143 j=(x - self.bbox.x.start) // self.cell_shape.x,

144 )

145

146 def bbox_of(self, cell: CellIJ) -> Box:

147 """Return the bounding box of the given cell.

148

149 Parameters

150 ----------

151 cell

152 Index of the cell whose bounding box is returned.

153 """

154 return Box.from_shape(

155 self.cell_shape,

156 start=YX(

157 y=cell.i * self.cell_shape.y + self.bbox.y.start,

158 x=cell.j * self.cell_shape.x + self.bbox.x.start,

159 ),

160 )

161

162 @staticmethod

163 def from_legacy(legacy: LegacyUniformGrid) -> CellGrid:

164 """Construct from a legacy `lsst.cell_coadds.UniformGrid` object.

165

166 Parameters

167 ----------

168 legacy

169 Legacy grid to convert.

170 """

171 if legacy.padding:

172 raise ValueError("Only cell grids with no padding are supported.")

173 bbox = Box.from_legacy(legacy.bbox)

174 cell_shape = YX(y=legacy.cell_size.y, x=legacy.cell_size.x)

175 return CellGrid(bbox=bbox, cell_shape=cell_shape)

176

177 def to_legacy(self) -> LegacyUniformGrid:

178 """Convert to a legacy `lsst.cell_coadds.UniformGrid` object."""

179 from lsst.cell_coadds import UniformGrid as LegacyUniformGrid

180

181 return LegacyUniformGrid(

182 self.cell_shape.to_legacy_int_extent(),

183 self.grid_size.to_legacy(),

184 min=self.bbox.min.to_legacy_int_point(),

185 )

186

187

188class CellGridBounds(pydantic.BaseModel, frozen=True):

189 """A region of pixels defined by a set of cells within a grid.

190

191 Notes

192 -----

193 This data structure is optimized for the case where a continguous

194 rectangular region of the grid (the `bbox` attribute) is populated with

195 only a few exceptions (the `missing` set).

196

197 Slicing a `CellGridBounds` with a `.Box` returns a new `CellGridBounds`

198 with just the cells that overlap that box. As always,

199 `CellGridBounds.bbox` will be snapped to the outer boundaries of those

200 cells, so it will contain (and not generally equal) the given box.

201 """

202

203 grid: CellGrid = pydantic.Field(description="Definition of the grid that defines the cells.")

204 bbox: Box = pydantic.Field(description="Pixel bounding box of the region (snapped to cell boundaries).")

205 missing: frozenset[CellIJ] = pydantic.Field(

206 default=frozenset(),

207 description=(

208 "Indices of cells that are missing, where (i=0, j=0) is the cell that starts at grid.bbox.start."

209 ),

210 )

211

212 def __str__(self) -> str:

213 text = f"{self.bbox} in grid ({self.grid})"

214 if self.missing:

215 missing = ", ".join(f"(i={c.i}, j={c.j})" for c in sorted(self.missing))

216 text = f"{text}, missing={{{missing}}}"

217 return text

218

219 @cached_property

220 def subgrid_start(self) -> CellIJ:

221 """The index of the first cell in this bounds' bounding box within

222 its grid.

223 """

224 return self.grid.index_of(y=self.bbox.y.start, x=self.bbox.x.start)

225

226 @cached_property

227 def subgrid_stop(self) -> CellIJ:

228 """One-past-the-last indices for the cells in these bounds, within

229 its grid.

230 """

231 return self.grid.index_of(y=self.bbox.y.stop, x=self.bbox.x.stop)

232

233 @cached_property

234 def subgrid_size(self) -> CellIJ:

235 """Number of cells within these bounds in both dimensions, not

236 accounting for `missing`.

237 """

238 return self.subgrid_stop - self.subgrid_start

239

240 @overload

241 def contains(self, point: XY[int | float] | YX[int | float], /) -> bool: ...

242

243 @overload

244 def contains(self, point: XY[npt.ArrayLike] | YX[npt.ArrayLike], /) -> np.ndarray: ...

245

246 @overload

247 def contains(self, *, x: int | float, y: int | float) -> bool: ...

248

249 @overload

250 def contains(self, *, x: npt.ArrayLike, y: npt.ArrayLike) -> np.ndarray: ...

251

252 def contains(self, point: XY[Any] | YX[Any] | None = None, /, *, x: Any = None, y: Any = None) -> Any: # type: ignore[misc]

253 """Test whether these bounds contain one or more points.

254

255 Parameters

256 ----------

257 point

258 An `.XY` or `.YX` coordinate pair to test for containment.

259 Mutually exclusive with ``x`` and ``y``.

260 x

261 One or more X coordinates to test for containment, as a scalar or

262 any array-like. Results are broadcast against ``y``.

263 Mutually exclusive with ``point``.

264 y

265 One or more Y coordinates to test for containment, as a scalar or

266 any array-like. Results are broadcast against ``x``.

267 Mutually exclusive with ``point``.

268

269 Returns

270 -------

271 `bool` | `numpy.ndarray`

272 If ``x`` and ``y`` are both scalars, a single `bool` value. If

273 ``x`` and ``y`` are array-like, a boolean array with their

274 broadcasted shape.

275 """

276 match point:

277 case None:

278 if x is None or y is None: 278 ↛ 279line 278 didn't jump to line 279 because the condition on line 278 was never true

279 raise TypeError("Pass either a point or both x= and y= to 'CellGridBounds.contains'.")

280 case XY() | YX(): 280 ↛ 286line 280 didn't jump to line 286 because the pattern on line 280 always matched

281 if x is not None or y is not None: 281 ↛ 282line 281 didn't jump to line 282 because the condition on line 281 was never true

282 raise TypeError(

283 "'CellGridBounds.contains' point argument is mutually exclusive with x= and y=."

284 )

285 x, y = point.x, point.y

286 case _:

287 raise TypeError(f"Unexpected positional argument type: {type(point)!r}.")

288 result = self.bbox.contains(x=x, y=y)

289 if not self.missing: 289 ↛ 290line 289 didn't jump to line 290 because the condition on line 289 was never true

290 return result

291 match result:

292 case False:

293 return False

294 case True:

295 return self.grid.index_of(x=x, y=y) not in self.missing

296 case np.ndarray(): 296 ↛ 299line 296 didn't jump to line 299 because the pattern on line 296 always matched

297 for box in self.missing_boxes():

298 result = np.logical_and(result, np.logical_not(box.contains(x=x, y=y)))

299 return result

300

301 def intersection(self, other: Bounds) -> Bounds:

302 """Compute the intersection of this bounds object with another.

303

304 Parameters

305 ----------

306 other

307 Bounds to intersect with this one.

308 """

309 from ._concrete_bounds import _intersect_cgb

310

311 return _intersect_cgb(self, other)

312

313 def contains_cell(self, index: CellIJ) -> bool:

314 """Test whether the given cell is in the bounds.

315

316 Parameters

317 ----------

318 index

319 Index of the cell to test.

320 """

321 return (

322 (index.i >= self.subgrid_start.i and index.i < self.subgrid_stop.i)

323 and (index.j >= self.subgrid_start.j and index.j < self.subgrid_stop.j)

324 and index not in self.missing

325 )

326

327 def missing_boxes(self) -> Iterator[Box]:

328 """Iterate over the bounding boxes of the missing cells."""

329 for index in sorted(self.missing):

330 yield self.grid.bbox_of(index)

331

332 def cell_indices(self) -> Iterator[CellIJ]:

333 """Iterate over the indices of the cells in these bounds."""

334 for i in range(self.subgrid_start.i, self.subgrid_stop.i):

335 for j in range(self.subgrid_start.j, self.subgrid_stop.j):

336 index = CellIJ(i=i, j=j)

337 if index not in self.missing:

338 yield index

339

340 def __getitem__(self, bbox: Box) -> CellGridBounds:

341 if not self.bbox.contains(bbox): 341 ↛ 342line 341 didn't jump to line 342 because the condition on line 341 was never true

342 raise ValueError(

343 f"Original grid bounding box {self.bbox} does not contain the subset bounding box {bbox}."

344 )

345 c = self.grid.cell_shape

346 s = self.grid.bbox.start

347 i1 = (bbox.y.start - s.y) // c.y

348 j1 = (bbox.x.start - s.x) // c.x

349 i2 = math.ceil((bbox.y.stop - s.y) / c.y)

350 j2 = math.ceil((bbox.x.stop - s.x) / c.x)

351 subset_bbox = Box.factory[i1 * c.y + s.y : i2 * c.y + s.y, j1 * c.x + s.x : j2 * c.x + s.x]

352 grid_as_box = Box.factory[i1:i2, j1:j2]

353 subset_missing = {index for index in self.missing if grid_as_box.contains(y=index.i, x=index.j)}

354 return CellGridBounds(grid=self.grid, bbox=subset_bbox, missing=frozenset(subset_missing))

355

356 def serialize(self) -> CellGridBounds:

357 """Convert a bounds instance into a serializable object."""

358 return self

359

360 def deserialize(self) -> CellGridBounds:

361 """Deserialize a bounds object on the assumption it is a

362 `CellGridBounds`.

363

364 This method just returns the `CellGridBounds` itself, since that

365 already provides Pydantic serialization hooks. It exists for

366 compatibility with the `.Bounds` protocol.

367 """

368 return self

369

370

371class PatchDefinition(pydantic.BaseModel, frozen=True):

372 """Identifiers and geometry for a full patch."""

373

374 id: int = pydantic.Field(description="ID for the patch.")

375 index: SerializableYX[int] = pydantic.Field(description="2-d index of this patch within the tract.")

376 inner_bbox: Box = pydantic.Field(description="Inner bounding box of this patch.")

377 cells: CellGrid = pydantic.Field(description="Cell grid for the full patch.")

378

379 def __str__(self) -> str:

380 return (

381 f"id={self.id}, index=(y={self.index.y}, x={self.index.x}), "

382 f"inner_bbox={self.inner_bbox}, cells=({self.cells})"

383 )

384

385 @property

386 def outer_bbox(self) -> Box:

387 """The outer bounding box of this patch (`.Box`)."""

388 return self.cells.bbox

389

390

391if TYPE_CHECKING:

392

393 def _test_types() -> None:

394 arr = np.zeros(3)

395 bbox = Box.from_shape((100, 200))

396 grid = CellGrid(bbox=bbox, cell_shape=YX(10, 20))

397 cgb = CellGridBounds(grid=grid, bbox=bbox)

398

399 # CellGridBounds satisfies the Bounds Protocol.

400 bounds: Bounds = cgb

401

402 # CellGridBounds.contains: XY/YX, scalar, array-like

403 assert_type(cgb.contains(x=1, y=2), bool)

404 assert_type(cgb.contains(x=1.0, y=2.0), bool)

405 assert_type(cgb.contains(x=arr, y=arr), np.ndarray)

406 assert_type(cgb.contains(XY(1, 2)), bool)

407 assert_type(cgb.contains(YX(2, 1)), bool)

408 assert_type(cgb.contains(XY(arr, arr)), np.ndarray)

409 assert_type(cgb.contains(YX(arr, arr)), np.ndarray)

410

411 # Via the Bounds Protocol view, same signatures hold.

412 assert_type(bounds.contains(x=1, y=1), bool)

413 assert_type(bounds.contains(x=1.0, y=1.0), bool)

414 assert_type(bounds.contains(x=arr, y=arr), np.ndarray)

415 assert_type(bounds.contains(XY(1, 1)), bool)

416 assert_type(bounds.contains(YX(1, 1)), bool)

417 assert_type(bounds.contains(XY(arr, arr)), np.ndarray)

418 assert_type(bounds.contains(YX(arr, arr)), np.ndarray)

Coverage for python/lsst/images/_cell_grid.py: 41%

142 statements