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_serialize_deserialize_differing_missing_cells() -> None: 

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

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

119 """ 

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

121 model, result = _roundtrip(ap) 

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

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

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

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

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

127 for cell in BASE_BOUNDS.cell_indices(): 

128 if cell not in field.bounds.missing: 

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

130 else: 

131 with pytest.raises(BoundsError): 

132 result[name].value_in_cell(cell) 

133 

134 

135def test_serialize_raises_on_inconsistent_grid() -> None: 

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

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

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

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

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

141 CellApertureCorrectionMapSerializationModel.serialize( 

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

143 ) 

144 

145 

146def test_deserialize_nan_column_becomes_missing() -> None: 

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

148 field. 

149 """ 

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

151 _, result = _roundtrip(ap) 

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

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

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

155 with pytest.raises(BoundsError): 

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