Coverage for tests/test_transforms.py: 75%

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 

14import dataclasses 

15import functools 

16import os 

17from typing import Any, ClassVar 

18 

19import astropy.units as u 

20import numpy as np 

21import pydantic 

22import pytest 

23 

24from lsst.images import ( 

25 ICRS, 

26 XY, 

27 YX, 

28 Box, 

29 CameraFrameSet, 

30 CameraFrameSetSerializationModel, 

31 DetectorFrame, 

32 FocalPlaneFrame, 

33 GeneralFrame, 

34 SkyProjection, 

35 Transform, 

36 TransformSerializationModel, 

37) 

38from lsst.images._transforms import _ast as astshim 

39from lsst.images.fits import PointerModel 

40from lsst.images.serialization import ArchiveTree, InputArchive, JsonRef, OutputArchive 

41from lsst.images.tests import ( 

42 DP2_VISIT_DETECTOR_DATA_ID, 

43 RoundtripFits, 

44 RoundtripJson, 

45 check_transform, 

46 compare_sky_projection_to_legacy_wcs, 

47 legacy_points_to_xy_array, 

48 make_random_sky_projection, 

49) 

50 

51EXTERNAL_DATA_DIR = os.environ.get("TESTDATA_IMAGES_DIR", None) 

52 

53 

54@pytest.fixture(scope="session") 

55def legacy_camera() -> Any: 

56 """Return a legacy Camera loaded from camera.fits. 

57 

58 Skips if TESTDATA_IMAGES_DIR is unset or lsst.afw.cameraGeom is 

59 unavailable. 

60 """ 

61 if EXTERNAL_DATA_DIR is None: 61 ↛ 63line 61 didn't jump to line 63 because the condition on line 61 was always true

62 pytest.skip("TESTDATA_IMAGES_DIR is not in the environment.") 

63 try: 

64 from lsst.afw.cameraGeom import Camera 

65 except ImportError: 

66 pytest.skip("'lsst.afw.cameraGeom' could not be imported.") 

67 filename = os.path.join(EXTERNAL_DATA_DIR, "dp2", "legacy", "camera.fits") 

68 return Camera.readFits(filename) 

69 

70 

71@pytest.fixture(scope="session") 

72def legacy_detector_wcs() -> dict[str, Any]: 

73 """Return WCS-related objects read from visit_image.fits. 

74 

75 Skips if TESTDATA_IMAGES_DIR is unset or lsst.afw.image is unavailable. 

76 """ 

77 if EXTERNAL_DATA_DIR is None: 77 ↛ 79line 77 didn't jump to line 79 because the condition on line 77 was always true

78 pytest.skip("TESTDATA_IMAGES_DIR is not in the environment.") 

79 try: 

80 from lsst.afw.image import ExposureFitsReader 

81 except ImportError: 

82 pytest.skip("'lsst.afw.image' could not be imported.") 

83 filename = os.path.join(EXTERNAL_DATA_DIR, "dp2", "legacy", "visit_image.fits") 

84 reader = ExposureFitsReader(filename) 

85 return { 

86 "legacy_wcs": reader.readWcs(), 

87 "wcs_bbox": Box.from_legacy(reader.readDetector().getBBox()), 

88 "subimage_bbox": Box.from_legacy(reader.readBBox()), 

89 } 

90 

91 

92def test_identity() -> None: 

93 """Test an identity transform.""" 

94 frame = DetectorFrame(**DP2_VISIT_DETECTOR_DATA_ID, bbox=Box.factory[:5, :4]) 

95 xy = frame.bbox.meshgrid().map(np.ravel) 

96 identity = Transform.identity(frame) 

97 check_transform(identity, xy, xy, frame, frame) 

98 assert identity.decompose() == [] 

99 with RoundtripJson(identity) as roundtrip: 

100 pass 

101 check_transform(roundtrip.result, xy, xy, frame, frame) 

102 

103 

104def test_transform_equality() -> None: 

105 """Test Transform.__eq__ across all of its comparison branches.""" 

106 pixel_frame = DetectorFrame(**DP2_VISIT_DETECTOR_DATA_ID, bbox=Box.factory[:5, :4]) 

107 focal_plane = FocalPlaneFrame(instrument="LSSTCam", visit=1, unit=u.mm) 

108 # A distinct frame for the in-frame and out-frame branches. 

109 alt_frame = DetectorFrame(instrument="LSSTCam", visit=1, detector=12, bbox=Box.factory[:5, :4]) 

110 in_bounds = Box.factory[:5, :4] 

111 out_bounds = Box.factory[:10, :8] 

112 

113 def make( 

114 *, 

115 in_frame: Any = pixel_frame, 

116 out_frame: Any = focal_plane, 

117 ast_mapping: astshim.Mapping | None = None, 

118 in_bounds_: Box | None = in_bounds, 

119 out_bounds_: Box | None = out_bounds, 

120 components: Any = (), 

121 ) -> Transform[Any, Any]: 

122 return Transform( 

123 in_frame, 

124 out_frame, 

125 ast_mapping if ast_mapping is not None else astshim.UnitMap(2), 

126 in_bounds=in_bounds_, 

127 out_bounds=out_bounds_, 

128 components=components, 

129 ) 

130 

131 base = make() 

132 

133 # Identity short-circuit: an object is always equal to itself. 

134 assert base == base 

135 

136 # Two independently constructed but equivalent transforms are equal, 

137 # and equality is symmetric. 

138 assert base == make() 

139 assert make() == base 

140 

141 # Comparison against a non-Transform yields NotImplemented, so Python 

142 # falls back to identity: the objects are unequal and != is True. 

143 assert not (base == "not a transform") 

144 assert base != "not a transform" 

145 assert base != None # noqa: E711 

146 assert base != 42 

147 

148 # Each remaining branch differs from base in exactly one attribute. 

149 assert base != make(ast_mapping=astshim.ShiftMap([1.0, 2.0])) 

150 assert base != make(in_bounds_=out_bounds) 

151 assert base != make(out_bounds_=in_bounds) 

152 assert base != make(in_frame=alt_frame) 

153 assert base != make(out_frame=alt_frame) 

154 assert base != make(components=[Transform.identity(alt_frame)]) 

155 

156 

157def test_sky_projection_equality() -> None: 

158 """Test SkyProjection.__eq__ across all of its comparison branches.""" 

159 pixel_frame = DetectorFrame(**DP2_VISIT_DETECTOR_DATA_ID, bbox=Box.factory[:5, :4]) 

160 

161 # Check the two failure modes. 

162 with pytest.raises(ValueError): 

163 SkyProjection(Transform(ICRS, ICRS, astshim.UnitMap(2))) 

164 

165 with pytest.raises(ValueError): 

166 SkyProjection(Transform(pixel_frame, pixel_frame, astshim.UnitMap(2))) 

167 

168 def make_pixel_to_sky(ast_mapping: astshim.Mapping | None = None) -> Transform[Any, Any]: 

169 mapping = ast_mapping if ast_mapping is not None else astshim.UnitMap(2) 

170 return Transform(pixel_frame, ICRS, mapping) 

171 

172 base = SkyProjection(make_pixel_to_sky()) 

173 

174 # Identity short-circuit: an object is always equal to itself. 

175 assert base == base 

176 

177 # Two independently constructed but equivalent projections are equal. 

178 assert base == SkyProjection(make_pixel_to_sky()) 

179 

180 # Comparison against a non-SkyProjection yields NotImplemented. 

181 assert not (base == "not a projection") 

182 assert base != "not a projection" 

183 assert base != None # noqa: E711 

184 

185 # Differ only in the pixel-to-sky transform. 

186 assert base != SkyProjection(make_pixel_to_sky(astshim.ShiftMap([1.0, 2.0]))) 

187 

188 # The fits_approximation branch: absent on base but present here. 

189 with_approx = SkyProjection( 

190 make_pixel_to_sky(), fits_approximation=make_pixel_to_sky(astshim.ShiftMap([0.1, 0.2])) 

191 ) 

192 assert base != with_approx 

193 

194 # Equal pixel-to-sky and equal fits_approximations are equal. 

195 with_approx_again = SkyProjection( 

196 make_pixel_to_sky(), fits_approximation=make_pixel_to_sky(astshim.ShiftMap([0.1, 0.2])) 

197 ) 

198 assert with_approx == with_approx_again 

199 

200 # Same pixel-to-sky transform but a different fits_approximation. 

201 other_approx = SkyProjection( 

202 make_pixel_to_sky(), fits_approximation=make_pixel_to_sky(astshim.ShiftMap([0.3, 0.4])) 

203 ) 

204 assert with_approx != other_approx 

205 

206 

207def test_affine_2x2() -> None: 

208 """Test an affine transform constructed from a 2x2 matrix.""" 

209 in_frame = DetectorFrame(**DP2_VISIT_DETECTOR_DATA_ID, bbox=Box.factory[:5, :4]) 

210 out_frame = GeneralFrame(unit=u.pix) 

211 transform_matrix = np.array([[2.0, 0.25], [-0.75, 0.8]]) 

212 in_xy = in_frame.bbox.meshgrid().map(np.ravel) 

213 in_matrix = np.array([in_xy.x, in_xy.y]) 

214 out_matrix = np.dot(transform_matrix, in_matrix) 

215 check_transform( 

216 Transform.affine(in_frame, out_frame, transform_matrix), 

217 in_xy, 

218 XY(x=out_matrix[0, :], y=out_matrix[1, :]), 

219 in_frame, 

220 out_frame, 

221 in_atol=1e-15 * u.pix, 

222 out_atol=1e-15 * u.pix, 

223 ) 

224 

225 

226def test_affine_3x3() -> None: 

227 """Test an affine transform constructed from a 3x3 matrix.""" 

228 in_frame = DetectorFrame(**DP2_VISIT_DETECTOR_DATA_ID, bbox=Box.factory[:5, :4]) 

229 out_frame = GeneralFrame(unit=u.pix) 

230 transform_matrix = np.array([[2.0, 0.25, -0.5], [-0.75, 0.8, 0.4], [0.0, 0.0, 1.0]]) 

231 in_xy = in_frame.bbox.meshgrid().map(np.ravel) 

232 in_matrix = np.array([in_xy.x, in_xy.y, np.ones(in_xy.x.shape)]) 

233 out_matrix = np.dot(transform_matrix, in_matrix) 

234 check_transform( 

235 Transform.affine(in_frame, out_frame, transform_matrix), 

236 in_xy, 

237 XY(x=out_matrix[0, :], y=out_matrix[1, :]), 

238 in_frame, 

239 out_frame, 

240 in_atol=1e-15 * u.pix, 

241 out_atol=1e-15 * u.pix, 

242 ) 

243 

244 

245def compare_to_legacy_camera(legacy_camera: Any, frame_set: CameraFrameSet) -> None: 

246 """Assert that transforms extracted from a CameraFrameSet match the 

247 legacy afw implementations. 

248 """ 

249 from lsst.afw.cameraGeom import FIELD_ANGLE, FOCAL_PLANE, PIXELS 

250 from lsst.geom import Point2D 

251 

252 legacy_detector = legacy_camera[16] 

253 pixel_legacy_points = [Point2D(50.0, 60.0), Point2D(801.2, 322.8), Point2D(33.5, 22.1)] 

254 fp_legacy_points = [legacy_detector.transform(p, PIXELS, FOCAL_PLANE) for p in pixel_legacy_points] 

255 fa_legacy_points = [legacy_detector.transform(p, PIXELS, FIELD_ANGLE) for p in pixel_legacy_points] 

256 pixel_xy_array = legacy_points_to_xy_array(pixel_legacy_points) 

257 fp_xy_array = legacy_points_to_xy_array(fp_legacy_points) 

258 fa_xy_array = legacy_points_to_xy_array(fa_legacy_points) 

259 # Test transforms extracted directly from the frame set. 

260 pixel_to_fp = frame_set[frame_set.detector(16), frame_set.focal_plane()] 

261 check_transform(pixel_to_fp, pixel_xy_array, fp_xy_array, frame_set.detector(16), frame_set.focal_plane()) 

262 pixel_to_fa = frame_set[frame_set.detector(16), frame_set.field_angle()] 

263 check_transform(pixel_to_fa, pixel_xy_array, fa_xy_array, frame_set.detector(16), frame_set.field_angle()) 

264 fp_to_fa = frame_set[frame_set.focal_plane(), frame_set.field_angle()] 

265 check_transform(fp_to_fa, fp_xy_array, fa_xy_array, frame_set.focal_plane(), frame_set.field_angle()) 

266 # Test a composition. 

267 pixel_to_fa_indirect = pixel_to_fp.then(fp_to_fa) 

268 check_transform( 

269 pixel_to_fa_indirect, 

270 pixel_xy_array, 

271 fa_xy_array, 

272 frame_set.detector(16), 

273 frame_set.field_angle(), 

274 ) 

275 pixel_to_fp_d, fp_to_fa_d = pixel_to_fa_indirect.decompose() 

276 check_transform( 

277 pixel_to_fp_d, pixel_xy_array, fp_xy_array, frame_set.detector(16), frame_set.focal_plane() 

278 ) 

279 check_transform(fp_to_fa_d, fp_xy_array, fa_xy_array, frame_set.focal_plane(), frame_set.field_angle()) 

280 fa_to_fp_d, fp_to_pixel_d = pixel_to_fa_indirect.inverted().decompose() 

281 check_transform(fa_to_fp_d, fa_xy_array, fp_xy_array, frame_set.field_angle(), frame_set.focal_plane()) 

282 check_transform( 

283 fp_to_pixel_d, fp_xy_array, pixel_xy_array, frame_set.focal_plane(), frame_set.detector(16) 

284 ) 

285 

286 

287def test_camera(legacy_camera: Any) -> None: 

288 """Test CameraFrameSet construction, transforms, and FITS/JSON 

289 serialization round-trips. 

290 

291 Also verifies the archive system's pointer and frame-set reference 

292 machinery. 

293 """ 

294 legacy_camera = legacy_camera 

295 frame_set = CameraFrameSet.from_legacy(legacy_camera) 

296 detector_id: int = DP2_VISIT_DETECTOR_DATA_ID["detector"] 

297 compare_to_legacy_camera(legacy_camera, frame_set) 

298 test_holder = FrameSetTestHolder( 

299 frames=frame_set, 

300 pixels_to_fp=frame_set[frame_set.detector(detector_id), frame_set.focal_plane()], 

301 ) 

302 with RoundtripFits(test_holder) as roundtrip1: 

303 assert len(roundtrip1.serialized.pixels_to_fp.frames) == 2 

304 assert len(roundtrip1.serialized.pixels_to_fp.bounds) == 2 

305 assert len(roundtrip1.serialized.pixels_to_fp.mappings) == 1 

306 # Instead of storing the AST mapping directly, we should have 

307 # stored a reference to the frame set: 

308 assert isinstance(roundtrip1.serialized.pixels_to_fp.mappings[0], PointerModel) 

309 compare_to_legacy_camera(legacy_camera, roundtrip1.result.frames) 

310 assert roundtrip1.result.pixels_to_fp.in_frame == frame_set.detector(detector_id) 

311 assert roundtrip1.result.pixels_to_fp.out_frame == frame_set.focal_plane() 

312 assert ( 

313 roundtrip1.result.pixels_to_fp._ast_mapping.simplified().show() 

314 == test_holder.pixels_to_fp._ast_mapping.simplified().show() 

315 ) 

316 with RoundtripJson(test_holder) as roundtrip2: 

317 assert len(roundtrip2.serialized.pixels_to_fp.frames) == 2 

318 assert len(roundtrip2.serialized.pixels_to_fp.bounds) == 2 

319 assert len(roundtrip2.serialized.pixels_to_fp.mappings) == 1 

320 # Instead of storing the AST mapping directly, we should have 

321 # stored a reference to the frame set: 

322 assert isinstance(roundtrip2.serialized.pixels_to_fp.mappings[0], JsonRef) 

323 raw_data = roundtrip2.inspect() 

324 assert len(raw_data["indirect"]) == 1 

325 assert raw_data["frames"] == {"$ref": "#/indirect/0"} 

326 compare_to_legacy_camera(legacy_camera, roundtrip2.result.frames) 

327 assert roundtrip2.result.pixels_to_fp.in_frame == frame_set.detector(detector_id) 

328 assert roundtrip2.result.pixels_to_fp.out_frame == frame_set.focal_plane() 

329 assert ( 

330 roundtrip2.result.pixels_to_fp._ast_mapping.simplified().show() 

331 == test_holder.pixels_to_fp._ast_mapping.simplified().show() 

332 ) 

333 

334 

335def test_fits_wcs_projection_to_legacy() -> None: 

336 """Verify that a projection created by from_fits_wcs can be converted 

337 to a legacy SkyWcs. 

338 

339 The AST pixel frame uses the domain PIXEL, while lsst.afw.geom.SkyWcs 

340 requires PIXELS, so to_legacy has to rename it. 

341 """ 

342 pytest.importorskip("lsst.afw.geom") 

343 rng = np.random.default_rng(43) 

344 bbox = Box.factory[75:275, 25:225] 

345 pixel_frame = GeneralFrame(unit=u.pix) 

346 sky_projection = make_random_sky_projection(rng, pixel_frame, bbox) 

347 legacy_wcs = sky_projection.to_legacy() 

348 compare_sky_projection_to_legacy_wcs(sky_projection, legacy_wcs, pixel_frame, bbox, is_fits=True) 

349 # The conversion must not modify the projection in place: its own AST 

350 # mapping keeps the PIXEL domain, and the conversion is repeatable. 

351 frame_set = sky_projection.pixel_to_sky_transform._ast_mapping 

352 assert isinstance(frame_set, astshim.FrameSet) 

353 domains = {frame_set.getFrame(i, copy=False).domain for i in range(1, frame_set.nFrame + 1)} 

354 assert "PIXEL" in domains 

355 assert "PIXELS" not in domains 

356 compare_sky_projection_to_legacy_wcs( 

357 sky_projection, sky_projection.to_legacy(), pixel_frame, bbox, is_fits=True 

358 ) 

359 

360 

361def test_detector_wcs(legacy_detector_wcs: dict[str, Any]) -> None: 

362 """Test the Transform/SkyProjection representation of a detector WCS.""" 

363 legacy_wcs = legacy_detector_wcs["legacy_wcs"] 

364 wcs_bbox = legacy_detector_wcs["wcs_bbox"] 

365 subimage_bbox = legacy_detector_wcs["subimage_bbox"] 

366 detector_frame = DetectorFrame(**DP2_VISIT_DETECTOR_DATA_ID, bbox=wcs_bbox) 

367 sky_projection = SkyProjection.from_legacy(legacy_wcs, detector_frame) 

368 assert sky_projection.fits_approximation is not None 

369 compare_sky_projection_to_legacy_wcs(sky_projection, legacy_wcs, detector_frame, subimage_bbox) 

370 # When we convert from a legacy SkyWcs, the internal AST Mapping needs 

371 # to really be an AST FrameSet in order to be able to convert back. 

372 assert "Begin FrameSet" in sky_projection.show() 

373 compare_sky_projection_to_legacy_wcs( 

374 sky_projection, sky_projection.to_legacy(), detector_frame, subimage_bbox 

375 ) 

376 assert "Begin FrameSet" in sky_projection.fits_approximation.show() 

377 compare_sky_projection_to_legacy_wcs( 

378 sky_projection.fits_approximation, 

379 sky_projection.fits_approximation.to_legacy(), 

380 detector_frame, 

381 subimage_bbox, 

382 is_fits=True, 

383 ) 

384 with RoundtripJson(sky_projection, "SkyProjection") as roundtrip: 

385 pass 

386 compare_sky_projection_to_legacy_wcs(roundtrip.result, legacy_wcs, detector_frame, subimage_bbox) 

387 # The AST FrameSet-ness needs to propagate through serialization. 

388 assert "Begin FrameSet" in roundtrip.result.show() 

389 compare_sky_projection_to_legacy_wcs( 

390 sky_projection, roundtrip.result.to_legacy(), detector_frame, subimage_bbox 

391 ) 

392 with RoundtripJson(sky_projection.fits_approximation, "SkyProjection") as roundtrip: 

393 pass 

394 compare_sky_projection_to_legacy_wcs( 

395 roundtrip.result, 

396 legacy_wcs.getFitsApproximation(), 

397 detector_frame, 

398 subimage_bbox, 

399 is_fits=True, 

400 ) 

401 assert "Begin FrameSet" in roundtrip.result.show() 

402 compare_sky_projection_to_legacy_wcs( 

403 sky_projection.fits_approximation, 

404 roundtrip.result.to_legacy(), 

405 detector_frame, 

406 subimage_bbox, 

407 is_fits=True, 

408 ) 

409 

410 

411@dataclasses.dataclass 

412class FrameSetTestHolder: 

413 """A top-level object that holds a CameraFrameSet and a transform 

414 extracted from it, for testing archive pointers and frame set references. 

415 """ 

416 

417 frames: CameraFrameSet 

418 pixels_to_fp: Transform[DetectorFrame, FocalPlaneFrame] 

419 

420 def serialize[P: pydantic.BaseModel](self, archive: OutputArchive[P]) -> FrameSetTestHolderModel[P]: 

421 frames_model = archive.serialize_frame_set( 

422 "frames", self.frames, self.frames.serialize, key=id(self.frames) 

423 ) 

424 pixels_to_fp_model = archive.serialize_direct( 

425 "pixels_to_fp", functools.partial(self.pixels_to_fp.serialize, use_frame_sets=True) 

426 ) 

427 return FrameSetTestHolderModel[P](frames=frames_model, pixels_to_fp=pixels_to_fp_model) 

428 

429 @staticmethod 

430 def _get_archive_tree_type[P: pydantic.BaseModel]( 

431 pointer_type: type[P], 

432 ) -> type[FrameSetTestHolderModel[P]]: 

433 return FrameSetTestHolderModel[pointer_type] # type: ignore 

434 

435 

436class FrameSetTestHolderModel[P: pydantic.BaseModel](ArchiveTree): 

437 """The serialization model for FrameSetTestHolder.""" 

438 

439 SCHEMA_NAME: ClassVar[str] = "_test_frame_set_holder" 

440 SCHEMA_VERSION: ClassVar[str] = "1.0.0" 

441 MIN_READ_VERSION: ClassVar[int] = 1 

442 PUBLIC_TYPE: ClassVar[type] = FrameSetTestHolder 

443 

444 frames: CameraFrameSetSerializationModel | P 

445 pixels_to_fp: TransformSerializationModel[P] 

446 

447 def deserialize(self, archive: InputArchive[Any]) -> FrameSetTestHolder: 

448 assert not isinstance(self.frames, CameraFrameSetSerializationModel), "Archive pointer expected." 

449 frames = archive.deserialize_pointer( 

450 self.frames, CameraFrameSetSerializationModel, CameraFrameSetSerializationModel.deserialize 

451 ) 

452 pixels_to_fp = self.pixels_to_fp.deserialize(archive) 

453 return FrameSetTestHolder(frames, pixels_to_fp) 

454 

455 

456@dataclasses.dataclass 

457class _BroadcastTestData: 

458 """Shared inputs for broadcasting/scalar tests on Transform and 

459 SkyProjection. 

460 """ 

461 

462 in_frame: DetectorFrame 

463 out_frame: GeneralFrame 

464 matrix: np.ndarray 

465 scalar_x: float 

466 scalar_y: float 

467 xv: list[int] 

468 yv: list[int] 

469 sky_proj: SkyProjection[DetectorFrame] 

470 

471 

472@pytest.fixture 

473def broadcast_test_data() -> _BroadcastTestData: 

474 """Return shared inputs for broadcasting/scalar tests.""" 

475 in_frame = DetectorFrame(instrument="Inst", visit=1, detector=1, bbox=Box.factory[0:20, 0:20]) 

476 return _BroadcastTestData( 

477 in_frame=in_frame, 

478 out_frame=GeneralFrame(unit=u.pix), 

479 matrix=np.array([[2.0, 0.5], [-0.3, 1.5]]), 

480 scalar_x=3.0, 

481 scalar_y=7.0, 

482 xv=[1, 2, 3], 

483 yv=[4, 5, 6], 

484 sky_proj=make_random_sky_projection(np.random.default_rng(42), in_frame, in_frame.bbox), 

485 ) 

486 

487 

488def test_apply_forward_scalar(broadcast_test_data: _BroadcastTestData) -> None: 

489 """Verify that apply_forward and apply_inverse accept scalar x/y and return 

490 scalar floats, and that the _q variants accept scalar Quantity inputs. 

491 """ 

492 t = broadcast_test_data.sky_proj.pixel_to_sky_transform 

493 # apply_forward with Python float scalars should return XY of floats. 

494 result_fwd = t.apply_forward(x=broadcast_test_data.scalar_x, y=broadcast_test_data.scalar_y) 

495 assert type(result_fwd.x) is float 

496 assert type(result_fwd.y) is float 

497 # Values must match the corresponding single-element array call. 

498 ref_fwd = t.apply_forward( 

499 x=np.array([broadcast_test_data.scalar_x]), y=np.array([broadcast_test_data.scalar_y]) 

500 ) 

501 assert result_fwd.x == ref_fwd.x[0] 

502 assert result_fwd.y == ref_fwd.y[0] 

503 # apply_inverse round-trips back to the original scalars. 

504 result_inv = t.apply_inverse(x=result_fwd.x, y=result_fwd.y) 

505 assert type(result_inv.x) is float 

506 assert type(result_inv.y) is float 

507 np.testing.assert_allclose(result_inv.x, broadcast_test_data.scalar_x, atol=1e-12) 

508 np.testing.assert_allclose(result_inv.y, broadcast_test_data.scalar_y, atol=1e-12) 

509 # apply_forward_q / apply_inverse_q with scalar Quantity inputs. 

510 x_q = broadcast_test_data.scalar_x * t.in_frame.unit 

511 y_q = broadcast_test_data.scalar_y * t.in_frame.unit 

512 result_fwd_q = t.apply_forward_q(x=x_q, y=y_q) 

513 assert result_fwd_q.x.shape == () 

514 assert result_fwd_q.y.shape == () 

515 np.testing.assert_allclose(result_fwd_q.x.to_value(t.out_frame.unit), result_fwd.x, atol=1e-12) 

516 result_inv_q = t.apply_inverse_q(x=result_fwd_q.x, y=result_fwd_q.y) 

517 assert result_inv_q.x.shape == () 

518 np.testing.assert_allclose( 

519 result_inv_q.x.to_value(t.in_frame.unit), broadcast_test_data.scalar_x, atol=1e-12 

520 ) 

521 

522 

523def test_apply_array_like_and_integer_input(broadcast_test_data: _BroadcastTestData) -> None: 

524 """Verify that apply_forward accepts Python lists and integer-dtype 

525 arrays, returning float64 ndarray results consistent with float64 

526 array input. 

527 """ 

528 t = broadcast_test_data.sky_proj.pixel_to_sky_transform 

529 # Python list input should return an ndarray. 

530 result_list = t.apply_forward(x=broadcast_test_data.xv, y=broadcast_test_data.yv) 

531 assert isinstance(result_list.x, np.ndarray) 

532 assert isinstance(result_list.y, np.ndarray) 

533 ref = t.apply_forward(x=np.array(broadcast_test_data.xv), y=np.array(broadcast_test_data.yv)) 

534 np.testing.assert_array_equal(result_list.x, ref.x) 

535 np.testing.assert_array_equal(result_list.y, ref.y) 

536 # Integer dtype arrays should not raise and should return float64. 

537 xi = np.array(broadcast_test_data.xv, dtype=np.int32) 

538 yi = np.array(broadcast_test_data.yv, dtype=np.int32) 

539 result_int = t.apply_forward(x=xi, y=yi) 

540 assert result_int.x.dtype == np.float64 

541 assert result_int.y.dtype == np.float64 

542 np.testing.assert_array_equal(result_int.x, ref.x) 

543 np.testing.assert_array_equal(result_int.y, ref.y) 

544 

545 

546def test_apply_broadcast(broadcast_test_data: _BroadcastTestData) -> None: 

547 """Verify that apply_forward and apply_inverse broadcast x and y like 

548 a NumPy ufunc, in both 1-D and 2-D cases. 

549 """ 

550 t = broadcast_test_data.sky_proj.pixel_to_sky_transform 

551 xv = np.array(broadcast_test_data.xv) 

552 yv = np.array(broadcast_test_data.yv + [7]) 

553 # 1-D broadcast: array x, scalar y. 

554 result_1d = t.apply_forward(x=xv, y=broadcast_test_data.scalar_y) 

555 assert isinstance(result_1d.x, np.ndarray) 

556 assert result_1d.x.shape == xv.shape 

557 ref_1d = t.apply_forward(x=xv, y=np.full_like(xv, broadcast_test_data.scalar_y)) 

558 np.testing.assert_array_equal(result_1d.x, ref_1d.x) 

559 np.testing.assert_array_equal(result_1d.y, ref_1d.y) 

560 # 2-D broadcast: column x (M,1) × row y (1,N) -> (M,N). 

561 x2d = xv[:, np.newaxis] # shape (3, 1) 

562 y2d = yv[np.newaxis, :] # shape (1, 4) 

563 result_2d = t.apply_forward(x=x2d, y=y2d) 

564 assert result_2d.x.shape == (3, 4) 

565 assert result_2d.y.shape == (3, 4) 

566 # Values must match the fully expanded meshgrid call. 

567 xmesh, ymesh = np.meshgrid(xv, yv, indexing="ij") 

568 ref_2d = t.apply_forward(x=xmesh, y=ymesh) 

569 np.testing.assert_array_equal(result_2d.x, ref_2d.x) 

570 np.testing.assert_array_equal(result_2d.y, ref_2d.y) 

571 # apply_inverse also broadcasts. 

572 result_inv_2d = t.apply_inverse(x=result_2d.x, y=result_2d.y) 

573 assert result_inv_2d.x.shape == (3, 4) 

574 np.testing.assert_allclose(result_inv_2d.x, xmesh, atol=1e-12) 

575 np.testing.assert_allclose(result_inv_2d.y, ymesh, atol=1e-12) 

576 

577 

578def test_sky_projection_broadcast(broadcast_test_data: _BroadcastTestData) -> None: 

579 """Verify that SkyProjection.pixel_to_sky, sky_to_pixel, and the 

580 Astropy view broadcast x and y like a NumPy ufunc. 

581 """ 

582 p = broadcast_test_data 

583 xv = np.array(p.xv) 

584 yv = np.array(p.yv + [7]) 

585 # 1-D broadcast: array x, scalar y. 

586 sc_1d = p.sky_proj.pixel_to_sky(x=xv, y=p.scalar_y) 

587 assert sc_1d.shape == xv.shape 

588 ref_1d = p.sky_proj.pixel_to_sky(x=xv, y=np.full_like(xv, p.scalar_y)) 

589 np.testing.assert_allclose(sc_1d.ra.rad, ref_1d.ra.rad, atol=1e-12) 

590 np.testing.assert_allclose(sc_1d.dec.rad, ref_1d.dec.rad, atol=1e-12) 

591 # 2-D broadcast: column x (M,1) × row y (1,N) -> (M,N). 

592 x2d = xv[:, np.newaxis] # shape (3, 1) 

593 y2d = yv[np.newaxis, :] # shape (1, 4) 

594 sc_2d = p.sky_proj.pixel_to_sky(x=x2d, y=y2d) 

595 assert sc_2d.shape == (3, 4) 

596 xmesh, ymesh = np.meshgrid(xv, yv, indexing="ij") 

597 ref_2d = p.sky_proj.pixel_to_sky(x=xmesh, y=ymesh) 

598 np.testing.assert_allclose(sc_2d.ra.rad, ref_2d.ra.rad, atol=1e-12) 

599 np.testing.assert_allclose(sc_2d.dec.rad, ref_2d.dec.rad, atol=1e-12) 

600 # sky_to_pixel round-trips back to the original grid. 

601 pix_2d = p.sky_proj.sky_to_pixel(sc_2d) 

602 assert pix_2d.x.shape == (3, 4) 

603 np.testing.assert_allclose(pix_2d.x, xmesh, atol=1e-9) 

604 np.testing.assert_allclose(pix_2d.y, ymesh, atol=1e-9) 

605 # SkyProjectionAstropyView.pixel_to_world_values also broadcasts. 

606 view = p.sky_proj.as_astropy() 

607 world_2d = view.pixel_to_world_values(x2d, y2d) 

608 assert world_2d[0].shape == (3, 4) 

609 assert world_2d[1].shape == (3, 4) 

610 np.testing.assert_allclose(world_2d[0], ref_2d.ra.rad, atol=1e-12) 

611 np.testing.assert_allclose(world_2d[1], ref_2d.dec.rad, atol=1e-12) 

612 # SkyProjectionAstropyView.world_to_pixel_values also broadcasts. 

613 ra_2d = ref_2d.ra.rad[:, np.newaxis, :] # (3, 1, 4) — over-broadcast to check 

614 dec_2d = ref_2d.dec.rad[np.newaxis, :, :] # (1, 3, 4) 

615 pix_world = view.world_to_pixel_values(ra_2d, dec_2d) 

616 assert pix_world[0].shape == (3, 3, 4) 

617 

618 

619def test_apply_xy_yx(broadcast_test_data: _BroadcastTestData) -> None: 

620 """Verify that apply_forward, apply_inverse, and the _q variants accept 

621 XY and YX positional arguments, producing results identical to the 

622 equivalent x=/y= keyword calls. 

623 """ 

624 p = broadcast_test_data 

625 t = p.sky_proj.pixel_to_sky_transform 

626 sx, sy = p.scalar_x, p.scalar_y 

627 xv, yv = np.array(p.xv, dtype=float), np.array(p.yv, dtype=float) 

628 

629 # --- apply_forward: scalar --- 

630 ref_fwd = t.apply_forward(x=sx, y=sy) 

631 assert t.apply_forward(XY(sx, sy)) == ref_fwd 

632 assert t.apply_forward(YX(sy, sx)) == ref_fwd 

633 

634 # --- apply_forward: array --- 

635 ref_fwd_arr = t.apply_forward(x=xv, y=yv) 

636 np.testing.assert_array_equal(t.apply_forward(XY(xv, yv)).x, ref_fwd_arr.x) 

637 np.testing.assert_array_equal(t.apply_forward(YX(yv, xv)).x, ref_fwd_arr.x) 

638 

639 # --- apply_inverse: scalar --- 

640 ref_inv = t.apply_inverse(x=ref_fwd.x, y=ref_fwd.y) 

641 assert t.apply_inverse(XY(ref_fwd.x, ref_fwd.y)) == ref_inv 

642 assert t.apply_inverse(YX(ref_fwd.y, ref_fwd.x)) == ref_inv 

643 

644 # --- apply_forward_q: scalar --- 

645 x_q = sx * t.in_frame.unit 

646 y_q = sy * t.in_frame.unit 

647 ref_fwd_q = t.apply_forward_q(x=x_q, y=y_q) 

648 result_q = t.apply_forward_q(XY(x_q, y_q)) 

649 np.testing.assert_allclose(result_q.x.value, ref_fwd_q.x.value, atol=1e-12) 

650 result_q_yx = t.apply_forward_q(YX(y_q, x_q)) 

651 np.testing.assert_allclose(result_q_yx.x.value, ref_fwd_q.x.value, atol=1e-12) 

652 

653 # --- apply_inverse_q: scalar --- 

654 ref_inv_q = t.apply_inverse_q(x=ref_fwd_q.x, y=ref_fwd_q.y) 

655 result_inv_q = t.apply_inverse_q(XY(ref_fwd_q.x, ref_fwd_q.y)) 

656 np.testing.assert_allclose(result_inv_q.x.value, ref_inv_q.x.value, atol=1e-12) 

657 

658 # --- TypeError on bad combinations --- 

659 with pytest.raises(TypeError): 

660 t.apply_forward(XY(sx, sy), x=sx) 

661 with pytest.raises(TypeError): 

662 t.apply_forward(YX(sy, sx), y=sy) 

663 with pytest.raises(TypeError): 

664 t.apply_forward()