Coverage for python/lsst/images/_polygon.py: 90%

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. 

11 

12from __future__ import annotations 

13 

14__all__ = ("Polygon", "Region", "RegionSerializationModel") 

15 

16from typing import TYPE_CHECKING, Any, Literal, assert_type, overload 

17 

18import numpy as np 

19import numpy.typing as npt 

20import pydantic 

21import pydantic_core.core_schema as pcs 

22import shapely 

23from pydantic.json_schema import JsonSchemaValue 

24 

25from ._geom import XY, YX, Bounds, Box 

26from .utils import round_half_down, round_half_up 

27 

28if TYPE_CHECKING: 

29 from ._geom import Bounds 

30 from ._transforms import Transform 

31 

32 try: 

33 from lsst.afw.geom import Polygon as LegacyPolygon 

34 except ImportError: 

35 type LegacyPolygon = Any # type: ignore[no-redef] 

36 

37 

38class Region: 

39 """A 2-d Euclidean region represented as one or more polygons with 

40 optional holes. 

41 

42 Parameters 

43 ---------- 

44 geometry 

45 A polygon or multi-polygon from the Shapely library. 

46 """ 

47 

48 def __init__(self, geometry: shapely.Polygon | shapely.MultiPolygon) -> None: 

49 self._impl = geometry 

50 

51 @property 

52 def area(self) -> float: 

53 """The area of the region (`float`).""" 

54 return self._impl.area 

55 

56 @property 

57 def bbox(self) -> Box: 

58 """The integer-coordinate bounding box of the region (`Box`). 

59 

60 Because a `Box` logically contains the entirety of the pixels on its 

61 boundary, but the centers of those pixels are the numerical values of 

62 its bounds, the region may contain vertices that are up to 0.5 beyond 

63 the integer box coordinates in either dimension. 

64 """ 

65 x_min, y_min, x_max, y_max = self._impl.bounds 

66 return Box.factory[ 

67 round_half_up(y_min) : round_half_down(y_max) + 1, 

68 round_half_up(x_min) : round_half_down(x_max) + 1, 

69 ] 

70 

71 @property 

72 def wkt(self) -> str: 

73 """The 'Well-Known Text' representation of this region (`str`).""" 

74 return self._impl.wkt 

75 

76 @staticmethod 

77 def from_wkt(wkt: str) -> Region: 

78 """Construct from a 'Well-Known Text' string. 

79 

80 Parameters 

81 ---------- 

82 wkt 

83 Well-Known Text representation of the region. 

84 """ 

85 impl = shapely.from_wkt(wkt) 

86 if not isinstance(impl, shapely.Polygon | shapely.MultiPolygon): 86 ↛ 87line 86 didn't jump to line 87 because the condition on line 86 was never true

87 raise ValueError("Only Polygon and MultiPolygon geometries can be converted to Regions.") 

88 return Region(impl).try_to_polygon() 

89 

90 def __str__(self) -> str: 

91 return self._impl.wkt 

92 

93 def __repr__(self) -> str: 

94 return f"Region.from_wkt({self._impl.wkt!r})" 

95 

96 def __eq__(self, other: object) -> bool: 

97 if isinstance(other, Region): 

98 return bool(shapely.equals(self._impl, other._impl)) 

99 return NotImplemented 

100 

101 @overload 

102 def contains(self, point: XY[int | float] | YX[int | float], /) -> bool: ... 102 ↛ exitline 102 didn't return from function 'contains' because

103 

104 @overload 

105 def contains(self, point: XY[npt.ArrayLike] | YX[npt.ArrayLike], /) -> np.ndarray: ... 105 ↛ exitline 105 didn't return from function 'contains' because

106 

107 @overload 

108 def contains(self, /, *, x: int | float, y: int | float) -> bool: ... 108 ↛ exitline 108 didn't return from function 'contains' because

109 

110 @overload 

111 def contains(self, /, *, x: npt.ArrayLike, y: npt.ArrayLike) -> np.ndarray: ... 111 ↛ exitline 111 didn't return from function 'contains' because

112 

113 @overload 

114 def contains(self, other: Polygon, /) -> bool: ... 114 ↛ exitline 114 didn't return from function 'contains' because

115 

116 def contains( 

117 self, 

118 other: Region | XY[Any] | YX[Any] | None = None, 

119 /, 

120 *, 

121 x: Any = None, 

122 y: Any = None, 

123 ) -> bool | np.ndarray: 

124 """Test whether the geometry contains the given points or another 

125 geometry. 

126 

127 Parameters 

128 ---------- 

129 other 

130 Another geometry, or an `XY` or `YX` coordinate pair, to compare 

131 to. Mutually exclusive with ``x`` and ``y``. 

132 x 

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

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

135 Mutually exclusive with ``other``. 

136 y 

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

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

139 Mutually exclusive with ``other``. 

140 """ 

141 match other: 

142 case None: 

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

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

145 case Region(): 

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

147 raise TypeError("'Region.contains' point argument is mutually exclusive with x= and y=.") 

148 return self._impl.contains(other._impl) 

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

150 if x is not None or y is not None: 

151 raise TypeError("'Region.contains' point argument is mutually exclusive with x= and y=.") 

152 x, y = other.x, other.y 

153 case _: 

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

155 return shapely.contains_xy(self._impl, x=x, y=y) 

156 

157 @overload 

158 def intersection(self, other: Region) -> Region: ... 158 ↛ exitline 158 didn't return from function 'intersection' because

159 

160 @overload 

161 def intersection(self, other: Box) -> Region | Box: ... 161 ↛ exitline 161 didn't return from function 'intersection' because

162 

163 @overload 

164 def intersection(self, other: Bounds) -> Bounds: ... 164 ↛ exitline 164 didn't return from function 'intersection' because

165 

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

167 """Compute the intersection of this region with a `Bounds` object. 

168 

169 Notes 

170 ----- 

171 Because `Region` implements the `Bounds` interface, its intersections 

172 need to support all other `Bounds` objects. This is not true of other 

173 `Region` point-set operations like `union` and `difference`. 

174 

175 Parameters 

176 ---------- 

177 other 

178 Bounds to intersect with this region. 

179 """ 

180 from ._concrete_bounds import _intersect_region 

181 

182 return _intersect_region(self, other) 

183 

184 def union(self, other: Region) -> Region: 

185 """Compute the point-set union of this region with another. 

186 

187 Parameters 

188 ---------- 

189 other 

190 Region to union with this one. 

191 """ 

192 impl = shapely.union(self._impl, other._impl) 

193 assert isinstance(impl, shapely.Polygon | shapely.MultiPolygon), ( 

194 "A union of Polygons and MultiPolygons should be one of those." 

195 ) 

196 return Region(impl).try_to_polygon() 

197 

198 def difference(self, other: Region) -> Region: 

199 """Compute the point-set difference of this region with another. 

200 

201 Parameters 

202 ---------- 

203 other 

204 Region to subtract from this one. 

205 """ 

206 impl = shapely.difference(self._impl, other._impl) 

207 assert isinstance(impl, shapely.Polygon | shapely.MultiPolygon), ( 

208 "A difference of Polygons and MultiPolygons should be one of those." 

209 ) 

210 return Region(impl).try_to_polygon() 

211 

212 def try_to_polygon(self) -> Region: 

213 """If the underlying geometry is a single polygon with no holes, 

214 return a `Polygon` instance holding it. 

215 

216 In all other cases ``self`` is returned. 

217 """ 

218 impl = self._impl 

219 if isinstance(impl, shapely.MultiPolygon) and len(impl.geoms) == 1: 219 ↛ 220line 219 didn't jump to line 220 because the condition on line 219 was never true

220 impl = impl.geoms[0] 

221 if isinstance(impl, shapely.Polygon) and not impl.interiors: 

222 vertices = np.array(impl.exterior.coords) 

223 return Polygon(x_vertices=vertices[:, 0], y_vertices=vertices[:, 1]) 

224 return self 

225 

226 def try_to_box(self) -> Region | Box: 

227 """If the underlying geometry is a rectangle that fully covers integer 

228 pixels (i.e. has all vertices at half-integer positions), return the 

229 equivalent `Box`. 

230 

231 In all other cases ``self`` is returned. 

232 """ 

233 if Polygon.from_box(self.bbox) == self: 

234 return self.bbox 

235 return self 

236 

237 def to_shapely(self) -> shapely.Polygon | shapely.MultiPolygon: 

238 """Convert to a `shapely.Polygon` or `shapely.MultiPolygon` object.""" 

239 return self._impl 

240 

241 def transform(self, transform: Transform[Any, Any]) -> Region: 

242 """Transform the coordinates of the region, returning a new one. 

243 

244 Parameters 

245 ---------- 

246 transform 

247 Coordinate transform to apply (in the forward direction). 

248 

249 Notes 

250 ----- 

251 This applies the transform to all vertices, assuming that the 

252 transform is close enough to affine that the topology of the geometry 

253 does not change and straight-line edges do not need to be subsampled. 

254 """ 

255 

256 def wrapper(x: np.ndarray, y: np.ndarray) -> XY[np.ndarray]: 

257 return transform.apply_forward(x=x, y=y) 

258 

259 return Region( 

260 # Shapely overloads don't seem to have been annotated rigorously 

261 shapely.transform(self._impl, wrapper, interleaved=False) # type: ignore[arg-type] 

262 ).try_to_polygon() 

263 

264 def serialize(self) -> RegionSerializationModel: 

265 """Serialize the region to a Pydantic model. 

266 

267 Region serialization uses a subset of the GeoJSON specification (IETF 

268 RFC 7946). 

269 """ 

270 return RegionSerializationModel.model_validate_json(shapely.to_geojson(self._impl)) 

271 

272 @classmethod 

273 def __get_pydantic_core_schema__( 

274 cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler 

275 ) -> pcs.CoreSchema: 

276 from_model_schema = pcs.chain_schema( 

277 [ 

278 handler(RegionSerializationModel), 

279 pcs.no_info_plain_validator_function(RegionSerializationModel.deserialize), 

280 ] 

281 ) 

282 return pcs.json_or_python_schema( 

283 json_schema=from_model_schema, 

284 python_schema=pcs.union_schema([pcs.is_instance_schema(cls), from_model_schema]), 

285 serialization=pcs.plain_serializer_function_ser_schema(cls.serialize), 

286 ) 

287 

288 @classmethod 

289 def __get_pydantic_json_schema__( 

290 cls, schema: pcs.CoreSchema, handler: pydantic.GetJsonSchemaHandler 

291 ) -> JsonSchemaValue: 

292 return handler(RegionSerializationModel.__pydantic_core_schema__) 

293 

294 

295class Polygon(Region): 

296 """A simple 2-d polygon in Euclidean coordinates, with no holes. 

297 

298 Parameters 

299 ---------- 

300 x_vertices 

301 The x coordinates of the vertices of the polygon. 

302 y_vertices 

303 The y coordinate of the vertices of the polygon. 

304 """ 

305 

306 def __init__(self, *, x_vertices: npt.ArrayLike, y_vertices: npt.ArrayLike) -> None: 

307 self._vertices = np.stack( 

308 [np.asarray(x_vertices).flat, np.asarray(y_vertices).flat], dtype=np.float64 

309 ).transpose() 

310 self._vertices.flags.writeable = False 

311 super().__init__(shapely.Polygon(self._vertices)) 

312 

313 @staticmethod 

314 def from_box(box: Box) -> Polygon: 

315 """Construct from an integer-coordinate box. 

316 

317 Parameters 

318 ---------- 

319 box 

320 Integer-coordinate box to convert to a polygon. 

321 

322 Notes 

323 ----- 

324 Because the integer min and max coordinates of the box are 

325 interpreted as pixel centers, these are expanded by 0.5 on all sides 

326 before using them to form the polygon vertices. 

327 """ 

328 return Polygon( 

329 x_vertices=[box.x.min - 0.5, box.x.min - 0.5, box.x.max + 0.5, box.x.max + 0.5], 

330 y_vertices=[box.y.min - 0.5, box.y.max + 0.5, box.y.max + 0.5, box.y.min - 0.5], 

331 ) 

332 

333 @property 

334 def n_vertices(self) -> int: 

335 """The number of vertices in the polygon.""" 

336 return self._vertices.shape[0] 

337 

338 @property 

339 def x_vertices(self) -> np.ndarray: 

340 """The x coordinates of the vertices of the polygon. 

341 

342 This is a read-only array; polygons are immutable. 

343 """ 

344 return self._vertices[:, 0] 

345 

346 @property 

347 def y_vertices(self) -> np.ndarray: 

348 """The y coordinates of the vertices of the polygon. 

349 

350 This is a read-only array; polygons are immutable. 

351 """ 

352 return self._vertices[:, 1] 

353 

354 @property 

355 def centroid(self) -> XY[float]: 

356 """The centroid of the polygon (`XY` [`float`]).""" 

357 c = self._impl.centroid 

358 return XY(x=c.x, y=c.y) 

359 

360 def __repr__(self) -> str: 

361 return f"Polygon(x_vertices={self.x_vertices!r}, y_vertices={self.y_vertices!r})" 

362 

363 def transform(self, transform: Transform[Any, Any]) -> Polygon: 

364 # Docstring inherited. 

365 result = super().transform(transform) 

366 assert isinstance(result, Polygon), "Transforming a polygon should not change its topology." 

367 return result 

368 

369 @staticmethod 

370 def from_legacy(legacy: LegacyPolygon) -> Polygon: 

371 """Convert from a legacy `lsst.afw.geom.Polygon` instance. 

372 

373 Parameters 

374 ---------- 

375 legacy 

376 Legacy `lsst.afw.geom.Polygon` to convert. 

377 """ 

378 vertices = legacy.getVertices() 

379 x_vertices = np.zeros(len(vertices), dtype=np.float64) 

380 y_vertices = np.zeros(len(vertices), dtype=np.float64) 

381 for n, point in enumerate(vertices): 

382 x_vertices[n] = point.x 

383 y_vertices[n] = point.y 

384 return Polygon(x_vertices=x_vertices, y_vertices=y_vertices) 

385 

386 def to_legacy(self) -> LegacyPolygon: 

387 """Convert to a legacy `lsst.afw.geom.Polygon` instance.""" 

388 from lsst.afw.geom import Polygon as LegacyPolygon 

389 from lsst.geom import Point2D 

390 

391 return LegacyPolygon([Point2D(x, y) for x, y in zip(self.x_vertices, self.y_vertices)]) 

392 

393 

394class RegionSerializationModel(pydantic.BaseModel): 

395 """Serialization model for `Region` and `Polygon`. 

396 

397 This model is a subset of the GeoJSON specification (IETF RFC 7946). 

398 """ 

399 

400 type: Literal["Polygon", "MultiPolygon"] = pydantic.Field(description="Geometry type.") 

401 

402 coordinates: list[list[tuple[float, float] | list[tuple[float, float]]]] = pydantic.Field( 

403 description="Vertices of the polygon or polygons." 

404 ) 

405 

406 def deserialize(self) -> Region: 

407 """Deserialize into a `Region` (a `Polygon`, if possible).""" 

408 region_impl = shapely.from_geojson(self.model_dump_json()) 

409 assert isinstance(region_impl, shapely.Polygon | shapely.MultiPolygon), ( 

410 "Other geometry types are not used." 

411 ) 

412 return Region(region_impl).try_to_polygon() 

413 

414 

415if TYPE_CHECKING: 

416 

417 def _test_types() -> None: 

418 region = Region(shapely.Point(0, 0).buffer(1.0)) 

419 arr = np.zeros(3) 

420 

421 # Region satisfies the Bounds Protocol. 

422 bounds: Bounds = region 

423 

424 # Region.contains: XY/YX, scalar, array-like 

425 assert_type(region.contains(x=1.0, y=2.0), bool) 

426 assert_type(region.contains(x=arr, y=arr), np.ndarray) 

427 assert_type(region.contains(XY(1.0, 2.0)), bool) 

428 assert_type(region.contains(YX(2.0, 1.0)), bool) 

429 assert_type(region.contains(XY(arr, arr)), np.ndarray) 

430 assert_type(region.contains(YX(arr, arr)), np.ndarray) 

431 

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

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

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

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

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

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

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