Coverage for tests / test_transforms.py: 29%

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 frame = DetectorFrame(**DP2_VISIT_DETECTOR_DATA_ID, bbox=Box.factory[:5, :4]) 

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

52 identity = Transform.identity(frame) 

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

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

55 

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

57 def test_camera(self) -> None: 

58 """Test that we can: 

59 

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

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

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

63 do all of that; 

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

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

66 

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

68 

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

70 frame-set reference machinery. 

71 """ 

72 try: 

73 from lsst.afw.cameraGeom import Camera 

74 except ImportError: 

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

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

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

78 legacy_camera = Camera.readFits(filename) 

79 frame_set = CameraFrameSet.from_legacy(legacy_camera) 

80 detector_id: int = DP2_VISIT_DETECTOR_DATA_ID["detector"] 

81 self.compare_to_legacy_camera(legacy_camera, frame_set) 

82 test_holder = FrameSetTestHolder( 

83 frames=frame_set, 

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

85 ) 

86 with RoundtripFits(self, test_holder) as roundtrip: 

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

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

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

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

91 # stored a reference to the frame set: 

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

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

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

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

96 self.assertEqual( 

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

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

99 ) 

100 

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

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

103 legacy lsst.afw.cameraGeom implementations. 

104 """ 

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

106 from lsst.geom import Point2D 

107 

108 legacy_detector = legacy_camera[16] 

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

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

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

112 pixel_xy_array = legacy_points_to_xy_array(pixel_legacy_points) 

113 fp_xy_array = legacy_points_to_xy_array(fp_legacy_points) 

114 fa_xy_array = legacy_points_to_xy_array(fa_legacy_points) 

115 # Test transforms extracted directly from the frame set. 

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

117 check_transform( 

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

119 ) 

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

121 check_transform( 

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

123 ) 

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

125 check_transform( 

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

127 ) 

128 # Test a composition. 

129 pixel_to_fa_indirect = pixel_to_fp.then(fp_to_fa) 

130 check_transform( 

131 self, 

132 pixel_to_fa_indirect, 

133 pixel_xy_array, 

134 fa_xy_array, 

135 frame_set.detector(16), 

136 frame_set.field_angle(), 

137 ) 

138 pixel_to_fp_d, fp_to_fa_d = pixel_to_fa_indirect.decompose() 

139 check_transform( 

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

141 ) 

142 check_transform( 

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

144 ) 

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

146 check_transform( 

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

148 ) 

149 check_transform( 

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

151 ) 

152 

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

154 def test_detector_wcs(self) -> None: 

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

156 try: 

157 from lsst.afw.image import ExposureFitsReader 

158 except ImportError: 

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

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

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

162 reader = ExposureFitsReader(filename) 

163 legacy_wcs = reader.readWcs() 

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

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

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

167 projection = Projection.from_legacy(legacy_wcs, detector_frame) 

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

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

170 pass 

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

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

173 pass 

174 compare_projection_to_legacy_wcs( 

175 self, 

176 roundtrip.result, 

177 legacy_wcs.getFitsApproximation(), 

178 detector_frame, 

179 subimage_bbox, 

180 is_fits=True, 

181 ) 

182 

183 

184@dataclasses.dataclass 

185class FrameSetTestHolder: 

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

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

188 """ 

189 

190 frames: CameraFrameSet 

191 pixels_to_fp: Transform[DetectorFrame, FocalPlaneFrame] 

192 

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

194 frames_model = archive.serialize_frame_set( 

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

196 ) 

197 pixels_to_fp_model = archive.serialize_direct( 

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

199 ) 

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

201 

202 @staticmethod 

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

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

205 frames = archive.deserialize_pointer( 

206 model.frames, CameraFrameSetSerializationModel, CameraFrameSet.deserialize 

207 ) 

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

209 return FrameSetTestHolder(frames, pixels_to_fp) 

210 

211 @staticmethod 

212 def _get_archive_tree_type[P: pydantic.BaseModel]( 

213 pointer_type: type[P], 

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

215 return FrameSetTestHolderModel[pointer_type] # type: ignore 

216 

217 

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

219 """The serialization model for FrameSetTestHolder.""" 

220 

221 frames: CameraFrameSetSerializationModel | P 

222 pixels_to_fp: TransformSerializationModel[P] 

223 

224 

225if __name__ == "__main__": 

226 unittest.main()