Coverage for python/lsst/dax/images/cutout/stencils.py: 86%

1# This file is part of dax_images_cutout.

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# This program is free software: you can redistribute it and/or modify

10# it under the terms of the GNU General Public License as published by

11# the Free Software Foundation, either version 3 of the License, or

12# (at your option) any later version.

13#

14# This program is distributed in the hope that it will be useful,

15# but WITHOUT ANY WARRANTY; without even the implied warranty of

16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the

17# GNU General Public License for more details.

18#

19# You should have received a copy of the GNU General Public License

20# along with this program. If not, see <https://www.gnu.org/licenses/>.

22from __future__ import annotations

24__all__ = (

25 "MaskBackend",

26 "PixelStencil",

27 "SkyCircle",

28 "SkyPolygon",

29 "SkyStencil",

30 "StencilNotContainedError",

31)

33import enum

34import struct

35from abc import ABC, abstractmethod

36from collections.abc import Iterable

37from hashlib import blake2b

39import astropy.coordinates

40import astropy.io.fits

41import astropy.units as u

42import numpy as np

43import starlink.Ast as Ast

44from astropy.coordinates import SkyCoord

46import lsst.sphgeom

47from lsst.images import Box, Mask, NoOverlapError, SkyProjection

48from lsst.sphgeom import Angle, LonLat, UnitVector3d

51class MaskBackend(enum.Enum):

52 """Selects the algorithm used to rasterize a stencil onto pixels."""

54 AST = enum.auto()

55 """Mask using starlink-pyast ``Region.mask`` on the true sky region."""

57 SPHGEOM = enum.auto()

58 """Mask by testing pixel centers against the ``lsst.sphgeom`` region."""

61def _skycoord_from_lonlat(lonlat: LonLat) -> SkyCoord:

62 """Return an ICRS `astropy.coordinates.SkyCoord` for a `sphgeom.LonLat`."""

63 return SkyCoord(

64 ra=lonlat.getLon().asRadians() * u.rad,

65 dec=lonlat.getLat().asRadians() * u.rad,

66 frame="icrs",

67 )

70class _AstLineSource:

71 """Feeds AST-native text lines to a `starlink.Ast.Channel`."""

73 def __init__(self, text: str) -> None:

74 self._lines = text.splitlines()

76 def astsource(self) -> str | None:

77 return self._lines.pop(0) if self._lines else None

80def _starlink_sky_to_pixel(projection: SkyProjection) -> Ast.Mapping:

81 """Return the sky->pixel mapping of ``projection`` as a starlink-pyast

82 Mapping.

84 ``lsst.images`` may wrap AST with either astshim or starlink-pyast

85 depending on the runtime environment. The transform is serialized to

86 AST's native text form via the public `~lsst.images.Transform.show`

87 method and re-read with starlink-pyast, so that all region masking

88 happens in starlink-pyast regardless of which wrapper ``lsst.images``

89 uses internally.

90 """

91 return Ast.Channel(_AstLineSource(projection.sky_to_pixel_transform.show())).read()

94class StencilNotContainedError(RuntimeError):

95 """Exception that may be raised when a stencil is not with a desired

96 bounding box.

97 """

100class PixelStencil(ABC):

101 """An image cutout stencil defined in pixel coordinates."""

102

103 @property

104 @abstractmethod

105 def bbox(self) -> Box:

106 """Bounding box of this stencil, as a `lsst.images.Box`."""

107 raise NotImplementedError()

108

109 @abstractmethod

110 def _coverage(self) -> np.ndarray:

111 """Boolean array over `bbox`, `True` for pixels the stencil covers.

112

113 The array has shape ``bbox.shape`` (``(ny, nx)``).

114 """

115 raise NotImplementedError()

116

117 def set_mask(self, mask: Mask, plane: str, *, covered: bool = True) -> None:

118 """Set a mask plane for pixels inside or outside the stencil.

119

120 Parameters

121 ----------

122 mask : `lsst.images.Mask`

123 Mask to modify in-place. Its schema must already define ``plane``

124 and its bounding box must contain `bbox`.

125 plane : `str`

126 Name of the mask plane to set.

127 covered : `bool`, optional

128 If `True` (default), set ``plane`` where the stencil covers a pixel

129 center. If `False`, set ``plane`` where the stencil does *not*

130 cover a pixel, including the region of ``mask`` that lies outside

131 `bbox`.

132 """

133 coverage = self._coverage()

134 if not covered:

135 coverage = np.logical_not(coverage)

136 # Pixels outside the stencil's bounding box are never covered, so they

137 # take the value assigned to uncovered pixels.

138 full = np.full(mask.bbox.shape, not covered, dtype=bool)

139 y_off = self.bbox.y.min - mask.bbox.y.min

140 x_off = self.bbox.x.min - mask.bbox.x.min

141 full[y_off : y_off + self.bbox.shape.y, x_off : x_off + self.bbox.shape.x] = coverage

142 mask.set(plane, full)

143

144

145class _AstPixelRegion(PixelStencil):

146 """Pixel-coordinate stencil backed by a starlink-pyast sky `Region`.

147

148 Parameters

149 ----------

150 sky_region : `starlink.Ast.Region`

151 The stencil region expressed in an ICRS sky frame.

152 sky_to_pixel : `starlink.Ast.Mapping`

153 Mapping whose forward direction transforms sky coordinates to pixels,

154 as required by ``Region.mask`` (region frame to grid).

155 bbox : `lsst.images.Box`

156 Bounding box the stencil is restricted to.

157 """

158

159 def __init__(self, sky_region: Ast.Region, sky_to_pixel: Ast.Mapping, bbox: Box) -> None:

160 self._sky_region = sky_region

161 self._sky_to_pixel = sky_to_pixel

162 self._bbox = bbox

163

164 @property

165 def bbox(self) -> Box:

166 # Docstring inherited.

167 return self._bbox

168

169 def _coverage(self) -> np.ndarray:

170 # Docstring inherited.

171 scratch = np.zeros(self._bbox.shape, dtype=np.int64)

172 self._sky_region.mask(

173 self._sky_to_pixel,

174 1,

175 [self._bbox.x.min, self._bbox.y.min],

176 [self._bbox.x.max, self._bbox.y.max],

177 scratch,

178 1,

179 )

180 return scratch != 0

181

182

183class _SphgeomPixelRegion(PixelStencil):

184 """Pixel-coordinate stencil that tests pixel centers against a sphgeom

185 region.

186

187 Parameters

188 ----------

189 region : `lsst.sphgeom.Region`

190 Sky region to test pixel centers against.

191 projection : `lsst.images.SkyProjection`

192 Mapping used to convert pixel centers to sky coordinates.

193 bbox : `lsst.images.Box`

194 Bounding box the stencil is restricted to.

195 """

196

197 def __init__(self, region: lsst.sphgeom.Region, projection: SkyProjection, bbox: Box) -> None:

198 self._region = region

199 self._projection = projection

200 self._bbox = bbox

201

202 @property

203 def bbox(self) -> Box:

204 # Docstring inherited.

205 return self._bbox

206

207 def _coverage(self) -> np.ndarray:

208 # Docstring inherited.

209 grid = self._bbox.meshgrid()

210 sky = self._projection.pixel_to_sky(x=grid.x.ravel(), y=grid.y.ravel())

211 return self._region.contains(sky.ra.radian, sky.dec.radian).reshape(self._bbox.shape)

212

213

214class SkyStencil(ABC):

215 """An image cutout stencil defined in sky (ICRS) coordinates."""

216

217 _clip: bool

218

219 def to_pixels(

220 self,

221 projection: SkyProjection,

222 bbox: Box,

223 *,

224 backend: MaskBackend = MaskBackend.AST,

225 ) -> PixelStencil:

226 """Transform to a pixel-coordinate stencil.

227

228 Parameters

229 ----------

230 projection : `lsst.images.SkyProjection`

231 Mapping from sky coordinates to pixel coordinates.

232 bbox : `lsst.images.Box`

233 Bounds that the returned stencil must lie within.

234 backend : `MaskBackend`, optional

235 Algorithm used to rasterize the stencil. Defaults to

236 `MaskBackend.AST`.

237

238 Returns

239 -------

240 pixels : `PixelStencil`

241 Pixel-coordinate stencil object. `PixelStencil.bbox` is guaranteed

242 to be contained by the given ``bbox``.

243

244 Raises

245 ------

246 StencilNotContainedError

247 Raised when ``clip`` is `False` and the pixel-coordinate stencil

248 does not lie within ``bbox``.

249 """

250 tight = self._pixel_bbox(projection)

251 final = self._resolve_box(tight, bbox)

252 if backend is MaskBackend.AST:

253 return _AstPixelRegion(self._ast_sky_region(), _starlink_sky_to_pixel(projection), final)

254 if backend is MaskBackend.SPHGEOM: 254 ↛ 256line 254 didn't jump to line 256 because the condition on line 254 was always true

255 return _SphgeomPixelRegion(self.region, projection, final)

256 raise ValueError(f"Unknown mask backend: {backend!r}.")

257

258 def _pixel_bbox(self, projection: SkyProjection) -> Box:

259 """Compute the tight pixel bounding box of this stencil.

260

261 The boundary is sampled on the sky and transformed to pixels, so both

262 mask backends share an identical bounding box.

263 """

264 xy = projection.sky_to_pixel(self._boundary_skycoord())

265 return Box.from_float_bounds(

266 x_min=float(np.min(xy.x)),

267 x_max=float(np.max(xy.x)),

268 y_min=float(np.min(xy.y)),

269 y_max=float(np.max(xy.y)),

270 )

271

272 def _resolve_box(self, tight: Box, box: Box) -> Box:

273 """Clip ``tight`` to ``box`` or raise if not contained.

274

275 Honors the stencil's ``clip`` flag: when clipping, returns the

276 intersection (raising `StencilNotContainedError` when disjoint); when

277 not clipping, returns ``tight`` only if ``box`` contains it.

278 """

279 if self._clip:

280 try:

281 return tight.intersection(box)

282 except NoOverlapError:

283 raise StencilNotContainedError(f"{self} does not overlap {box}.") from None

284 if not box.contains(tight):

285 raise StencilNotContainedError(f"{self} has pixel bbox {tight}, which is not within {box}.")

286 return tight

287

288 @abstractmethod

289 def _ast_sky_region(self) -> Ast.Region:

290 """Return a starlink-pyast `Region` in an ICRS sky frame."""

291 raise NotImplementedError()

292

293 @abstractmethod

294 def _boundary_skycoord(self) -> SkyCoord:

295 """Return sky coordinates sampling the stencil boundary.

296

297 Used to size the pixel bounding box.

298 """

299 raise NotImplementedError()

300

301 @property

302 @abstractmethod

303 def region(self) -> lsst.sphgeom.Region:

304 """A `lsst.sphgeom.Region` that bounds this stencil on the sky."""

305 raise NotImplementedError()

306

307 @abstractmethod

308 def to_fits_metadata(self) -> astropy.io.fits.Header:

309 """Return FITS header cards that describe the stencil.

310

311 The cards carry per-keyword comments and are merged into the cutout

312 provenance header.

313 """

314 raise NotImplementedError()

315

316 @property

317 @abstractmethod

318 def fingerprint(self) -> bytes:

319 """A 16-byte blob that is unique to this stencil."""

320 raise NotImplementedError()

321

322

323class SkyCircle(SkyStencil):

324 """A sky-coordinate circular stencil.

325

326 Parameters

327 ----------

328 center : `lsst.sphgeom.LonLat`

329 The center of the circle, in ICRS (longitude, latitude).

330 radius : `lsst.sphgeom.Angle`

331 Radius of the circle.

332 clip : `bool`, optional

333 If `True` (`False` is default), clip pixel stencils returned by

334 `to_pixels` instead of raising `StencilNotContainedError`.

335 """

336

337 #: Number of points used to sample the circle boundary when sizing the

338 #: pixel bounding box.

339 BOUNDARY_SAMPLES = 64

340

341 def __init__(self, center: LonLat, radius: Angle, clip: bool = False):

342 self._center = center

343 self._radius = radius

344 self._clip = clip

345

346 def __repr__(self) -> str:

347 return (

348 f"SkyCircle(LonLat.fromRadians({self._center.getLon().asRadians()!r}, "

349 f"{self._center.getLat().asRadians()!r}), "

350 f"Angle({self._radius.asRadians()!r}), clip={self._clip!r})"

351 )

352

353 @classmethod

354 def from_astropy(

355 cls, center: astropy.coordinates.SkyCoord, radius: astropy.coordinates.Angle, clip: bool = False

356 ) -> SkyCircle:

357 """Construct from `astropy.coordinates` arguments.

358

359 Parameters

360 ----------

361 center : `astropy.coordinates.SkyCoord`

362 The center of the circle, in ICRS (ra, dec). Must be scalar.

363 radius : `astropy.coordinates.Angle`

364 Radius of the circle. Must be scalar.

365 clip : `bool`, optional

366 If `True` (`False` is default), clip pixel stencils returned by

367 `to_pixels` instead of raising `StencilNotContainedError`.

368

369 Returns

370 -------

371 stencil : `SkyCircle`

372 Circular stencil.

373 """

374 return cls(

375 center=_lonlat_from_astropy(center),

376 radius=_angle_from_astropy(radius),

377 clip=clip,

378 )

379

380 @classmethod

381 def from_sphgeom(cls, circle: lsst.sphgeom.Circle, clip: bool = False) -> SkyCircle:

382 """Construct from a `lsst.sphgeom.Circle` instance."""

383 return cls(LonLat(circle.getCenter()), circle.getOpeningAngle(), clip=clip)

384

385 def to_polygon(self, n_vertices: int = 16) -> SkyPolygon:

386 """Return a polygon sky stencil that approximates this circle.

387

388 Parameters

389 ----------

390 n_vertices : `int`, optional

391 Number of polygon vertices in the approximation.

392

393 Returns

394 -------

395 polygon : `SkyPolygon`

396 Polygon approximation.

397

398 Notes

399 -----

400 This helper is retained for callers that want a polygon approximation;

401 it is no longer used by `to_pixels`, which masks the true circle.

402 """

403 center = _skycoord_from_lonlat(self._center)

404 position_angle = (np.arange(n_vertices) / n_vertices * 2.0 * np.pi) * u.rad

405 radius = astropy.coordinates.Angle(self._radius.asRadians() * u.rad)

406 points = center.directional_offset_by(position_angle, radius)

407 vertices = [LonLat.fromRadians(float(p.ra.rad), float(p.dec.rad)) for p in points]

408 return SkyPolygon(vertices, clip=self._clip)

409

410 def _ast_sky_region(self) -> Ast.Region:

411 # Docstring inherited.

412 return Ast.Circle(

413 Ast.SkyFrame("System=ICRS"),

414 1,

415 [self._center.getLon().asRadians(), self._center.getLat().asRadians()],

416 [self._radius.asRadians()],

417 )

418

419 def _boundary_skycoord(self) -> SkyCoord:

420 # Docstring inherited.

421 center = _skycoord_from_lonlat(self._center)

422 position_angle = (np.arange(self.BOUNDARY_SAMPLES) / self.BOUNDARY_SAMPLES * 2.0 * np.pi) * u.rad

423 radius = astropy.coordinates.Angle(self._radius.asRadians() * u.rad)

424 return center.directional_offset_by(position_angle, radius)

425

426 @property

427 def region(self) -> lsst.sphgeom.Region:

428 # Docstring inherited.

429 return lsst.sphgeom.Circle(UnitVector3d(self._center), self._radius)

430

431 def to_fits_metadata(self) -> astropy.io.fits.Header:

432 # Docstring inherited.

433 header = astropy.io.fits.Header()

434 header.set("ST_TYPE", "CIRCLE", "Type of stencil used to create this cutout")

435 header.set("ST_RA", self._center.getLon().asDegrees(), "[deg] Circle center Right Ascension")

436 header.set("ST_DEC", self._center.getLat().asDegrees(), "[deg] Circle center Declination")

437 header.set("ST_RAD", self._radius.asDegrees(), "[deg] Circle radius")

438 return header

439

440 @property

441 def fingerprint(self) -> bytes:

442 # Docstring inherited.

443 hasher = blake2b(digest_size=16)

444 hasher.update(b"CIRCLE")

445 hasher.update(struct.pack("!d", self._center.getLon().asRadians()))

446 hasher.update(struct.pack("!d", self._center.getLat().asRadians()))

447 hasher.update(struct.pack("!d", self._radius.asRadians()))

448 return hasher.digest()

449

450

451class SkyPolygon(SkyStencil):

452 """A sky-coordinate stencil in the shape of a great-circle polygon.

453

454 Parameters

455 ----------

456 vertices : `Iterable` [ `lsst.sphgeom.LonLat` ]

457 Vertices of the polygon, CCW when looking out from the origin.

458 Implicitly closed (the first vertex should not be duplicated as the

459 last).

460 clip : `bool`, optional

461 If `True` (`False` is default), clip pixel stencils returned by

462 `to_pixels` instead of raising `StencilNotContainedError`.

463

464 Notes

465 -----

466 Vertex orientation is not checked at construction, and incorrect

467 orientation may result in unspecified failures in `to_pixels`.

468 """

469

470 def __init__(self, vertices: Iterable[LonLat], clip: bool = False):

471 self._vertices = tuple(vertices)

472 self._clip = clip

473

474 @classmethod

475 def from_astropy(cls, vertices: astropy.coordinates.SkyCoord, clip: bool = False) -> SkyPolygon:

476 """Construct from an array-valued `astropy.coordinates.SkyCoord`.

477

478 Parameters

479 ----------

480 vertices : `astropy.coordinates.SkyCoord`

481 Array of vertices, CCW when looking out from the origin.

482 Implicitly closed (the first vertex should not be duplicated as the

483 last).

484 clip : `bool`, optional

485 If `True` (`False` is default), clip pixel stencils returned by

486 `to_pixels` instead of raising `StencilNotContainedError`.

487

488 Returns

489 -------

490 stencil : `SkyPolygon`

491 Polygon stencil.

492 """

493 return cls((_lonlat_from_astropy(v) for v in vertices), clip=clip)

494

495 def _ast_sky_region(self) -> Ast.Region:

496 # Docstring inherited.

497 sky_frame = Ast.SkyFrame("System=ICRS")

498 ra = [v.getLon().asRadians() for v in self._vertices]

499 dec = [v.getLat().asRadians() for v in self._vertices]

500 centroid = lsst.sphgeom.LonLat(self.region.getCentroid())

501 probe = [centroid.getLon().asRadians(), centroid.getLat().asRadians()]

502 polygon = Ast.Polygon(sky_frame, np.array([ra, dec]))

503 # AST's bounded interior depends on vertex winding: with the wrong

504 # winding the polygon represents its own complement. ``negate`` flips

505 # ``pointinregion`` but not the ``mask`` polarity, so reverse the

506 # vertices instead to obtain a region whose interior is the polygon.

507 if not polygon.pointinregion(probe): 507 ↛ 508line 507 didn't jump to line 508 because the condition on line 507 was never true

508 polygon = Ast.Polygon(sky_frame, np.array([ra[::-1], dec[::-1]]))

509 # ``Region.mask`` rasterizes by simplifying the region into the pixel

510 # frame. By default AST re-fits the polygon to straight-edged

511 # pixel-space vertices, discarding the great-circle curvature of the

512 # edges whenever the sky-to-pixel projection is non-gnomonic.

513 # ``SimpVertices=0`` makes AST keep the curved edges unless they match

514 # the straight approximation to within the region's uncertainty.

515 polygon.set("SimpVertices=0")

516 return polygon

517

518 def _boundary_skycoord(self) -> SkyCoord:

519 # Docstring inherited.

520 return SkyCoord(

521 ra=[v.getLon().asRadians() for v in self._vertices] * u.rad,

522 dec=[v.getLat().asRadians() for v in self._vertices] * u.rad,

523 frame="icrs",

524 )

525

526 @property

527 def region(self) -> lsst.sphgeom.Region:

528 # Docstring inherited.

529 return lsst.sphgeom.ConvexPolygon([UnitVector3d(v) for v in self._vertices])

530

531 def to_fits_metadata(self) -> astropy.io.fits.Header:

532 # Docstring inherited.

533 header = astropy.io.fits.Header()

534 header.set("ST_TYPE", "POLYGON", "Type of stencil used to create this cutout")

535 if len(self._vertices) > 100: 535 ↛ 536line 535 didn't jump to line 536 because the condition on line 535 was never true

536 raise NotImplementedError(

537 "TODO: FITS limitations make it difficult to serialize big stencils to the header."

538 )

539 for n, v in enumerate(self._vertices):

540 header.set(f"ST_RA{n:02d}", v.getLon().asDegrees(), f"[deg] Vertex {n} Right Ascension")

541 header.set(f"ST_DEC{n:02d}", v.getLat().asDegrees(), f"[deg] Vertex {n} Declination")

542 return header

543

544 @property

545 def fingerprint(self) -> bytes:

546 # Docstring inherited.

547 hasher = blake2b(digest_size=16)

548 hasher.update(b"POLYGON")

549 for v in self._vertices:

550 hasher.update(struct.pack("!dd", v.getLon().asRadians(), v.getLat().asRadians()))

551 return hasher.digest()

552

553

554def _angle_from_astropy(angle: astropy.coordinates.Angle) -> Angle:

555 """Convert an `astropy.coordinates.Angle` to a `lsst.sphgeom.Angle`.

556

557 Parameters

558 ----------

559 angle : `astropy.coordinates.Angle`

560 Astropy Angle to convert. Must be a scalar.

561

562 Returns

563 -------

564 angle : `lsst.sphgeom.Angle`

565 Equivalent sphgeom angle.

566 """

567 if not angle.isscalar: 567 ↛ 568line 567 didn't jump to line 568 because the condition on line 567 was never true

568 raise ValueError("Only scalar angles are supported.")

569 return Angle(angle.to_value(u.rad))

570

571

572def _lonlat_from_astropy(skycoord: astropy.coordinates.SkyCoord) -> LonLat:

573 """Convert an `astropy.coordinates.SkyCoord` to a `lsst.sphgeom.LonLat`.

574

575 Parameters

576 ----------

577 skycoord : `astropy.coordinates.SkyCoord`

578 Astropy coordinates to convert. Must be a scalar.

579

580 Returns

581 -------

582 lonlat : `lsst.sphgeom.LonLat`

583 Equivalent spherical point.

584 """

585 if not skycoord.isscalar: 585 ↛ 586line 585 didn't jump to line 586 because the condition on line 585 was never true

586 raise ValueError("Only scalar coordinates are supported.")

587 icrs = skycoord.transform_to("icrs")

588 return LonLat.fromRadians(float(icrs.ra.rad), float(icrs.dec.rad))