diff --git a/doc/lsst.scarlet.lite/getting_started.rst b/doc/lsst.scarlet.lite/getting_started.rst index 192bdc14..044f79db 100644 --- a/doc/lsst.scarlet.lite/getting_started.rst +++ b/doc/lsst.scarlet.lite/getting_started.rst @@ -110,18 +110,10 @@ Creating an initial model for each component is a non-trivial task for a number 1. The models exist in a partially deconvolved space. This means that if the component morphologies are initialized using the observed images, convolution by the difference kernel will cause all of the initial models to be "puffy." 2. There is no analytic way to initialize the bulge and disk components of a two component source. This means that some approximation or iterative proceedure must be used to initialize the bulge and disk. -The scarlet developers attempted several different initialization schemes over the years, and certainly there is a lot of theoretical space we have yet to explore, but scarlet lite has implemented two of the most useful (and fastest) methods that we've found so far. -An important note is that for factorized components, initializing the morphology is the most important and difficult algorithmic challenge, as a least squares algorithm can be (and is) used to fit the initial spectra once a set of morphology templates is generated (see :ref:`lsst.scarlet.lite-fit-spectra`). +For factorized components, initializing the morphology is the most important and difficult algorithmic challenge, as a least squares algorithm can be (and is) used to fit the initial spectra once a set of morphology templates is generated (see :ref:`lsst.scarlet.lite-fit-spectra`). -In general the chi^2 initializaation appears to work better with the Adaprox optimizer, which usually gives the best overall fit. -However, the wavelet initialization tends to work better with the PGM optmizer, and also generally starts with more compact sources. -So there is benefit to using the wavelet initialization in crowded fields, or fields with heavy blending like galaxy clusters. - -Chi^2 Initialization -^^^^^^^^^^^^^^^^^^^^ - -The idea behind the chi^2 initialization algorithm is that it initializes the morphologies on the chi^2 (variance weighted single band) coadd, referred to in the remainder of this section as the *detection image*. -For detection, chi^2 coadds have been shown to be optimal for multiband point source detection (see `Szalay et. al 1998 `_) +:py:class:`~lsst.scarlet.lite.initialization.FactorizedInitialization` initializes the morphologies on a detection image, by default the chi^2 (variance weighted single band) coadd. +For detection, chi^2 coadds have been shown to be optimal for multiband point source detection (see `Szalay et. al 1998 `_). Initializing the morphologies this way is almost trivial when each source has a single component, in which case sources are initialized by making the detection image symmetric about the center of the source by taking the minimum of each pixel and its symmetric counterpart, similar to the SDSS deblender. This helps prevent the source from growing too large and swallowing too much flux from neighbors. Next a monotonicity operator is used to smooth over non-monotonic regions that are usually due to neighboring sources. @@ -129,7 +121,7 @@ For sources with low signal to noise, there are not a sufficient number of pixel For sources with sufficiently large signal to noise, two separate components are initialized by specifying the fraction of the overall flux to attribute to the disk. The "bulge" component contains all of the flux above the disk threshold, and the "disk" component contains all of the flux below it. -To initialize a set of sources using chi^2 initialization simply use +To initialize a set of sources use >>> # Initialize the monotonicity operator. >>> # This is only done once and effectively creates a lookup table that all of the monotonic components use. @@ -137,23 +129,10 @@ To initialize a set of sources using chi^2 initialization simply use >>> # If a source is larger, the operator will be modified to the new value. >>> monotonicity = scl.operators.Monotonicity((51, 51)) >>> # Initialize the sources ->>> chi2init = scl.initialization.FactorizedChi2Initialization(observation, centers, monotonicity=monotonicity) ->>> sources = chi2init.sources - -Wavelet Initialization -^^^^^^^^^^^^^^^^^^^^^^ +>>> initialization = scl.initialization.FactorizedInitialization(observation, centers, monotonicity=monotonicity) +>>> sources = initialization.sources -The idea behind wavelet initialization is that much of the flux at lower scales in the image is due to the PSF, so using only the higher frequency scales prevents much of the observed PSF from getting into the initial source models. -The other key difference is that the bulge and disk are initialized with the detection image wavelets at different scales, where the high frequency wavelets are used for the bulge and the lower frequency wavelets are used for the disk. -To initialize sources using wavelet coefficients use - ->>> # Initialize the sources ->>> wavelet_init = scl.initialization.FactorizedWaveletInitialization( -... observation, -... centers, -... monotonicity=monotonicity -... ) ->>> wavelet_sources = wavelet_init.sources +To initialize on a different detection image, for example one built from wavelet coefficients, pass it as the ``detect`` argument. .. _lsst.scarlet.lite-fit-model: diff --git a/python/lsst/scarlet/lite/detect.py b/python/lsst/scarlet/lite/detect.py index af988211..4a7eeae6 100644 --- a/python/lsst/scarlet/lite/detect.py +++ b/python/lsst/scarlet/lite/detect.py @@ -22,10 +22,18 @@ from __future__ import annotations import logging -from typing import Sequence +from dataclasses import dataclass +from math import comb, log +from typing import Sequence, cast import numpy as np -from lsst.scarlet.lite.detect_pybind11 import Footprint, get_footprints # type: ignore +from deprecated.sphinx import deprecated # type: ignore +from lsst.scarlet.lite.detect_pybind11 import Footprint, Peak, get_footprints # type: ignore +from scipy import special +from scipy.ndimage import binary_dilation, label +from scipy.sparse import csr_array +from scipy.sparse.csgraph import connected_components +from scipy.spatial import cKDTree from .bbox import Box, overlapped_slices from .image import Image @@ -34,6 +42,7 @@ get_multiresolution_support, get_starlet_scales, multiband_starlet_reconstruction, + multiband_starlet_transform, starlet_transform, ) @@ -165,6 +174,12 @@ def footprints_to_image(footprints: Sequence[Footprint], bbox: Box) -> Image: return result +@deprecated( + reason="get_wavelets is only used by the deprecated detect_footprints " + "and will be removed after v31.0.", + version="v31.0", + category=FutureWarning, +) def get_wavelets( images: np.ndarray, variance: np.ndarray, @@ -209,6 +224,11 @@ def get_wavelets( return coeffs +@deprecated( + reason="get_detect_wavelets is superseded by detect_peaks and will be " "removed after v31.0.", + version="v31.0", + category=FutureWarning, +) def get_detect_wavelets(images: np.ndarray, variance: np.ndarray, scales: int = 3) -> np.ndarray: """Get an array of wavelet coefficents to use for detection @@ -246,6 +266,12 @@ def get_detect_wavelets(images: np.ndarray, variance: np.ndarray, scales: int = return (support.support * _coeffs).astype(images.dtype) +@deprecated( + reason="detect_footprints is replaced by detect_peaks and will be removed after v31.0. " + "Use detect_peaks instead.", + version="v31.0", + category=FutureWarning, +) def detect_footprints( images: np.ndarray, variance: np.ndarray, @@ -359,3 +385,1986 @@ def detect_footprints( ) return footprints + + +# Record type for a peak candidate: one local maximum in one plane +# (band, scale) of the significance map that survived the plane's +# contrast-limited watershed. +CANDIDATE_DTYPE = np.dtype( + [ + ("y", int), + ("x", int), + ("band", int), + ("scale", int), + ("flux", float), + ("saddle", float), + ("position", int), + ] +) + +# Record type for a unique candidate position, the intermediate between the +# candidates and the final peaks. +POSITION_DTYPE = np.dtype( + [ + ("y", int), + ("x", int), + ("scale", int), + ("flux", float), + ("peak_sigma", float), + ("band_flags", int), + ("scale_flags", int), + ("plane_flags", np.int64), + ("n_candidates", int), + ("n_linked", int), + ("peak", int), + ("component", int), + ("ambiguous", bool), + ("max_linked_fraction", float), + ("max_linked_peak", int), + ] +) + +# Record type for a final detection, one row per source. +DETECTION_DTYPE = np.dtype( + [ + ("y", int), + ("x", int), + ("scale", int), + ("flux", float), + ("peak_sigma", float), + ("band_flags", int), + ("scale_flags", int), + ("n_candidates", int), + ("n_positions", int), + ("n_ambiguous", int), + ("footprint", int), + ("location_fallback", bool), + ("search_capped", bool), + ("nearest_distance", float), + ("nearest_distance_fwhm", float), + ("stat_err_y", float), + ("stat_err_x", float), + ("moment_yy", float), + ("moment_xx", float), + ("moment_xy", float), + ("scatter_y", float), + ("scatter_x", float), + ("err_y", float), + ("err_x", float), + ("structure_peaks", int), + ("structure_id", int), + ] +) + +# Record type for a pair of positions within the largest link radius, linked +# or not. +PAIR_DTYPE = np.dtype( + [ + ("i", int), + ("j", int), + ("distance", float), + ("linked", bool), + ("shares_plane", bool), + ("within_min", bool), + ("within_max", bool), + ] +) + +# Record type for a connected component of the position graph. +COMPONENT_DTYPE = np.dtype( + [ + ("n_nodes", int), + ("lower_bound", int), + ("upper_bound", int), + ("k", int), + ("search_ran", bool), + ("n_visited", int), + ("cap_hit", bool), + ("n_min_covers", int), + ] +) + +# The PAIR_DTYPE field that decides whether a pair is close enough to link, +# for each value of ``radius_rule``. +RADIUS_RULES = {"min": "within_min", "max": "within_max"} + +# Minimum covers are only counted for components this small, and at most this +# many of them. +COUNT_COVERS_MAX_NODES = 12 +COUNT_COVERS_MAX = 100 +SEARCH_CAP = 100_000 + +SIGMA_TO_FWHM = 2 * np.sqrt(2 * np.log(2)) +MAD_TO_SIGMA = 1.4826 + + +@dataclass +class PeakDetectionResult: + """Detected peaks and the products used to detect them. + + Attributes + ---------- + peaks : + Structured array of detections with dtype `DETECTION_DTYPE`, one row + per source. Row ids link back to ``positions`` via + ``positions["peak"]``; ``footprint`` is the index into + ``footprints`` of the footprint containing the peak. The location, + flags and errors are described in ``_build_detections``. + footprints : + The detected footprints, each holding the `Peak` objects of the + detections inside it, brightest first. Every peak lies in exactly + one footprint and every footprint holds at least one peak (see + ``_build_footprints``). + positions : + Structured array of unique candidate positions with dtype + `POSITION_DTYPE`, the intermediate between candidates and peaks. Row + ids link back to ``candidates`` via ``candidates["position"]``. + ``plane_flags`` is a bitmask of the planes (band, scale) the position + was a distinct peak in. ``peak`` and ``component`` index ``peaks`` + and ``components``; ``n_linked``, ``ambiguous`` and the + ``max_linked_*`` fields describe the links to other peaks (see + ``_link_fractions``). + candidates : + Structured array of peak candidates with dtype `CANDIDATE_DTYPE`. + The same source is expected to appear multiple times, once per band + and scale it is significant in. ``saddle`` is the level at which the + candidate's basin joined a brighter candidate's basin in its plane + (NaN for the brightest candidate in a footprint). + significance_map : + The per-scale detection map from ``_build_significance_map``, with + shape ``(n_scales, n_bands + 1, Ny, Nx)``. Every plane is in units of + Gaussian sigma; the final plane is the clipped chi coadd mapped to + sigma. + starlets : + The multiband starlet coefficients, with shape + ``(scales + 1, n_bands, Ny, Nx)``. + sigma : + The coefficient noise std, shaped to broadcast against ``starlets``. + See ``_build_detection_starlets`` for the two shapes it takes. + pairs : + Structured array of position pairs with dtype `PAIR_DTYPE`. ``i`` and + ``j`` index ``positions``. + components : + Structured array of the connected components of the position graph + with dtype `COMPONENT_DTYPE` (see ``_clique_cover``). + metadata : + The settings the peaks were grouped with: the per-band + ``psf_fwhm``, the ``link_radius`` at each of ``link_radius_scales``, + ``min_separation``, ``radius_rule`` and ``search_cap``, and the + ``structure_thresh`` the structures were found with. + """ + + peaks: np.ndarray + footprints: list[Footprint] + positions: np.ndarray + candidates: np.ndarray + significance_map: np.ndarray + starlets: np.ndarray + sigma: np.ndarray + pairs: np.ndarray + components: np.ndarray + metadata: dict + + +def _starlet_scale_factors(scales: int, generation: int = 2) -> np.ndarray: + """Sum of squared effective-kernel weights at each starlet scale. + + Parameters + ---------- + scales : + The number of wavelet scales. The result has ``scales + 1`` entries + to account for the coarse scale. + generation : + The generation of the starlet transform, either ``1`` or ``2``. + + Returns + ------- + factors : + The factor ``F_j = sum(K_j**2)`` for each scale, where ``K_j`` is the + effective a trous kernel. White noise of variance ``v`` transforms to + coefficients of variance ``F_j * v`` at scale ``j``. + + Notes + ----- + Because starlets change the scale of the data, we have to change the + scale of the variance by a related factor in order to properly calculate + the significance of detections. This algorithm computes the sum of squared + effective-kernel weights at each starlet scale. + + The factors are read off a delta image transformed through the same fast + a trous transform, so they match its generation and boundary handling. + """ + size = 2 ** (scales + 3) + 1 + delta = np.zeros((size, size)) + delta[size // 2, size // 2] = 1.0 + coeffs = starlet_transform(delta, scales=scales, generation=generation) + return np.sum(coeffs**2, axis=(1, 2)) + + +def _build_detection_starlets( + images: np.ndarray, + variance: np.ndarray, + scales: int = 3, + generation: int = 2, + variance_mode: str = "median", +) -> tuple[np.ndarray, np.ndarray]: + """Transform a multiband image and estimate its per-scale noise. + + Parameters + ---------- + images : + The multiband image with shape ``(n_bands, Ny, Nx)``. + variance : + The per-pixel variance, with the same shape as ``images``. + scales : + The maximum number of wavelet scales to use. + generation : + The generation of the starlet transform, either ``1`` or ``2``. + variance_mode : + How ``variance`` is reduced to a coefficient noise. ``"median"`` takes + one value per band, ``"pixel"`` keeps the variance plane per pixel. + + Returns + ------- + starlets : + The multiband starlet coefficients with shape + ``(scales + 1, n_bands, Ny, Nx)``. + sigma : + The coefficient noise std at each scale and band, shaped to broadcast + against ``starlets``: ``(scales + 1, n_bands, 1, 1)`` for + ``variance_mode="median"`` and ``(scales + 1, n_bands, Ny, Nx)`` for + ``variance_mode="pixel"``. + + Raises + ------ + ValueError + Raised if ``images`` and ``variance`` differ in shape, if ``images`` is + not 3D, or if ``variance_mode`` is not one of the two allowed values. + + Notes + ----- + The coefficient noise is ``sigma_{j,b} = sqrt(F_j * var_b)``, where ``F_j`` + is the per-scale factor from ``_starlet_scale_factors``. + + With ``variance_mode="median"`` the variance is a single + ``nanmedian(var_b)`` per band, which treats the noise as stationary within + a band. With ``"pixel"`` the variance plane is used as it stands, which + follows a varying depth at the memory cost of a ``sigma`` that is as + large as ``starlets``. + + Neither is the exact coefficient noise, which is ``var_b`` convolved with + the square of the effective kernel at each scale. Both approximate that by + the scalar ``F_j``, so ``"pixel"`` is the better estimate only where the + variance is smooth across a kernel, and the kernel doubles in width with + every scale. + """ + if images.shape != variance.shape: + raise ValueError("images and variance must have the same shape") + if images.ndim != 3: + raise ValueError("images and variance must be 3D (bands, Ny, Nx)") + if variance_mode not in ("median", "pixel"): + raise ValueError(f"variance_mode must be 'median' or 'pixel', not {variance_mode!r}") + + starlets = multiband_starlet_transform(images, scales=scales, generation=generation) + # The transform caps the scale count at the image size, so read the + # realized number of scales back off the coefficients. + realized_scales = starlets.shape[0] - 1 + factors = _starlet_scale_factors(realized_scales, generation=generation) + if variance_mode == "median": + band_variance = np.nanmedian(variance, axis=(1, 2))[:, None, None] + else: + band_variance = variance + sigma = np.sqrt(factors[:, None, None, None] * band_variance[None]).astype(starlets.dtype) + return starlets, sigma + + +def _log_gammaincc(a: float, z: np.ndarray) -> np.ndarray: + """Log of the regularized upper incomplete gamma function ``Q(a, z)``. + + Parameters + ---------- + a : + The shape parameter, ``a > 0``. + z : + The lower limit(s) of integration, ``z >= 0``. + + Returns + ------- + log_q : + ``log(Q(a, z))``, with the same shape as ``z``. + + Notes + ----- + ``scipy.special.gammaincc`` underflows to zero for ``z`` above a few + hundred, which would map every sufficiently bright pixel to the same + saturated significance. Above ``z = 500`` this uses the asymptotic + expansion + ``Q(a, z) ~ z**(a-1) exp(-z) / Gamma(a) * sum_m prod_{i<=m}(a-i) / z**m``, + whose truncation error at ``z >= 500`` is far below double precision for + any ``a`` of interest here (``a = k/2`` for ``k`` bands). + """ + z = np.asarray(z, dtype=float) + out = np.empty(z.shape, dtype=float) + small = z < 500 + out[small] = np.log(special.gammaincc(a, z[small])) + zl = z[~small] + if zl.size: + series = np.ones_like(zl) + term = np.ones_like(zl) + for m in range(1, 8): + term = term * (a - m) / zl + series += term + out[~small] = (a - 1) * np.log(zl) - zl - special.gammaln(a) + np.log(series) + return out + + +def _chi2_log_survival(chi2: np.ndarray | float, n_bands: int) -> np.ndarray: + """Log survival function of the clipped chi-squared coadd under noise. + + Parameters + ---------- + chi2 : + The clipped chi-squared coadd value(s); scalar or array. + n_bands : + The number of bands combined into ``chi2``. + + Returns + ------- + log_survival : + ``log P(C > chi2)`` for pure noise, with the shape of ``chi2``. + + Notes + ----- + See ``_chi2_to_sigma`` for the derivation of this mixture of ``chi^2_k`` + survival functions weighted by the binomial coefficients ``C(n, k)/2**n``. + The mixture is evaluated in log space so it stays finite for arbitrarily + bright pixels; ``chi2_k.sf(x) = Q(k/2, x/2)``. + """ + z = np.asarray(chi2, dtype=float) / 2 + log_norm = -n_bands * log(2) + terms = np.stack( + [log(comb(n_bands, k)) + log_norm + _log_gammaincc(k / 2, z) for k in range(1, n_bands + 1)] + ) + return special.logsumexp(terms, axis=0) + + +def _chi2_to_sigma(chi2: np.ndarray | float, n_bands: int) -> np.ndarray: + """Convert a clipped chi-squared coadd to Gaussian sigma. + + Parameters + ---------- + chi2 : + The chi-squared coadd ``sum_b max(w_b, 0)**2`` of per-band + coefficients standardized to unit noise variance. This can also be + a single pixel, like the value of a peak, to detect its significance. + n_bands : + The number of bands combined into ``chi2``. + + Returns + ------- + sigma : + The equivalent Gaussian significance (upper-tail), with the same + shape as ``chi2``. + + Notes + ----- + For an unclipped chi-squared coadd with ``n_bands`` degrees of freedom, + ``scipy.stats.chi2.sf(chi2, n)`` gives the probability that a pure-noise + pixel exceeds ``chi2``. However, ``_build_significance_map`` clips each + band's coefficients at zero before squaring to supress false detections + from negative structure (wavelet sidelobes, oversubtracted sky, etc). + This changes the noise distribution: each standardized coeffficent is in + a Gaussian with mean 0 and unit variance, so in any band a noise pixel is + clipped to zero with probability 1/2, and the number of bands that + contribute to the sum is binomially distributed. If ``k`` bands contribute, + the sum is a ``chi^2_k`` deviate, so the probability that a noise pixel + exceeds ``chi2`` is the sum over ``k`` of ``chi2_k.sf(chi2)`` weighted by + the binomial probability ``C(n, k) / 2**n``. That probability is converted + to an equivalent Gaussian sigma so that thresholds are directly comparable + to a single-band ``n``-sigma cut. + + Due to divergence for large ``chi2``, instead of using scipy.stats.chi2 + function directly, we calculate in log space (``_chi2_log_survival`` and + ``scipy.special.ndtri_exp``), so the mapping is strictly monotonic with no + saturation, however bright the pixel. + """ + return -special.ndtri_exp(_chi2_log_survival(chi2, n_bands)) + + +def _chi_to_sigma(chi: np.ndarray, n_bands: int) -> np.ndarray: + """Map a clipped chi coadd plane to Gaussian sigma through a lookup table. + + Parameters + ---------- + chi : + The chi coadd ``sqrt(sum_b max(w_b, 0)**2)``; any shape. + n_bands : + The number of bands combined into ``chi``. + + Returns + ------- + sigma : + ``_chi2_to_sigma(chi**2, n_bands)`` with the shape and dtype of + ``chi``, evaluated by linear interpolation of a table. + + Notes + ----- + ``chi -> sigma`` is a smooth, monotonic, one-dimensional function for a + fixed band count, so evaluating it on a few thousand grid points and + interpolating is far cheaper than evaluating the survival-function mixture + per pixel, and accurate to well below 1e-4 sigma. The grid is dense where + the function has curvature (``chi < 64``) and geometric beyond, where + ``sigma`` approaches ``chi`` minus a slowly varying offset. + """ + chi_max = float(np.nanmax(chi)) if chi.size else 1.0 + n_dense = 4096 + dense_max = 64.0 + grid = np.linspace(0.0, dense_max, n_dense + 1) + table = _chi2_to_sigma(grid**2, n_bands) + + # The dense grid is uniform, so interpolate by direct indexing (a bin + # search per pixel with np.interp is ~10x slower on a full image). + invalid = np.isnan(chi) + t = np.clip(chi, 0, dense_max) * (n_dense / dense_max) + if invalid.any(): + t[invalid] = 0.0 + idx = t.astype(np.intp) + np.minimum(idx, n_dense - 1, out=idx) + frac = t - idx + sigma = table[idx] + sigma += frac * (table[idx + 1] - sigma) + if invalid.any(): + sigma[invalid] = np.nan + + if chi_max > dense_max: + # Bright pixels are rare; a geometric grid and np.interp are fine here. + bright = chi > dense_max + grid_hi = np.geomspace(dense_max, chi_max * 1.01, 513) + table_hi = _chi2_to_sigma(grid_hi**2, n_bands) + sigma[bright] = np.interp(chi[bright], grid_hi, table_hi) + return sigma.astype(chi.dtype, copy=False) + + +def build_chi2_significance( + standardized: np.ndarray, +) -> np.ndarray: + """Coadd standardized single-band planes into a map of Gaussian sigma. + + Parameters + ---------- + standardized : + The single-band planes with shape ``(n_bands, Ny, Nx)``, each + standardized to unit noise variance. These are starlet coefficients + divided by their per-scale noise std, or an image divided by the + square root of its variance. + + Returns + ------- + significance : + The coadd with shape ``(Ny, Nx)``, in units of Gaussian sigma. + + Notes + ----- + The standardized single-band planes are clipped at zero before computing + the chi coadd, to avoid negative contributions. This changes the noise + distribution of the coadd (see ``_chi2_to_sigma``), which is why the coadd + is mapped to sigma rather than used directly. + """ + chi = np.sqrt(np.sum(np.clip(standardized, 0, None) ** 2, axis=0)) + return _chi_to_sigma(chi, standardized.shape[0]) + + +def _build_significance_map( + starlets: np.ndarray, + sigma: np.ndarray, + first_scale: int = 1, +) -> np.ndarray: + """Build a per-scale detection map from multiband starlet coefficients. + + Parameters + ---------- + starlets : + The multiband starlet coefficients with shape + ``(scales + 1, n_bands, Ny, Nx)``. + sigma : + The coefficient noise std from ``_build_detection_starlets``, shaped + to broadcast against ``starlets``. + first_scale : + The first starlet scale to keep. Scales below this (the highest + frequencies) and the final residual scale are dropped. + + Returns + ------- + significance_map : + The detection map with shape ``(n_scales, n_bands + 1, Ny, Nx)``, + where ``n_scales = scales - first_scale``. The first ``n_bands`` + planes are the standardized single-band coefficients, in units of + Gaussian sigma. The last plane is the clipped chi coadd + ``sqrt(sum_b max(w_b, 0)**2)`` mapped to Gaussian sigma with + ``_chi_to_sigma``, so every plane is on the same footing and a + single threshold or contrast applies to all of them. + """ + _, n_bands, height, width = starlets.shape + scale_slice = slice(first_scale, -1) + coeffs = starlets[scale_slice] + scale_sigma = sigma[scale_slice] + n_scales = coeffs.shape[0] + + significance_map = np.zeros((n_scales, n_bands + 1, height, width), dtype=np.float32) + standardized = coeffs / scale_sigma + significance_map[:, :n_bands] = standardized + for i in range(n_scales): + significance_map[i, n_bands] = build_chi2_significance(standardized[i]) + return significance_map + + +def _dilate_bad_mask(mask: np.ndarray, radius: float) -> np.ndarray: + """Grow a bad-pixel mask by a circular structuring element. + + Parameters + ---------- + mask : + A 2D boolean mask, `True` where a pixel is bad. + radius : + The dilation radius in pixels. + + Returns + ------- + dilated : + The dilated mask. ``mask`` is returned unchanged when ``radius`` is + not positive or no pixel is set. + + Notes + ----- + Starlet ringing around a bad region reaches beyond the flagged pixels, so + the peaks it seeds have to be suppressed over a slightly wider region than + the mask itself. + """ + if radius <= 0 or not mask.any(): + return mask + r = int(np.ceil(radius)) + yy, xx = np.mgrid[-r : r + 1, -r : r + 1] + element = yy**2 + xx**2 <= radius**2 + return binary_dilation(mask, structure=element) + + +def _find_peak_candidates( + significance_map: np.ndarray, + min_separation: float = 0, + min_area: int = 4, + peak_thresh: float = 3, + footprint_thresh: float = 2, + kappa: float = 3, + first_scale: int = 1, + origin: tuple[int, int] = (0, 0), + bad_pixel_mask: np.ndarray | None = None, + dilation_radius: float = 0, +) -> tuple[np.ndarray, np.ndarray]: + """Find peak candidates in each plane of a significance map. + + Parameters + ---------- + significance_map : + A detection map from ``_build_significance_map``, with every plane + in sigma. + min_separation : + A hard floor on the separation between peaks in a plane, in pixels. + Peaks closer than this to a brighter peak are dropped regardless of + the contrast test; ``0`` disables it. + min_area : + The minimum area of a footprint in pixels. + peak_thresh : + The peak detection threshold, in sigma. + footprint_thresh : + The footprint detection threshold, in sigma. + kappa : + The minimum prominence of a peak, in sigma: the peak must rise at + least ``kappa`` above the saddle connecting it to any brighter peak + in its footprint, otherwise it is a bump on that peak's flank and is + culled. This is what makes two candidates in the same plane distinct + sources (see ``_build_pairs``). + first_scale : + The first starlet scale in ``significance_map``, used to label the + ``scale`` field of each candidate. + origin : + The ``(y, x)`` location of the lower corner of the image. + bad_pixel_mask : + Per-band boolean mask with shape ``(n_bands, Ny, Nx)``, `True` where a + band is bad. Peaks inside the dilated mask of their band are dropped; + the chi coadd plane (``band == n_bands``) uses the union across bands. + `None` disables the suppression. + dilation_radius : + The base radius in pixels by which ``bad_pixel_mask`` is grown before + it suppresses peaks, to catch the starlet ringing just outside a bad + region. The starlet kernel doubles in width with each scale, so the + radius applied at a plane is ``dilation_radius * 2 ** (scale - + first_scale)``. + + Returns + ------- + candidates : + Structured array of candidates with dtype `CANDIDATE_DTYPE`. The + ``flux`` field is the peak significance in sigma for every plane. + footprint_mask : + Boolean image with the spatial shape of ``significance_map`` that is + `True` in every pixel of any plane's footprint that has a peak. + """ + y0, x0 = origin + height, width = significance_map.shape[-2:] + n_bands = significance_map.shape[1] - 1 + footprint_mask = np.zeros((height, width), dtype=bool) + # Dilated bad mask per (band, scale), filled lazily. + # The ring widens with scale, so each scale grows the mask by its own + # radius. The chi coadd plane is the union of all bands. + dilated_masks: dict[tuple[int, int], np.ndarray] = {} + candidates = [] + for scale_index, scale_planes in enumerate(significance_map): + scale = scale_index + first_scale + radius = dilation_radius * 2 ** (scale - first_scale) + for band, plane in enumerate(scale_planes): + bad = None + if bad_pixel_mask is not None: + if (band, scale) not in dilated_masks: + raw = bad_pixel_mask.any(axis=0) if band == n_bands else bad_pixel_mask[band] + dilated_masks[(band, scale)] = _dilate_bad_mask(raw, radius) + bad = dilated_masks[(band, scale)] + footprints = get_footprints( + plane, + min_separation, + min_area, + peak_thresh, + footprint_thresh, + True, + y0, + x0, + kappa, + ) + for footprint in footprints: + bottom, top, left, right = footprint.bounds + footprint_mask[bottom - y0 : top - y0 + 1, left - x0 : right - x0 + 1] |= footprint.data + for peak in footprint.peaks: + # Do not add peaks that fall on bad pixels. + if bad is not None and bad[peak.y - y0, peak.x - x0]: + continue + # position is -1 until _collapse_positions fills it. + candidates.append((peak.y, peak.x, band, scale, peak.flux, peak.saddle, -1)) + return np.array(candidates, dtype=CANDIDATE_DTYPE), footprint_mask + + +def _plane_flags(candidates: np.ndarray, n_bands: int, first_scale: int) -> np.ndarray: + """One bit per plane ``(band, scale)`` for each candidate. + + Parameters + ---------- + candidates : + Structured array of candidates with dtype `CANDIDATE_DTYPE`. + n_bands : + The number of single-band planes (the chi plane is ``band == + n_bands``). + first_scale : + The first starlet scale in the significance map. + + Returns + ------- + flags : + ``1 << plane`` for each candidate, where + ``plane = (scale - first_scale) * (n_bands + 1) + band``. + + Raises + ------ + ValueError + Raised if there are more than 63 planes. + """ + n_planes = int(candidates["scale"].max() - first_scale + 1) * (n_bands + 1) if len(candidates) else 0 + if n_planes > 63: + raise ValueError(f"plane_flags supports at most 63 planes, got {n_planes}") + plane = (candidates["scale"] - first_scale) * (n_bands + 1) + candidates["band"] + return np.left_shift(np.int64(1), plane.astype(np.int64)) + + +def _collapse_positions(candidates: np.ndarray, n_bands: int, first_scale: int) -> np.ndarray: + """Collapse candidates to unique positions, filling their ``position``. + + Parameters + ---------- + candidates : + Structured array of candidates with dtype `CANDIDATE_DTYPE`. The + ``position`` field is filled in place. + n_bands : + The number of single-band planes. + first_scale : + The first starlet scale in the significance map. + + Returns + ------- + positions : + The unique positions with dtype `POSITION_DTYPE`, sorted by + ``(y, x)``. The grouping fields are left for ``_group_peaks``. + + Notes + ----- + Everything here is a grouped reduction over the candidates, done with + ``np.ufunc.at`` and one lexsort rather than a loop over positions, so + the cost is linear in the number of candidates. + """ + yx = np.column_stack([candidates["y"], candidates["x"]]) + unique_yx, inverse, counts = np.unique(yx, axis=0, return_inverse=True, return_counts=True) + inverse = np.asarray(inverse).ravel() + n = len(unique_yx) + candidates["position"] = inverse + + positions = np.zeros(n, dtype=POSITION_DTYPE) + positions["y"] = unique_yx[:, 0] + positions["x"] = unique_yx[:, 1] + positions["n_candidates"] = counts + positions["peak"] = -1 + + band_flags = np.zeros(n, dtype=np.int64) + np.bitwise_or.at(band_flags, inverse, np.left_shift(np.int64(1), candidates["band"].astype(np.int64))) + scale_flags = np.zeros(n, dtype=np.int64) + np.bitwise_or.at(scale_flags, inverse, np.left_shift(np.int64(1), candidates["scale"].astype(np.int64))) + plane_flags = np.zeros(n, dtype=np.int64) + np.bitwise_or.at(plane_flags, inverse, _plane_flags(candidates, n_bands, first_scale)) + positions["band_flags"] = band_flags + positions["scale_flags"] = scale_flags + positions["plane_flags"] = plane_flags + + scale = np.full(n, np.iinfo(int).max, dtype=int) + np.minimum.at(scale, inverse, candidates["scale"]) + positions["scale"] = scale + peak_sigma = np.full(n, -np.inf) + np.maximum.at(peak_sigma, inverse, candidates["flux"]) + positions["peak_sigma"] = peak_sigma + + # flux is the brightest candidate at the finest scale of each position: + # sort by (position, scale, -flux) and take the first row per position. + order = np.lexsort((-candidates["flux"], candidates["scale"], inverse)) + sorted_inverse = inverse[order] + first = np.ones(len(order), dtype=bool) + first[1:] = sorted_inverse[1:] != sorted_inverse[:-1] + positions["flux"][sorted_inverse[first]] = candidates["flux"][order][first] + return positions + + +def _link_radius(scale: int | np.ndarray, psf_fwhm: float) -> float | np.ndarray: + """Linking radius for a position at a given starlet scale. + + Parameters + ---------- + scale : + The absolute starlet scale of the position (scalar or array). + psf_fwhm : + The PSF FWHM in pixels, used as a floor at the finest scales. + + Returns + ------- + radius : + The radius in pixels, ``max(psf_fwhm, 2**scale)``. + """ + return np.maximum(psf_fwhm, 2.0 ** np.asarray(scale, dtype=float)) + + +def _starlet_kernel_fwhm(psf_fwhm: float | np.ndarray, scale: int | np.ndarray) -> float | np.ndarray: + """Effective FWHM of a point source at a starlet scale. + + Parameters + ---------- + psf_fwhm : + The PSF FWHM in pixels (scalar or array). + scale : + The absolute starlet scale (scalar or array). + + Returns + ------- + fwhm : + The FWHM in pixels of the PSF convolved with the smoothing kernel that + precedes ``scale``. + + Notes + ----- + Scale ``j`` is the difference between the image smoothed ``j`` times and + its further smoothing. Pass ``i`` of the B3 spline has a variance of + ``4**i`` per axis, so the smoothing before scale ``j`` has a variance of + ``(4**j - 1) / 3``. Both kernels are treated as Gaussians and added in + quadrature. + """ + kernel_variance = (4.0 ** np.asarray(scale, dtype=float) - 1) / 3 + return np.sqrt(np.asarray(psf_fwhm, dtype=float) ** 2 + SIGMA_TO_FWHM**2 * kernel_variance) + + +def _build_pairs(positions: np.ndarray, psf_fwhm: float, radius_rule: str) -> np.ndarray: + """Find every pair of positions within the largest link radius. + + Parameters + ---------- + positions : + The unique positions with dtype `POSITION_DTYPE`. + psf_fwhm : + The PSF FWHM in pixels, passed to ``_link_radius``. + radius_rule : + Which link radius of a pair bounds its separation, a key of + `RADIUS_RULES`. + + Returns + ------- + pairs : + The pairs with dtype `PAIR_DTYPE`, sorted by ``(i, j)`` with + ``i < j``. + + Notes + ----- + Two positions are linked if they are within the link radius and share no + plane (starlet scale and band). Positions that were both peaks in one + plane were separated by that plane's contrast-limited watershed, + so they are distinct sources however close they are. + """ + if len(positions) < 2: + return np.zeros(0, dtype=PAIR_DTYPE) + yx = np.column_stack([positions["y"], positions["x"]]).astype(float) + radii = cast(np.ndarray, _link_radius(positions["scale"], psf_fwhm)) + ij = cKDTree(yx).query_pairs(float(radii.max()), output_type="ndarray") + ij = ij[np.lexsort((ij[:, 1], ij[:, 0]))] + i, j = ij[:, 0], ij[:, 1] + flags = positions["plane_flags"] + + pairs = np.zeros(len(ij), dtype=PAIR_DTYPE) + pairs["i"] = i + pairs["j"] = j + pairs["distance"] = np.hypot(yx[i, 0] - yx[j, 0], yx[i, 1] - yx[j, 1]) + pairs["shares_plane"] = (flags[i] & flags[j]) != 0 + pairs["within_min"] = pairs["distance"] <= np.minimum(radii[i], radii[j]) + pairs["within_max"] = pairs["distance"] <= np.maximum(radii[i], radii[j]) + pairs["linked"] = pairs[RADIUS_RULES[radius_rule]] & ~pairs["shares_plane"] + return pairs + + +def _full_cliques(neighbors: list[int], clique: list[int], sizes: list[int]) -> list[int]: + """Cliques that a node is adjacent to every member of, in ascending order. + + Parameters + ---------- + neighbors : + The neighbors of the node. + clique : + The clique of every node, ``-1`` if unassigned. + sizes : + The number of members of each clique. + + Returns + ------- + cliques : + The cliques the node could join. + """ + counts: dict[int, int] = {} + for u in neighbors: + g = clique[u] + if g >= 0: + counts[g] = counts.get(g, 0) + 1 + return sorted(g for g, count in counts.items() if count == sizes[g]) + + +def _greedy_independent_set(adjacency: list[list[int]], order: list[int]) -> int: + """Size of a greedy independent set, visiting nodes in ``order``. + + Parameters + ---------- + adjacency : + The neighbors of each node. + order : + The order in which nodes are offered to the set. + + Returns + ------- + size : + The number of nodes in the set. + + Notes + ----- + Two adjacent nodes (positions) can never be in the same clique (peak), + so the size of any independent set is a lower bound on the number of + cliques. + """ + chosen = [False] * len(adjacency) + for v in order: + if not any(chosen[u] for u in adjacency[v]): + chosen[v] = True + return sum(chosen) + + +def _greedy_cover(adjacency: list[list[int]], order: list[int]) -> list[int]: + """Cover a graph with cliques greedily, visiting nodes in ``order``. + + Parameters + ---------- + adjacency : + The neighbors of each node. + order : + The order in which nodes are assigned. + + Returns + ------- + clique : + The clique of each node. Each node joins the first clique it is + adjacent to every member of, or opens a new one. + """ + clique = [-1] * len(adjacency) + sizes: list[int] = [] + for v in order: + eligible = _full_cliques(adjacency[v], clique, sizes) + if eligible: + clique[v] = eligible[0] + sizes[eligible[0]] += 1 + else: + clique[v] = len(sizes) + sizes.append(1) + return clique + + +def _search_covers( + adjacency: list[list[int]], + order: list[int], + max_cliques: int, + max_visits: float, + max_covers: int, +) -> tuple[list[list[int]], int, bool]: + """Search for covers of a graph by at most ``k`` cliques. + + Parameters + ---------- + adjacency : + The neighbors of each node. + order : + The order in which nodes are assigned. + max_cliques : + The maximum number of cliques. + max_visits : + Stop after visiting this many nodes. + max_covers : + Stop after finding this many covers. + + Returns + ------- + covers : + The clique of each node, for each cover found. + n_visited : + The number of nodes visited. + cap_hit : + `True` if the search stopped at ``max_visits`` before finishing. + + Notes + ----- + A depth-first search that assigns the nodes in ``order``: each node may + join any clique it is adjacent to every member of, or open a new one + while fewer than ``max_cliques`` exist. Cliques are numbered in the order + they are opened, so every partition is reached by exactly one path. + The stack is explicit because components can be far deeper than the + recursion limit. + """ + n = len(order) + clique = [-1] * n + sizes: list[int] = [] + covers: list[list[int]] = [] + + def options(v: int) -> list[int]: + eligible = _full_cliques(adjacency[v], clique, sizes) + if len(sizes) < max_cliques: + eligible.append(len(sizes)) + # Options are popped from the end, so the oldest clique is tried first. + return eligible[::-1] + + def unassign(v: int) -> None: + g = clique[v] + clique[v] = -1 + sizes[g] -= 1 + if sizes[g] == 0: + # Only the newest clique can empty, since later ones were opened + # deeper in the search and have already been unwound. + sizes.pop() + + stack = [options(order[0])] + n_visited = 1 + while stack: + depth = len(stack) - 1 + if not stack[-1]: + stack.pop() + if depth > 0: + unassign(order[depth - 1]) + continue + v = order[depth] + g = stack[-1].pop() + if g == len(sizes): + sizes.append(1) + else: + sizes[g] += 1 + clique[v] = g + if depth + 1 == n: + covers.append(clique.copy()) + if len(covers) >= max_covers: + return covers, n_visited, False + unassign(v) + continue + if n_visited >= max_visits: + return covers, n_visited, True + n_visited += 1 + stack.append(options(order[depth + 1])) + return covers, n_visited, False + + +def _cover_component( + adjacency: list[list[int]], + plane_bound: int, + search_cap: int, + count_covers: bool, +) -> tuple[list[int], dict[str, int | bool]]: + """Find a minimum clique cover of one connected component. + + Parameters + ---------- + adjacency : + The neighbors of each node. + plane_bound : + The largest number of nodes that share a plane. They are pairwise + unlinked, so this is a lower bound on the cover. + search_cap : + The number of nodes the exact search may visit, summed over every + ``k`` it tries. + count_covers : + Count the distinct minimum covers of a small component. + This is a diagnostic that should always be ``False`` in production. + + Returns + ------- + clique : + The clique of each node. + stats : + The values of the `COMPONENT_DTYPE` fields. + + Notes + ----- + Finding the true clique cover is an NP-hard problem, so this function uses + a combination of greedy heuristics and exact search within a limited search + cap to find a minimum cover for small components. + """ + n = len(adjacency) + order = sorted(range(n), key=lambda v: (len(adjacency[v]), v)) + # The independent set is a lower bound on the number of cliques, + # as no two nodes in an independent set can be in the same clique. + # However, it is dependent on the order of nodes considered and is + # not guaranteed to be the maximum. + # It's a cheap way to get a lower bound that might be higher than the + # number of components. + lower = max(_greedy_independent_set(adjacency, order), plane_bound) + + # Estimate the upper bound on the number of cliques using a greedy cover. + clique = _greedy_cover(adjacency, order) + upper = max(clique) + 1 + max_cliques = upper + n_visited = 0 + cap_hit = False + for trial in range(lower, upper): + if n_visited >= search_cap: + cap_hit = True + break + # Try an exact search for a minimum clique cover + covers, visited, cap_hit = _search_covers(adjacency, order, trial, search_cap - n_visited, 1) + n_visited += visited + if covers: + # A valid cover was found, so update the clique and the maximum + # number of cliques. + clique = covers[0] + max_cliques = trial + break + if cap_hit: + break + n_min_covers = -1 + if count_covers and n <= COUNT_COVERS_MAX_NODES and not cap_hit: + + covers, _, _ = _search_covers(adjacency, order, max_cliques, np.inf, COUNT_COVERS_MAX) + n_min_covers = len(covers) + stats: dict[str, int | bool] = { + "lower_bound": lower, + "upper_bound": upper, + "k": max_cliques, + "search_ran": lower < upper, + "n_visited": n_visited, + "cap_hit": cap_hit, + "n_min_covers": n_min_covers, + } + return clique, stats + + +def _max_shared_plane(plane_flags: np.ndarray, component: np.ndarray, n_components: int) -> np.ndarray: + """Largest number of positions in each component that share a plane. + + Parameters + ---------- + plane_flags : + The plane bitmask of each position. + component : + The component of each position. + n_components : + The number of components. + + Returns + ------- + count : + The count for each component. + """ + flags = plane_flags.astype(np.int64) + count = np.zeros(n_components, dtype=int) + n_bits = int(flags.max()).bit_length() if len(flags) else 0 + for bit in range(n_bits): + has_plane = ((flags >> bit) & 1).astype(bool) + count = np.maximum(count, np.bincount(component[has_plane], minlength=n_components)) + return count + + +def _clique_cover( + plane_flags: np.ndarray, + edges: np.ndarray, + search_cap: int = 100_000, + count_covers: bool = False, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Partition the position graph into a minimum number of cliques. + + Parameters + ---------- + plane_flags : + The plane bitmask of each position, used for a lower bound. + edges : + The linked pairs, with shape ``(n_edges, 2)``. + search_cap : + The number of nodes the exact search may visit in a component. + count_covers : + Count the distinct minimum covers of components with at most + `COUNT_COVERS_MAX_NODES` nodes, up to `COUNT_COVERS_MAX`. + + Returns + ------- + clique : + The clique (peak) of each position. Cliques are numbered in the + order of their first member. + component : + The connected component of each position, numbered in the + order of their first member. A component is a set of linked positions + that is split into one or more peaks (cliques). + components : + One row per component with dtype `COMPONENT_DTYPE`. + + Notes + ----- + A missing edge is evidence of two sources, so every peak is a clique and + the number of peaks is the size of the minimum clique cover. Each + connected component is covered on its own. A complete component is a + single clique. Otherwise nodes are visited fewest neighbors first (ties by + index) to build a greedy cover, an upper bound, and a greedy independent + set, which together with ``plane_flags`` gives a lower bound. When the + bounds differ, ``_search_covers`` tries each ``k`` from the lower bound + up. If it hits ``search_cap`` the greedy cover is kept and ``cap_hit`` is + set. + """ + n = len(plane_flags) + i, j = edges[:, 0], edges[:, 1] + # The adjacency array needs both directions for each edge, meaning we need + # i, j and j, i. We do this by concatenating them and switching the + # order. + # The format of a CSR array has three attributes: + # - `data`: the non-zero values of the array. In this case all True. + # - `indices`: the column indices of the non-zero elements. These are in + # the same order as the `data` array, arranged row by row, and by + # increasing column within each row (after `sort_indices`). + # - `indptr`: points to the start of each row in the `indices` and `data` + # arrays. It has one entry per position plus a final entry marking the + # end of the last row, so the neighbors of position ``v`` are + # ``indices[indptr[v]:indptr[v + 1]]``. A row with no non-zero + # elements has ``indptr[v] == indptr[v + 1]``. + adjacency = csr_array( + (np.ones(2 * len(edges), dtype=bool), (np.concatenate([i, j]), np.concatenate([j, i]))), + shape=(n, n), + ) + adjacency.sort_indices() + indptr, indices = adjacency.indptr, adjacency.indices + + # Split the graph into components. + # Each component is at a minimum a single peak. + n_components, component = connected_components(adjacency, directed=False) + # Count the number of nodes and edges in each component. + sizes = np.bincount(component, minlength=n_components) + n_edges = np.bincount(component[i], minlength=n_components) + # The largest number of positions in each component that share a single + # (scale, band) plane. Positions in the same plane are always distinct, + # so this gives us a lower bound on the number of peaks (cliques) in the + # component. + plane_bound = _max_shared_plane(plane_flags, component, n_components) + + components = np.zeros(n_components, dtype=COMPONENT_DTYPE) + components["n_nodes"] = sizes + complete = n_edges == sizes * (sizes - 1) // 2 + components["lower_bound"] = 1 + components["upper_bound"] = 1 + # Number of cliques (peaks) that each component is divided into. + components["k"] = 1 + components["n_min_covers"] = -1 + if count_covers: + components["n_min_covers"][complete & (sizes <= COUNT_COVERS_MAX_NODES)] = 1 + + # Nodes (positions) grouped by component, in index order within each one. + by_component = np.argsort(component, kind="stable") + # Starts gives the starting index of each component in the sorted array + # of nodes, so component `c`'s nodes are + # `by_component[starts[c] : starts[c + 1]]`. + starts = np.concatenate([[0], np.cumsum(sizes)]) + local = np.empty(n, dtype=int) + # The local index of each node (candidate) within its component. + local[by_component] = np.arange(n) - starts[component[by_component]] + # The local clique index of each node within its component. + local_clique = np.zeros(n, dtype=int) + for c in np.flatnonzero(~complete): + nodes = by_component[starts[c] : starts[c + 1]] + # The local indices of the neighbors of each node + # These are the other positions that can potentially be merged + # with the node. + neighbors = [local[indices[indptr[v] : indptr[v + 1]]].tolist() for v in nodes] + # Subdivide the component into cliques (peaks), where each clique is + # a set of nodes all connected to each other. + component_clique, stats = _cover_component(neighbors, int(plane_bound[c]), search_cap, count_covers) + local_clique[nodes] = component_clique + for name, value in stats.items(): + components[name][c] = value + + # The starting offset of each component's cliques in the raw numbering, + # where cliques are numbered component by component. + offsets = np.cumsum(components["k"]) - components["k"] + # A unique but not yet ordered clique label for each position + raw_position_cliques = offsets[component] + local_clique + n_cliques = int(components["k"].sum()) + # Determine the first occurrence of each clique in the global ordering. + first = np.full(n_cliques, n, dtype=int) + np.minimum.at(first, raw_position_cliques, np.arange(n)) + # The global index of each clique, determined by the first occurrence + # of any of its nodes. + global_clique = np.empty(n_cliques, dtype=int) + global_clique[np.argsort(first)] = np.arange(n_cliques) + # The clique that each position (node) belongs to. + clique = global_clique[raw_position_cliques] + return clique, component, components + + +def _link_fractions( + clique: np.ndarray, edges: np.ndarray, n_cliques: int +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Measure how strongly each position links to the other cliques. + + Parameters + ---------- + clique : + The clique of each position. + edges : + The linked pairs, with shape ``(n_edges, 2)``. + n_cliques : + The number of cliques. + + Returns + ------- + n_linked : + The number of cliques the position is adjacent to every member of, + counting its own. + ambiguous : + `True` where ``n_linked > 1``: the position could move to another + clique without breaking it. + max_fraction : + The largest fraction of another clique's members the position is + adjacent to, ``0`` if none. + max_clique : + The clique of ``max_fraction`` (the lowest id on a tie), ``-1`` if + none. + """ + n = len(clique) + n_linked = np.ones(n, dtype=int) + max_fraction = np.zeros(n) + max_clique = np.full(n, -1, dtype=int) + # Edges are not directed and only stored once. + # We concatenate to give us both directions. + # Note that positions with multiple edges will appear multiple times. + node = np.concatenate([edges[:, 0], edges[:, 1]]).astype(np.int64) + # We swap the edge order in the concatenate, so each node is paired + # with its linked partner. `other` is the clique of the partner node, + # which tells us which cliques each node's partners belong to. + other = clique[np.concatenate([edges[:, 1], edges[:, 0]])] + # If the node and its partner are in the same clique, + # the link carries no information so we ignore those. + keep = other != clique[node] + # Pack each (position, foreign clique) pair into one integer. + # This works because other < n_cliques. + # np.unique counts how many members of that clique the position is + # linked to. + key, count = np.unique(node[keep] * n_cliques + other[keep], return_counts=True) + if len(key) == 0: + # This happens if each component has only a single clique (peak) + # or if there are no edges at all. + return n_linked, n_linked > 1, max_fraction, max_clique + # Unpack the position and its partner's clique from the key. + node, other = np.divmod(key, n_cliques) + # The number of members in the partner's clique of each pair. + size = np.bincount(clique, minlength=n_cliques)[other] + fraction = count / size + # A position linked to every member of a foreign clique could join + # that clique. + n_linked += np.bincount(node[count == size], minlength=n) + + # lexsort uses its last key first: group by position, then descending + # fraction, then the lowest clique id on a tie. + order = np.lexsort((other, -fraction, node)) + # The first row of each position's run is its strongest foreign clique. + first = np.ones(len(order), dtype=bool) + first[1:] = node[order[1:]] != node[order[:-1]] + best = order[first] + # Positions with no foreign links are absent from node and keep the + # defaults. + max_fraction[node[best]] = fraction[best] + max_clique[node[best]] = other[best] + return n_linked, n_linked > 1, max_fraction, max_clique + + +def _clique_median(values: np.ndarray, cliques: np.ndarray, n_cliques: int) -> np.ndarray: + """Median of ``values`` within each clique. + + Parameters + ---------- + values : + The values. + cliques : + The clique of each value. + n_cliques : + The number of cliques. + + Returns + ------- + median : + The median of each clique, NaN for an empty clique. + """ + # Sort by clique first, then by value, so the values of each clique + # form one contiguous sorted run. + ordered = values[np.lexsort((values, cliques))].astype(float) + counts = np.bincount(cliques, minlength=n_cliques) + # The offset in `ordered` where each clique's run begins. + starts = np.cumsum(counts) - counts + median = np.full(n_cliques, np.nan) + has = counts > 0 + # The two middle elements of each run. For an odd count they are the + # same element, so averaging them is the usual median either way. + lower = starts[has] + (counts[has] - 1) // 2 + upper = starts[has] + counts[has] // 2 + median[has] = 0.5 * (ordered[lower] + ordered[upper]) + return median + + +def _unambiguous_members(positions: np.ndarray, n_peaks: int) -> tuple[np.ndarray, np.ndarray]: + """Select the positions that describe each peak. + + Parameters + ---------- + positions : + The positions with ``peak`` and ``ambiguous`` filled. + n_peaks : + The number of peaks. + + Returns + ------- + use : + `True` for the unambiguous positions, and for every position of a + peak that has none. + fallback : + `True` for each peak that has no unambiguous position. + """ + peak = positions["peak"] + clear = ~positions["ambiguous"] + # Peaks that have no clear positions get a fallback. + fallback = np.bincount(peak[clear], minlength=n_peaks) == 0 + return clear | fallback[peak], fallback + + +def _locate_peaks( + positions: np.ndarray, + candidates: np.ndarray, + use: np.ndarray, + n_peaks: int, + psf_fwhm: np.ndarray, +) -> np.ndarray: + """Choose the position that locates each peak. + + Parameters + ---------- + positions : + The positions with ``peak`` filled. + candidates : + The candidates with ``position`` filled. + use : + The positions a peak may be located at, from + ``_unambiguous_members``. + n_peaks : + The number of peaks. + psf_fwhm : + The PSF FWHM of each single band, in pixels. + + Returns + ------- + location : + The index of the chosen position for each peak. + + Notes + ----- + The chosen position is the one closest to the median ``x`` and ``y`` of + the peak's candidates, so a position counts once per candidate. Ties go to + the finest scale, then the smallest PSF FWHM among the bands of the + position's finest-scale candidates (the chi^2 plane ranks last), then the + lowest position index. + """ + position_of = candidates["position"] + # Work at the candidate level so a position detected in several planes + # pulls the median toward itself once per plane. + in_use = use[position_of] + candidate_peak = positions["peak"][position_of[in_use]] + median_y = _clique_median(candidates["y"][in_use], candidate_peak, n_peaks) + median_x = _clique_median(candidates["x"][in_use], candidate_peak, n_peaks) + + # A position's scale is the finest of its candidates, so this selects the + # candidates at that scale. + finest = candidates["scale"] == positions["scale"][position_of] + # Map a candidate's band to its PSF FWHM. The chi^2 plane is band index + # n_bands and maps to inf, so it ranks after every single band. + band_rank = np.append(psf_fwhm, np.inf) + # The smallest PSF FWHM among the finest-scale candidates of each + # position. + finest_fwhm = np.full(len(positions), np.inf) + np.minimum.at(finest_fwhm, position_of[finest], band_rank[candidates["band"][finest]]) + + index = np.flatnonzero(use) + peak = positions["peak"][index] + # The medians are multiples of 0.5, so distances that tie compare equal. + distance2 = (positions["y"][index] - median_y[peak]) ** 2 + (positions["x"][index] - median_x[peak]) ** 2 + # lexsort uses its last key first: group by peak, then the tie breakers + # in the order listed in the Notes. + order = np.lexsort((index, finest_fwhm[index], positions["scale"][index], distance2, peak)) + # The first row of each peak's run is its chosen position. Every peak + # has at least one row because _unambiguous_members never leaves a peak + # without a usable position. + sorted_peak = peak[order] + first = np.ones(len(order), dtype=bool) + first[1:] = sorted_peak[1:] != sorted_peak[:-1] + location = np.empty(n_peaks, dtype=int) + location[sorted_peak[first]] = index[order][first] + return location + + +def _fill_position_errors( + peaks: np.ndarray, + positions: np.ndarray, + candidates: np.ndarray, + use: np.ndarray, + psf_fwhm: np.ndarray, +) -> None: + """Fill the position error fields of each peak. + + Parameters + ---------- + peaks : + The detections with ``y`` and ``x`` set. The error fields are filled + in place. + positions : + The positions with ``peak`` filled. + candidates : + The candidates with ``position`` filled. + use : + The positions that describe each peak, from ``_unambiguous_members``. + psf_fwhm : + The PSF FWHM of each single band, in pixels. + + Notes + ----- + Every term is computed over the candidates of the ``use`` positions. The + statistical error of a candidate is its effective width from + ``_starlet_kernel_fwhm`` over ``2.355 * significance``. Scales within a + band are not independent, so each band contributes its best candidate and + the bands are combined by inverse variance. The chi^2 plane is a + combination of the bands and is used only for a peak with no single-band + candidate. The provisional error is the largest of the statistical error, + the robust scatter, and the ``1 / sqrt(12)`` quantization floor. + """ + n_peaks = len(peaks) + n_bands = len(psf_fwhm) + position_of = candidates["position"] + keep = use[position_of] + peak = positions["peak"][position_of[keep]] + band = candidates["band"][keep] + scale = candidates["scale"][keep] + y = candidates["y"][keep].astype(float) + x = candidates["x"][keep].astype(float) + # The chi^2 plane takes the widest band's PSF. + plane_fwhm = np.append(psf_fwhm, psf_fwhm.max()) + + width = _starlet_kernel_fwhm(plane_fwhm[band], scale) + sigma = width / (SIGMA_TO_FWHM * candidates["flux"][keep]) + best = np.full((n_peaks, n_bands + 1), np.inf) + np.minimum.at(best, (peak, band), sigma) + inverse_variance = np.sum(best[:, :n_bands] ** -2.0, axis=1) + stat = best[:, n_bands].copy() + has_band = inverse_variance > 0 + stat[has_band] = inverse_variance[has_band] ** -0.5 + peaks["stat_err_y"] = stat + peaks["stat_err_x"] = stat + + count = np.bincount(peak, minlength=n_peaks) + dy = y - peaks["y"][peak] + dx = x - peaks["x"][peak] + peaks["moment_yy"] = np.bincount(peak, weights=dy * dy, minlength=n_peaks) / count + peaks["moment_xx"] = np.bincount(peak, weights=dx * dx, minlength=n_peaks) / count + peaks["moment_xy"] = np.bincount(peak, weights=dx * dy, minlength=n_peaks) / count + + for axis, values in (("y", y), ("x", x)): + median = _clique_median(values, peak, n_peaks) + scatter = MAD_TO_SIGMA * _clique_median(np.abs(values - median[peak]), peak, n_peaks) + scatter[count < 3] = np.nan + peaks[f"scatter_{axis}"] = scatter + # fmax ignores the NaN scatter of a peak with too few candidates. + peaks[f"err_{axis}"] = np.fmax(np.fmax(stat, scatter), 1 / np.sqrt(12)) + + nearest = np.full(n_peaks, np.inf) + if n_peaks > 1: + yx = np.column_stack([peaks["y"], peaks["x"]]).astype(float) + nearest = cKDTree(yx).query(yx, k=2)[0][:, 1] + min_fwhm = np.full(n_peaks, np.inf) + np.minimum.at(min_fwhm, peak, plane_fwhm[band]) + peaks["nearest_distance"] = nearest + peaks["nearest_distance_fwhm"] = nearest / min_fwhm + + +def _build_detections( + positions: np.ndarray, + candidates: np.ndarray, + components: np.ndarray, + psf_fwhm: np.ndarray, +) -> np.ndarray: + """Reduce grouped positions to one detection record per peak. + + Parameters + ---------- + positions : + The positions with the grouping fields filled by ``_group_peaks``. + candidates : + The candidates with ``position`` filled. + components : + The components from ``_clique_cover``. + psf_fwhm : + The PSF FWHM of each single band, in pixels. + + Returns + ------- + peaks : + The detections with dtype `DETECTION_DTYPE`. The location and + ``flux`` come from the position chosen by ``_locate_peaks``; + ``scale``, ``peak_sigma`` and the flags are reductions over the + unambiguous positions (all positions for a peak with none), while the + counts include every position. + """ + n_peaks = int(positions["peak"].max()) + 1 if len(positions) else 0 + peaks = np.zeros(n_peaks, dtype=DETECTION_DTYPE) + if n_peaks == 0: + return peaks + peak_of = positions["peak"] + # To determine the location of a peak we only use positions that are + # unambiguously contained in its clique. + use, fallback = _unambiguous_members(positions, n_peaks) + location = _locate_peaks(positions, candidates, use, n_peaks, psf_fwhm) + members = positions[use] + member_peak = peak_of[use] + + peaks["y"] = positions["y"][location] + peaks["x"] = positions["x"][location] + peaks["flux"] = positions["flux"][location] + scale = np.full(n_peaks, np.iinfo(int).max, dtype=int) + np.minimum.at(scale, member_peak, members["scale"]) + peaks["scale"] = scale + peak_sigma = np.full(n_peaks, -np.inf) + np.maximum.at(peak_sigma, member_peak, members["peak_sigma"]) + peaks["peak_sigma"] = peak_sigma + band_flags = np.zeros(n_peaks, dtype=np.int64) + np.bitwise_or.at(band_flags, member_peak, members["band_flags"].astype(np.int64)) + peaks["band_flags"] = band_flags + scale_flags = np.zeros(n_peaks, dtype=np.int64) + np.bitwise_or.at(scale_flags, member_peak, members["scale_flags"].astype(np.int64)) + peaks["scale_flags"] = scale_flags + + peaks["n_candidates"] = np.bincount(peak_of, weights=positions["n_candidates"], minlength=n_peaks) + peaks["n_positions"] = np.bincount(peak_of, minlength=n_peaks) + peaks["n_ambiguous"] = np.bincount(peak_of, weights=positions["ambiguous"], minlength=n_peaks) + peaks["location_fallback"] = fallback + peaks["search_capped"] = components["cap_hit"][positions["component"][location]] + _fill_position_errors(peaks, positions, candidates, use, psf_fwhm) + return peaks + + +def _build_footprints( + footprint_mask: np.ndarray, + peaks: np.ndarray, + origin: tuple[int, int] = (0, 0), +) -> list[Footprint]: + """Split the footprint mask into footprints and place the peaks in them. + + Parameters + ---------- + footprint_mask : + The union of the per-plane footprints from ``_find_peak_candidates``. + peaks : + The detections with dtype `DETECTION_DTYPE`. The ``footprint`` field + is filled in place. + origin : + The ``(y, x)`` location of the lower corner of ``footprint_mask``. + + Returns + ------- + footprints : + The connected components of ``footprint_mask`` that contain at least + one peak, each holding a `Peak` per detection inside it, brightest + first. The `Peak` flux is the detection's ``peak_sigma``. + + Raises + ------ + RuntimeError + Raised if a peak lies outside every footprint. Every peak is located + at a candidate inside one of the per-plane footprints, so this + indicates a bug upstream. + + Notes + ----- + Per-plane footprints of the same source differ in extent (and may be + split or merged) from one plane to the next, so a footprint is identified + with a source only after the peaks are grouped: the mask is relabeled into + 4-connected components and each peak is looked up in the component that + contains it. A component whose candidates were all grouped with a peak + located in another component ends up with no peak and is dropped, so + ``footprints`` never holds an empty footprint. + """ + y0, x0 = origin + if len(peaks) == 0: + return [] + components = get_footprints( + footprint_mask.astype(np.float32), + min_separation=0, + min_area=1, + peak_thresh=0, + footprint_thresh=0.5, + find_peaks=False, + y0=y0, + x0=x0, + ) + bbox = Box(footprint_mask.shape, origin=origin) + # Labels are the component index plus one, zero outside every footprint. + labels = footprints_to_image(components, bbox).data + component_of = labels[peaks["y"] - y0, peaks["x"] - x0] - 1 + if np.any(component_of < 0): + raise RuntimeError("A detected peak lies outside every footprint") + + # Brightest first within a footprint, matching get_footprints. + order = np.lexsort((-peaks["peak_sigma"], component_of)) + footprints: list[Footprint] = [] + footprint_of = np.empty(len(peaks), dtype=int) + last = -1 + for row in order: + component = component_of[row] + if component != last: + footprints.append(components[component]) + last = component + footprint_of[row] = len(footprints) - 1 + footprints[-1].add_peak( + Peak(int(peaks["y"][row]), int(peaks["x"][row]), float(peaks["peak_sigma"][row])) + ) + peaks["footprint"] = footprint_of + return footprints + + +def _group_peaks( + candidates: np.ndarray, + n_bands: int, + first_scale: int, + psf_fwhm: np.ndarray, + radius_rule: str = "min", + search_cap: int = 100_000, + count_covers: bool = False, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Collapse candidates to positions and group positions into detections. + + Parameters + ---------- + candidates : + Structured array of candidates with dtype `CANDIDATE_DTYPE`. The + ``position`` field is filled in place. + n_bands : + The number of single-band planes. + first_scale : + The first starlet scale in the significance map. + psf_fwhm : + The PSF FWHM of each single band, in pixels. The largest sets the link + radius. + radius_rule : + Passed to ``_build_pairs``. + search_cap : + Passed to ``_clique_cover``. + count_covers : + Passed to ``_clique_cover``. + + Returns + ------- + positions : + The unique positions with dtype `POSITION_DTYPE`. + peaks : + The detections with dtype `DETECTION_DTYPE`. + pairs : + The position pairs with dtype `PAIR_DTYPE`. + components : + The connected components of the position graph with dtype + `COMPONENT_DTYPE`. + """ + if len(candidates) == 0: + return ( + np.empty(0, dtype=POSITION_DTYPE), + np.empty(0, dtype=DETECTION_DTYPE), + np.empty(0, dtype=PAIR_DTYPE), + np.empty(0, dtype=COMPONENT_DTYPE), + ) + positions = _collapse_positions(candidates, n_bands, first_scale) + pairs = _build_pairs(positions, float(np.max(psf_fwhm)), radius_rule) + linked = pairs[pairs["linked"]] + edges = np.column_stack([linked["i"], linked["j"]]) + peak, component, components = _clique_cover(positions["plane_flags"], edges, search_cap, count_covers) + positions["peak"] = peak + positions["component"] = component + n_linked, ambiguous, max_fraction, max_peak = _link_fractions(peak, edges, int(peak.max()) + 1) + positions["n_linked"] = n_linked + positions["ambiguous"] = ambiguous + positions["max_linked_fraction"] = max_fraction + positions["max_linked_peak"] = max_peak + peaks = _build_detections(positions, candidates, components, psf_fwhm) + return positions, peaks, pairs, components + + +def find_structures( + peaks: np.ndarray, + significance: np.ndarray, + origin: tuple[int, int] = (0, 0), + structure_thresh: float = 1.0, +) -> None: + """Label peaks by the connected region of the finest scale they share. + + Parameters + ---------- + peaks : + The detections with dtype `DETECTION_DTYPE`, all inside + ``significance``. The ``structure_peaks`` and ``structure_id`` + fields are filled in place. + significance : + The chi coadd plane of the finest starlet scale searched, with shape + ``(Ny, Nx)`` and in units of Gaussian sigma. + origin : + The ``(y, x)`` location of the lower corner of ``significance``. + structure_thresh : + The threshold in sigma above which pixels of ``significance`` are + connected into structures. + + Notes + ----- + Spiral arms and star forming regions of a resolved galaxy are detected + as separate peaks, which the deblender wants but a catalog user may not. + In the chi coadd of the finest starlet scale the pieces of a galaxy are + connected at a low threshold into one region, so the peaks of a region + are labeled with a common ``structure_id`` (counting from one, in raster + order of the region's first pixel) and the number of peaks it holds in + ``structure_peaks``. A region of three or more peaks is nearly always a + shredded galaxy, whereas a region of two is usually a pair of blended + stars, so ``structure_peaks`` is left for the user to cut on. + + A peak found only in one band, or only at a coarser scale, can sit on a + pixel below ``structure_thresh`` and belongs to no region: its + ``structure_peaks`` and ``structure_id`` are both zero. + """ + y0, x0 = origin + peaks["structure_peaks"] = 0 + peaks["structure_id"] = 0 + if len(peaks) == 0: + return + labels, n_labels = label(significance >= structure_thresh) + label_of = labels[peaks["y"] - y0, peaks["x"] - x0] + used = np.unique(label_of[label_of > 0]) + # Dense ids for the labels that hold a peak, zero for the rest. + structure_of = np.zeros(n_labels + 1, dtype=int) + structure_of[used] = np.arange(1, len(used) + 1) + structure_id = structure_of[label_of] + counts = np.bincount(structure_id, minlength=len(used) + 1) + counts[0] = 0 + peaks["structure_id"] = structure_id + peaks["structure_peaks"] = counts[structure_id] + + +def detect_peaks( + images: np.ndarray, + variance: np.ndarray, + scales: int = 3, + generation: int = 2, + first_scale: int = 1, + origin: tuple[int, int] | None = None, + min_separation: float | None = None, + min_area: int = 4, + peak_thresh: float = 3, + footprint_thresh: float = 2, + psf_fwhm: float | Sequence[float] = 3.5, + kappa: float = 3.0, + variance_mode: str = "median", + bad_pixel_mask: np.ndarray | None = None, + dilation_radius: float = 0, + radius_rule: str = "min", + search_cap: int = SEARCH_CAP, + count_covers: bool = False, + structure_thresh: float = 1.0, +) -> PeakDetectionResult: + """Detect peaks across bands and starlet scales. + + Parameters + ---------- + images : + The multiband image with shape ``(n_bands, Ny, Nx)``. + variance : + The per-pixel variance, with the same shape as ``images``. + scales : + The maximum number of wavelet scales to use. + generation : + The generation of the starlet transform, either ``1`` or ``2``. + first_scale : + The first starlet scale to search for peaks. + origin : + The ``(y, x)`` location of the lower corner of the image. + min_separation : + A hard floor on the separation between peaks within a plane, in + pixels; ``0`` disables it and relies on ``kappa`` alone. + `None` uses the smallest PSF FWHM to calculate the minimum separation + based on two Gaussians of equal amplitude. + min_area : + The minimum area of a footprint in pixels. + peak_thresh : + The peak detection threshold, in sigma. + footprint_thresh : + The footprint detection threshold, in sigma. + psf_fwhm : + The PSF FWHM in pixels, either one value for every band or one per + band. The largest is the floor on the linking radius; the per-band + values break ties in the peak location and set the position errors. + kappa : + The minimum prominence in sigma of a peak above its saddle to a + brighter peak in the same plane; smaller bumps are culled. + variance_mode : + How ``variance`` is reduced to a coefficient noise, either ``"median"`` + for one value per band or ``"pixel"`` to follow the variance plane; see + ``_build_detection_starlets``. + bad_pixel_mask : + Per-band boolean mask with the same shape as ``images``, `True` for + pixels that are bad in that band. Candidate peaks that fall within + ``dilation_radius`` of a bad pixel in a given band are suppressed + in that band only. This allows eg. saturated sources in one band to + be detected in other, clean, bands. + `None` disables the suppression. + dilation_radius : + The base radius in pixels by which ``bad_pixel_mask`` is grown before + it suppresses peaks, to catch the starlet ringing just outside a bad + region. The radius doubles with each starlet scale, since the kernel + does; see ``_find_peak_candidates``. + radius_rule : + Whether two positions link within the smaller (``"min"``) or larger + (``"max"``) of their link radii; see ``_build_pairs``. + search_cap : + The number of nodes the exact clique cover search may visit in a + component before it falls back to the greedy cover; see + ``_clique_cover``. + count_covers : + Count the distinct minimum covers of small components; see + ``_clique_cover``. This is a diagnostic that should always be + ``False`` in production. + structure_thresh : + The threshold in sigma that connects the chi coadd of the first + scale into structures; see ``find_structures``. + + Returns + ------- + result : + The detected peaks, their footprints, and the intermediate detection + products. + + Raises + ------ + ValueError + Raised if ``bad_pixel_mask`` is given and does not match the shape of + ``images``, if ``psf_fwhm`` is a sequence without one value per + band, or if ``radius_rule`` is not a key of `RADIUS_RULES`. + """ + band_fwhm = np.asarray(psf_fwhm, dtype=float) + if band_fwhm.ndim == 0: + band_fwhm = np.full(len(images), float(band_fwhm)) + elif band_fwhm.shape != (len(images),): + raise ValueError(f"psf_fwhm must be a scalar or have one value per band, got {band_fwhm.shape}") + if radius_rule not in RADIUS_RULES: + raise ValueError(f"radius_rule must be one of {list(RADIUS_RULES)}, got {radius_rule!r}") + link_fwhm = float(band_fwhm.max()) + if min_separation is None: + min_separation = band_fwhm.min() * 0.84932 + if origin is None: + origin = (0, 0) + if bad_pixel_mask is not None and bad_pixel_mask.shape != images.shape: + raise ValueError("bad_pixel_mask must have the same shape as images") + starlets, sigma = _build_detection_starlets( + images, + variance, + scales=scales, + generation=generation, + variance_mode=variance_mode, + ) + significance_map = _build_significance_map(starlets, sigma, first_scale=first_scale) + candidates, footprint_mask = _find_peak_candidates( + significance_map, + min_separation, + min_area, + peak_thresh, + footprint_thresh, + kappa, + first_scale, + origin, + bad_pixel_mask, + dilation_radius, + ) + n_bands = significance_map.shape[1] - 1 + positions, peaks, pairs, components = _group_peaks( + candidates, n_bands, first_scale, band_fwhm, radius_rule, search_cap, count_covers + ) + footprints = _build_footprints(footprint_mask, peaks, origin) + # Only the finest scale searched separates the pieces of a galaxy from + # the sky while keeping them connected; coarser scales merge neighbors. + find_structures(peaks, significance_map[0, -1], origin, structure_thresh) + link_scales = np.arange(first_scale, first_scale + significance_map.shape[0]) + metadata = { + "psf_fwhm": band_fwhm.tolist(), + "link_radius_scales": link_scales.tolist(), + "link_radius": np.atleast_1d(_link_radius(link_scales, link_fwhm)).tolist(), + "min_separation": float(min_separation), + "radius_rule": radius_rule, + "search_cap": int(search_cap), + "structure_thresh": float(structure_thresh), + } + return PeakDetectionResult( + peaks=peaks, + footprints=footprints, + positions=positions, + candidates=candidates, + significance_map=significance_map, + starlets=starlets, + sigma=sigma, + pairs=pairs, + components=components, + metadata=metadata, + ) diff --git a/python/lsst/scarlet/lite/detect_pybind11.cc b/python/lsst/scarlet/lite/detect_pybind11.cc index e5fb0b5b..70ed106c 100644 --- a/python/lsst/scarlet/lite/detect_pybind11.cc +++ b/python/lsst/scarlet/lite/detect_pybind11.cc @@ -4,6 +4,7 @@ #include #include #include +#include #include #include #include @@ -129,33 +130,43 @@ MatrixB get_connected_multipeak( /// A Peak in a Footprint -/// This class is meant to keep track of both the position and -/// flux at the location of a maximum in a Footprint +/// This class keeps track of the position and flux at the location of a +/// local maximum in a Footprint, together with the `saddle`: the highest +/// level at which the peak's basin connects to the basin of a brighter peak +/// (NaN for the brightest peak in a footprint, which never joins a brighter +/// basin). `flux - saddle` is the peak's prominence, the quantity that +/// `kappa` thresholds in `get_peaks`. class Peak { public: - Peak(int y, int x, double flux){ + Peak(int y, int x, double flux, double saddle=std::numeric_limits::quiet_NaN()){ _y = y; _x = x; _flux = flux; + _saddle = saddle; } - int getY(){ + int getY() const { return _y; } - int getX(){ + int getX() const { return _x; } - double getFlux(){ + double getFlux() const { return _flux; } + double getSaddle() const { + return _saddle; + } + private: int _y; int _x; double _flux; + double _saddle; }; @@ -165,79 +176,182 @@ bool sortBrightness(Peak a, Peak b){ } -// Get a list of peaks found in an image. -// To make ut easier to cull peaks that are too close together -// and ensure that every footprint has at least one peak, -// this algorithm is meant to be run on a single footprint -// created by `get_connected_pixels`. +/// Find the root of a union-find tree with path compression. +inline int union_find(std::vector& parent, int i){ + int root = i; + while (parent[root] != root) { + root = parent[root]; + } + while (parent[i] != root) { + int next = parent[i]; + parent[i] = root; + i = next; + } + return root; +} + + +/** + * Find the peaks in a masked image with a contrast-limited watershed. + * + * Pixels inside `mask` are visited in decreasing order of value. A pixel + * with no visited 8-neighbor is a strict local maximum and seeds a new basin + * (and a peak, if it is at least `peak_thresh`). Otherwise it joins the + * basins of its visited neighbors. When a pixel connects two basins for the + * first time, its value is the saddle level between them: the basin with the + * lower maximum is absorbed into the other, and its peak is culled if its + * prominence above that saddle, `flux - saddle`, is less than `kappa`. Because + * pixels are processed in decreasing order every later saddle is lower, so a + * peak that survives its first test survives all of them and each peak is + * tested exactly once. + * + * With `kappa <= 0` nothing is culled and the result is the set of strict + * 8-connected local maxima above `peak_thresh`. + * + * The cost is O(N log N) in the number of masked pixels, independent of the + * number of peaks. + */ template -std::vector get_peaks( - M& image, +std::vector watershed_peaks( + const M& image, + const MatrixB& mask, const double min_separation, const double peak_thresh, + const double kappa, const int y0, const int x0 ){ const int height = image.rows(); const int width = image.cols(); - - std::vector peaks; - - for(int i=0; i> pixels; + pixels.reserve(mask.count()); + for (int i = 0; i < height; ++i) { + for (int j = 0; j < width; ++j) { + if (mask(i, j)) { + pixels.emplace_back(static_cast(image(i, j)), i * width + j); } - if(i > 0 && image(i, j) <= image(i-1, j)){ - continue; - } - if(i < height-1 && image(i,j) <= image(i+1, j)){ - continue; - } - if(j > 0 && image(i, j) <= image(i, j-1)){ - continue; - } - if(j < width-1 && image(i,j) <= image(i, j+1)){ - continue; + } + } + std::vector peaks; + if (pixels.empty()) { + return peaks; + } + std::sort(pixels.begin(), pixels.end(), [](const std::pair& a, const std::pair& b) { + return a.first > b.first || (a.first == b.first && a.second < b.second); + }); + + // Union-find over pixels. parent == -1 marks an unvisited pixel. + std::vector parent(n, -1); + // For each root, the pixel index of the basin's maximum. + std::vector comp_max(n, -1); + // For each pixel, the index into `candidates` of the peak it seeded, or -1. + std::vector peak_id(n, -1); + + struct Candidate { + int pixel; + double flux; + double saddle; + bool culled; + }; + std::vector candidates; + + auto value = [&](int idx) { return static_cast(image(idx / width, idx % width)); }; + + // Merge the basins rooted at idx_a and idx_b at saddle level pixel_value; returns the new root. + auto merge = [&](int idx_a, int idx_b, double pixel_value) { + const double value_a = value(comp_max[idx_a]); + const double value_b = value(comp_max[idx_b]); + int hi = idx_a, lo = idx_b; + if (value_b > value_a) { + hi = idx_b; + lo = idx_a; + } + const int pk = peak_id[comp_max[lo]]; + if (pk >= 0) { + Candidate& c = candidates[pk]; + c.saddle = pixel_value; + if (c.flux - pixel_value < kappa) { + c.culled = true; } - - if(i > 0 && j > 0 && image(i, j) <= image(i-1, j-1)){ + } + parent[lo] = hi; + return hi; + }; + + const int di[8] = {-1, -1, -1, 0, 0, 1, 1, 1}; + const int dj[8] = {-1, 0, 1, -1, 1, -1, 0, 1}; + + for (const auto& pv : pixels) { + const double pixel_value = pv.first; + const int pixel_index = pv.second; + const int i = pixel_index / width; + const int j = pixel_index % width; + + int root = -1; + for (int k = 0; k < 8; ++k) { + const int ni = i + di[k]; + const int nj = j + dj[k]; + if (ni < 0 || ni >= height || nj < 0 || nj >= width) { continue; } - if(i < height-1 && j < width-1 && image(i,j) <= image(i+1, j+1)){ - continue; + const int neighbor_index = ni * width + nj; + if (parent[neighbor_index] < 0) { + continue; // unvisited, i.e. lower than pixel_index (or masked out) } - if(i < height-1 && j > 0 && image(i, j) <= image(i+1, j-1)){ - continue; + const int root_index = union_find(parent, neighbor_index); + if (root < 0) { + // First visited neighbor: pixel_index joins its basin. + parent[pixel_index] = root_index; + root = root_index; + } else if (root_index != root) { + root = merge(root, root_index, pixel_value); } - if(i > 0 && j < width-1 && image(i,j) <= image(i-1, j+1)){ - continue; + } + if (root < 0) { + // Strict local maximum: seed a new basin. + parent[pixel_index] = pixel_index; + comp_max[pixel_index] = pixel_index; + if (pixel_value >= peak_thresh) { + peak_id[pixel_index] = static_cast(candidates.size()); + candidates.push_back({pixel_index, pixel_value, std::numeric_limits::quiet_NaN(), false}); } - - peaks.push_back(Peak(i+y0, j+x0, static_cast(image(i, j)))); } } - if(peaks.empty()){ + for (const Candidate& c : candidates) { + if (!c.culled) { + peaks.push_back(Peak(c.pixel / width + y0, c.pixel % width + x0, c.flux, c.saddle)); + } + } + if (peaks.empty()) { return peaks; } /// Sort the peaks in the footprint so that the brightest are first - std::sort (peaks.begin(), peaks.end(), sortBrightness); - - // Remove peaks within min_separation - double min_separation2 = min_separation * min_separation; - for (size_t i = 0; i < peaks.size() - 1; ++i) { - for (size_t j = i + 1; j < peaks.size();) { - Peak *p1 = &peaks[i]; - Peak *p2 = &peaks[j]; - double dy = p1->getY() - p2->getY(); - double dx = p1->getX() - p2->getX(); - double separation2 = dy*dy + dx*dx; - if (separation2 < min_separation2) { - peaks.erase(peaks.begin() + j); - } else { - ++j; + std::sort(peaks.begin(), peaks.end(), sortBrightness); + + // Remove peaks within min_separation of a brighter peak. This is a hard + // floor on top of the contrast test and is skipped when min_separation <= 0. + if (min_separation > 0) { + double min_separation2 = min_separation * min_separation; + for (size_t i = 0; i < peaks.size() - 1; ++i) { + for (size_t j = i + 1; j < peaks.size();) { + Peak *p1 = &peaks[i]; + Peak *p2 = &peaks[j]; + double dy = p1->getY() - p2->getY(); + double dx = p1->getX() - p2->getX(); + double separation2 = dy*dy + dx*dx; + if (separation2 < min_separation2) { + peaks.erase(peaks.begin() + j); + } else { + ++j; + } } } } @@ -245,6 +359,32 @@ std::vector get_peaks( } +// Get a list of peaks found in an image. +// This is meant to be run on a single footprint created by +// `get_connected_pixels`. Only pixels above `footprint_thresh` take part in +// the watershed; with the default (-inf) every pixel does. +template +std::vector get_peaks( + const M& image, + const double min_separation, + const double peak_thresh, + const int y0, + const int x0, + const double kappa=0.0, + const double footprint_thresh=-std::numeric_limits::infinity() +){ + const int height = image.rows(); + const int width = image.cols(); + MatrixB mask(height, width); + for (int i = 0; i < height; ++i) { + for (int j = 0; j < width; ++j) { + mask(i, j) = static_cast(image(i, j)) > footprint_thresh; + } + } + return watershed_peaks(image, mask, min_separation, peak_thresh, kappa, y0, x0); +} + + // A detected footprint class Footprint { public: @@ -274,23 +414,6 @@ class Footprint { }; -template -void maskImage( - py::EigenDRef image, - py::EigenDRef footprint -){ - const int height = image.rows(); - const int width = image.cols(); - - for(int i=0; i get_footprints( const double footprint_thresh, const bool find_peaks=true, const int y0=0, - const int x0=0 + const int x0=0, + const double kappa=0.0 ){ const int height = image.rows(); const int width = image.cols(); @@ -333,14 +460,17 @@ std::vector get_footprints( MatrixB subFootprint = footprint.block(bounds[0], bounds[2], subHeight, subWidth); int area = subFootprint.count(); if(area >= min_area){ - M patch = image.block(bounds[0], bounds[2], subHeight, subWidth); - maskImage(patch, subFootprint); std::vector _peaks; if(find_peaks){ - _peaks = get_peaks( + // The watershed only visits masked pixels, so the + // patch does not need to be zeroed outside the footprint. + M patch = image.block(bounds[0], bounds[2], subHeight, subWidth); + _peaks = watershed_peaks( patch, + subFootprint, min_separation, peak_thresh, + kappa, bounds[0] + y0, bounds[2] + x0 ); @@ -361,6 +491,84 @@ std::vector get_footprints( } +/** + * Build footprints around a supplied set of peaks + * + * Each peak seeds the connected component of pixels above + * `footprint_thresh` that contains it (4-connectivity, as in + * `get_connected_pixels`). Components not reached by any peak are + * ignored, and a component reached by several peaks yields one + * footprint holding all of them. The input peaks carry no watershed + * information, so each resulting `Peak` records the image value at its + * pixel and a NaN saddle. + * + * @param image: The image the peaks were detected on + * @param peaks: Peak positions as (y, x) pairs in image coordinates + * @param min_area: The minimum area of a footprint in pixels + * @param footprint_thresh: The minimum flux of a pixel to be included in a footprint + * @param y0: The y-coordinate of the top-left corner of the image + * @param x0: The x-coordinate of the top-left corner of the image + * + * @return: A list of Footprints, each containing at least one peak + */ +template +std::vector peaks_to_footprints( + py::EigenDRef image, + const std::vector>& peaks, + const int min_area, + const double footprint_thresh, + const int y0=0, + const int x0=0 +){ + const int height = image.rows(); + const int width = image.cols(); + + std::vector footprints; + MatrixB unchecked = MatrixB::Ones(height, width); + MatrixB footprint = MatrixB::Zero(height, width); + + for(const auto& peak : peaks){ + // The catalog is in image coordinates; shift into the local array. + const int i = peak[0] - y0; + const int j = peak[1] - x0; + + if (i < 0 || i >= height || j < 0 || j >= width) { + throw std::out_of_range("Peak coordinates (" + std::to_string(peak[0]) + ", " + + std::to_string(peak[1]) + ") are out of image bounds"); + } + + // A peak already swept into an earlier footprint seeds no new one. + if(!unchecked(i, j)){ + continue; + } + + Bounds bounds; bounds << i, i, j, j; + get_connected_pixels(i, j, image, unchecked, footprint, bounds, footprint_thresh); + int subHeight = bounds[1]-bounds[0]+1; + int subWidth = bounds[3]-bounds[2]+1; + MatrixB subFootprint = footprint.block(bounds[0], bounds[2], subHeight, subWidth); + int area = subFootprint.count(); + if(area >= min_area){ + // Gather every supplied peak that lands on this footprint. + std::vector _peaks; + for(const auto& other : peaks){ + const int py = other[0] - y0; + const int px = other[1] - x0; + if(py >= bounds[0] && py <= bounds[1] && px >= bounds[2] && px <= bounds[3] + && subFootprint(py - bounds[0], px - bounds[2])){ + _peaks.push_back(Peak(other[0], other[1], static_cast(image(py, px)))); + } + } + Bounds trueBounds; trueBounds << bounds[0] + y0, + bounds[1] + y0, bounds[2] + x0, bounds[3] + x0; + footprints.push_back(Footprint(subFootprint, _peaks, trueBounds)); + } + footprint.block(bounds[0], bounds[2], subHeight, subWidth) = MatrixB::Zero(subHeight, subWidth); + } + return footprints; +} + + PYBIND11_MODULE(detect_pybind11, mod) { typedef Eigen::Matrix MatrixF; @@ -383,18 +591,29 @@ PYBIND11_MODULE(detect_pybind11, mod) { "Trim pixels not conencted to a center from a list of centers"); mod.def("get_peaks", &get_peaks, - "Get a list of peaks in a footprint created by get_connected_pixels"); + "Get a list of peaks in a footprint with a contrast-limited watershed", + "image"_a, "min_separation"_a, "peak_thresh"_a, "y0"_a, "x0"_a, + "kappa"_a=0.0, "footprint_thresh"_a=-std::numeric_limits::infinity()); mod.def("get_peaks", &get_peaks, - "Get a list of peaks in a footprint created by get_connected_pixels"); + "Get a list of peaks in a footprint with a contrast-limited watershed", + "image"_a, "min_separation"_a, "peak_thresh"_a, "y0"_a, "x0"_a, + "kappa"_a=0.0, "footprint_thresh"_a=-std::numeric_limits::infinity()); mod.def("get_footprints", &get_footprints, "Create a list of all of the footprints in an image, with their peaks", "image"_a, "min_separation"_a, "min_area"_a, "peak_thresh"_a, "footprint_thresh"_a, - "find_peaks"_a=true, "y0"_a=0, "x0"_a=0); + "find_peaks"_a=true, "y0"_a=0, "x0"_a=0, "kappa"_a=0.0); mod.def("get_footprints", &get_footprints, "Create a list of all of the footprints in an image, with their peaks", "image"_a, "min_separation"_a, "min_area"_a, "peak_thresh"_a, "footprint_thresh"_a, - "find_peaks"_a=true, "y0"_a=0, "x0"_a=0); + "find_peaks"_a=true, "y0"_a=0, "x0"_a=0, "kappa"_a=0.0); + + mod.def("peaks_to_footprints", &peaks_to_footprints, + "Build footprints around a supplied set of peaks", + "image"_a, "peaks"_a, "min_area"_a, "footprint_thresh"_a, "y0"_a=0, "x0"_a=0); + mod.def("peaks_to_footprints", &peaks_to_footprints, + "Build footprints around a supplied set of peaks", + "image"_a, "peaks"_a, "min_area"_a, "footprint_thresh"_a, "y0"_a=0, "x0"_a=0); py::class_(mod, "Footprint") .def(py::init, Bounds>(), @@ -405,9 +624,10 @@ PYBIND11_MODULE(detect_pybind11, mod) { .def("add_peak", &Footprint::addPeak); py::class_(mod, "Peak") - .def(py::init(), - "y"_a, "x"_a, "flux"_a) + .def(py::init(), + "y"_a, "x"_a, "flux"_a, "saddle"_a=std::numeric_limits::quiet_NaN()) .def_property_readonly("y", &Peak::getY) .def_property_readonly("x", &Peak::getX) - .def_property_readonly("flux", &Peak::getFlux); + .def_property_readonly("flux", &Peak::getFlux) + .def_property_readonly("saddle", &Peak::getSaddle); } diff --git a/python/lsst/scarlet/lite/fft.py b/python/lsst/scarlet/lite/fft.py index 7b9096d9..232d9abd 100644 --- a/python/lsst/scarlet/lite/fft.py +++ b/python/lsst/scarlet/lite/fft.py @@ -214,10 +214,9 @@ def get_fft_shape( if use_max: shape[n] = np.max([shape1[ax], shape2[ax]]) - shape += padding + shape = np.asarray(shape, dtype=int) + padding # Use the next fastest shape in each dimension - shape = [fftpack.next_fast_len(s) for s in shape] - return tuple(shape) + return tuple(fftpack.next_fast_len(s) for s in shape) class Fourier: diff --git a/python/lsst/scarlet/lite/initialization.py b/python/lsst/scarlet/lite/initialization.py index d524f838..bc8090c2 100644 --- a/python/lsst/scarlet/lite/initialization.py +++ b/python/lsst/scarlet/lite/initialization.py @@ -23,11 +23,10 @@ from typing import Sequence, cast import numpy as np -from deprecated.sphinx import deprecated # type: ignore from .bbox import Box -from .component import Component, FactorizedComponent -from .detect import bounds_to_bbox, get_detect_wavelets +from .component import FactorizedComponent +from .detect import bounds_to_bbox from .image import Image from .measure import calculate_snr from .observation import Observation @@ -604,237 +603,3 @@ def init_source(self, center: tuple[int, int]) -> Source: components = [component] return Source(components) # type: ignore - - -@deprecated( - reason="This class is replaced by FactorizedInitialization and will be removed after v29.0", - version="v29.0", - category=FutureWarning, -) -class FactorizedChi2Initialization(FactorizedInitialization): - """Initialize all sources with chi^2 detections - - There are a large number of parameters that are universal for all of the - sources being initialized from the same set of observed images. - To simplify the API those parameters are all initialized by this class - and passed to `init_main_source` for each source. - It also creates temporary objects that only need to be created once for - all of the sources in a blend. - - Parameters - ---------- - observation: - The observation containing the blend - centers: - The center of each source to initialize. - detect: - The array that contains a 2D image used for detection. - min_snr: - The minimum SNR required per component. - So a 2-component source requires at least `2*min_snr` while sources - with SNR < `min_snr` will be initialized with the PSF. - monotonicity: - When `monotonicity` is `None`, - the component is initialized with only the - monotonic pixels, otherwise the monotonicity operator is used to - project the morphology to a monotonic solution. - disk_percentile: - The percentage of the overall flux to attribute to the disk. - thresh: - The threshold used to trim the morphology, - so all pixels below `thresh * bg_rms` are set to zero. - padding: - The amount to pad the morphology to allow for extra flux - in the first few iterations before resizing. - """ - - pass - - -@deprecated( - reason="FactorizedWaveletInitialization will be removed after v29.0 " - "since it does not appear to offer any advantages over " - "FactorizedChi2Initialization. Consider switching to " - "FactorizedInitialization now.", - version="v29.0", - category=FutureWarning, -) -class FactorizedWaveletInitialization(FactorizedInitialization): - """Parameters used to initialize all sources with wavelet detections - - There are a large number of parameters that are universal for all of the - sources being initialized from the same set of wavelet coefficients. - To simplify the API those parameters are all initialized by this class - and passed to `init_wavelet_source` for each source. - - Parameters - ---------- - observation: - The multiband observation of the blend. - centers: - The center of each source to initialize. - bulge_slice, disk_slice: - The slice used to select the wavelet scales used for the - bulge/disk. - bulge_padding, disk_padding: - The number of pixels to grow the bounding box of the bulge/disk - to leave extra room for growth in the first few iterations. - use_psf: - Whether or not to use the PSF for single component sources. - If `use_psf` is `False` then only sources with low signal - at all scales are initialized with the PSF morphology. - scales: - Number of wavelet scales to use. - wavelets: - The array of wavelet coefficients `(scale, y, x)` - used for detection. - monotonicity: - When `monotonicity` is `None`, - the component is initialized with only the - monotonic pixels, otherwise the monotonicity operator is used to - project the morphology to a monotonic solution. - min_snr: - The minimum SNR required per component. - So a 2-component source requires at least `2*min_snr` while sources - with SNR < `min_snr` will be initialized with the PSF. - """ - - def __init__( - self, - observation: Observation, - centers: Sequence[tuple[int, int]], - bulge_slice: slice = slice(None, 2), - disk_slice: slice = slice(2, -1), - bulge_padding: int = 5, - disk_padding: int = 5, - use_psf: bool = True, - scales: int = 5, - wavelets: np.ndarray | None = None, - monotonicity: Monotonicity | None = None, - min_snr: float = 50, - use_sparse_init: bool = True, - ): - if wavelets is None: - wavelets = get_detect_wavelets( - observation.images.data, - observation.variance.data, - scales=scales, - ) - wavelets[wavelets < 0] = 0 - # The detection coadd for single component sources - detectlets = np.sum(wavelets[:-1], axis=0) - # The detection coadd for the bulge - bulgelets = np.sum(wavelets[bulge_slice], axis=0) - # The detection coadd for the disk - disklets = np.sum(wavelets[disk_slice], axis=0) - - self.detectlets = detectlets - self.bulgelets = bulgelets - self.disklets = disklets - self.bulge_grow = bulge_padding - self.disk_grow = disk_padding - self.use_psf = use_psf - - # Initialize the sources - super().__init__( - observation=observation, - centers=centers, - detect=detectlets, - min_snr=min_snr, - monotonicity=monotonicity, - use_sparse_init=use_sparse_init, - ) - - def init_source(self, center: tuple[int, int]) -> Source: - """Initialize a source from a chi^2 detection. - - Parameters - ---------- - center: - The center of the source. - """ - local_center = ( - center[0] - self.observation.bbox.origin[0], - center[1] - self.observation.bbox.origin[1], - ) - nbr_components = self.get_snr(center) - observation = self.observation - - components: list[Component] | None = None - - if (nbr_components < 1 and self.use_psf) or self.detectlets[local_center[0], local_center[1]] <= 0: - # Initialize the source as an PSF source - components = [self.get_psf_component(center)] - elif nbr_components < 2: - # Inititialize with a single component - component = self.get_single_component(center, self.detectlets, 0, self.disk_grow) - if component is not None: - components = [component] - else: - # Initialize with a 2 component model - bulge_box, bulge_morph = init_monotonic_morph( - self.bulgelets, center, observation.bbox, self.bulge_grow - ) - disk_box, disk_morph = init_monotonic_morph( - self.disklets, center, observation.bbox, self.disk_grow - ) - if bulge_morph is None or disk_morph is None: - if bulge_morph is None: - if disk_morph is None: - raise RuntimeError("Both components are None") - # One of the components was null, - # so initialize as a single component - component = self.get_single_component(center, self.detectlets, 0, self.disk_grow) - if component is not None: - components = [component] - else: - local_bulge_box = bulge_box - self.observation.bbox.origin - local_disk_box = disk_box - self.observation.bbox.origin - bulge_morph = bulge_morph[local_bulge_box.slices] - disk_morph = disk_morph[local_disk_box.slices] - - bulge_spectrum, disk_spectrum = multifit_spectra( - observation, - [ - Image(bulge_morph, yx0=cast(tuple[int, int], bulge_box.origin)), - Image(disk_morph, yx0=cast(tuple[int, int], disk_box.origin)), - ], - ) - - components = [] - if np.any(bulge_spectrum != 0): - components.append( - FactorizedComponent( - observation.bands, - bulge_spectrum, - bulge_morph, - bulge_box, - center, - monotonicity=self.monotonicity, - ) - ) - else: - logger.debug("cut bulge") - if np.any(disk_spectrum != 0): - components.append( - FactorizedComponent( - observation.bands, - disk_spectrum, - disk_morph, - disk_box, - center, - monotonicity=self.monotonicity, - ) - ) - else: - logger.debug("cut disk") - - # If every init path above either failed (left components as - # None) or produced no components (empty list when bulge and - # disk spectra were both cut), fall back to a PSF source -- a - # point-source model is the most conservative non-empty model - # we can return. - if not components: - logger.debug("fall back to PSF source") - components = [self.get_psf_component(center)] - return Source(components) diff --git a/python/lsst/scarlet/lite/io/blend.py b/python/lsst/scarlet/lite/io/blend.py index 7b2d353f..ea9bf9f0 100644 --- a/python/lsst/scarlet/lite/io/blend.py +++ b/python/lsst/scarlet/lite/io/blend.py @@ -177,8 +177,9 @@ def to_blend(self, observation: Observation) -> Blend: @staticmethod @deprecated( - reason="ScarletBlendData.from_blend is deprecated. Use blend.to_data() instead.", - version="v30.0", + reason="ScarletBlendData.from_blend is deprecated and will be removed " + "after v31.0. Use blend.to_data() instead.", + version="v31.0", category=FutureWarning, ) def from_blend(blend: Blend) -> ScarletBlendData: diff --git a/python/lsst/scarlet/lite/io/cube_component.py b/python/lsst/scarlet/lite/io/cube_component.py index 482497b1..20506238 100644 --- a/python/lsst/scarlet/lite/io/cube_component.py +++ b/python/lsst/scarlet/lite/io/cube_component.py @@ -24,9 +24,9 @@ @deprecated( - reason="ComponentCube is deprecated and will be removed after scarlet_lite v30.0. " + reason="ComponentCube is deprecated and will be removed after v31.0. " "Please use CubeComponent instead.", - version="scarlet_lite v30.0", + version="v31.0", category=FutureWarning, ) class ComponentCube(CubeComponent): diff --git a/python/lsst/scarlet/lite/io/source.py b/python/lsst/scarlet/lite/io/source.py index 596b9a92..41c90d91 100644 --- a/python/lsst/scarlet/lite/io/source.py +++ b/python/lsst/scarlet/lite/io/source.py @@ -106,8 +106,8 @@ def to_source(self, observation: Observation) -> Source: @staticmethod @deprecated( - reason="from_source is deprecated and will be removed in a future release.", - version="v30.0", + reason="from_source is deprecated and will be removed after v31.0.", + version="v31.0", category=FutureWarning, ) def from_source(source: Source) -> ScarletSourceData: diff --git a/python/lsst/scarlet/lite/observation.py b/python/lsst/scarlet/lite/observation.py index 6b2dfaac..f8c2b5fd 100644 --- a/python/lsst/scarlet/lite/observation.py +++ b/python/lsst/scarlet/lite/observation.py @@ -73,7 +73,7 @@ def get_filter_coords(filter_values: np.ndarray, center: tuple[int, int] | None return coords -def get_filter_bounds(coords: np.ndarray) -> tuple[int, int, int, int]: +def get_filter_bounds(coords: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """Get the slices in x and y to apply a filter Parameters @@ -97,7 +97,11 @@ def get_filter_bounds(coords: np.ndarray) -> tuple[int, int, int, int]: return y_start, y_end, x_start, x_end -def convolve(image: np.ndarray, psf: np.ndarray, bounds: tuple[int, int, int, int]): +def convolve( + image: np.ndarray, + psf: np.ndarray, + bounds: tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray], +): """Convolve an image with a PSF in real space Parameters @@ -262,7 +266,7 @@ def __init__( self.diff_kernel = None self.grad_kernel = None - self._convolution_bounds: tuple[int, int, int, int] | None = None + self._convolution_bounds: tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray] | None = None @property def bands(self) -> tuple: @@ -518,7 +522,7 @@ def dtype(self) -> npt.DTypeLike: return self.images.dtype @property - def convolution_bounds(self) -> tuple[int, int, int, int]: + def convolution_bounds(self) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """Build the slices needed for convolution in real space""" if self._convolution_bounds is None: coords = get_filter_coords(cast(Fourier, self.diff_kernel).image[0]) diff --git a/python/lsst/scarlet/lite/parameters.py b/python/lsst/scarlet/lite/parameters.py index 5fc7894e..d5b3d77c 100644 --- a/python/lsst/scarlet/lite/parameters.py +++ b/python/lsst/scarlet/lite/parameters.py @@ -703,4 +703,4 @@ def relative_step( axis: int | Sequence[int] | None = None, ): """Step size set at `factor` times the mean of `X` in direction `axis`""" - return np.maximum(minimum, factor * x.mean(axis=axis)) + return np.maximum(minimum, factor * x.mean(axis=cast(Any, axis))) diff --git a/python/lsst/scarlet/lite/wavelet.py b/python/lsst/scarlet/lite/wavelet.py index b75ec42f..55aeefda 100644 --- a/python/lsst/scarlet/lite/wavelet.py +++ b/python/lsst/scarlet/lite/wavelet.py @@ -25,6 +25,7 @@ "multiband_starlet_transform", "multiband_starlet_reconstruction", "get_multiresolution_support", + "get_multiband_multiresolution_support", ] from dataclasses import dataclass @@ -245,12 +246,28 @@ class MultiResolutionSupport: support: A per-scale mask, with the shape of the starlet coefficients, that is non-zero where a coefficient is considered significant. + Shape ``(scales+1, Ny, Nx)``, or ``(scales+1, bands, Ny, Nx)`` + when ``multiband`` is `True`. sigma: The noise standard deviation estimated at each scale. + Shape ``(scales+1,)``, or ``(scales+1, bands)`` when + ``multiband`` is `True`. + multiband: + Whether the support spans multiple bands, adding a band axis to + both ``support`` and ``sigma``. """ support: np.ndarray sigma: np.ndarray + multiband: bool = False + + def __post_init__(self) -> None: + support_ndim = 4 if self.multiband else 3 + sigma_ndim = 2 if self.multiband else 1 + if self.support.ndim != support_ndim: + raise ValueError(f"Expected support with {support_ndim} dimensions, got {self.support.ndim}") + if self.sigma.ndim != sigma_ndim: + raise ValueError(f"Expected sigma with {sigma_ndim} dimensions, got {self.sigma.ndim}") def get_multiresolution_support( @@ -262,6 +279,7 @@ def get_multiresolution_support( max_iter: int = 20, image_type: str = "ground", rng: np.random.Generator | None = None, + generation: int = 2, ) -> MultiResolutionSupport: """Calculate the multi-resolution support for a dictionary of starlet coefficients. @@ -301,6 +319,9 @@ def get_multiresolution_support( realization that calibrates ``sigma_je`` in the ``space`` branch. Defaults to ``np.random.default_rng(0)`` so repeated calls with the same input return the same support. + generation: + The generation of the starlet transform used for the noise + realization in the ``space`` branch. Returns ------- @@ -316,7 +337,7 @@ def get_multiresolution_support( if rng is None: rng = np.random.default_rng(0) noise_img = rng.normal(size=image.shape) - noise_starlet = starlet_transform(noise_img, generation=1, scales=len(starlets) - 1) + noise_starlet = starlet_transform(noise_img, generation=generation, scales=len(starlets) - 1) sigma_je = np.zeros((len(noise_starlet),)) for j, star in enumerate(noise_starlet): sigma_je[j] = np.std(star) @@ -330,7 +351,7 @@ def get_multiresolution_support( if np.abs(sigma_i - last_sigma_i) / sigma_i < epsilon: break last_sigma_i = sigma_i - sigma_j = sigma_je + sigma_j = sigma_je * sigma_i else: # Sigma to use for significance at each scale # Initially we use the input `sigma` @@ -363,6 +384,63 @@ def get_multiresolution_support( return MultiResolutionSupport(support=m.astype(int), sigma=sigma_j) +def get_multiband_multiresolution_support( + image: np.ndarray, + starlets: np.ndarray, + sigma: np.floating | Sequence[np.floating] | np.ndarray, + sigma_scaling: float = 3, + epsilon: float = 1e-1, + max_iter: int = 20, + image_type: str = "ground", + rng: np.random.Generator | None = None, + generation: int = 2, +) -> MultiResolutionSupport: + """Calculate the multi-resolution support of a multiband image. + + See `get_multiresolution_support` for a description of the + remainder of the parameters. + + Parameters + ---------- + image: + The multiband image with dimension (bands, Ny, Nx). + starlets: + The multiband starlet dictionary with dimension + (scales+1, bands, Ny, Nx). + sigma: + The standard deviation of the image, either a single value used + for every band or a per-band sequence of length ``bands``. + + Returns + ------- + support: + The multiband support, with ``support`` shaped + (scales+1, bands, Ny, Nx) and ``sigma`` shaped (scales+1, bands). + """ + _, bands, height, width = starlets.shape + # Broadcast a scalar sigma to one value per band so the loop always + # indexes a per-band array. + band_sigma = np.broadcast_to(np.asarray(sigma), (bands,)) + + support = np.zeros(starlets.shape, dtype=int) + sigma_out = np.zeros((len(starlets), bands), dtype=starlets.dtype) + for band in range(bands): + result = get_multiresolution_support( + image=image[band], + starlets=starlets[:, band], + sigma=band_sigma[band], + sigma_scaling=sigma_scaling, + epsilon=epsilon, + max_iter=max_iter, + image_type=image_type, + rng=rng, + generation=generation, + ) + support[:, band] = result.support + sigma_out[:, band] = result.sigma + return MultiResolutionSupport(support=support, sigma=sigma_out, multiband=True) + + def apply_wavelet_denoising( image: np.ndarray, sigma: np.floating | None = None, diff --git a/tests/test_detect.py b/tests/test_detect.py index 02846084..c70d5428 100644 --- a/tests/test_detect.py +++ b/tests/test_detect.py @@ -24,9 +24,38 @@ import numpy as np from lsst.scarlet.lite import Box, Image from lsst.scarlet.lite.detect import ( + CANDIDATE_DTYPE, + COMPONENT_DTYPE, + COUNT_COVERS_MAX, + DETECTION_DTYPE, + PAIR_DTYPE, + POSITION_DTYPE, + SIGMA_TO_FWHM, + _build_detection_starlets, + _build_footprints, + _build_pairs, + _build_significance_map, + _chi2_log_survival, + _chi2_to_sigma, + _chi_to_sigma, + _clique_cover, + _collapse_positions, + _dilate_bad_mask, + _find_peak_candidates, + _group_peaks, + _link_fractions, + _link_radius, + _locate_peaks, + _max_shared_plane, + _plane_flags, + _starlet_kernel_fwhm, + _starlet_scale_factors, + _unambiguous_members, bbox_to_bounds, bounds_to_bbox, detect_footprints, + detect_peaks, + find_structures, footprints_to_image, get_detect_wavelets, get_wavelets, @@ -37,9 +66,12 @@ get_connected_multipeak, get_connected_pixels, get_footprints, + peaks_to_footprints, ) from lsst.scarlet.lite.utils import integrated_circular_gaussian -from numpy.testing import assert_array_equal +from lsst.scarlet.lite.wavelet import starlet_transform +from numpy.testing import assert_allclose, assert_array_equal +from scipy import stats from utils import ScarletTestCase @@ -184,6 +216,76 @@ def test_get_footprints_min_area_boundary(self): footprints = get_footprints(img, 1, 5, 1e-15, 1e-15, False) self.assertEqual(len(footprints), 0) + def test_peaks_to_footprints(self): + # Sources 0 and 1 overlap into one footprint, so passing all four + # centers yields three footprints. Footprints follow peak-input + # order, and peaks within a footprint keep their input order (unlike + # get_footprints, which emits raster order sorted brightest first). + peaks = [list(center) for center in self.centers] + footprints = peaks_to_footprints(self.image.data, peaks, 4, 1e-15) + self.assertEqual(len(footprints), 3) + + # First seeded by center 0; absorbs center 1 as a second peak. + merged = self.sources[0] + self.sources[1] + assert_array_equal(footprints[0].data, merged.data > 1e-15) + self.assertBoxEqual(footprints[0].bbox, merged.bbox) + self.assertEqual( + [(peak.y, peak.x) for peak in footprints[0].peaks], + [self.centers[0], self.centers[1]], + ) + + # Then the two single-peak footprints, in input order. + assert_array_equal(footprints[1].data, self.sources[2].data > 1e-15) + self.assertBoxEqual(footprints[1].bbox, self.sources[2].bbox) + self.assertEqual([(peak.y, peak.x) for peak in footprints[1].peaks], [self.centers[2]]) + + assert_array_equal(footprints[2].data, self.sources[3].data > 1e-15) + self.assertBoxEqual(footprints[2].bbox, self.sources[3].bbox) + self.assertEqual([(peak.y, peak.x) for peak in footprints[2].peaks], [self.centers[3]]) + + def test_peaks_to_footprints_subset(self): + # A single peak in the merged (source 0 + source 1) footprint keeps + # that footprint but attaches only the supplied peak, and the three + # unseeded footprints are dropped. + footprints = peaks_to_footprints(self.image.data, [list(self.centers[1])], 4, 1e-15) + self.assertEqual(len(footprints), 1) + truth = self.sources[0] + self.sources[1] + assert_array_equal(footprints[0].data, truth.data > 1e-15) + self.assertBoxEqual(footprints[0].bbox, truth.bbox) + self.assertEqual(len(footprints[0].peaks), 1) + self.assertEqual(footprints[0].peaks[0].y, self.centers[1][0]) + self.assertEqual(footprints[0].peaks[0].x, self.centers[1][1]) + + def test_peaks_to_footprints_min_area(self): + # min_area drops a footprint smaller than the threshold. + img = np.zeros((10, 10), dtype=np.float32) + img[3:5, 5:7] = 1.0 # 2x2 filled square, area 4 + + footprints = peaks_to_footprints(img, [[3, 5]], 4, 1e-15) + self.assertEqual(len(footprints), 1) + footprints = peaks_to_footprints(img, [[3, 5]], 5, 1e-15) + self.assertEqual(len(footprints), 0) + + def test_peaks_to_footprints_offset(self): + # Peaks are in image coordinates; y0/x0 shift them onto the local + # array, and the output stays in image coordinates. + img = np.zeros((10, 10), dtype=np.float32) + img[3:5, 5:7] = 1.0 + y0, x0 = 100, 200 + + footprints = peaks_to_footprints(img, [[3 + y0, 5 + x0]], 4, 1e-15, y0, x0) + self.assertEqual(len(footprints), 1) + bounds = list(footprints[0].bounds) + self.assertEqual(bounds, [3 + y0, 4 + y0, 5 + x0, 6 + x0]) + self.assertEqual(footprints[0].peaks[0].y, 3 + y0) + self.assertEqual(footprints[0].peaks[0].x, 5 + x0) + + def test_peaks_to_footprints_out_of_bounds(self): + img = np.zeros((10, 10), dtype=np.float32) + img[3:5, 5:7] = 1.0 + with self.assertRaises(IndexError): + peaks_to_footprints(img, [[20, 20]], 4, 1e-15) + def _check_peaks(self, peaks): matched_peaks = [] for center in self.centers: @@ -196,42 +298,50 @@ def _check_peaks(self, peaks): def test_detect_footprints(self): # This method doesn't test for accurracy, since # there is no variance, so we set it to ones. + # detect_footprints is deprecated in favor of detect_peaks; + # the calls assert the FutureWarning while still covering the + # legacy behavior. variance = np.ones(self.image.shape, dtype=self.image.dtype) - footprints = detect_footprints( - self.image.data[None, :, :], - variance[None, :, :], - scales=1, - generation=2, - origin=(0, 0), - min_separation=1, - min_area=4, - peak_thresh=1e-15, - footprint_thresh=1e-15, - find_peaks=True, - remove_high_freq=False, - min_pixel_detect=1, - ) + with self.assertWarns(FutureWarning): + footprints = detect_footprints( + self.image.data[None, :, :], + variance[None, :, :], + scales=1, + generation=2, + origin=(0, 0), + min_separation=1, + min_area=4, + peak_thresh=1e-15, + footprint_thresh=1e-15, + find_peaks=True, + remove_high_freq=False, + min_pixel_detect=1, + ) self.assertEqual(len(footprints), 3) peaks = [peak for footprint in footprints for peak in footprint.peaks] self._check_peaks(peaks) - footprints = detect_footprints( - self.image.data[None, :, :], - variance[None, :, :], - scales=1, - generation=1, - min_separation=1, - min_area=4, - peak_thresh=1e-15, - footprint_thresh=1e-15, - find_peaks=True, - remove_high_freq=True, - min_pixel_detect=1, - ) + with self.assertWarns(FutureWarning): + footprints = detect_footprints( + self.image.data[None, :, :], + variance[None, :, :], + scales=1, + generation=1, + min_separation=1, + min_area=4, + peak_thresh=1e-15, + footprint_thresh=1e-15, + find_peaks=True, + remove_high_freq=True, + min_pixel_detect=1, + ) - self.assertEqual(len(footprints), 2) + # Three footprints: the two source blends plus the flat 0.5 patch at + # [30:32, 40]. The watershed seeds a peak on the patch's plateau, so + # the footprint is kept rather than dropped for having no peak. + self.assertEqual(len(footprints), 3) peaks = [peak for footprint in footprints for peak in footprint.peaks] self._check_peaks(peaks) @@ -245,20 +355,21 @@ def test_detect_footprints_min_pixel_detect(self): # Single-band: with min_pixel_detect=2 every pixel fails the # "at least 2 bands above 0" check, so nothing survives. - footprints = detect_footprints( - self.image.data[None, :, :], - variance[None, :, :], - scales=1, - generation=2, - origin=(0, 0), - min_separation=1, - min_area=4, - peak_thresh=1e-15, - footprint_thresh=1e-15, - find_peaks=True, - remove_high_freq=False, - min_pixel_detect=2, - ) + with self.assertWarns(FutureWarning): + footprints = detect_footprints( + self.image.data[None, :, :], + variance[None, :, :], + scales=1, + generation=2, + origin=(0, 0), + min_separation=1, + min_area=4, + peak_thresh=1e-15, + footprint_thresh=1e-15, + find_peaks=True, + remove_high_freq=False, + min_pixel_detect=2, + ) self.assertEqual(len(footprints), 0) # Two-band: band 0 only contains sources 0+1, band 1 only @@ -271,38 +382,40 @@ def test_detect_footprints_min_pixel_detect(self): b1 = (full + band1).data images = np.stack([b0, b1]) variance2 = np.ones(images.shape, dtype=images.dtype) - footprints = detect_footprints( - images, - variance2, - scales=1, - generation=2, - origin=(0, 0), - min_separation=1, - min_area=4, - peak_thresh=1e-15, - footprint_thresh=1e-15, - find_peaks=True, - remove_high_freq=False, - min_pixel_detect=2, - ) + with self.assertWarns(FutureWarning): + footprints = detect_footprints( + images, + variance2, + scales=1, + generation=2, + origin=(0, 0), + min_separation=1, + min_area=4, + peak_thresh=1e-15, + footprint_thresh=1e-15, + find_peaks=True, + remove_high_freq=False, + min_pixel_detect=2, + ) self.assertEqual(len(footprints), 0) # Sanity: with min_pixel_detect=1 the same multi-band input # produces the union of both bands' footprints. - footprints = detect_footprints( - images, - variance2, - scales=1, - generation=2, - origin=(0, 0), - min_separation=1, - min_area=4, - peak_thresh=1e-15, - footprint_thresh=1e-15, - find_peaks=True, - remove_high_freq=False, - min_pixel_detect=1, - ) + with self.assertWarns(FutureWarning): + footprints = detect_footprints( + images, + variance2, + scales=1, + generation=2, + origin=(0, 0), + min_separation=1, + min_area=4, + peak_thresh=1e-15, + footprint_thresh=1e-15, + find_peaks=True, + remove_high_freq=False, + min_pixel_detect=1, + ) self.assertGreater(len(footprints), 0) def test_bounds_to_bbox(self): @@ -357,7 +470,8 @@ def test_footprint(self): def test_get_wavelets(self): images = self.hsc_data["images"] variance = self.hsc_data["variance"] - wavelets = get_wavelets(images, variance) + with self.assertWarns(FutureWarning): + wavelets = get_wavelets(images, variance) self.assertTupleEqual(wavelets.shape, (5, 5, 58, 48)) self.assertEqual(wavelets.dtype, np.float32) @@ -365,6 +479,1065 @@ def test_get_wavelets(self): def test_get_detect_wavelets(self): images = self.hsc_data["images"] variance = self.hsc_data["variance"] - wavelets = get_detect_wavelets(images, variance) + with self.assertWarns(FutureWarning): + wavelets = get_detect_wavelets(images, variance) self.assertTupleEqual(wavelets.shape, (4, 58, 48)) + + +class TestPeakDetection(ScarletTestCase): + def setUp(self): + # Three isotropic Gaussian sources present in every band on top of + # per-band white noise. The bands have different noise levels so the + # per-band standardization is actually exercised. + rng = np.random.default_rng(42) + self.n_bands = 3 + self.shape = (64, 64) + self.centers = [(20, 15), (40, 45), (12, 50)] + band_std = np.array([1.5, 2.0, 3.0], dtype=np.float32) + + yy, xx = np.mgrid[0 : self.shape[0], 0 : self.shape[1]] + images = rng.standard_normal((self.n_bands,) + self.shape) * band_std[:, None, None] + for cy, cx in self.centers: + bump = np.exp(-((yy - cy) ** 2 + (xx - cx) ** 2) / (2 * 2.5**2)) + images += 40.0 * bump + self.images = images.astype(np.float32) + self.variance = np.tile((band_std**2)[:, None, None], (1,) + self.shape) + self.band_std = band_std + + def test_starlet_scale_factors(self): + factors = _starlet_scale_factors(3, generation=2) + self.assertEqual(factors.shape, (4,)) + self.assertTrue(np.all(factors > 0)) + # Coarser scales spread the kernel wider, so each successive + # sum(K_j**2) is smaller. + self.assertTrue(np.all(np.diff(factors) < 0)) + + # White noise of variance ``v`` produces coefficients of variance + # ``F_j * v`` at scale ``j``; this is the whole reason the factors + # exist. Measure it on the interior to avoid boundary effects. + rng = np.random.default_rng(0) + v = 4.0 + noise = rng.standard_normal((512, 512)) * np.sqrt(v) + coeffs = starlet_transform(noise, scales=3, generation=2) + measured = np.array([c[50:-50, 50:-50].var() for c in coeffs]) + assert_allclose(measured, factors * v, rtol=0.05) + + def test_build_detection_starlets(self): + starlets, sigma = _build_detection_starlets(self.images, self.variance, scales=3) + self.assertEqual(starlets.shape, (4, self.n_bands) + self.shape) + self.assertEqual(sigma.shape, (4, self.n_bands, 1, 1)) + + # sigma_{j,b} = sqrt(F_j * nanmedian(var_b)). + factors = _starlet_scale_factors(3) + band_var = np.nanmedian(self.variance, axis=(1, 2)) + expected = np.sqrt(factors[:, None] * band_var[None, :]) + assert_allclose(sigma[..., 0, 0], expected, rtol=1e-4) + + def test_build_detection_starlets_per_pixel_variance(self): + starlets, sigma = _build_detection_starlets( + self.images, self.variance, scales=3, variance_mode="pixel" + ) + self.assertEqual(sigma.shape, (4, self.n_bands) + self.shape) + + # sigma_{j,b}(x) = sqrt(F_j * var_b(x)), the same factors applied to + # the variance plane rather than to one value per band. + factors = _starlet_scale_factors(3) + expected = np.sqrt(factors[:, None, None, None] * self.variance[None]) + assert_allclose(sigma, expected, rtol=1e-4) + + # A flat variance plane makes the two modes agree. + flat = np.full_like(self.variance, 4.0) + _, by_pixel = _build_detection_starlets(self.images, flat, scales=3, variance_mode="pixel") + _, by_median = _build_detection_starlets(self.images, flat, scales=3, variance_mode="median") + assert_allclose(by_pixel, np.broadcast_to(by_median, by_pixel.shape), rtol=1e-6) + + def test_build_detection_starlets_ignores_nan_variance(self): + variance = self.variance.copy() + variance[:, :10, :10] = np.nan + _, sigma = _build_detection_starlets(self.images, variance, scales=3) + self.assertFalse(np.any(np.isnan(sigma))) + + def test_build_detection_starlets_validation(self): + with self.assertRaises(ValueError): + _build_detection_starlets(self.images, self.variance[:, :10, :10]) + with self.assertRaises(ValueError): + _build_detection_starlets(self.images[0], self.variance[0]) + with self.assertRaises(ValueError): + _build_detection_starlets(self.images, self.variance, variance_mode="rms") + + def test_chi2_log_survival(self): + c = np.array([0.5, 2.0, 8.0, 20.0]) + # For one band the coadd is zero half the time and a chi^2_1 deviate + # otherwise, so the survival function is 0.5 * chi2_1. + assert_allclose(np.exp(_chi2_log_survival(c, 1)), 0.5 * stats.chi2.sf(c, 1)) + # Two bands: binomial mixture of chi^2_1 and chi^2_2. + expected = 0.5 * stats.chi2.sf(c, 1) + 0.25 * stats.chi2.sf(c, 2) + assert_allclose(np.exp(_chi2_log_survival(c, 2)), expected) + + # At zero the survival equals the probability that at least one band + # is positive, 1 - 2**-n. + for n in (1, 2, 3, 6): + self.assertAlmostEqual(float(np.exp(_chi2_log_survival(0.0, n))), 1 - 2.0**-n) + + # Strictly decreasing, and finite in log space where the survival + # function itself has underflowed to zero. + grid = np.linspace(0, 30, 200) + self.assertTrue(np.all(np.diff(_chi2_log_survival(grid, 3)) < 0)) + self.assertTrue(np.isfinite(float(_chi2_log_survival(1e6, 3)))) + + def test_chi2_log_survival_matches_simulation(self): + rng = np.random.default_rng(1) + n = 3 + y = rng.standard_normal((2_000_000, n)) + coadd = np.sum(np.clip(y, 0, None) ** 2, axis=1) + for thr in (2.0, 8.0, 12.0): + empirical = np.mean(coadd > thr) + assert_allclose(empirical, float(np.exp(_chi2_log_survival(thr, n))), rtol=0.06) + + def test_chi2_to_sigma(self): + n = 3 + c = np.array([1.0, 5.0, 20.0]) + assert_allclose(_chi2_to_sigma(c, n), stats.norm.isf(np.exp(_chi2_log_survival(c, n)))) + # Works on a scalar as well as an array. + self.assertAlmostEqual(float(_chi2_to_sigma(5.0, n)), float(_chi2_to_sigma(c, n)[1])) + # Monotonically increasing. + grid = np.linspace(0.1, 50, 100) + self.assertTrue(np.all(np.diff(_chi2_to_sigma(grid, n)) > 0)) + # A bright core keeps climbing rather than saturating: the log-space + # mapping stays finite and monotonic however bright the pixel. + bright = float(_chi2_to_sigma(1e6, n)) + self.assertTrue(np.isfinite(bright)) + self.assertGreater(bright, 900.0) + + def test_chi_to_sigma(self): + n = 3 + chi = np.array([0.0, 1.0, 3.0, 10.0, 80.0], dtype=np.float32) + # The lookup table reproduces the exact chi**2 -> sigma mapping, in + # both the dense (< 64) and the geometric (> 64) regimes. + assert_allclose(_chi_to_sigma(chi, n), _chi2_to_sigma(chi.astype(float) ** 2, n), atol=1e-4) + # The dtype of the input is preserved. + self.assertEqual(_chi_to_sigma(chi, n).dtype, np.float32) + # NaN in, NaN out; finite pixels are untouched. + chi_nan = chi.copy() + chi_nan[2] = np.nan + out = _chi_to_sigma(chi_nan, n) + self.assertTrue(np.isnan(out[2])) + self.assertFalse(np.any(np.isnan(out[[0, 1, 3, 4]]))) + + def test_build_significance_map(self): + starlets, sigma = _build_detection_starlets(self.images, self.variance, scales=3) + smap = _build_significance_map(starlets, sigma, first_scale=1) + + # The finest ``first_scale`` scales and the coarse residual are + # dropped, and a chi coadd plane is appended to the bands. + self.assertEqual(smap.shape, (2, self.n_bands + 1) + self.shape) + self.assertEqual(smap.dtype, np.float32) + + # The single-band planes are the standardized coefficients. + standardized = starlets[1:-1] / sigma[1:-1] + assert_allclose(smap[:, : self.n_bands], standardized, rtol=1e-5) + + # The final plane is the clipped chi coadd mapped to sigma, so it is + # on the same footing as the single-band planes. + chi = np.sqrt(np.sum(np.clip(standardized, 0, None) ** 2, axis=1)) + assert_allclose(smap[:, self.n_bands], _chi_to_sigma(chi, self.n_bands), rtol=1e-5) + + def test_build_significance_map_per_pixel_variance(self): + # A per-pixel noise divides the same coefficients, so the map keeps + # its shape and standardizes against the variance plane. + variance = self.variance.copy() + variance[:, :, :32] *= 4.0 + starlets, sigma = _build_detection_starlets(self.images, variance, scales=3, variance_mode="pixel") + smap = _build_significance_map(starlets, sigma, first_scale=1) + + self.assertEqual(smap.shape, (2, self.n_bands + 1) + self.shape) + standardized = starlets[1:-1] / sigma[1:-1] + assert_allclose(smap[:, : self.n_bands], standardized, rtol=1e-5) + + # The deeper half is divided by a smaller noise, so it is not the same + # map the stationary noise would have given. + _, flat_sigma = _build_detection_starlets(self.images, variance, scales=3) + flat_map = _build_significance_map(starlets, flat_sigma, first_scale=1) + self.assertFalse(np.allclose(smap, flat_map)) + + def test_find_peak_candidates(self): + ny, nx = 40, 40 + n_bands = 2 + # Every plane of a significance map is already in sigma, so a bump + # is read back at its own amplitude. + significance_map = np.zeros((1, n_bands + 1, ny, nx), dtype=np.float32) + yy, xx = np.mgrid[0:ny, 0:nx] + + # A 10-sigma bump in band 0. + band0_center = (25, 12) + significance_map[0, 0] = 10.0 * np.exp( + -((yy - band0_center[0]) ** 2 + (xx - band0_center[1]) ** 2) / (2 * 2.0**2) + ) + # A 9-sigma bump in the chi plane. + chi_center = (8, 30) + significance_map[0, n_bands] = 9.0 * np.exp( + -((yy - chi_center[0]) ** 2 + (xx - chi_center[1]) ** 2) / (2 * 2.0**2) + ) + + candidates, footprint_mask = _find_peak_candidates( + significance_map, + min_separation=1, + min_area=1, + peak_thresh=5, + footprint_thresh=3, + first_scale=1, + ) + self.assertEqual(candidates.dtype, CANDIDATE_DTYPE) + # The mask is the union of both planes' footprints above + # footprint_thresh. + self.assertEqual(footprint_mask.shape, (ny, nx)) + assert_array_equal(footprint_mask, np.any(significance_map[0] > 3, axis=0)) + # The empty middle band contributes nothing. + self.assertEqual(np.sum(candidates["band"] == 1), 0) + # Candidates are labeled with the offset starlet scale. + assert_array_equal(np.unique(candidates["scale"]), [1]) + + band0 = candidates[candidates["band"] == 0] + self.assertEqual(len(band0), 1) + self.assertEqual((int(band0["y"][0]), int(band0["x"][0])), band0_center) + self.assertAlmostEqual(float(band0["flux"][0]), 10.0, places=4) + # A single peak in a footprint never joins a brighter basin, so its + # saddle is NaN. + self.assertTrue(np.isnan(band0["saddle"][0])) + + chi = candidates[candidates["band"] == n_bands] + self.assertEqual(len(chi), 1) + self.assertEqual((int(chi["y"][0]), int(chi["x"][0])), chi_center) + self.assertAlmostEqual(float(chi["flux"][0]), 9.0, places=3) + + def test_find_peak_candidates_kappa_culls_flank(self): + # A bright bump with a low-prominence bump on its flank, in one plane. + ny, nx = 5, 40 + x = np.arange(nx) + profile = 10.0 * np.exp(-((x - 8) ** 2) / (2 * 2.0**2)) + 4.0 * np.exp( + -((x - 14) ** 2) / (2 * 2.0**2) + ) + significance_map = np.zeros((1, 2, ny, nx), dtype=np.float32) + significance_map[0, 0, 2] = profile + + # kappa=0 keeps both local maxima. + low, _ = _find_peak_candidates( + significance_map, min_separation=0, min_area=1, peak_thresh=1, footprint_thresh=0.5, kappa=0.0 + ) + self.assertEqual(np.sum(low["band"] == 0), 2) + + # kappa=3 culls the flank bump; it rises less than 3 sigma above the + # saddle to the brighter peak. + high, _ = _find_peak_candidates( + significance_map, min_separation=0, min_area=1, peak_thresh=1, footprint_thresh=0.5, kappa=3.0 + ) + band0 = high[high["band"] == 0] + self.assertEqual(len(band0), 1) + self.assertEqual(int(band0["x"][0]), 8) + + def test_find_peak_candidates_empty(self): + significance_map = np.zeros((2, 3, 32, 32), dtype=np.float32) + candidates, footprint_mask = _find_peak_candidates( + significance_map, peak_thresh=5, footprint_thresh=3 + ) + self.assertEqual(len(candidates), 0) + self.assertEqual(footprint_mask.shape, (32, 32)) + self.assertFalse(footprint_mask.any()) + + def test_find_peak_candidates_mask_excludes_peakless(self): + # A significant blob with no pixel above peak_thresh yields no + # candidate and contributes nothing to the mask. + significance_map = np.zeros((1, 1, 20, 20), dtype=np.float32) + significance_map[0, 0, 5:9, 5:9] = 4.0 + significance_map[0, 0, 12:16, 12:16] = 8.0 + candidates, footprint_mask = _find_peak_candidates( + significance_map, min_area=1, peak_thresh=5, footprint_thresh=3 + ) + self.assertEqual(len(candidates), 1) + truth = np.zeros((20, 20), dtype=bool) + truth[12:16, 12:16] = True + assert_array_equal(footprint_mask, truth) + + def _make_peaks(self, rows): + """Build detections from ``(y, x, peak_sigma)`` rows.""" + peaks = np.zeros(len(rows), dtype=DETECTION_DTYPE) + for i, (y, x, peak_sigma) in enumerate(rows): + peaks[i]["y"] = y + peaks[i]["x"] = x + peaks[i]["peak_sigma"] = peak_sigma + peaks[i]["footprint"] = -1 + return peaks + + def test_build_footprints(self): + origin = (100, 200) + mask = np.zeros((20, 30), dtype=bool) + # Two components: one holding two peaks, one holding a single peak, + # and a third with no peak at all. + mask[2:6, 3:12] = True + mask[10:14, 20:25] = True + mask[15:18, 2:5] = True + peaks = self._make_peaks( + [ + (100 + 3, 200 + 5, 6.0), + (100 + 11, 200 + 22, 9.0), + (100 + 4, 200 + 10, 7.5), + ] + ) + + footprints = _build_footprints(mask, peaks, origin) + + # The peakless component is dropped. + self.assertEqual(len(footprints), 2) + # Every peak is placed in the footprint that contains it. + for row, peak in enumerate(peaks): + footprint = footprints[peak["footprint"]] + self.assertTrue(footprint.bbox.contains((peak["y"], peak["x"]))) + self.assertTrue(footprint.data[peak["y"] - footprint.yx0[0], peak["x"] - footprint.yx0[1]]) + self.assertEqual(peaks["footprint"][0], peaks["footprint"][2]) + self.assertNotEqual(peaks["footprint"][0], peaks["footprint"][1]) + + # Peaks within a footprint are brightest first with the detection's + # peak_sigma as their flux. + pair = footprints[peaks["footprint"][0]] + self.assertEqual([(p.y, p.x) for p in pair.peaks], [(104, 210), (103, 205)]) + assert_allclose([p.flux for p in pair.peaks], [7.5, 6.0]) + single = footprints[peaks["footprint"][1]] + self.assertEqual([(p.y, p.x) for p in single.peaks], [(111, 222)]) + + # The footprint pixels are the component of the mask, in absolute + # coordinates. + self.assertBoxEqual(pair.bbox, Box((4, 9), origin=(102, 203))) + assert_array_equal(pair.data, mask[2:6, 3:12]) + self.assertBoxEqual(single.bbox, Box((4, 5), origin=(110, 220))) + + def test_build_footprints_empty(self): + mask = np.zeros((8, 8), dtype=bool) + mask[2:4, 2:4] = True + self.assertEqual(_build_footprints(mask, np.empty(0, dtype=DETECTION_DTYPE)), []) + + def test_build_footprints_peak_outside_mask(self): + mask = np.zeros((8, 8), dtype=bool) + mask[2:4, 2:4] = True + peaks = self._make_peaks([(6, 6, 5.0)]) + with self.assertRaises(RuntimeError): + _build_footprints(mask, peaks) + + def _make_candidates(self, rows): + """Build a candidate array from ``(y, x, band, scale, flux)`` rows.""" + candidates = np.zeros(len(rows), dtype=CANDIDATE_DTYPE) + for i, (y, x, band, scale, flux) in enumerate(rows): + candidates[i] = (y, x, band, scale, flux, np.nan, -1) + return candidates + + def _make_positions(self, rows): + """Build positions from ``(y, x, scale, peak_sigma, plane_flags)``.""" + positions = np.zeros(len(rows), dtype=POSITION_DTYPE) + for i, (y, x, scale, peak_sigma, plane_flags) in enumerate(rows): + positions[i]["y"] = y + positions[i]["x"] = x + positions[i]["scale"] = scale + positions[i]["peak_sigma"] = peak_sigma + positions[i]["flux"] = peak_sigma + positions[i]["plane_flags"] = plane_flags + positions[i]["peak"] = -1 + return positions + + def test_link_radius(self): + # The PSF FWHM floors the radius at fine scales. + self.assertEqual(_link_radius(1, 3.5), 3.5) + # The 2**scale term dominates at coarse scales. + self.assertEqual(_link_radius(3, 3.5), 8.0) + # Vectorized over an array of scales. + assert_array_equal(_link_radius(np.array([1, 3]), 3.5), [3.5, 8.0]) + + def test_plane_flags(self): + # plane = (scale - first_scale) * (n_bands + 1) + band, so band 0 at + # first_scale is plane 0, the chi plane (band 3) is plane 3, and band 0 + # at the next scale is plane 4. + candidates = self._make_candidates( + [ + (0, 0, 0, 1, 1.0), + (0, 0, 3, 1, 1.0), + (0, 0, 0, 2, 1.0), + ] + ) + flags = _plane_flags(candidates, n_bands=3, first_scale=1) + assert_array_equal(flags, [1 << 0, 1 << 3, 1 << 4]) + + # More than 63 planes cannot fit in an int64 bitmask. + too_many = self._make_candidates([(0, 0, 0, 20, 1.0)]) + with self.assertRaisesRegex(ValueError, "at most 63 planes"): + _plane_flags(too_many, n_bands=3, first_scale=1) + + def test_collapse_positions(self): + # Two candidates at the same pixel in different bands and scales + # collapse to one position; a third at another pixel stays separate. + candidates = self._make_candidates( + [ + (10, 10, 0, 1, 8.0), + (10, 10, 2, 2, 5.0), + (10, 14, 1, 1, 6.0), + ] + ) + positions = _collapse_positions(candidates, n_bands=3, first_scale=1) + self.assertEqual(positions.dtype, POSITION_DTYPE) + self.assertEqual(len(positions), 2) + + merged = positions[(positions["y"] == 10) & (positions["x"] == 10)][0] + self.assertEqual(merged["band_flags"], (1 << 0) | (1 << 2)) + self.assertEqual(merged["scale_flags"], (1 << 1) | (1 << 2)) + self.assertEqual(merged["n_candidates"], 2) + # The finest scale wins the position flux; peak_sigma is the brightest. + self.assertEqual(merged["scale"], 1) + self.assertAlmostEqual(float(merged["flux"]), 8.0) + self.assertAlmostEqual(float(merged["peak_sigma"]), 8.0) + # ``position`` is filled in place, sending the two co-located + # candidates to one position and the third to another. + self.assertEqual(candidates["position"][0], candidates["position"][1]) + self.assertNotEqual(candidates["position"][0], candidates["position"][2]) + + def test_build_pairs(self): + positions = self._make_positions( + [ + (10, 10, 1, 8.0, 0b1), + # Shares a plane with the first: within reach but not linked. + (10, 12, 1, 7.0, 0b1), + # Shares no plane with the first: linked. + (12, 10, 1, 7.0, 0b10), + # Beyond every link radius: no pair at all. + (40, 40, 1, 7.0, 0b10), + ] + ) + pairs = _build_pairs(positions, psf_fwhm=3.5, radius_rule="min") + self.assertEqual(pairs.dtype, PAIR_DTYPE) + assert_array_equal(pairs["i"], [0, 0, 1]) + assert_array_equal(pairs["j"], [1, 2, 2]) + assert_allclose(pairs["distance"], [2, 2, np.sqrt(8)]) + assert_array_equal(pairs["shares_plane"], [True, False, False]) + assert_array_equal(pairs["linked"], [False, True, True]) + + # A fine position (radius 2) and a coarse one (radius 8) five pixels + # apart link only under the max rule. + positions = self._make_positions([(10, 10, 1, 8.0, 0b1), (10, 15, 3, 7.0, 0b10)]) + for rule, linked in (("min", False), ("max", True)): + pairs = _build_pairs(positions, psf_fwhm=1.0, radius_rule=rule) + self.assertEqual(len(pairs), 1) + self.assertFalse(pairs["within_min"][0]) + self.assertTrue(pairs["within_max"][0]) + self.assertEqual(pairs["linked"][0], linked) + + self.assertEqual(len(_build_pairs(positions[:1], psf_fwhm=1.0, radius_rule="min")), 0) + + @staticmethod + def _cover(n, edges, **kwargs): + edges = np.array(edges, dtype=int).reshape(-1, 2) + return _clique_cover(np.zeros(n, dtype=np.int64), edges, **kwargs) + + def test_clique_cover_clique(self): + edges = [(0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3)] + group, component, components = self._cover(4, edges, count_covers=True) + assert_array_equal(group, 0) + assert_array_equal(component, 0) + self.assertEqual(components.dtype, COMPONENT_DTYPE) + self.assertEqual(len(components), 1) + self.assertEqual(components["k"][0], 1) + self.assertFalse(components["search_ran"][0]) + self.assertEqual(components["n_min_covers"][0], 1) + + def test_clique_cover_chain(self): + # A-B-C needs two cliques, and B could belong to either. + edges = np.array([(0, 1), (1, 2)]) + group, _, components = self._cover(3, edges, count_covers=True) + self.assertEqual(components["k"][0], 2) + self.assertEqual(components["n_min_covers"][0], 2) + self.assertNotEqual(group[0], group[2]) + + n_linked, ambiguous, max_fraction, max_group = _link_fractions(group, edges, 2) + assert_array_equal(ambiguous, [False, True, False]) + assert_array_equal(n_linked, [1, 2, 1]) + # The end that shares B's clique has no other clique to link to; the + # other end is adjacent to half of B's clique. + lone = 0 if group[0] != group[1] else 2 + paired = 2 - lone + self.assertEqual(max_fraction[lone], 0.5) + self.assertEqual(max_group[lone], group[1]) + self.assertEqual(max_fraction[paired], 0) + self.assertEqual(max_group[paired], -1) + self.assertEqual(max_fraction[1], 1) + self.assertEqual(max_group[1], group[lone]) + + def test_clique_cover_components(self): + # A triangle and a separate edge, plus an isolated node. + edges = [(0, 1), (1, 2), (0, 2), (3, 5)] + group, component, components = self._cover(6, edges) + assert_array_equal(component, [0, 0, 0, 1, 2, 1]) + assert_array_equal(components["n_nodes"], [3, 2, 1]) + assert_array_equal(components["k"], 1) + # Peaks are numbered in the order of their first position. + assert_array_equal(group, [0, 0, 0, 1, 2, 1]) + + def test_clique_cover_ring(self): + # A five-node ring: a greedy independent set gives 2, but no two + # edges cover five nodes, so three cliques are needed. + edges = [(0, 1), (1, 2), (2, 3), (3, 4), (4, 0)] + group, _, components = self._cover(5, edges, count_covers=True) + row = components[0] + self.assertEqual(row["lower_bound"], 2) + self.assertEqual(row["upper_bound"], 3) + self.assertEqual(row["k"], 3) + self.assertTrue(row["search_ran"]) + self.assertFalse(row["cap_hit"]) + self.assertGreater(row["n_visited"], 0) + # One node is alone and the rest pair up around the ring. + self.assertEqual(row["n_min_covers"], 5) + self.assertEqual(sorted(np.bincount(group)), [1, 2, 2]) + + def test_clique_cover_cap(self): + edges = [(0, 1), (1, 2), (2, 3), (3, 4), (4, 0)] + group, _, components = self._cover(5, edges, search_cap=1, count_covers=True) + row = components[0] + self.assertTrue(row["cap_hit"]) + self.assertEqual(row["n_visited"], 1) + # The greedy cover is kept, and no covers are counted. + self.assertEqual(row["k"], row["upper_bound"]) + self.assertEqual(len(np.unique(group)), row["upper_bound"]) + self.assertEqual(row["n_min_covers"], -1) + + def test_clique_cover_plane_bound(self): + # Three positions that share a plane bound the cover from below. + flags = np.array([0b1, 0b1, 0b11, 0b100], dtype=np.int64) + component = np.array([0, 0, 0, 1]) + assert_array_equal(_max_shared_plane(flags, component, 2), [3, 1]) + + @staticmethod + def _brute_force_cover(n, adjacent): + """Return the minimum clique cover size and how many reach it.""" + best = [n + 1, 0] + + def partitions(node, blocks): + if node == n: + yield blocks + return + for block in blocks: + block.append(node) + yield from partitions(node + 1, blocks) + block.pop() + blocks.append([node]) + yield from partitions(node + 1, blocks) + blocks.pop() + + for blocks in partitions(0, []): + if all(adjacent[u, v] for block in blocks for u in block for v in block if u < v): + if len(blocks) < best[0]: + best = [len(blocks), 1] + elif len(blocks) == best[0]: + best[1] += 1 + return best + + def test_clique_cover_brute_force(self): + rng = np.random.default_rng(52721) + searched = 0 + for _ in range(300): + n = int(rng.integers(1, 9)) + adjacent = np.triu(rng.random((n, n)) < rng.uniform(0.2, 0.9), 1) + edges = np.argwhere(adjacent) + adjacent = adjacent | adjacent.T + group, component, components = self._cover(n, edges, count_covers=True) + + k, n_covers = self._brute_force_cover(n, adjacent) + self.assertEqual(components["k"].sum(), k) + self.assertEqual(len(np.unique(group)), k) + # Every group is a clique within one component. + for g in np.unique(group): + members = np.flatnonzero(group == g) + self.assertTrue(all(adjacent[u, v] for u in members for v in members if u < v)) + self.assertEqual(len(np.unique(component[members])), 1) + # Covers multiply across components, and each count stops at the + # cap. + counted = np.prod(components["n_min_covers"]) + if np.any(components["n_min_covers"] == COUNT_COVERS_MAX): + self.assertGreaterEqual(n_covers, counted) + else: + self.assertEqual(counted, n_covers) + searched += int(np.sum(components["search_ran"] & (components["k"] < components["upper_bound"]))) + # Some graphs must need the exact search to beat the greedy cover. + self.assertGreater(searched, 0) + + def _locate(self, rows, n_bands, psf_fwhm, ambiguous=None): + """Locate one peak made of every position of ``rows``.""" + candidates = self._make_candidates(rows) + positions = _collapse_positions(candidates, n_bands=n_bands, first_scale=1) + positions["peak"] = 0 + if ambiguous is not None: + positions["ambiguous"] = ambiguous + use, fallback = _unambiguous_members(positions, 1) + location = _locate_peaks(positions, candidates, use, 1, np.array(psf_fwhm)) + return (int(positions["y"][location[0]]), int(positions["x"][location[0]])), bool(fallback[0]) + + def test_locate_peaks_median(self): + # Every candidate counts, so the two at (10, 14) pull the median to it. + rows = [(10, 10, 0, 1, 9.0), (10, 14, 0, 2, 5.0), (10, 14, 1, 2, 5.0), (11, 13, 1, 1, 50.0)] + self.assertEqual(self._locate(rows, 2, [1.0, 1.0]), ((10, 14), False)) + + def test_locate_peaks_tie_breaks(self): + # The median x is 11, one pixel from both positions. + # Finest scale wins. + rows = [(10, 10, 0, 2, 50.0), (10, 12, 0, 1, 5.0)] + self.assertEqual(self._locate(rows, 2, [1.0, 1.0])[0], (10, 12)) + # Then the smallest PSF FWHM among the finest-scale bands. + rows = [(10, 10, 1, 1, 5.0), (10, 12, 0, 1, 50.0)] + self.assertEqual(self._locate(rows, 2, [2.0, 1.0])[0], (10, 10)) + # The chi^2 plane (band 2 here) ranks after every band. + rows = [(10, 10, 2, 1, 50.0), (10, 12, 1, 1, 5.0)] + self.assertEqual(self._locate(rows, 2, [1.0, 9.0])[0], (10, 12)) + # Then the position index. + rows = [(10, 12, 1, 1, 50.0), (10, 10, 0, 1, 5.0)] + self.assertEqual(self._locate(rows, 2, [1.0, 1.0])[0], (10, 10)) + + def test_locate_peaks_ambiguous(self): + # The ambiguous position is ignored for the median and the choice. + rows = [(10, 10, 0, 1, 5.0), (10, 13, 0, 2, 50.0), (10, 13, 1, 2, 50.0), (10, 13, 2, 2, 50.0)] + self.assertEqual(self._locate(rows, 2, [1.0, 1.0], ambiguous=[False, True]), ((10, 10), False)) + # With every position ambiguous they are all used and the peak flagged. + self.assertEqual(self._locate(rows, 2, [1.0, 1.0], ambiguous=[True, True]), ((10, 13), True)) + + def test_starlet_kernel_fwhm(self): + # The smoothing before scale j matches a delta through the transform. + size = 129 + delta = np.zeros((size, size)) + delta[size // 2, size // 2] = 1 + offset = np.arange(size) - size // 2 + for scale in (1, 2, 3): + profile = starlet_transform(delta, scales=scale)[-1].sum(axis=0) + variance = np.sum(profile * offset**2) / np.sum(profile) + fwhm = _starlet_kernel_fwhm(0.0, scale) + self.assertAlmostEqual(fwhm / SIGMA_TO_FWHM, np.sqrt(variance)) + self.assertAlmostEqual(_starlet_kernel_fwhm(3.0, 0), 3.0) + + def _group(self, rows, n_bands, psf_fwhm, **kwargs): + candidates = self._make_candidates(rows) + return candidates, *_group_peaks(candidates, n_bands, 1, np.array(psf_fwhm, dtype=float), **kwargs) + + def test_group_peaks_absorbs_coarse_midpoint(self): + # Two sources resolved in the fine planes, 6 pixels apart, with a + # bright coarse-scale blend detection between them. The midpoint links + # to both, so it is absorbed by one and flagged, and never becomes a + # peak of its own. + rows = [ + (10, 10, 0, 1, 10.0), + (10, 10, 1, 1, 9.0), + (10, 16, 0, 1, 8.0), + (10, 16, 1, 1, 7.0), + (10, 13, 0, 2, 50.0), + ] + _, positions, peaks, pairs, components = self._group(rows, 2, [3.5, 3.5]) + self.assertEqual(len(peaks), 2) + assert_array_equal(peaks["y"], 10) + assert_array_equal(peaks["x"], [10, 16]) + midpoint = positions[positions["x"] == 13][0] + self.assertTrue(midpoint["ambiguous"]) + self.assertEqual(midpoint["max_linked_fraction"], 1) + assert_array_equal(peaks["n_positions"].sum(), 3) + assert_array_equal(peaks["n_ambiguous"].sum(), 1) + # The midpoint is counted but does not set the peak's significance. + assert_allclose(peaks["peak_sigma"], [10.0, 8.0]) + self.assertEqual(peaks["n_candidates"].sum(), len(rows)) + self.assertEqual(len(components), 1) + self.assertEqual(components["k"][0], 2) + assert_allclose(peaks["nearest_distance"], 6) + + def test_group_peaks_order_independent(self): + rng = np.random.default_rng(3) + rows = [ + (int(y), int(x), int(band), int(scale), float(flux)) + for y, x, band, scale, flux in zip( + rng.integers(0, 30, 80), + rng.integers(0, 30, 80), + rng.integers(0, 3, 80), + rng.integers(1, 3, 80), + rng.uniform(5, 50, 80), + ) + ] + # A plane holds at most one candidate per pixel. + rows = list({row[:4]: row for row in rows}.values()) + _, positions, peaks, pairs, components = self._group(rows, 2, [2.0, 3.0], count_covers=True) + shuffled = [rows[i] for i in rng.permutation(len(rows))] + _, positions2, peaks2, pairs2, components2 = self._group(shuffled, 2, [2.0, 3.0], count_covers=True) + self.assertEqual(peaks.tobytes(), peaks2.tobytes()) + self.assertEqual(positions.tobytes(), positions2.tobytes()) + self.assertEqual(pairs.tobytes(), pairs2.tobytes()) + self.assertEqual(components.tobytes(), components2.tobytes()) + + def test_group_peaks_search_cap_flags_peaks(self): + # A pentagon of positions, one band each, linked only to its two + # neighbors. + rows = [ + (8, 0, 0, 1, 10.0), + (2, 7, 1, 1, 10.0), + (-6, 5, 2, 1, 10.0), + (-6, -4, 3, 1, 10.0), + (2, -7, 4, 1, 10.0), + ] + rows = [(y + 20, x + 20, band, scale, flux) for y, x, band, scale, flux in rows] + _, _, peaks, pairs, components = self._group(rows, 5, [10.0] * 5) + self.assertEqual(int(pairs["linked"].sum()), 5) + self.assertEqual(len(peaks), 3) + self.assertFalse(peaks["search_capped"].any()) + + _, _, peaks, _, components = self._group(rows, 5, [10.0] * 5, search_cap=1) + self.assertEqual(len(peaks), 3) + self.assertTrue(components["cap_hit"][0]) + self.assertTrue(peaks["search_capped"].all()) + + def test_group_peaks_errors(self): + rows = [(10, 10, 0, 1, 10.0), (10, 10, 0, 2, 20.0), (10, 11, 1, 1, 10.0)] + psf_fwhm = [2.0, 3.0] + _, positions, peaks, _, _ = self._group(rows, 2, psf_fwhm) + self.assertEqual(len(peaks), 1) + peak = peaks[0] + self.assertEqual((peak["y"], peak["x"]), (10, 10)) + + def sigma(band, scale, flux): + return _starlet_kernel_fwhm(psf_fwhm[band], scale) / (SIGMA_TO_FWHM * flux) + + # Band 0 contributes its better scale only; the bands are combined by + # inverse variance. + band0 = min(sigma(0, 1, 10.0), sigma(0, 2, 20.0)) + stat = (band0**-2 + sigma(1, 1, 10.0) ** -2) ** -0.5 + self.assertAlmostEqual(peak["stat_err_x"], stat) + self.assertAlmostEqual(peak["stat_err_y"], stat) + # Moments about the peak, one entry per candidate. + self.assertAlmostEqual(peak["moment_xx"], 1 / 3) + self.assertEqual(peak["moment_yy"], 0) + self.assertEqual(peak["moment_xy"], 0) + self.assertEqual(peak["scatter_x"], 0) + self.assertEqual(peak["scatter_y"], 0) + self.assertAlmostEqual(peak["err_x"], max(stat, 1 / np.sqrt(12))) + self.assertEqual(peak["nearest_distance"], np.inf) + + # Two candidates are too few for a robust scatter, which leaves the + # provisional error to the other terms. + _, _, peaks, _, _ = self._group(rows[:2], 2, psf_fwhm) + self.assertTrue(np.isnan(peaks["scatter_x"][0])) + self.assertAlmostEqual(peaks["err_x"][0], max(peaks["stat_err_x"][0], 1 / np.sqrt(12))) + + # A peak seen only in the chi^2 plane falls back to it, with the widest + # band's PSF. + _, _, peaks, _, _ = self._group([(10, 10, 2, 1, 10.0)], 2, psf_fwhm) + self.assertAlmostEqual(peaks["stat_err_x"][0], sigma(1, 1, 10.0)) + assert_allclose(peaks["nearest_distance_fwhm"], np.inf) + + def test_group_peaks_nearest_distance_fwhm(self): + rows = [(10, 10, 0, 1, 10.0), (10, 10, 1, 1, 10.0), (10, 30, 1, 1, 10.0)] + _, _, peaks, _, _ = self._group(rows, 2, [2.0, 4.0]) + assert_allclose(peaks["nearest_distance"], [20, 20]) + # The smallest member PSF sets the unit. + assert_allclose(peaks["nearest_distance_fwhm"], [10, 5]) + + def test_detect_peaks_groups_sources(self): + result = detect_peaks(self.images, self.variance, scales=3, peak_thresh=5, footprint_thresh=3) + self.assertEqual(result.peaks.dtype, DETECTION_DTYPE) + self.assertEqual(result.positions.dtype, POSITION_DTYPE) + + # Every injected source is recovered exactly once, within a pixel. + self.assertEqual(len(result.peaks), len(self.centers)) + for cy, cx in self.centers: + matched = [p for p in result.peaks if abs(p["y"] - cy) <= 1 and abs(p["x"] - cx) <= 1] + self.assertEqual(len(matched), 1) + + # The many per-band, per-scale candidates collapse into each source. + peak = result.peaks[0] + self.assertGreater(peak["n_candidates"], 1) + self.assertGreaterEqual(peak["peak_sigma"], peak["flux"]) + # These bright sources are seen in all 3 bands and the chi coadd. + expected_bands = (1 << self.n_bands + 1) - 1 + self.assertEqual(peak["band_flags"], expected_bands) + self.assertNotEqual(peak["scale_flags"], 0) + + # Each well-separated source gets its own footprint holding its peak. + self.assertEqual(len(result.footprints), len(self.centers)) + assert_array_equal(np.sort(result.peaks["footprint"]), np.arange(len(self.centers))) + for peak in result.peaks: + footprint = result.footprints[peak["footprint"]] + self.assertEqual(len(footprint.peaks), 1) + self.assertEqual((footprint.peaks[0].y, footprint.peaks[0].x), (peak["y"], peak["x"])) + self.assertEqual(footprint.peaks[0].flux, peak["peak_sigma"]) + self.assertTrue(footprint.data[peak["y"] - footprint.yx0[0], peak["x"] - footprint.yx0[1]]) + + def test_detect_peaks_footprints_offset_origin(self): + origin = (1000, 2000) + result = detect_peaks( + self.images, self.variance, scales=3, peak_thresh=5, footprint_thresh=3, origin=origin + ) + reference = detect_peaks(self.images, self.variance, scales=3, peak_thresh=5, footprint_thresh=3) + self.assertEqual(len(result.footprints), len(reference.footprints)) + for footprint, ref in zip(result.footprints, reference.footprints): + self.assertBoxEqual(footprint.bbox, ref.bbox + origin) + assert_array_equal(footprint.data, ref.data) + self.assertEqual( + [(p.y, p.x) for p in footprint.peaks], + [(p.y + origin[0], p.x + origin[1]) for p in ref.peaks], + ) + + def test_detect_peaks_blend_shares_footprint(self): + # Two sources close enough that their footprints touch but far + # enough apart to be resolved land in one footprint with two peaks. + rng = np.random.default_rng(7) + shape = (48, 48) + yy, xx = np.mgrid[0 : shape[0], 0 : shape[1]] + images = rng.standard_normal((self.n_bands,) + shape).astype(np.float32) + for cy, cx in [(24, 18), (24, 30)]: + images += 40.0 * np.exp(-((yy - cy) ** 2 + (xx - cx) ** 2) / (2 * 2.5**2)).astype(np.float32) + variance = np.ones_like(images) + result = detect_peaks(images, variance, scales=3, peak_thresh=5, footprint_thresh=3) + self.assertEqual(len(result.peaks), 2) + self.assertEqual(len(result.footprints), 1) + assert_array_equal(result.peaks["footprint"], 0) + self.assertEqual(len(result.footprints[0].peaks), 2) + + def test_detect_peaks_retraceable(self): + result = detect_peaks(self.images, self.variance, scales=3, peak_thresh=5, footprint_thresh=3) + # Every position is assigned to a peak in this clean blend. + assert_array_equal(result.candidates["position"] >= 0, True) + for peak_id, peak in enumerate(result.peaks): + member_positions = np.nonzero(result.positions["peak"] == peak_id)[0] + self.assertEqual(len(member_positions), peak["n_positions"]) + + member_candidates = np.concatenate( + [np.nonzero(result.candidates["position"] == pid)[0] for pid in member_positions] + ) + self.assertEqual(len(member_candidates), peak["n_candidates"]) + + def test_detect_peaks_graph_products(self): + psf_fwhm = [3.0, 3.5, 4.0] + result = detect_peaks( + self.images, self.variance, scales=3, peak_thresh=5, footprint_thresh=3, psf_fwhm=psf_fwhm + ) + self.assertEqual(result.pairs.dtype, PAIR_DTYPE) + self.assertEqual(result.components.dtype, COMPONENT_DTYPE) + self.assertEqual(result.components["n_nodes"].sum(), len(result.positions)) + self.assertEqual(result.components["k"].sum(), len(result.peaks)) + self.assertEqual(result.positions["component"].max() + 1, len(result.components)) + self.assertEqual(result.metadata["psf_fwhm"], psf_fwhm) + self.assertEqual(result.metadata["radius_rule"], "min") + self.assertEqual(result.metadata["link_radius_scales"], [1, 2]) + assert_allclose(result.metadata["link_radius"], [4.0, 4.0]) + self.assertAlmostEqual(result.metadata["min_separation"], 3.0 * 0.84932) + + def _make_structured_peaks(self): + """Lay out regions on a significance plane and the peaks inside them. + + Region A holds three peaks, B one, and C two; A and B are separated + by a gap. Peak D sits on a pixel below the threshold. + """ + significance = np.zeros((40, 80), dtype=np.float32) + regions = { + "A": ((10, 15, 10, 31), [(12, 12), (12, 20), (12, 28)]), + "B": ((10, 15, 33, 36), [(12, 34)]), + "C": ((30, 33, 5, 15), [(31, 6), (31, 13)]), + } + centers = [] + for (y0, y1, x0, x1), members in regions.values(): + significance[y0:y1, x0:x1] = 2.0 + centers.extend(members) + centers.append((5, 70)) + peaks = np.zeros(len(centers), dtype=DETECTION_DTYPE) + peaks["y"] = [y for y, _ in centers] + peaks["x"] = [x for _, x in centers] + return significance, peaks + + def test_find_structures(self): + significance, peaks = self._make_structured_peaks() + find_structures(peaks, significance) + # Structures are numbered in raster order: A, B, C. + assert_array_equal(peaks["structure_id"], [1, 1, 1, 2, 3, 3, 0]) + assert_array_equal(peaks["structure_peaks"], [3, 3, 3, 1, 2, 2, 0]) + + # Raising the threshold above the plane leaves every peak unlabeled. + find_structures(peaks, significance, structure_thresh=5.0) + assert_array_equal(peaks["structure_id"], 0) + assert_array_equal(peaks["structure_peaks"], 0) + + def test_find_structures_origin(self): + significance, peaks = self._make_structured_peaks() + shifted = peaks.copy() + find_structures(peaks, significance) + origin = (100, 200) + shifted["y"] += origin[0] + shifted["x"] += origin[1] + find_structures(shifted, significance, origin=origin) + for field in ("structure_peaks", "structure_id"): + assert_array_equal(shifted[field], peaks[field]) + + def test_find_structures_empty(self): + significance = np.zeros((8, 8), dtype=np.float32) + peaks = np.zeros(0, dtype=DETECTION_DTYPE) + find_structures(peaks, significance) + self.assertEqual(len(peaks), 0) + + def test_detect_peaks_structures(self): + # Three sources in a row, close enough to connect at the first + # scale but far enough apart to resolve, share one structure. + rng = np.random.default_rng(7) + shape = (48, 64) + yy, xx = np.mgrid[0 : shape[0], 0 : shape[1]] + images = rng.standard_normal((self.n_bands,) + shape).astype(np.float32) + for cy, cx in [(24, 24), (24, 32), (24, 40)]: + images += 40.0 * np.exp(-((yy - cy) ** 2 + (xx - cx) ** 2) / (2 * 2.5**2)).astype(np.float32) + variance = np.ones_like(images) + result = detect_peaks(images, variance, scales=3, peak_thresh=5, footprint_thresh=3) + self.assertEqual(len(result.peaks), 3) + assert_array_equal(result.peaks["structure_id"], 1) + assert_array_equal(result.peaks["structure_peaks"], 3) + self.assertEqual(result.metadata["structure_thresh"], 1.0) + + # Well separated sources each get a structure of their own. + result = detect_peaks(self.images, self.variance, scales=3, peak_thresh=5, footprint_thresh=3) + assert_array_equal(np.sort(result.peaks["structure_id"]), [1, 2, 3]) + assert_array_equal(result.peaks["structure_peaks"], 1) + + def test_detect_peaks_validation(self): + with self.assertRaisesRegex(ValueError, "one value per band"): + detect_peaks(self.images, self.variance, psf_fwhm=[3.0, 3.0]) + with self.assertRaisesRegex(ValueError, "radius_rule"): + detect_peaks(self.images, self.variance, radius_rule="mean") + + def test_detect_peaks_empty(self): + # A pure-noise floor with a threshold nothing clears yields no peaks + # but still well-formed tables. + images = np.zeros((self.n_bands,) + self.shape, dtype=np.float32) + variance = np.ones_like(images) + result = detect_peaks(images, variance, scales=3, peak_thresh=5, footprint_thresh=3) + self.assertEqual(len(result.candidates), 0) + self.assertEqual(len(result.positions), 0) + self.assertEqual(len(result.peaks), 0) + self.assertEqual(result.footprints, []) + self.assertEqual(result.peaks.dtype, DETECTION_DTYPE) + self.assertEqual(result.positions.dtype, POSITION_DTYPE) + + def test_dilate_bad_mask(self): + mask = np.zeros((21, 21), dtype=bool) + mask[10, 10] = True + + # A single pixel grows to the disk of the given radius. + dilated = _dilate_bad_mask(mask, 3.0) + yy, xx = np.mgrid[0:21, 0:21] + expected = (yy - 10) ** 2 + (xx - 10) ** 2 <= 3.0**2 + assert_array_equal(dilated, expected) + + # A non-positive radius and an empty mask are both no-ops. + assert_array_equal(_dilate_bad_mask(mask, 0.0), mask) + empty = np.zeros((5, 5), dtype=bool) + assert_array_equal(_dilate_bad_mask(empty, 3.0), empty) + + def test_find_peak_candidates_bad_pixel_mask(self): + ny, nx = 40, 40 + n_bands = 2 + significance_map = np.zeros((1, n_bands + 1, ny, nx), dtype=np.float32) + yy, xx = np.mgrid[0:ny, 0:nx] + # The same 10-sigma bump in both single-band planes and the chi plane. + center = (25, 12) + bump = 10.0 * np.exp(-((yy - center[0]) ** 2 + (xx - center[1]) ** 2) / (2 * 2.0**2)) + significance_map[0, 0] = bump + significance_map[0, 1] = bump + significance_map[0, n_bands] = bump + + # Band 0 is bad at the peak; band 1 is clean. + bad = np.zeros((n_bands, ny, nx), dtype=bool) + bad[0, center[0], center[1]] = True + + candidates, _ = _find_peak_candidates( + significance_map, + min_separation=1, + min_area=1, + peak_thresh=5, + footprint_thresh=3, + first_scale=1, + bad_pixel_mask=bad, + dilation_radius=2.0, + ) + # The candidate is dropped in the bad band, kept in the clean band. + self.assertEqual(np.sum(candidates["band"] == 0), 0) + self.assertEqual(np.sum(candidates["band"] == 1), 1) + # The chi plane takes the union of the bands, so it is dropped too. + self.assertEqual(np.sum(candidates["band"] == n_bands), 0) + + def test_find_peak_candidates_dilation_radius(self): + ny, nx = 40, 40 + n_bands = 1 + significance_map = np.zeros((1, n_bands + 1, ny, nx), dtype=np.float32) + yy, xx = np.mgrid[0:ny, 0:nx] + center = (20, 20) + significance_map[0, 0] = 10.0 * np.exp( + -((yy - center[0]) ** 2 + (xx - center[1]) ** 2) / (2 * 2.0**2) + ) + # A bad pixel five pixels from the peak. + bad = np.zeros((n_bands, ny, nx), dtype=bool) + bad[0, center[0], center[1] - 5] = True + + kw = dict(min_separation=1, min_area=1, peak_thresh=5, footprint_thresh=3, first_scale=1) + # A radius short of the peak leaves it; one that reaches it drops it. + near, _ = _find_peak_candidates(significance_map, bad_pixel_mask=bad, dilation_radius=2.0, **kw) + far, _ = _find_peak_candidates(significance_map, bad_pixel_mask=bad, dilation_radius=5.0, **kw) + self.assertEqual(np.sum(near["band"] == 0), 1) + self.assertEqual(np.sum(far["band"] == 0), 0) + + def test_find_peak_candidates_dilation_grows_with_scale(self): + # The same peak and bad pixel in two scales. The bad pixel is beyond + # the base radius but within its doubling at the next scale, so the + # peak survives at the finer scale and is suppressed at the coarser. + ny, nx = 40, 40 + n_bands = 1 + significance_map = np.zeros((2, n_bands + 1, ny, nx), dtype=np.float32) + yy, xx = np.mgrid[0:ny, 0:nx] + center = (20, 20) + bump = 10.0 * np.exp(-((yy - center[0]) ** 2 + (xx - center[1]) ** 2) / (2 * 2.0**2)) + significance_map[0, 0] = bump + significance_map[1, 0] = bump + # Six pixels away: outside a base radius of 4, inside its double of 8. + bad = np.zeros((n_bands, ny, nx), dtype=bool) + bad[0, center[0], center[1] - 6] = True + + candidates, _ = _find_peak_candidates( + significance_map, + min_separation=1, + min_area=1, + peak_thresh=5, + footprint_thresh=3, + first_scale=1, + bad_pixel_mask=bad, + dilation_radius=4.0, + ) + band0 = candidates[candidates["band"] == 0] + assert_array_equal(np.sort(band0["scale"]), [1]) + + def test_detect_peaks_bad_pixel_mask_keeps_clean_band(self): + # Mark the first source bad in every band but one; it stays detected, + # seeded by the clean band, and the others are unaffected. + bad = np.zeros_like(self.images, dtype=bool) + cy, cx = self.centers[0] + bad[1:, cy - 3 : cy + 4, cx - 3 : cx + 4] = True + + result = detect_peaks( + self.images, + self.variance, + scales=3, + peak_thresh=5, + footprint_thresh=3, + bad_pixel_mask=bad, + dilation_radius=3.0, + ) + self.assertEqual(len(result.peaks), len(self.centers)) + matched = [p for p in result.peaks if abs(p["y"] - cy) <= 1 and abs(p["x"] - cx) <= 1] + self.assertEqual(len(matched), 1) + # Only the clean band and no chi plane seed the suppressed source. + self.assertEqual(matched[0]["band_flags"], 1 << 0) + + def test_detect_peaks_bad_pixel_mask_validation(self): + with self.assertRaisesRegex(ValueError, "same shape as images"): + detect_peaks( + self.images, + self.variance, + bad_pixel_mask=np.zeros(self.shape, dtype=bool), + ) diff --git a/tests/test_initialization.py b/tests/test_initialization.py index 6493cfa2..471cfbaf 100644 --- a/tests/test_initialization.py +++ b/tests/test_initialization.py @@ -20,14 +20,11 @@ # along with this program. If not, see . import os -from unittest.mock import patch import numpy as np -from deprecated.sphinx import deprecated from lsst.scarlet.lite import Box, Image, Observation from lsst.scarlet.lite.initialization import ( FactorizedInitialization, - FactorizedWaveletInitialization, init_monotonic_morph, multifit_spectra, trim_morphology, @@ -344,57 +341,3 @@ def test_factorized_chi2_init(self): self.assertEqual(len(init.sources), 8) for src in init.sources: self.assertEqual(src.get_model().dtype, np.float32) - - @deprecated( - version="v29.0", - reason="FactorizedWaveletInitialization is deprecated and will be removed after v29.0", - ) - def test_wavelet_init_source_falls_back_to_psf(self): - """init_source must always return a Source with at least one - component, even when individual init paths fail. - - Audit finding I-1: when get_single_component returned None or - the two-component path produced all-zero spectra, ``components`` - was either left unbound (UnboundLocalError) or set to an empty - list. Both cases must now fall back to a PSF component. - """ - init = FactorizedWaveletInitialization(self.observation, self.centers) - int_centers = [(int(round(c[0])), int(round(c[1]))) for c in self.centers] - - # Failure mode 1: get_single_component always returns None. - # Centers hitting the single-component branch or the - # two-component fallback must fall back to PSF rather than - # raising or producing an empty Source. - with patch.object(FactorizedWaveletInitialization, "get_single_component", return_value=None): - for center in int_centers: - source = init.init_source(center) - self.assertGreater(len(source.components), 0) - - # Failure mode 2: two-component path returns all-zero spectra - # for both bulge and disk -> empty components list, also a fall - # back case. - n_bands = len(self.observation.bands) - with patch( - "lsst.scarlet.lite.initialization.multifit_spectra", - return_value=np.zeros((2, n_bands), dtype=np.float32), - ): - for center in int_centers: - source = init.init_source(center) - self.assertGreater(len(source.components), 0) - - @deprecated( - version="v29.0", - reason="FactorizedWaveletInitialization is deprecated and will be removed after v29.0", - ) - def test_factorized_wavelet_init(self): - # Test default parameters - init = FactorizedWaveletInitialization(self.observation, self.centers) - self.assertEqual(init.observation, self.observation) - self.assertEqual(init.min_snr, 50) - self.assertIsNone(init.monotonicity) - self.assertTupleEqual((init.py, init.px), (7, 7)) - self.assertEqual(len(init.sources), 7) - components = np.sum([len(src.components) for src in init.sources]) - self.assertEqual(components, 8) - for src in init.sources: - self.assertEqual(src.get_model().dtype, np.float32) diff --git a/tests/test_wavelet.py b/tests/test_wavelet.py index 7b59d2cc..2fc6a047 100644 --- a/tests/test_wavelet.py +++ b/tests/test_wavelet.py @@ -23,14 +23,16 @@ import numpy as np from lsst.scarlet.lite.wavelet import ( + MultiResolutionSupport, apply_wavelet_denoising, + get_multiband_multiresolution_support, get_multiresolution_support, multiband_starlet_reconstruction, multiband_starlet_transform, starlet_reconstruction, starlet_transform, ) -from numpy.testing import assert_almost_equal +from numpy.testing import assert_allclose, assert_almost_equal from utils import ScarletTestCase @@ -91,6 +93,38 @@ def test_extras(self): get_multiresolution_support(image, starlets, 0.1, image_type="space") apply_wavelet_denoising(image) + def test_multiband_multiresolution_support(self): + image = self.data["images"].astype(float) + bands = image.shape[0] + starlets = multiband_starlet_transform(image, scales=3) + + support = get_multiband_multiresolution_support(image, starlets, 0.1) + self.assertTrue(support.multiband) + self.assertTupleEqual(support.support.shape, starlets.shape) + self.assertTupleEqual(support.sigma.shape, (len(starlets), bands)) + + # Each band must match a direct single-band call. + for band in range(bands): + single = get_multiresolution_support(image[band], starlets[:, band], 0.1) + np.testing.assert_array_equal(support.support[:, band], single.support) + np.testing.assert_array_equal(support.sigma[:, band], single.sigma) + + # A per-band sigma is accepted as an alternative to a scalar. + per_band = get_multiband_multiresolution_support(image, starlets, np.full(bands, 0.1)) + np.testing.assert_array_equal(per_band.support, support.support) + + def test_multiresolution_support_dimension_check(self): + # The band axis makes multiband support 4D and sigma 2D; the + # single-band case is one dimension lower on each. A mismatch + # against the ``multiband`` flag is rejected. + support = np.zeros((3, 2, 8, 8), dtype=int) + sigma = np.zeros((2, 3)) + MultiResolutionSupport(support=support, sigma=sigma, multiband=True) + with self.assertRaises(ValueError): + MultiResolutionSupport(support=support, sigma=sigma, multiband=False) + with self.assertRaises(ValueError): + MultiResolutionSupport(support=support[0], sigma=sigma, multiband=True) + def test_ground_branch_unbiased_sigma(self): """Audit finding D-5: the per-scale noise estimate in the ``image_type='ground'`` branch must compute std over the @@ -202,3 +236,19 @@ def test_space_branch_iterates_sigma(self): overestimate_count = result_overestimate.support.sum() self.assertGreater(overestimate_count, 0) self.assertLess(abs(overestimate_count - correct_count), correct_count) + + def test_space_branch_sigma_scaling(self): + """The ``image_type='space'`` branch scales the per-scale noise + by the estimated image noise. + """ + rng = np.random.default_rng(3) + image_sigma = 5.0 + image = rng.standard_normal((256, 256)) * image_sigma + starlets = starlet_transform(image, scales=4, generation=2) + + result = get_multiresolution_support(image, starlets, np.std(image), image_type="space") + + # The transform is linear, so for pure noise the true per-scale + # coefficient noise is the std of each scale's coefficients. + true_sigma = np.array([scale.std() for scale in starlets]) + assert_allclose(result.sigma, true_sigma, rtol=0.1)