Coverage for tests / test_transforms.py: 28%

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 

17import unittest 

18from typing import Any 

19 

20import numpy as np 

21import pydantic 

22 

23from lsst.images import ( 

24 Box, 

25 CameraFrameSet, 

26 CameraFrameSetSerializationModel, 

27 DetectorFrame, 

28 FocalPlaneFrame, 

29 Projection, 

30 Transform, 

31 TransformSerializationModel, 

32) 

33from lsst.images.serialization import ArchiveTree, InputArchive, OutputArchive, TableCellReferenceModel 

34from lsst.images.tests import ( 

35 DP2_VISIT_DETECTOR_DATA_ID, 

36 RoundtripFits, 

37 check_transform, 

38 compare_projection_to_legacy_wcs, 

39 legacy_points_to_xy_array, 

40) 

41 

42DATA_DIR = os.environ.get("TESTDATA_IMAGES_DIR", None) 

43 

44 

45class TransformTestCase(unittest.TestCase): 

46 """Tests for the Transform, Projection, and FrameSet classes.""" 

47 

48 def test_identity(self) -> None: 

49 """Test an identity transform.""" 

50 try: 

51 import astshim # noqa: F401 

52 except ImportError: 

53 raise unittest.SkipTest("astshim could not be imported.") from None 

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

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

56 identity = Transform.identity(frame) 

57 check_transform(self, identity, xy, xy, frame, frame) 

58 self.assertEqual(identity.decompose(), []) 

59 

60 @unittest.skipUnless(DATA_DIR is not None, "TESTDATA_IMAGES_DIR is not in the environment.") 

61 def test_camera(self) -> None: 

62 """Test that we can: 

63 

64 - make a CameraFrameSet from the AST representation returned by afw; 

65 - transform points and get the same result as afw; 

66 - round-trip the CameraFrameSet through FITS serialization and still 

67 do all of that; 

68 - also roundtrip a Transform that can be obtained from the 

69 CameraFrameSet, by referencing the mappings in the frame set. 

70 

71 This test is skipped if legacy modules cannot be imported. 

72 

73 This test provides coverage for the archive system's pointer and 

74 frame-set reference machinery. 

75 """ 

76 try: 

77 from lsst.afw.cameraGeom import Camera 

78 except ImportError: 

79 raise unittest.SkipTest("'lsst.afw.cameraGeom' could not be imported.") from None 

80 assert DATA_DIR is not None, "Guaranteed by decorator." 

81 filename = os.path.join(DATA_DIR, "dp2", "legacy", "camera.fits") 

82 legacy_camera = Camera.readFits(filename) 

83 frame_set = CameraFrameSet.from_legacy(legacy_camera) 

84 detector_id: int = DP2_VISIT_DETECTOR_DATA_ID["detector"] 

85 self.compare_to_legacy_camera(legacy_camera, frame_set) 

86 test_holder = FrameSetTestHolder( 

87 frames=frame_set, 

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

89 ) 

90 with RoundtripFits(self, test_holder) as roundtrip: 

91 self.assertEqual(len(roundtrip.serialized.pixels_to_fp.frames), 2) 

92 self.assertEqual(len(roundtrip.serialized.pixels_to_fp.bounds), 2) 

93 self.assertEqual(len(roundtrip.serialized.pixels_to_fp.mappings), 1) 

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

95 # stored a reference to the frame set: 

96 self.assertIsInstance(roundtrip.serialized.pixels_to_fp.mappings[0], TableCellReferenceModel) 

97 self.compare_to_legacy_camera(legacy_camera, roundtrip.result.frames) 

98 self.assertEqual(roundtrip.result.pixels_to_fp.in_frame, frame_set.detector(detector_id)) 

99 self.assertEqual(roundtrip.result.pixels_to_fp.out_frame, frame_set.focal_plane()) 

100 self.assertEqual( 

101 roundtrip.result.pixels_to_fp._ast_mapping.simplified().show(), 

102 test_holder.pixels_to_fp._ast_mapping.simplified().show(), 

103 ) 

104 

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

106 """Test the transforms extracted from a CameraFrameSet against the 

107 legacy lsst.afw.cameraGeom implementations. 

108 """ 

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

110 from lsst.geom import Point2D 

111 

112 legacy_detector = legacy_camera[16] 

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

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

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

116 pixel_xy_array = legacy_points_to_xy_array(pixel_legacy_points) 

117 fp_xy_array = legacy_points_to_xy_array(fp_legacy_points) 

118 fa_xy_array = legacy_points_to_xy_array(fa_legacy_points) 

119 # Test transforms extracted directly from the frame set. 

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

121 check_transform( 

122 self, pixel_to_fp, pixel_xy_array, fp_xy_array, frame_set.detector(16), frame_set.focal_plane() 

123 ) 

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

125 check_transform( 

126 self, pixel_to_fa, pixel_xy_array, fa_xy_array, frame_set.detector(16), frame_set.field_angle() 

127 ) 

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

129 check_transform( 

130 self, fp_to_fa, fp_xy_array, fa_xy_array, frame_set.focal_plane(), frame_set.field_angle() 

131 ) 

132 # Test a composition. 

133 pixel_to_fa_indirect = pixel_to_fp.then(fp_to_fa) 

134 check_transform( 

135 self, 

136 pixel_to_fa_indirect, 

137 pixel_xy_array, 

138 fa_xy_array, 

139 frame_set.detector(16), 

140 frame_set.field_angle(), 

141 ) 

142 pixel_to_fp_d, fp_to_fa_d = pixel_to_fa_indirect.decompose() 

143 check_transform( 

144 self, pixel_to_fp_d, pixel_xy_array, fp_xy_array, frame_set.detector(16), frame_set.focal_plane() 

145 ) 

146 check_transform( 

147 self, fp_to_fa_d, fp_xy_array, fa_xy_array, frame_set.focal_plane(), frame_set.field_angle() 

148 ) 

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

150 check_transform( 

151 self, fa_to_fp_d, fa_xy_array, fp_xy_array, frame_set.field_angle(), frame_set.focal_plane() 

152 ) 

153 check_transform( 

154 self, fp_to_pixel_d, fp_xy_array, pixel_xy_array, frame_set.focal_plane(), frame_set.detector(16) 

155 ) 

156 

157 @unittest.skipUnless(DATA_DIR is not None, "TESTDATA_IMAGES_DIR is not in the environment.") 

158 def test_detector_wcs(self) -> None: 

159 """Test the Transform/Projection representation of a detector WCS.""" 

160 try: 

161 from lsst.afw.image import ExposureFitsReader 

162 except ImportError: 

163 raise unittest.SkipTest("'lsst.afw.image' could not be imported.") from None 

164 assert DATA_DIR is not None, "Guaranteed by decorator." 

165 filename = os.path.join(DATA_DIR, "dp2", "legacy", "visit_image.fits") 

166 reader = ExposureFitsReader(filename) 

167 legacy_wcs = reader.readWcs() 

168 wcs_bbox = Box.from_legacy(reader.readDetector().getBBox()) 

169 subimage_bbox = Box.from_legacy(reader.readBBox()) 

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

171 projection = Projection.from_legacy(legacy_wcs, detector_frame) 

172 compare_projection_to_legacy_wcs(self, projection, legacy_wcs, detector_frame, subimage_bbox) 

173 with RoundtripFits(self, projection, "Projection") as roundtrip: 

174 pass 

175 compare_projection_to_legacy_wcs(self, roundtrip.result, legacy_wcs, detector_frame, subimage_bbox) 

176 with RoundtripFits(self, projection.fits_approximation, "Projection") as roundtrip: 

177 pass 

178 compare_projection_to_legacy_wcs( 

179 self, 

180 roundtrip.result, 

181 legacy_wcs.getFitsApproximation(), 

182 detector_frame, 

183 subimage_bbox, 

184 is_fits=True, 

185 ) 

186 

187 

188@dataclasses.dataclass 

189class FrameSetTestHolder: 

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

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

192 """ 

193 

194 frames: CameraFrameSet 

195 pixels_to_fp: Transform[DetectorFrame, FocalPlaneFrame] 

196 

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

198 frames_model = archive.serialize_frame_set( 

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

200 ) 

201 pixels_to_fp_model = archive.serialize_direct( 

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

203 ) 

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

205 

206 @staticmethod 

207 def deserialize(model: FrameSetTestHolderModel[Any], archive: InputArchive[Any]) -> FrameSetTestHolder: 

208 assert not isinstance(model.frames, CameraFrameSetSerializationModel), "Archive pointer expected." 

209 frames = archive.deserialize_pointer( 

210 model.frames, CameraFrameSetSerializationModel, CameraFrameSet.deserialize 

211 ) 

212 pixels_to_fp = Transform.deserialize(model.pixels_to_fp, archive) 

213 return FrameSetTestHolder(frames, pixels_to_fp) 

214 

215 @staticmethod 

216 def _get_archive_tree_type[P: pydantic.BaseModel]( 

217 pointer_type: type[P], 

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

219 return FrameSetTestHolderModel[pointer_type] # type: ignore 

220 

221 

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

223 """The serialization model for FrameSetTestHolder.""" 

224 

225 frames: CameraFrameSetSerializationModel | P 

226 pixels_to_fp: TransformSerializationModel[P] 

227 

228 

229if __name__ == "__main__": 

230 unittest.main()