Coverage for tests/test_cell_aperture_corrections.py: 95%

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 

14from typing import Any 

15 

16import numpy as np 

17import pytest 

18 

19from lsst.images import YX, BoundsError, Box 

20from lsst.images.cells import ( 

21 CellApertureCorrectionMapSerializationModel, 

22 CellField, 

23 CellGrid, 

24 CellGridBounds, 

25 CellIJ, 

26) 

27from lsst.images.json import JsonInputArchive, JsonOutputArchive 

28 

29try: 

30 from lsst.cell_coadds import StitchedApertureCorrection as LegacyStitchedApertureCorrection 

31 from lsst.cell_coadds import UniformGrid as LegacyUniformGrid 

32 from lsst.skymap import Index2D as LegacyIndex2D 

33 

34 HAVE_LEGACY = True 

35except ImportError: 

36 HAVE_LEGACY = False 

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

38 type LegacyUniformGrid = Any # type: ignore[no-redef] 

39 type LegacyIndex2D = Any # type: ignore[no-redef] 

40 

41skip_no_legacy = pytest.mark.skipif(not HAVE_LEGACY, reason="lsst.cell_coadds could not be imported.") 

42 

43# A 3x3 grid with 10x10-pixel cells. 

44GRID = CellGrid(bbox=Box.from_shape((30, 30)), cell_shape=YX(10, 10)) 

45BASE_BOUNDS = CellGridBounds(grid=GRID, bbox=Box.factory[0:30, 0:30]) 

46CELL_X = CellIJ(i=0, j=0) 

47CELL_Y = CellIJ(i=1, j=1) 

48 

49 

50def _make_field(missing: frozenset[CellIJ]) -> CellField: 

51 """Build a CellField over the shared 3x3 grid with ``missing`` cells set 

52 to NaN and all other cells set to 0.5. 

53 """ 

54 bounds = CellGridBounds(grid=GRID, bbox=Box.factory[0:30, 0:30], missing=missing) 

55 array = np.full((3, 3), 0.5) 

56 for cell in missing: 

57 index = cell - bounds.subgrid_start 

58 array[index.i, index.j] = np.nan 

59 return CellField(bounds, array) 

60 

61 

62def _make_legacy(gc: dict[LegacyIndex2D, float]) -> LegacyStitchedApertureCorrection: 

63 """Construct a legacy `StitchedApertureCorrection` on the shared 3x3 grid 

64 from the given per-cell values. 

65 """ 

66 ugrid = LegacyUniformGrid(cell_size=YX(10, 10).to_legacy_int_extent(), shape=LegacyIndex2D(x=3, y=3)) 

67 return LegacyStitchedApertureCorrection(ugrid, gc) 

68 

69 

70def _roundtrip( 

71 aperture_corrections: dict[CellIJ, CellField], 

72) -> tuple[CellApertureCorrectionMapSerializationModel, dict[str, CellField]]: 

73 """Serialize a CellField map, persist it to JSON, and read it back. 

74 

75 Returns the serialized model and the deserialized map, so tests can assert 

76 both the shared serialized bounds and the per-field result. 

77 

78 We can't use RoundtripJson for this because the in-memory type is just 

79 `dict` and hence doesn't have a ``serialize`` method. 

80 """ 

81 output_archive = JsonOutputArchive() 

82 model = CellApertureCorrectionMapSerializationModel.serialize(aperture_corrections, output_archive) 

83 output_tree = output_archive.finish(model) 

84 input_tree = CellApertureCorrectionMapSerializationModel.model_validate_json( 

85 output_tree.model_dump_json() 

86 ) 

87 input_archive = JsonInputArchive() 

88 return input_tree, input_tree.deserialize(input_archive) 

89 

90 

91@skip_no_legacy 

92def test_from_legacy_records_missing_cells() -> None: 

93 """Test that cells absent from the legacy map are recorded as missing.""" 

94 missing = {CELL_X, CELL_Y} 

95 present = [c for c in BASE_BOUNDS.cell_indices() if c not in missing] 

96 gc = {c.to_legacy(): 0.5 + 0.1 * c.i + 0.01 * c.j for c in present} 

97 legacy = _make_legacy(gc) 

98 field = CellField.from_legacy_aperture_correction(legacy, BASE_BOUNDS) 

99 assert field.bounds.missing == missing 

100 for cell in present: 

101 np.testing.assert_allclose(field.value_in_cell(cell), gc[cell.to_legacy()]) 

102 for cell in missing: 

103 with pytest.raises(BoundsError): 

104 field.value_in_cell(cell) 

105 

106 

107@skip_no_legacy 

108def test_from_legacy_all_cells_present() -> None: 

109 """Test that when every legacy cell has an entry, bounds are unchanged.""" 

110 gc = {c.to_legacy(): 0.5 for c in BASE_BOUNDS.cell_indices()} 

111 legacy = _make_legacy(gc) 

112 field = CellField.from_legacy_aperture_correction(legacy, BASE_BOUNDS) 

113 assert field.bounds.missing == frozenset() 

114 

115 

116def test_value_in_cell_rejects_index_below_bounds() -> None: 

117 """An index below the bounds must not wrap around the NumPy array.""" 

118 field = _make_field(frozenset()) 

119 with pytest.raises(BoundsError, match="out of bounds"): 

120 field.value_in_cell(CellIJ(i=-1, j=0)) 

121 

122 

123def test_serialize_deserialize_differing_missing_cells() -> None: 

124 """Test that a map whose fields differ in missing cells round-trips, 

125 restoring each field's per-cell missing set. 

126 """ 

127 ap = {"A": _make_field(frozenset({CELL_X})), "B": _make_field(frozenset({CELL_X, CELL_Y}))} 

128 model, result = _roundtrip(ap) 

129 # Shared serialized bounds keep only cells missing from every field. 

130 assert model.bounds.missing == frozenset({CELL_X}) 

131 assert result["A"].bounds.missing == frozenset({CELL_X}) 

132 assert result["B"].bounds.missing == frozenset({CELL_X, CELL_Y}) 

133 for name, field in ap.items(): 

134 for cell in BASE_BOUNDS.cell_indices(): 

135 if cell not in field.bounds.missing: 

136 np.testing.assert_allclose(result[name].value_in_cell(cell), field.value_in_cell(cell)) 

137 else: 

138 with pytest.raises(BoundsError): 

139 result[name].value_in_cell(cell) 

140 

141 

142def test_serialize_raises_on_inconsistent_grid() -> None: 

143 """Test that fields on different grids cannot be serialized together.""" 

144 other_grid = CellGrid(bbox=Box.from_shape((40, 40)), cell_shape=YX(10, 10)) 

145 other_bounds = CellGridBounds(grid=other_grid, bbox=Box.factory[0:30, 0:30]) 

146 other_field = CellField(other_bounds, np.full((3, 3), 0.25)) 

147 with pytest.raises(ValueError, match="do not have a consistent grid"): 

148 CellApertureCorrectionMapSerializationModel.serialize( 

149 {"A": _make_field(frozenset()), "B": other_field}, JsonOutputArchive() 

150 ) 

151 

152 

153def test_serialize_raises_on_inconsistent_bbox() -> None: 

154 """Test that fields covering different parts of one grid are rejected.""" 

155 other_bounds = CellGridBounds(grid=GRID, bbox=Box.factory[0:20, 10:30]) 

156 other_field = CellField(other_bounds, np.full((2, 2), 0.25)) 

157 with pytest.raises(ValueError, match="do not have a consistent bounding box"): 

158 CellApertureCorrectionMapSerializationModel.serialize( 

159 {"A": _make_field(frozenset()), "B": other_field}, JsonOutputArchive() 

160 ) 

161 

162 

163def test_deserialize_nan_column_becomes_missing() -> None: 

164 """Test that a NaN in one field column makes that cell missing for that 

165 field. 

166 """ 

167 ap = {"A": _make_field(frozenset()), "B": _make_field(frozenset({CELL_Y}))} 

168 _, result = _roundtrip(ap) 

169 assert result["A"].bounds.missing == frozenset() 

170 assert result["B"].bounds.missing == frozenset({CELL_Y}) 

171 np.testing.assert_allclose(result["A"].value_in_cell(CELL_Y), ap["A"].value_in_cell(CELL_Y)) 

172 with pytest.raises(BoundsError): 

173 result["B"].value_in_cell(CELL_Y)