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

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 @property 

151 def arbitrary(self) -> T: 

152 """An arbitrary cell from the container. 

153 

154 It is typically the first cell, but in the case where it might not be 

155 available, it is a well-behaved cell. 

156 

157 Raises 

158 ------ 

159 RuntimeError: 

160 Raised if the grid container is empty. 

161 """ 

162 if not self._cells: 

163 raise RuntimeError("GridContainer is empty.") 

164 

165 cell = next(iter(self._cells.values())) 

166 return cell 

167 

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

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

170 

171 Parameters 

172 ---------- 

173 grid : `~lsst.cell_coadds.UniformGrid` 

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

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

176 container's cells. 

177 bbox : `~lsst.geom.Box2I` 

178 Bounding box that returned cells must overlap. 

179 

180 Returns 

181 ------- 

182 grid_container : `GridContainer` 

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

184 """ 

185 bbox = bbox.clippedTo(grid.bbox) 

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

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

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

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

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

191 for index in gc.indices(): 

192 gc[index] = self[index] 

193 return gc 

194 

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

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

197 

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

199 cell value in this object. 

200 

201 Parameters 

202 ---------- 

203 transform : Callable[[T], T] 

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

205 """ 

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

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

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

209 return gc