Coverage for python/lsst/scarlet/lite/observation.py: 25%
130 statements
« prev ^ index » next coverage.py v7.4.1, created at 2024-02-07 11:26 +0000
« prev ^ index » next coverage.py v7.4.1, created at 2024-02-07 11:26 +0000
1# This file is part of scarlet_lite.
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
24__all__ = ["Observation", "convolve"]
26from typing import Any, cast
28import numpy as np
29import numpy.typing as npt
31from .bbox import Box
32from .fft import Fourier, _pad, centered
33from .fft import convolve as fft_convolve
34from .fft import match_kernel
35from .image import Image
38def get_filter_coords(filter_values: np.ndarray, center: tuple[int, int] | None = None) -> np.ndarray:
39 """Create filter coordinate grid needed for the apply filter function
41 Parameters
42 ----------
43 filter_values:
44 The 2D array of the filter to apply.
45 center:
46 The center (y,x) of the filter. If `center` is `None` then
47 `filter_values` must have an odd number of rows and columns
48 and the center will be set to the center of `filter_values`.
50 Returns
51 -------
52 coords:
53 The coordinates of the pixels in `filter_values`,
54 where the coordinates of the `center` pixel are `(0,0)`.
55 """
56 if filter_values.ndim != 2:
57 raise ValueError("`filter_values` must be 2D")
58 if center is None:
59 if filter_values.shape[0] % 2 == 0 or filter_values.shape[1] % 2 == 0:
60 msg = """Ambiguous center of the `filter_values` array,
61 you must use a `filter_values` array
62 with an odd number of rows and columns or
63 calculate `coords` on your own."""
64 raise ValueError(msg)
65 center = tuple([filter_values.shape[0] // 2, filter_values.shape[1] // 2]) # type: ignore
66 center = cast(tuple[int, int], center)
67 x = np.arange(filter_values.shape[1])
68 y = np.arange(filter_values.shape[0])
69 x, y = np.meshgrid(x, y)
70 x -= center[1]
71 y -= center[0]
72 coords = np.dstack([y, x])
73 return coords
76def get_filter_bounds(coords: np.ndarray) -> tuple[int, int, int, int]:
77 """Get the slices in x and y to apply a filter
79 Parameters
80 ----------
81 coords:
82 The coordinates of the filter,
83 defined by `get_filter_coords`.
85 Returns
86 -------
87 y_start, y_end, x_start, x_end:
88 The start and end of each slice that is passed to `apply_filter`.
89 """
90 z = np.zeros((len(coords),), dtype=int)
91 # Set the y slices
92 y_start = np.max([z, coords[:, 0]], axis=0)
93 y_end = -np.min([z, coords[:, 0]], axis=0)
94 # Set the x slices
95 x_start = np.max([z, coords[:, 1]], axis=0)
96 x_end = -np.min([z, coords[:, 1]], axis=0)
97 return y_start, y_end, x_start, x_end
100def convolve(image: np.ndarray, psf: np.ndarray, bounds: tuple[int, int, int, int]):
101 """Convolve an image with a PSF in real space
103 Parameters
104 ----------
105 image:
106 The multi-band image to convolve.
107 psf:
108 The psf to convolve the image with.
109 bounds:
110 The filter bounds required by the ``apply_filter`` C++ method,
111 usually obtained by calling `get_filter_bounds`.
112 """
113 from lsst.scarlet.lite.operators_pybind11 import apply_filter # type: ignore
115 result = np.empty(image.shape, dtype=image.dtype)
116 for band in range(len(image)):
117 img = image[band]
119 apply_filter(
120 img,
121 psf[band].reshape(-1),
122 bounds[0],
123 bounds[1],
124 bounds[2],
125 bounds[3],
126 result[band],
127 )
128 return result
131def _set_image_like(images: np.ndarray | Image, bands: tuple | None = None, bbox: Box | None = None) -> Image:
132 """Ensure that an image-like array is cast appropriately as an image
134 Parameters
135 ----------
136 images:
137 The multiband image-like array to cast as an Image.
138 If it already has `bands` and `bbox` properties then it is returned
139 with no modifications.
140 bands:
141 The bands for the multiband-image.
142 If `images` is a numpy array, this parameter is mandatory.
143 If `images` is an `Image` and `bands` is not `None`,
144 then `bands` is ignored.
145 bbox:
146 Bounding box containing the image.
147 If `images` is a numpy array, this parameter is mandatory.
148 If `images` is an `Image` and `bbox` is not `None`,
149 then `bbox` is ignored.
151 Returns
152 -------
153 images: Image
154 The input images converted into an image.
155 """
156 if isinstance(images, Image):
157 # This is already an image
158 if bbox is not None and images.bbox != bbox:
159 raise ValueError(f"Bounding boxes {images.bbox} and {bbox} do not agree")
160 return images
162 if bbox is None:
163 bbox = Box(images.shape[-2:])
164 return Image(images, bands=bands, yx0=cast(tuple[int, int], bbox.origin))
167class Observation:
168 """A single observation
170 This class contains all of the observed images and derived
171 properties, like PSFs, variance map, and weight maps,
172 required for most optimizers.
173 This includes methods to match a scarlet model PSF to the oberved PSF
174 in each band.
176 Notes
177 -----
178 This is effectively a combination of the `Observation` and
179 `Renderer` class from scarlet main, greatly simplified due
180 to the assumptions that the observations are all resampled
181 onto the same pixel grid and that the `images` contain all
182 of the information for all of the model bands.
184 Parameters
185 ----------
186 images:
187 (bands, y, x) array of observed images.
188 variance:
189 (bands, y, x) array of variance for each image pixel.
190 weights:
191 (bands, y, x) array of weights to use when calculate the
192 likelihood of each pixel.
193 psfs:
194 (bands, y, x) array of the PSF image in each band.
195 model_psf:
196 (bands, y, x) array of the model PSF image in each band.
197 If `model_psf` is `None` then convolution is performed,
198 which should only be done when the observation is a
199 PSF matched coadd, and the scarlet model has the same PSF.
200 noise_rms:
201 Per-band average noise RMS. If `noise_rms` is `None` then the mean
202 of the sqrt of the variance is used.
203 bbox:
204 The bounding box containing the model. If `bbox` is `None` then
205 a `Box` is created that is the shape of `images` with an origin
206 at `(0, 0)`.
207 padding:
208 Padding to use when performing an FFT convolution.
209 convolution_mode:
210 The method of convolution. This should be either "fft" or "real".
211 """
213 def __init__(
214 self,
215 images: np.ndarray | Image,
216 variance: np.ndarray | Image,
217 weights: np.ndarray | Image,
218 psfs: np.ndarray,
219 model_psf: np.ndarray | None = None,
220 noise_rms: np.ndarray | None = None,
221 bbox: Box | None = None,
222 bands: tuple | None = None,
223 padding: int = 3,
224 convolution_mode: str = "fft",
225 ):
226 # Convert the images to a multi-band `Image` and use the resulting
227 # bbox and bands.
228 images = _set_image_like(images, bands, bbox)
229 bands = images.bands
230 bbox = images.bbox
231 self.images = images
232 self.variance = _set_image_like(variance, bands, bbox)
233 self.weights = _set_image_like(weights, bands, bbox)
234 # make sure that the images and psfs have the same dtype
235 if psfs.dtype != images.dtype:
236 psfs = psfs.astype(images.dtype)
237 self.psfs = psfs
239 if convolution_mode not in [
240 "fft",
241 "real",
242 ]:
243 raise ValueError("convolution_mode must be either 'fft' or 'real'")
244 self.mode = convolution_mode
245 if noise_rms is None:
246 noise_rms = np.array(np.mean(np.sqrt(variance.data), axis=(1, 2)))
247 self.noise_rms = noise_rms
249 # Create a difference kernel to convolve the model to the PSF
250 # in each band
251 self.model_psf = model_psf
252 self.padding = padding
253 if model_psf is not None:
254 if model_psf.dtype != images.dtype:
255 self.model_psf = model_psf.astype(images.dtype)
256 self.diff_kernel: Fourier | None = cast(Fourier, match_kernel(psfs, model_psf, padding=padding))
257 # The gradient of a convolution is another convolution,
258 # but with the flipped and transposed kernel.
259 diff_img = self.diff_kernel.image
260 self.grad_kernel: Fourier | None = Fourier(diff_img[:, ::-1, ::-1])
261 else:
262 self.diff_kernel = None
263 self.grad_kernel = None
265 self._convolution_bounds: tuple[int, int, int, int] | None = None
267 @property
268 def bands(self) -> tuple:
269 """The bands in the observations."""
270 return self.images.bands
272 @property
273 def bbox(self) -> Box:
274 """The bounding box for the full observation."""
275 return self.images.bbox
277 def convolve(self, image: Image, mode: str | None = None, grad: bool = False) -> Image:
278 """Convolve the model into the observed seeing in each band.
280 Parameters
281 ----------
282 image:
283 The 3D image to convolve.
284 mode:
285 The convolution mode to use.
286 This should be "real" or "fft" or `None`,
287 where `None` will use the default `convolution_mode`
288 specified during init.
289 grad:
290 Whether this is a backward gradient convolution
291 (`grad==True`) or a pure convolution with the PSF.
293 Returns
294 -------
295 result:
296 The convolved image.
297 """
298 if grad:
299 kernel = self.grad_kernel
300 else:
301 kernel = self.diff_kernel
303 if kernel is None:
304 return image
306 if mode is None:
307 mode = self.mode
308 if mode == "fft":
309 result = fft_convolve(
310 Fourier(image.data),
311 kernel,
312 axes=(1, 2),
313 return_fourier=False,
314 )
315 elif mode == "real":
316 dy = image.shape[1] - kernel.image.shape[1]
317 dx = image.shape[2] - kernel.image.shape[2]
318 if dy < 0 or dx < 0:
319 # The image needs to be padded because it is smaller than
320 # the psf kernel
321 _image = image.data
322 newshape = list(_image.shape)
323 if dy < 0:
324 newshape[1] += kernel.image.shape[1] - image.shape[1]
325 if dx < 0:
326 newshape[2] += kernel.image.shape[2] - image.shape[2]
327 _image = _pad(_image, newshape)
328 result = convolve(_image, kernel.image, self.convolution_bounds)
329 result = centered(result, image.data.shape) # type: ignore
330 else:
331 result = convolve(image.data, kernel.image, self.convolution_bounds)
332 else:
333 raise ValueError(f"mode must be either 'fft' or 'real', got {mode}")
334 return Image(cast(np.ndarray, result), bands=image.bands, yx0=image.yx0)
336 def log_likelihood(self, model: Image) -> float:
337 """Calculate the log likelihood of the given model
339 Parameters
340 ----------
341 model:
342 Model to compare with the observed images.
344 Returns
345 -------
346 result:
347 The log-likelihood of the given model.
348 """
349 result = 0.5 * -np.sum((self.weights * (self.images - model) ** 2).data)
350 return result
352 @property
353 def shape(self) -> tuple[int, int, int]:
354 """The shape of the images, variance, etc."""
355 return cast(tuple[int, int, int], self.images.shape)
357 @property
358 def n_bands(self) -> int:
359 """The number of bands in the observation"""
360 return self.images.shape[0]
362 @property
363 def dtype(self) -> npt.DTypeLike:
364 """The dtype of the observation is the dtype of the images"""
365 return self.images.dtype
367 @property
368 def convolution_bounds(self) -> tuple[int, int, int, int]:
369 """Build the slices needed for convolution in real space"""
370 if self._convolution_bounds is None:
371 coords = get_filter_coords(cast(Fourier, self.diff_kernel).image[0])
372 self._convolution_bounds = get_filter_bounds(coords.reshape(-1, 2))
373 return self._convolution_bounds
375 @staticmethod
376 def empty(
377 bands: tuple[Any], psfs: np.ndarray, model_psf: np.ndarray, bbox: Box, dtype: npt.DTypeLike
378 ) -> Observation:
379 dummy_image = np.zeros((len(bands),) + bbox.shape, dtype=dtype)
381 return Observation(
382 images=dummy_image,
383 variance=dummy_image,
384 weights=dummy_image,
385 psfs=psfs,
386 model_psf=model_psf,
387 noise_rms=np.zeros((len(bands),), dtype=dtype),
388 bbox=bbox,
389 bands=bands,
390 convolution_mode="real",
391 )