Coverage for python/lsst/cell_coadds/_grid_container.py: 46%
70 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-13 11:23 +0000
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-13 11:23 +0000
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/>.
22from __future__ import annotations
24from collections.abc import Callable, Iterable, Iterator, MutableMapping
25from itertools import product
26from typing import TYPE_CHECKING, TypeVar
28from lsst.skymap import Index2D
30from ._uniform_grid import UniformGrid
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
36__all__ = ("GridContainer",)
38T = TypeVar("T")
41class GridContainer(MutableMapping[Index2D, T]):
42 """A container whose elements form a 2-d grid.
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 """
53 def __init__(self, shape: Index2D, offset: Index2D | None = None) -> None:
54 super().__init__()
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
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 )
74 def __repr__(self) -> str:
75 return f"{type(self).__name__}(shape={self.shape}, offset={self.offset})"
77 # Implement the necessary methods for MutableMapping.
78 def __len__(self) -> int:
79 return len(self._cells)
81 def __getitem__(self, index: Index2D) -> T:
82 return self._cells[index]
84 def __setitem__(self, index: Index2D, value: T) -> None:
85 self._check(index)
86 self._cells[index] = value
88 def __delitem__(self, index: Index2D) -> None:
89 self._check(index)
90 del self._cells[index]
92 def __iter__(self) -> Iterator[T]:
93 return iter(self._cells)
95 # def keys(self) -> Iterable[Index2D]: # populated_indices
96 # """Return an iterator over the indices in the container with values.
98 # See also
99 # --------
100 # indices
101 # """
102 # yield from self._cells.keys()
104 def indices(self) -> Iterable[Index2D]:
105 """Return an iterator over all possible indices for the container.
107 Unlike `keys`, this method returns an iterator over all valid indices,
108 whether the corresponding value is set or not.
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)
121 @property
122 def shape(self) -> Index2D:
123 """Number of cells in the container in each dimension."""
124 return self._shape
126 @property
127 def offset(self) -> Index2D:
128 """Index of the first cell in the container."""
129 return self._offset
131 @property
132 def size(self) -> int:
133 """The number of cells expected in the container.
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
140 @property
141 def first(self) -> T:
142 """The cell at the lower left corner of the container."""
143 return self._cells[self.offset]
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)]
150 def subset_overlapping(self, grid: UniformGrid, bbox: geom.Box2I) -> GridContainer:
151 """Return a new GridContainer with cells that overlap a bounding box.
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.
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
177 def rebuild_transformed(self, transform: Callable[[T], T]) -> GridContainer:
178 """Return a GridContainer with the same shape and offset.
180 The cell values are created by applying a callback function to each
181 cell value in this object.
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