Coverage for python / lsst / scarlet / lite / observation.py: 17%
163 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-04-17 08:40 +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 copy import deepcopy
27from typing import Any, cast
29import numpy as np
30import numpy.typing as npt
32from .bbox import Box
33from .fft import Fourier, _pad, centered
34from .fft import convolve as fft_convolve
35from .fft import match_kernel
36from .image import Image
39def get_filter_coords(filter_values: np.ndarray, center: tuple[int, int] | None = None) -> np.ndarray:
40 """Create filter coordinate grid needed for the apply filter function
42 Parameters
43 ----------
44 filter_values:
45 The 2D array of the filter to apply.
46 center:
47 The center (y,x) of the filter. If `center` is `None` then
48 `filter_values` must have an odd number of rows and columns
49 and the center will be set to the center of `filter_values`.
51 Returns
52 -------
53 coords:
54 The coordinates of the pixels in `filter_values`,
55 where the coordinates of the `center` pixel are `(0,0)`.
56 """
57 if filter_values.ndim != 2:
58 raise ValueError("`filter_values` must be 2D")
59 if center is None:
60 if filter_values.shape[0] % 2 == 0 or filter_values.shape[1] % 2 == 0:
61 msg = """Ambiguous center of the `filter_values` array,
62 you must use a `filter_values` array
63 with an odd number of rows and columns or
64 calculate `coords` on your own."""
65 raise ValueError(msg)
66 center = tuple([filter_values.shape[0] // 2, filter_values.shape[1] // 2]) # type: ignore
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(v[np.isfinite(v)])) for v in self.variance.data])
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 def __getitem__(self, indices: Any) -> Observation:
353 """Get a view for the subset of an image
355 Parameters
356 ----------
357 indices:
358 The indices to select a subsection of the image.
360 Returns
361 -------
362 result:
363 The resulting image obtained by selecting subsets of the iamge
364 based on the `indices`.
365 """
366 new_image = self.images[indices]
367 new_variance = self.variance[indices]
368 new_weights = self.weights[indices]
370 # If the indices is a single band, make sure to keep the band axis
371 if new_image.ndim == 2:
372 if indices in self.bands:
373 new_bands = (indices,)
374 else:
375 # The indices contain spatial and band indices
376 new_bands = (indices[0],)
377 new_image = Image(
378 new_image.data[None, :, :],
379 yx0=new_image.yx0,
380 bands=new_bands,
381 )
382 new_variance = Image(
383 new_variance.data[None, :, :],
384 yx0=new_variance.yx0,
385 bands=new_bands,
386 )
387 new_weights = Image(
388 new_weights.data[None, :, :],
389 yx0=new_weights.yx0,
390 bands=new_bands,
391 )
393 # Extract the appropriate bands from the PSF
394 bands = self.images.bands
395 new_bands = new_image.bands
396 if bands != new_bands:
397 band_indices = self.images.spectral_indices(new_bands)
398 psfs = self.psfs[band_indices,]
399 noise_rms = self.noise_rms[band_indices,]
400 else:
401 psfs = self.psfs
402 noise_rms = self.noise_rms
404 return Observation(
405 images=new_image,
406 variance=new_variance,
407 weights=new_weights,
408 psfs=psfs,
409 model_psf=self.model_psf,
410 noise_rms=noise_rms,
411 bbox=new_image.bbox,
412 bands=new_bands,
413 padding=self.padding,
414 convolution_mode=self.mode,
415 )
417 def __copy__(self) -> Observation:
418 """Create a copy of the observation
420 Returns
421 -------
422 result:
423 The copy of the observation.
424 """
425 return Observation(
426 images=self.images,
427 variance=self.variance,
428 weights=self.weights,
429 psfs=self.psfs,
430 model_psf=self.model_psf,
431 noise_rms=self.noise_rms,
432 bands=self.bands,
433 padding=self.padding,
434 convolution_mode=self.mode,
435 )
437 def __deepcopy__(self, memo: dict[int, Any]) -> Observation:
438 """Create a deep copy of the observation
440 Parameters
441 ----------
442 memo: dict[int, Any]
443 The memoization dictionary used by `copy.deepcopy`.
445 Returns
446 -------
447 result:
448 The deep copy of the observation.
449 """
450 # Check if already copied
451 if id(self) in memo:
452 return memo[id(self)]
454 # Create placeholder and add to memo FIRST
455 result = Observation.__new__(Observation)
456 memo[id(self)] = result
458 # Now safely initialize the placeholder with deepcopied arguments
459 result.__init__( # type: ignore[misc]
460 images=deepcopy(self.images, memo),
461 variance=deepcopy(self.variance, memo),
462 weights=deepcopy(self.weights, memo),
463 psfs=deepcopy(self.psfs, memo),
464 model_psf=deepcopy(self.model_psf, memo),
465 noise_rms=deepcopy(self.noise_rms, memo),
466 bands=deepcopy(self.bands, memo),
467 padding=deepcopy(self.padding, memo),
468 convolution_mode=self.mode,
469 )
471 return result
473 def copy(self, deep: bool = False) -> Observation:
474 """Create a copy of the observation
476 Parameters
477 ----------
478 deep:
479 Whether to perform a deep copy or not.
481 Returns
482 -------
483 result:
484 The copy of the observation.
485 """
486 if deep:
487 return self.__deepcopy__({})
488 return self.__copy__()
490 @property
491 def shape(self) -> tuple[int, int, int]:
492 """The shape of the images, variance, etc."""
493 return cast(tuple[int, int, int], self.images.shape)
495 @property
496 def n_bands(self) -> int:
497 """The number of bands in the observation"""
498 return self.images.shape[0]
500 @property
501 def dtype(self) -> npt.DTypeLike:
502 """The dtype of the observation is the dtype of the images"""
503 return self.images.dtype
505 @property
506 def convolution_bounds(self) -> tuple[int, int, int, int]:
507 """Build the slices needed for convolution in real space"""
508 if self._convolution_bounds is None:
509 coords = get_filter_coords(cast(Fourier, self.diff_kernel).image[0])
510 self._convolution_bounds = get_filter_bounds(coords.reshape(-1, 2))
511 return self._convolution_bounds
513 @staticmethod
514 def empty(
515 bands: tuple[Any], psfs: np.ndarray, model_psf: np.ndarray, bbox: Box, dtype: npt.DTypeLike
516 ) -> Observation:
517 dummy_image = np.zeros((len(bands),) + bbox.shape, dtype=dtype)
519 return Observation(
520 images=dummy_image,
521 variance=dummy_image,
522 weights=dummy_image,
523 psfs=psfs,
524 model_psf=model_psf,
525 noise_rms=np.zeros((len(bands),), dtype=dtype),
526 bbox=bbox,
527 bands=bands,
528 convolution_mode="real",
529 )