Coverage for tests/test_transforms.py: 83%

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 astropy.wcs 

21import numpy as np 

22import pydantic 

23import pytest 

24from astropy.coordinates import Longitude 

25 

26from lsst.images import ( 

27 ICRS, 

28 XY, 

29 YX, 

30 Box, 

31 DetectorFrame, 

32 FocalPlaneFrame, 

33 GeneralFrame, 

34 SkyProjection, 

35 Transform, 

36 TransformSerializationModel, 

37) 

38from lsst.images._transforms import _ast as astshim 

39from lsst.images.cameras import CameraFrameSet, CameraFrameSetSerializationModel 

40from lsst.images.describe import FieldRole, Report 

41from lsst.images.fits import PointerModel 

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

43from lsst.images.tests import ( 

44 DP2_VISIT_DETECTOR_DATA_ID, 

45 RoundtripFits, 

46 RoundtripJson, 

47 check_transform, 

48 compare_sky_projection_to_legacy_wcs, 

49 legacy_points_to_xy_array, 

50 make_random_sky_projection, 

51) 

52 

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

54 

55 

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

57def legacy_camera() -> Any: 

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

59 

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

61 unavailable. 

62 """ 

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

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

65 try: 

66 from lsst.afw.cameraGeom import Camera 

67 except ImportError: 

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

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

70 return Camera.readFits(filename) 

71 

72 

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

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

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

76 

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

78 """ 

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

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

81 try: 

82 from lsst.afw.image import ExposureFitsReader 

83 except ImportError: 

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

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

86 reader = ExposureFitsReader(filename) 

87 return { 

88 "legacy_wcs": reader.readWcs(), 

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

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

91 } 

92 

93 

94def test_identity() -> None: 

95 """Test an identity transform.""" 

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

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

98 identity = Transform.identity(frame) 

99 check_transform(identity, xy, xy, frame, frame) 

100 assert identity.decompose() == [] 

101 with RoundtripJson(identity) as roundtrip: 

102 pass 

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

104 

105 

106def test_transform_equality() -> None: 

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

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

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

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

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

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

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

114 

115 def make( 

116 *, 

117 in_frame: Any = pixel_frame, 

118 out_frame: Any = focal_plane, 

119 ast_mapping: astshim.Mapping | None = None, 

120 in_bounds_: Box | None = in_bounds, 

121 out_bounds_: Box | None = out_bounds, 

122 components: Any = (), 

123 ) -> Transform[Any, Any]: 

124 return Transform( 

125 in_frame, 

126 out_frame, 

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

128 in_bounds=in_bounds_, 

129 out_bounds=out_bounds_, 

130 components=components, 

131 ) 

132 

133 base = make() 

134 

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

136 assert base == base 

137 

138 # Two independently constructed but equivalent transforms are equal, 

139 # and equality is symmetric. 

140 assert base == make() 

141 assert make() == base 

142 

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

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

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

146 assert base != "not a transform" 

147 assert base != None # noqa: E711 

148 assert base != 42 

149 

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

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

152 assert base != make(in_bounds_=out_bounds) 

153 assert base != make(out_bounds_=in_bounds) 

154 assert base != make(in_frame=alt_frame) 

155 assert base != make(out_frame=alt_frame) 

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

157 

158 

159def test_sky_projection_equality() -> None: 

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

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

162 

163 # Check the two failure modes. 

164 with pytest.raises(ValueError, match="not a mapping from pixel coordinates"): 

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

166 

167 with pytest.raises(ValueError, match="not a mapping to ICRS"): 

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

169 

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

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

172 return Transform(pixel_frame, ICRS, mapping) 

173 

174 base = SkyProjection(make_pixel_to_sky()) 

175 

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

177 assert base == base 

178 

179 # Two independently constructed but equivalent projections are equal. 

180 assert base == SkyProjection(make_pixel_to_sky()) 

181 

182 # Comparison against a non-SkyProjection yields NotImplemented. 

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

184 assert base != "not a projection" 

185 assert base != None # noqa: E711 

186 

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

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

189 

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

191 with_approx = SkyProjection( 

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

193 ) 

194 assert base != with_approx 

195 

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

197 with_approx_again = SkyProjection( 

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

199 ) 

200 assert with_approx == with_approx_again 

201 

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

203 other_approx = SkyProjection( 

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

205 ) 

206 assert with_approx != other_approx 

207 

208 

209def test_affine_2x2() -> None: 

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

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

212 out_frame = GeneralFrame(unit=u.pix) 

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

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

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

216 out_matrix = np.dot(transform_matrix, in_matrix) 

217 check_transform( 

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

219 in_xy, 

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

221 in_frame, 

222 out_frame, 

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

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

225 ) 

226 

227 

228def test_affine_3x3() -> None: 

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

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

231 out_frame = GeneralFrame(unit=u.pix) 

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

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

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

235 out_matrix = np.dot(transform_matrix, in_matrix) 

236 check_transform( 

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

238 in_xy, 

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

240 in_frame, 

241 out_frame, 

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

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

244 ) 

245 

246 

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

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

249 legacy afw implementations. 

250 """ 

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

252 from lsst.geom import Point2D 

253 

254 legacy_detector = legacy_camera[16] 

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

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

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

258 pixel_xy_array = legacy_points_to_xy_array(pixel_legacy_points) 

259 fp_xy_array = legacy_points_to_xy_array(fp_legacy_points) 

260 fa_xy_array = legacy_points_to_xy_array(fa_legacy_points) 

261 # Test transforms extracted directly from the frame set. 

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

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

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

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

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

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

268 # Test a composition. 

269 pixel_to_fa_indirect = pixel_to_fp.then(fp_to_fa) 

270 check_transform( 

271 pixel_to_fa_indirect, 

272 pixel_xy_array, 

273 fa_xy_array, 

274 frame_set.detector(16), 

275 frame_set.field_angle(), 

276 ) 

277 pixel_to_fp_d, fp_to_fa_d = pixel_to_fa_indirect.decompose() 

278 check_transform( 

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

280 ) 

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

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

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

284 check_transform( 

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

286 ) 

287 

288 

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

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

291 serialization round-trips. 

292 

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

294 machinery. 

295 """ 

296 legacy_camera = legacy_camera 

297 frame_set = CameraFrameSet.from_legacy(legacy_camera) 

298 detector_id: int = DP2_VISIT_DETECTOR_DATA_ID["detector"] 

299 compare_to_legacy_camera(legacy_camera, frame_set) 

300 test_holder = FrameSetTestHolder( 

301 frames=frame_set, 

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

303 ) 

304 with RoundtripFits(test_holder) as roundtrip1: 

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

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

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

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

309 # stored a reference to the frame set: 

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

311 compare_to_legacy_camera(legacy_camera, roundtrip1.result.frames) 

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

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

314 assert ( 

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

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

317 ) 

318 with RoundtripJson(test_holder) as roundtrip2: 

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

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

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

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

323 # stored a reference to the frame set: 

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

325 raw_data = roundtrip2.inspect() 

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

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

328 compare_to_legacy_camera(legacy_camera, roundtrip2.result.frames) 

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

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

331 assert ( 

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

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

334 ) 

335 

336 

337def test_fits_wcs_projection_to_legacy() -> None: 

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

339 to a legacy SkyWcs. 

340 

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

342 requires PIXELS, so to_legacy has to rename it. 

343 """ 

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

345 rng = np.random.default_rng(43) 

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

347 pixel_frame = GeneralFrame(unit=u.pix) 

348 sky_projection = make_random_sky_projection(rng, pixel_frame, bbox) 

349 legacy_wcs = sky_projection.to_legacy() 

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

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

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

353 frame_set = sky_projection.pixel_to_sky_transform._ast_mapping 

354 assert isinstance(frame_set, astshim.FrameSet) 

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

356 assert "PIXEL" in domains 

357 assert "PIXELS" not in domains 

358 compare_sky_projection_to_legacy_wcs( 

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

360 ) 

361 

362 

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

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

365 legacy_wcs = legacy_detector_wcs["legacy_wcs"] 

366 wcs_bbox = legacy_detector_wcs["wcs_bbox"] 

367 subimage_bbox = legacy_detector_wcs["subimage_bbox"] 

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

369 sky_projection = SkyProjection.from_legacy(legacy_wcs, detector_frame) 

370 assert sky_projection.fits_approximation is not None 

371 compare_sky_projection_to_legacy_wcs(sky_projection, legacy_wcs, detector_frame, subimage_bbox) 

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

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

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

375 compare_sky_projection_to_legacy_wcs( 

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

377 ) 

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

379 compare_sky_projection_to_legacy_wcs( 

380 sky_projection.fits_approximation, 

381 sky_projection.fits_approximation.to_legacy(), 

382 detector_frame, 

383 subimage_bbox, 

384 is_fits=True, 

385 ) 

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

387 pass 

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

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

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

391 compare_sky_projection_to_legacy_wcs( 

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

393 ) 

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

395 pass 

396 compare_sky_projection_to_legacy_wcs( 

397 roundtrip.result, 

398 legacy_wcs.getFitsApproximation(), 

399 detector_frame, 

400 subimage_bbox, 

401 is_fits=True, 

402 ) 

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

404 compare_sky_projection_to_legacy_wcs( 

405 sky_projection.fits_approximation, 

406 roundtrip.result.to_legacy(), 

407 detector_frame, 

408 subimage_bbox, 

409 is_fits=True, 

410 ) 

411 

412 

413@dataclasses.dataclass 

414class FrameSetTestHolder: 

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

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

417 """ 

418 

419 frames: CameraFrameSet 

420 pixels_to_fp: Transform[DetectorFrame, FocalPlaneFrame] 

421 

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

423 frames_model = archive.serialize_frame_set( 

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

425 ) 

426 pixels_to_fp_model = archive.serialize_direct( 

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

428 ) 

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

430 

431 @staticmethod 

432 def _get_archive_tree_type[P: pydantic.BaseModel]( 

433 pointer_type: type[P], 

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

435 return FrameSetTestHolderModel[pointer_type] # type: ignore 

436 

437 

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

439 """The serialization model for FrameSetTestHolder.""" 

440 

441 SCHEMA_NAME: ClassVar[str] = "_test_frame_set_holder" 

442 SCHEMA_VERSION: ClassVar[str] = "1.0.0" 

443 MIN_READ_VERSION: ClassVar[int] = 1 

444 PUBLIC_TYPE: ClassVar[type] = FrameSetTestHolder 

445 

446 frames: CameraFrameSetSerializationModel | P 

447 pixels_to_fp: TransformSerializationModel[P] 

448 

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

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

451 frames = archive.deserialize_pointer( 

452 self.frames, CameraFrameSetSerializationModel, CameraFrameSetSerializationModel.deserialize 

453 ) 

454 pixels_to_fp = self.pixels_to_fp.deserialize(archive) 

455 return FrameSetTestHolder(frames, pixels_to_fp) 

456 

457 

458@dataclasses.dataclass 

459class _BroadcastTestData: 

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

461 SkyProjection. 

462 """ 

463 

464 in_frame: DetectorFrame 

465 out_frame: GeneralFrame 

466 matrix: np.ndarray 

467 scalar_x: float 

468 scalar_y: float 

469 xv: list[int] 

470 yv: list[int] 

471 sky_proj: SkyProjection[DetectorFrame] 

472 

473 

474@pytest.fixture 

475def broadcast_test_data() -> _BroadcastTestData: 

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

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

478 return _BroadcastTestData( 

479 in_frame=in_frame, 

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

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

482 scalar_x=3.0, 

483 scalar_y=7.0, 

484 xv=[1, 2, 3], 

485 yv=[4, 5, 6], 

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

487 ) 

488 

489 

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

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

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

493 """ 

494 t = broadcast_test_data.sky_proj.pixel_to_sky_transform 

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

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

497 assert type(result_fwd.x) is float 

498 assert type(result_fwd.y) is float 

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

500 ref_fwd = t.apply_forward( 

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

502 ) 

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

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

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

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

507 assert type(result_inv.x) is float 

508 assert type(result_inv.y) is float 

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

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

511 # apply_forward_q / apply_inverse_q with scalar Quantity inputs. 

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

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

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

515 assert result_fwd_q.x.shape == () 

516 assert result_fwd_q.y.shape == () 

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

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

519 assert result_inv_q.x.shape == () 

520 np.testing.assert_allclose( 

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

522 ) 

523 

524 

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

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

527 arrays, returning float64 ndarray results consistent with float64 

528 array input. 

529 """ 

530 t = broadcast_test_data.sky_proj.pixel_to_sky_transform 

531 # Python list input should return an ndarray. 

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

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

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

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

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

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

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

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

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

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

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

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

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

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

546 

547 

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

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

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

551 """ 

552 t = broadcast_test_data.sky_proj.pixel_to_sky_transform 

553 xv = np.array(broadcast_test_data.xv) 

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

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

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

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

558 assert result_1d.x.shape == xv.shape 

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

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

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

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

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

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

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

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

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

568 # Values must match the fully expanded meshgrid call. 

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

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

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

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

573 # apply_inverse also broadcasts. 

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

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

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

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

578 

579 

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

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

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

583 """ 

584 p = broadcast_test_data 

585 xv = np.array(p.xv) 

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

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

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

589 assert sc_1d.shape == xv.shape 

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

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

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

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

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

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

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

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

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

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

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

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

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

603 pix_2d = p.sky_proj.sky_to_pixel(sc_2d) 

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

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

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

607 # SkyProjectionAstropyView.pixel_to_world_values also broadcasts. 

608 view = p.sky_proj.as_astropy() 

609 world_2d = view.pixel_to_world_values(x2d, y2d) 

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

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

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

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

614 # SkyProjectionAstropyView.world_to_pixel_values also broadcasts. 

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

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

617 pix_world = view.world_to_pixel_values(ra_2d, dec_2d) 

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

619 

620 

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

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

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

624 equivalent x=/y= keyword calls. 

625 """ 

626 p = broadcast_test_data 

627 t = p.sky_proj.pixel_to_sky_transform 

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

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

630 

631 # --- apply_forward: scalar --- 

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

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

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

635 

636 # --- apply_forward: array --- 

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

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

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

640 

641 # --- apply_inverse: scalar --- 

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

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

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

645 

646 # --- apply_forward_q: scalar --- 

647 x_q = sx * t.in_frame.unit 

648 y_q = sy * t.in_frame.unit 

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

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

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

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

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

654 

655 # --- apply_inverse_q: scalar --- 

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

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

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

659 

660 # --- TypeError on bad combinations --- 

661 with pytest.raises(TypeError): 

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

663 with pytest.raises(TypeError):