Coverage for python/lsst/cell_coadds/_grid_container.py: 46%

70 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-23 02:17 -0700

1# This file is part of cell_coadds. 

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# This program is free software: you can redistribute it and/or modify 

10# it under the terms of the GNU General Public License as published by 

11# the Free Software Foundation, either version 3 of the License, or 

12# (at your option) any later version. 

13# 

14# This program is distributed in the hope that it will be useful, 

15# but WITHOUT ANY WARRANTY; without even the implied warranty of 

16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

17# GNU General Public License for more details. 

18# 

19# You should have received a copy of the GNU General Public License 

20# along with this program. If not, see <https://www.gnu.org/licenses/>. 

21 

22from __future__ import annotations 

23 

24from collections.abc import Callable, Iterable, Iterator, MutableMapping 

25from itertools import product 

26from typing import TYPE_CHECKING, TypeVar 

27 

28from lsst.skymap import Index2D 

29 

30from ._uniform_grid import UniformGrid 

31 

32if TYPE_CHECKING: 32 ↛ 33line 32 didn't jump to line 33, because the condition on line 32 was never true

33 import lsst.geom as geom 

34 

35 

36__all__ = ("GridContainer",) 

37 

38T = TypeVar("T") 

39 

40 

41class GridContainer(MutableMapping[Index2D, T]): 

42 """A container whose elements form a 2-d grid. 

43 

44 Parameters 

45 ---------- 

46 shape : `~lsst.skymap.Index2D` 

47 The number of cells in the grid in each dimension. 

48 offset : `~lsst.skymap.Index2D` or None, optional 

49 The integer offset of the grid in each dimension. If `None`, the offset 

50 is Index2D(0, 0). 

51 """ 

52 

53 def __init__(self, shape: Index2D, offset: Index2D | None = None) -> None: 

54 super().__init__() 

55 

56 self._offset = offset if offset else Index2D(0, 0) 

57 self._shape = shape 

58 self._cells: dict[Index2D, T] = {} 

59 self._mapping = self._cells 

60 

61 def _check(self, index: Index2D) -> None: 

62 """Check if a given index belongs to the container or not.""" 

63 if index.x < self.offset.x or index.x >= self.offset.x + self.shape.x: 

64 raise ValueError( 

65 f"x index {index.x} out of range; expected a value between {self.offset.x} and " 

66 f"{self.offset.x + self.shape.x -1}." 

67 ) 

68 if index.y < self.offset.y or index.y >= self.offset.y + self.shape.y: 

69 raise ValueError( 

70 f"y index {index.y} out of range; expected a value between {self.offset.y} and " 

71 f"{self.offset.y + self.shape.y -1}." 

72 ) 

73 

74 def __repr__(self) -> str: 

75 return f"{type(self).__name__}(shape={self.shape}, offset={self.offset})" 

76 

77 # Implement the necessary methods for MutableMapping. 

78 def __len__(self) -> int: 

79 return len(self._cells) 

80 

81 def __getitem__(self, index: Index2D) -> T: 

82 return self._cells[index] 

83 

84 def __setitem__(self, index: Index2D, value: T) -> None: 

85 self._check(index) 

86 self._cells[index] = value 

87 

88 def __delitem__(self, index: Index2D) -> None: 

89 self._check(index) 

90 del self._cells[index] 

91 

92 def __iter__(self) -> Iterator[T]: 

93 return iter(self._cells) 

94 

95 # def keys(self) -> Iterable[Index2D]: # populated_indices 

96 # """Return an iterator over the indices in the container with values. 

97 

98 # See also 

99 # -------- 

100 # indices 

101 # """ 

102 # yield from self._cells.keys() 

103 

104 def indices(self) -> Iterable[Index2D]: 

105 """Return an iterator over all possible indices for the container. 

106 

107 Unlike `keys`, this method returns an iterator over all valid indices, 

108 whether the corresponding value is set or not. 

109 

110 See Also 

111 -------- 

112 keys 

113 """ 

114 iterator = product( 

115 range(self.offset.y, self.offset.y + self.shape.y), 

116 range(self.offset.x, self.offset.x + self.shape.x), 

117 ) 

118 for y, x in iterator: 

119 yield Index2D(x, y) 

120 

121 @property 

122 def shape(self) -> Index2D: 

123 """Number of cells in the container in each dimension.""" 

124 return self._shape 

125 

126 @property 

127 def offset(self) -> Index2D: 

128 """Index of the first cell in the container.""" 

129 return self._offset 

130 

131 @property 

132 def size(self) -> int: 

133 """The number of cells expected in the container. 

134 

135 This does not indicate the number of cells that have been filled. 

136 Use len() instead. 

137 """ 

138 return self._shape.x * self._shape.y 

139 

140 @property 

141 def first(self) -> T: 

142 """The cell at the lower left corner of the container.""" 

143 return self._cells[self.offset] 

144 

145 @property 

146 def last(self) -> T: 

147 """The cell at the upper right corner of the container.""" 

148 return self._cells[Index2D(x=self.offset.x + self.shape.x - 1, y=self.offset.y + self.shape.y - 1)] 

149 

150 def subset_overlapping(self, grid: UniformGrid, bbox: geom.Box2I) -> GridContainer: 

151 """Return a new GridContainer with cells that overlap a bounding box. 

152 

153 Parameters 

154 ---------- 

155 grid : `~lsst.cell_coadds.UniformGrid` 

156 Grid that maps the container's cells to the coordinates used to 

157 define the bounding box. May define a grid that is a super of the 

158 container's cells. 

159 bbox : `~lsst.geom.Box2I` 

160 Bounding box that returned cells must overlap. 

161 

162 Returns 

163 ------- 

164 grid_container : `GridContainer` 

165 GridContainer with just the cells that overlap the bounding box. 

166 """ 

167 bbox = bbox.clippedTo(grid.bbox) 

168 offset = grid.index(bbox.getBegin()) 

169 last = grid.index(bbox.getMax()) 

170 end = Index2D(last.x + 1, last.y + 1) 

171 shape = Index2D(end.x - offset.x, end.y - offset.y) 

172 gc = GridContainer[T](shape, offset) 

173 for index in gc.indices(): 

174 gc[index] = self[index] 

175 return gc 

176 

177 def rebuild_transformed(self, transform: Callable[[T], T]) -> GridContainer: 

178 """Return a GridContainer with the same shape and offset. 

179 

180 The cell values are created by applying a callback function to each 

181 cell value in this object. 

182 

183 Parameters 

184 ---------- 

185 transform : Callable[[T], T] 

186 A callable function that takes a cell value and returns a new 

187 """ 

188 gc = GridContainer[T](self.shape, self.offset) 

189 for key in self.keys(): # noqa: SIM118 

190 gc[key] = transform(self[key]) 

191 return gc