Coverage for python/lsst/images/_transforms/_sky_projection.py: 69%

176 statements  

« prev     ^ index     » next       coverage.py v7.14.1, created at 2026-06-24 01:50 -0700

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__ = ("SkyProjection", "SkyProjectionAstropyView", "SkyProjectionSerializationModel") 

15 

16import functools 

17from typing import TYPE_CHECKING, Any, ClassVar, Self, TypeVar, final 

18 

19import astropy.units as u 

20import astropy.wcs 

21import numpy as np 

22import pydantic 

23from astropy.coordinates import ICRS, Latitude, Longitude, SkyCoord 

24from astropy.wcs.wcsapi import BaseLowLevelWCS, HighLevelWCSMixin 

25 

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

27from ..serialization import ArchiveTree, InputArchive, InvalidParameterError, OutputArchive 

28from ..utils import is_none 

29from . import _ast as astshim 

30from ._frames import Frame, SkyFrame 

31from ._transform import Transform, TransformSerializationModel, _ast_apply 

32 

33if TYPE_CHECKING: 

34 try: 

35 from lsst.afw.geom import SkyWcs as LegacySkyWcs 

36 except ImportError: 

37 type LegacySkyWcs = Any # type: ignore[no-redef] 

38 

39 

40# This pre-python-3.12 declaration is needed by Sphinx (probably the 

41# autodoc-typehints plugin. 

42F = TypeVar("F", bound=Frame) 

43P = TypeVar("P", bound=pydantic.BaseModel) 

44 

45 

46def _set_ast_skyframe_system(frame: astshim.SkyFrame, system: str) -> None: 

47 """Set an AST SkyFrame coordinate system across supported wrappers.""" 

48 if hasattr(frame, "_impl"): 48 ↛ 49line 48 didn't jump to line 49 because the condition on line 48 was never true

49 frame._impl.System = system 

50 else: 

51 setattr(frame, "system", system) 

52 

53 

54@final 

55class SkyProjection[F: Frame]: 

56 """A transform from pixel coordinates to sky coordinates. 

57 

58 Parameters 

59 ---------- 

60 pixel_to_sky 

61 A low-level transform that maps pixel coordinates to sky coordinates. 

62 fits_approximation 

63 An approximation to ``pixel_to_sky`` that is guaranteed to have a 

64 `~Transform.as_fits_wcs` method that does not return `None`. This 

65 should not be provided if ``pixel_to_sky`` is itself representable 

66 as a FITS WCS. 

67 

68 Notes 

69 ----- 

70 `Transform` is conceptually immutable (the internal AST Mapping should 

71 never be modified in-place after construction), and hence does not need to 

72 be copied when any object that holds it is copied. 

73 """ 

74 

75 def __init__( 

76 self, pixel_to_sky: Transform[F, SkyFrame], fits_approximation: Transform[F, SkyFrame] | None = None 

77 ) -> None: 

78 self._pixel_to_sky = pixel_to_sky 

79 if pixel_to_sky.in_frame.unit != u.pix: 

80 raise ValueError("Transform is not a mapping from pixel coordinates.") 

81 if pixel_to_sky.out_frame != SkyFrame.ICRS: 

82 raise ValueError("Transform is not a mapping to ICRS.") 

83 self._fits_approximation = fits_approximation 

84 

85 def __eq__(self, other: Any) -> bool: 

86 if self is other: 

87 return True 

88 if not isinstance(other, SkyProjection): 

89 return NotImplemented 

90 # Even though two approximations could be different and yet consistent 

91 # with the primary mapping (for example using different tolerances 

92 # on construction) we require them to be equal to declare that the 

93 # two objects are equal. 

94 if self._fits_approximation != other._fits_approximation: 

95 return False 

96 return self._pixel_to_sky == other._pixel_to_sky 

97 

98 @staticmethod 

99 def from_fits_wcs( 

100 fits_wcs: astropy.wcs.WCS, 

101 pixel_frame: F, 

102 pixel_bounds: Bounds | None = None, 

103 x0: int = 0, 

104 y0: int = 0, 

105 ) -> SkyProjection[F]: 

106 """Construct a transform from a FITS WCS. 

107 

108 Parameters 

109 ---------- 

110 fits_wcs 

111 FITS WCS to convert. 

112 pixel_frame 

113 Coordinate frame for the pixel grid. 

114 pixel_bounds 

115 The region that bounds valid pixels for this transform. 

116 x0 

117 Logical coordinate of the first column in the array this WCS 

118 relates to world coordinates. 

119 y0 

120 Logical coordinate of the first column in the array this WCS 

121 relates to world coordinates. 

122 

123 Notes 

124 ----- 

125 The ``x0`` and ``y0`` parameters reflect the fact that for FITS, the 

126 first row and column are always labeled ``(1, 1)``, while in Astropy 

127 and most other Python libraries, they are ``(0, 0)``. The `types` in 

128 this package (e.g. `Image`, `Mask`) allow them to be any pair of 

129 integers. 

130 

131 See Also 

132 -------- 

133 Transform.from_fits_wcs 

134 """ 

135 return SkyProjection( 

136 Transform.from_fits_wcs( 

137 fits_wcs, pixel_frame, SkyFrame.ICRS, in_bounds=pixel_bounds, x0=x0, y0=y0 

138 ) 

139 ) 

140 

141 @staticmethod 

142 def from_ast_frame_set( 

143 ast_frame_set: astshim.FrameSet, 

144 pixel_frame: F, 

145 pixel_bounds: Bounds | None = None, 

146 ) -> SkyProjection[F]: 

147 """Construct a sky projection from an AST FrameSet. 

148 

149 The current frame of the FrameSet must be an AST SkyFrame. Its 

150 coordinate system is forced to ICRS (AST adjusts the mapping 

151 automatically) so the resulting Projection is always in ICRS 

152 regardless of the original sky system. 

153 

154 Parameters 

155 ---------- 

156 ast_frame_set 

157 An AST FrameSet whose base frame is pixel coordinates and 

158 whose current frame is a SkyFrame (in any supported sky 

159 coordinate system). 

160 pixel_frame 

161 Coordinate frame for the pixel grid. 

162 pixel_bounds 

163 The region that bounds valid pixels for this transform. 

164 

165 Raises 

166 ------ 

167 ValueError 

168 If the current frame of the FrameSet is not a SkyFrame. 

169 """ 

170 current_frame = ast_frame_set.getFrame(ast_frame_set.current, copy=False) 

171 if not isinstance(current_frame, astshim.SkyFrame): 171 ↛ 172line 171 didn't jump to line 172 because the condition on line 171 was never true

172 raise ValueError( 

173 "The current frame of the AST FrameSet is not a SkyFrame " 

174 f"(got {type(current_frame).__name__})." 

175 ) 

176 _set_ast_skyframe_system(current_frame, "ICRS") 

177 return SkyProjection(Transform(pixel_frame, SkyFrame.ICRS, ast_frame_set, in_bounds=pixel_bounds)) 

178 

179 @property 

180 def pixel_frame(self) -> F: 

181 """Coordinate frame for the pixel grid.""" 

182 return self._pixel_to_sky.in_frame 

183 

184 @property 

185 def sky_frame(self) -> SkyFrame: 

186 """Coordinate frame for the sky.""" 

187 return self._pixel_to_sky.out_frame 

188 

189 @property 

190 def pixel_bounds(self) -> Bounds | None: 

191 """The region that bounds valid pixel points (`Bounds` | `None`).""" 

192 return self._pixel_to_sky.in_bounds 

193 

194 @property 

195 def pixel_to_sky_transform(self) -> Transform[F, SkyFrame]: 

196 """Low-level transform from pixel to sky coordinates (`Transform`).""" 

197 return self._pixel_to_sky 

198 

199 @property 

200 def sky_to_pixel_transform(self) -> Transform[SkyFrame, F]: 

201 """Low-level transform from sky to pixel coordinates (`Transform`).""" 

202 return self._pixel_to_sky.inverted() 

203 

204 @property 

205 def fits_approximation(self) -> SkyProjection[F] | None: 

206 """An approximation to this projection that is guaranteed to have an 

207 `as_fits_wcs` method that does not return `None`. 

208 """ 

209 return SkyProjection(self._fits_approximation) if self._fits_approximation is not None else None 

210 

211 def show(self, simplified: bool = False, comments: bool = False) -> str: 

212 """Return the AST native representation of the transform. 

213 

214 Parameters 

215 ---------- 

216 simplified 

217 Whether to ask AST to simplify the mapping before showing it. 

218 This will make it much more likely that two equivalent transforms 

219 have the same `show` result. 

220 comments 

221 Whether to include descriptive comments. 

222 """ 

223 return self._pixel_to_sky.show(simplified=simplified, comments=comments) 

224 

225 def pixel_to_sky[T: np.ndarray | float](self, *, x: T, y: T) -> SkyCoord: 

226 """Transform one or more pixel points to sky coordinates. 

227 

228 Parameters 

229 ---------- 

230 x : `numpy.ndarray` | `float` 

231 ``x`` values of the pixel points to transform. 

232 y : `numpy.ndarray` | `float` 

233 ``y`` values of the pixel points to transform. 

234 

235 Returns 

236 ------- 

237 astropy.coordinates.SkyCoord 

238 Transformed sky coordinates. 

239 """ 

240 sky_rad = self._pixel_to_sky.apply_forward(x=x, y=y) 

241 return SkyCoord(ra=sky_rad.x, dec=sky_rad.y, unit=u.rad) 

242 

243 def sky_to_pixel(self, sky: SkyCoord) -> XY[np.ndarray | float]: 

244 """Transform one or more sky coordinates to pixels 

245 

246 Parameters 

247 ---------- 

248 sky 

249 Sky coordinates to transform. 

250 

251 Returns 

252 ------- 

253 `XY` [`numpy.ndarray` | `float`] 

254 Transformed pixel coordinates. 

255 """ 

256 if sky.frame.name != "icrs": 

257 sky = sky.transform_to("icrs") 

258 ra: Longitude = sky.ra 

259 dec: Latitude = sky.dec 

260 return self._pixel_to_sky.apply_inverse( 

261 x=ra.to_value(u.rad), 

262 y=dec.to_value(u.rad), 

263 ) 

264 

265 def as_astropy(self, bbox: Box | None = None) -> SkyProjectionAstropyView: 

266 """Return an `astropy.wcs` view of this `SkyProjection`. 

267 

268 Parameters 

269 ---------- 

270 bbox 

271 Bounding box of the array the view will describe. This 

272 projection object is assumed to work on the same coordinate system 

273 in which ``bbox`` is defined, while the Astropy view will consider 

274 the first row and column in that box to be ``(0, 0)``. 

275 

276 Notes 

277 ----- 

278 This returns an object that satisfies the 

279 `astropy.wcs.wcsapi.BaseHighLevelWCS` and 

280 `astropy.wcs.wcsapi.BaseLowLevelWCS` interfaces while evaluating the 

281 underlying `SkyProjection` itself. It is *not* an `astropy.wcs.WCS` 

282 instance, which is a type that also satisfies those interfaces but 

283 only supports FITS WCS representations (see `as_fits_wcs`). 

284 """ 

285 return SkyProjectionAstropyView(self._pixel_to_sky._ast_mapping, bbox) 

286 

287 def as_fits_wcs(self, bbox: Box, allow_approximation: bool = False) -> astropy.wcs.WCS | None: 

288 """Return a FITS WCS representation of this projection, if possible. 

289 

290 Parameters 

291 ---------- 

292 bbox 

293 Bounding box of the array the FITS WCS will describe. This 

294 transform object is assumed to work on the same coordinate system 

295 in which ``bbox`` is defined, while the FITS WCS will consider the 

296 first row and column in that box to be ``(0, 0)`` (in Astropy 

297 interfaces) or ``(1, 1)`` (in the FITS representation itself). 

298 allow_approximation 

299 If `True` and this `SkyProjection` holds a FITS approximation to 

300 itself, return that approximation. 

301 """ 

302 if allow_approximation and self._fits_approximation: 302 ↛ 303line 302 didn't jump to line 303 because the condition on line 302 was never true

303 return self._fits_approximation.as_fits_wcs(bbox) 

304 return self._pixel_to_sky.as_fits_wcs(bbox) 

305 

306 def serialize[P: pydantic.BaseModel]( 

307 self, archive: OutputArchive[P], *, use_frame_sets: bool = False 

308 ) -> SkyProjectionSerializationModel[P]: 

309 """Serialize a projection to an archive. 

310 

311 Parameters 

312 ---------- 

313 archive 

314 Archive to serialize to. 

315 use_frame_sets 

316 If `True`, decompose the underlying transform and try to reference 

317 component mappings that were already serialized into a `FrameSet` 

318 in the archive. The FITS approximation transform is never 

319 decomposed. 

320 

321 Returns 

322 ------- 

323 `SkyProjectionSerializationModel` 

324 Serialized form of the projection. 

325 """ 

326 pixel_to_sky = archive.serialize_direct( 

327 "pixel_to_sky", functools.partial(self._pixel_to_sky.serialize, use_frame_sets=use_frame_sets) 

328 ) 

329 fits_approximation = ( 

330 archive.serialize_direct("fits_approximation", self._fits_approximation.serialize) 

331 if self._fits_approximation is not None 

332 else None 

333 ) 

334 return SkyProjectionSerializationModel( 

335 pixel_to_sky=pixel_to_sky, fits_approximation=fits_approximation 

336 ) 

337 

338 @staticmethod 

339 def _get_archive_tree_type[P: pydantic.BaseModel]( 

340 pointer_type: type[P], 

341 ) -> type[SkyProjectionSerializationModel[P]]: 

342 """Return the serialization model type for this object for an archive 

343 type that uses the given pointer type. 

344 """ 

345 return SkyProjectionSerializationModel[pointer_type] # type: ignore 

346 

347 @staticmethod 

348 def from_legacy( 

349 sky_wcs: LegacySkyWcs, pixel_frame: F, pixel_bounds: Bounds | None = None 

350 ) -> SkyProjection[F]: 

351 """Construct a transform from a legacy `lsst.afw.geom.SkyWcs`. 

352 

353 Parameters 

354 ---------- 

355 sky_wcs : `lsst.afw.geom.SkyWcs` 

356 Legacy WCS object. 

357 pixel_frame 

358 Coordinate frame for the pixel grid. 

359 pixel_bounds 

360 The region that bounds valid pixels for this transform. 

361 """ 

362 fits_approximation: Transform[F, SkyFrame] | None = None 

363 if (legacy_fits_approximation := sky_wcs.getFitsApproximation()) is not None: 

364 fits_approximation = Transform( 

365 pixel_frame, 

366 SkyFrame.ICRS, 

367 legacy_fits_approximation.getFrameDict(), 

368 pixel_bounds, 

369 ) 

370 return SkyProjection( 

371 Transform(pixel_frame, SkyFrame.ICRS, sky_wcs.getFrameDict(), pixel_bounds), 

372 fits_approximation=fits_approximation, 

373 ) 

374 

375 def to_legacy(self) -> LegacySkyWcs: 

376 """Convert to a legacy `lsst.afw.geom.SkyWcs` instance.""" 

377 from lsst.afw.geom import SkyWcs as LegacySkyWcs 

378 

379 try: 

380 ast_mapping = astshim.FrameDict(self._pixel_to_sky._ast_mapping) 

381 except TypeError as err: 

382 err.add_note( 

383 "Only Projections created by from_legacy and from_fits_wcs " 

384 "are guaranteed to be convertible to SkyWcs." 

385 ) 

386 raise 

387 legacy_wcs = LegacySkyWcs(ast_mapping) 

388 if self.fits_approximation is not None: 

389 legacy_wcs = legacy_wcs.copyWithFitsApproximation(self.fits_approximation.to_legacy()) 

390 return legacy_wcs 

391 

392 

393class SkyProjectionAstropyView(BaseLowLevelWCS, HighLevelWCSMixin): 

394 """An Astropy-interface view of a `SkyProjection`. 

395 

396 Notes 

397 ----- 

398 The constructor of this classe is considered a private implementation 

399 detail; use `SkyProjection.as_astropy` instead. 

400 

401 This object satisfies the `astropy.wcs.wcsapi.BaseHighLevelWCS` and 

402 `astropy.wcs.wcsapi.BaseLowLevelWCS` interfaces while evaluating the 

403 underlying `SkyProjection` itself. It is *not* an `astropy.wcs.WCS` 

404 subclass, which is a type that also satisfies those interfaces but 

405 only supports FITS WCS representations (see `SkyProjection.as_fits_wcs`). 

406 """ 

407 

408 def __init__(self, ast_pixel_to_sky: astshim.Mapping, bbox: Box | None) -> None: 

409 self._bbox = bbox 

410 if bbox is not None: 410 ↛ 412line 410 didn't jump to line 412 because the condition on line 410 was always true

411 ast_pixel_to_sky = astshim.ShiftMap(list(bbox.start.xy)).then(ast_pixel_to_sky) 

412 self._ast_pixel_to_sky = ast_pixel_to_sky 

413 

414 @property 

415 def low_level_wcs(self) -> Self: 

416 return self 

417 

418 @property 

419 def array_shape(self) -> YX[int] | None: 

420 return self._bbox.shape if self._bbox is not None else None 

421 

422 @property 

423 def axis_correlation_matrix(self) -> np.ndarray: 

424 return np.array([[True, True], [True, True]]) 

425 

426 @property 

427 def pixel_axis_names(self) -> XY[str]: 

428 return XY("x", "y") 

429 

430 @property 

431 def pixel_bounds(self) -> XY[tuple[int, int]] | None: 

432 if self._bbox is None: 

433 return None 

434 return XY((self._bbox.x.min, self._bbox.x.max), (self._bbox.y.min, self._bbox.y.max)) 

435 

436 @property 

437 def pixel_n_dim(self) -> int: 

438 return 2 

439 

440 @property 

441 def pixel_shape(self) -> XY[int] | None: 

442 array_shape = self.array_shape 

443 return array_shape.xy if array_shape is not None else None 

444 

445 @property 

446 def serialized_classes(self) -> bool: 

447 return False 

448 

449 @property 

450 def world_axis_names(self) -> tuple[str, str]: 

451 return ("ra", "dec") 

452 

453 @property 

454 def world_axis_object_classes(self) -> dict[str, tuple[type[SkyCoord], tuple[()], dict[str, Any]]]: 

455 return {"celestial": (SkyCoord, (), {"frame": ICRS, "unit": (u.rad, u.rad)})} 

456 

457 @property 

458 def world_axis_object_components(self) -> list[tuple[str, int, str]]: 

459 return [("celestial", 0, "spherical.lon.radian"), ("celestial", 1, "spherical.lat.radian")] 

460 

461 @property 

462 def world_axis_physical_types(self) -> tuple[str, str]: 

463 return ("pos.eq.ra", "pos.eq.dec") 

464 

465 @property 

466 def world_axis_units(self) -> tuple[str, str]: 

467 return ("rad", "rad") 

468 

469 @property 

470 def world_n_dim(self) -> int: 

471 return 2 

472 

473 def pixel_to_world_values(self, x: np.ndarray, y: np.ndarray) -> XY[np.ndarray]: 

474 return _ast_apply(self._ast_pixel_to_sky.applyForward, x=x, y=y) 

475 

476 def world_to_pixel_values(self, ra: np.ndarray, dec: np.ndarray) -> XY[np.ndarray]: 

477 return _ast_apply(self._ast_pixel_to_sky.applyInverse, x=ra, y=dec) 

478 

479 

480class SkyProjectionSerializationModel[P: pydantic.BaseModel](ArchiveTree): 

481 """Serialization model for projetions.""" 

482 

483 SCHEMA_NAME: ClassVar[str] = "sky_projection" 

484 SCHEMA_VERSION: ClassVar[str] = "1.0.0" 

485 MIN_READ_VERSION: ClassVar[int] = 1 

486 PUBLIC_TYPE: ClassVar[type] = SkyProjection 

487 

488 pixel_to_sky: TransformSerializationModel[P] = pydantic.Field( 

489 description="The transform that maps pixel coordinates to the sky." 

490 ) 

491 fits_approximation: TransformSerializationModel[P] | None = pydantic.Field( 

492 default=None, 

493 description=( 

494 "An approximation of the pixel-to-sky transform that is exactly representable as a FITS WCS." 

495 ), 

496 exclude_if=is_none, 

497 ) 

498 

499 def deserialize(self, archive: InputArchive[P], **kwargs: Any) -> SkyProjection[Any]: 

500 """Deserialize a projection from an archive. 

501 

502 Parameters 

503 ---------- 

504 archive 

505 Archive to read from. 

506 **kwargs 

507 Unsupported keyword arguments are accepted only to provide better 

508 error messages (raising `serialization.InvalidParameterError`). 

509 """ 

510 if kwargs: 510 ↛ 511line 510 didn't jump to line 511 because the condition on line 510 was never true

511 raise InvalidParameterError(f"Unrecognized parameters for Projection: {set(kwargs.keys())}.") 

512 pixel_to_sky = self.pixel_to_sky.deserialize(archive) 

513 fits_approximation = ( 

514 self.fits_approximation.deserialize(archive) if self.fits_approximation is not None else None 

515 ) 

516 return SkyProjection(pixel_to_sky, fits_approximation=fits_approximation)