diff --git a/CHANGES.rst b/CHANGES.rst index 8de2e392b..29d0a1a8b 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -295,20 +295,19 @@ New Features results identical to the single-threaded computation. [#2407] - Significantly improved the performance of source deblending in - ``deblend_sources`` and ``SourceFinder``, producing identical results. - The multithreshold watershed markers are now built by compiled code. - Deblending is typically ~7-30 times faster, both for fields of many - small blended sources and for large segments with many markers. - [#2408] + ``deblend_sources`` and ``SourceFinder``, producing identical + results. The multithreshold watershed markers are now built by + compiled code. Deblending is typically ~7-30 times faster, both for + fields of many small blended sources and for large segments with + many markers. [#2408, #2413] - Added an ``n_threads`` keyword to ``deblend_sources`` and ``SourceFinder`` to deblend the sources using multiple threads. The sources are divided into chunks and processed concurrently, - producing results identical to the single-threaded computation. - The marker-building kernels release the GIL, as do the - watershed and most of the array operations, so multithreading - can significantly speed up deblending, especially for large - sources. [#2409] + producing results identical to the single-threaded computation. Each + chunk is deblended by a few compiled calls that release the GIL, so + multithreading speeds up deblending for fields of many small sources + as well as for large sources. [#2409, #2413] - Added a ``contrast_method`` keyword to ``deblend_sources`` and ``SourceFinder`` to select the flux used by the deblending contrast diff --git a/benchmarks/bench_deblend.py b/benchmarks/bench_deblend.py index f910bda41..80301886e 100755 --- a/benchmarks/bench_deblend.py +++ b/benchmarks/bench_deblend.py @@ -324,8 +324,8 @@ def bench_stages(*, size=1000, n_peaks=25, mode='exponential', * make_markers: the multithreshold levels, the level quantization, and the compiled component-tree kernel * watershed: a single watershed call over the cutout - * apply_watershed: the watershed contrast loop (one watershed - call per removed marker) + * apply_watershed: the reference watershed contrast loop (one + watershed call per removed marker) * deblend_source: the full pipeline Parameters diff --git a/docs/whats_new/3.1.rst b/docs/whats_new/3.1.rst index 6a4609860..94c918222 100644 --- a/docs/whats_new/3.1.rst +++ b/docs/whats_new/3.1.rst @@ -653,7 +653,7 @@ Source Deblending Performance Improvements ========================================== Source deblending with :func:`~photutils.segmentation.deblend_sources` -and `~photutils.segmentation.SourceFinder` is now typically 7--30 times +and `~photutils.segmentation.SourceFinder` is now typically 7--50 times faster, producing identical results. The whole watershed step now runs in compiled code. @@ -667,8 +667,9 @@ A new ``n_threads`` keyword in `~photutils.segmentation.SourceFinder` deblends the sources using multiple threads. The sources are divided into chunks and processed concurrently, producing results identical to the single-threaded -computation. The compiled kernels release the GIL, so multithreading can -significantly speed up deblending, especially for large sources. +computation. Each chunk of sources is deblended by a few compiled calls +that release the GIL, so multithreading also speeds up fields of many +small sources. A new ``contrast_method`` keyword selects the flux used by the deblending contrast criterion. The ``'basin'`` method preserves the diff --git a/photutils/segmentation/_deblend_markers.pyx b/photutils/segmentation/_deblend_markers.pyx index 23be97b3a..717c22a94 100644 --- a/photutils/segmentation/_deblend_markers.pyx +++ b/photutils/segmentation/_deblend_markers.pyx @@ -28,9 +28,8 @@ The multithreshold levels themselves are computed by the caller (vectorized in NumPy over all the sources of a chunk, see ``photutils.segmentation.deblend``) and passed in, so that they are bitwise identical to the pure-Python reference implementation on every -platform. The kernels here compute the per-source data extrema that the -levels depend on, quantize each cutout against its levels, and build the -markers. +platform. The kernels here compute the per-source data extrema and +flux, quantize each cutout against its levels, and build the markers. The kernels run without the GIL and use no global mutable state, so this module is safe to use from multiple threads, including on @@ -42,7 +41,7 @@ import numpy as np from libc.math cimport NAN, isnan from libc.stdlib cimport free, malloc, realloc -__all__ = ['deblend_markers_chunk', 'deblend_source_extrema', +__all__ = ['deblend_markers_chunk', 'deblend_source_stats', 'make_deblend_markers'] ctypedef fused data_t: @@ -708,16 +707,17 @@ cdef inline int _count_below(const double* thresholds, int n_levels, return lo -cdef void _source_extrema(const data_t* data, const segm_t* segm, - Py_ssize_t img_nx, long long label, - Py_ssize_t y0, Py_ssize_t y1, Py_ssize_t x0, - Py_ssize_t x1, double* smin, - double* smax) noexcept nogil: +cdef void _source_stats(const data_t* data, const segm_t* segm, + Py_ssize_t img_nx, long long label, + Py_ssize_t y0, Py_ssize_t y1, Py_ssize_t x0, + Py_ssize_t x1, double* smin, double* smax, + double* ssum) noexcept nogil: """ - Compute the minimum and maximum data value of one source segment. + Compute the minimum, maximum, and flux of one source segment. - NaN pixels are excluded. Both outputs are NaN if the segment has no - finite pixel. + NaN pixels are excluded. The flux is accumulated sequentially in + float64 in raster order. The minimum and maximum are NaN, and the + flux is 0, if the segment has no finite pixel. """ cdef Py_ssize_t iy, ix, idx cdef double value @@ -725,6 +725,7 @@ cdef void _source_extrema(const data_t* data, const segm_t* segm, smin[0] = NAN smax[0] = NAN + ssum[0] = 0.0 for iy in range(y1 - y0): for ix in range(x1 - x0): idx = (y0 + iy) * img_nx + x0 + ix @@ -733,6 +734,7 @@ cdef void _source_extrema(const data_t* data, const segm_t* segm, value = data[idx] if isnan(value): continue + ssum[0] += value if not has_value: smin[0] = value smax[0] = value @@ -743,18 +745,20 @@ cdef void _source_extrema(const data_t* data, const segm_t* segm, smax[0] = value -def deblend_source_extrema(const data_t[:, ::1] data, - const segm_t[:, ::1] segm_data, - const long long[::1] labels, - const long long[::1] y0, - const long long[::1] y1, - const long long[::1] x0, - const long long[::1] x1): +def deblend_source_stats(const data_t[:, ::1] data, + const segm_t[:, ::1] segm_data, + const long long[::1] labels, + const long long[::1] y0, + const long long[::1] y1, + const long long[::1] x0, + const long long[::1] x1): """ - Compute the minimum and maximum data value of each source segment. + Compute the minimum, maximum, and flux of each source segment. - NaN pixels are excluded, so the results are identical to the - ``nanmin`` and ``nanmax`` reductions over the segment pixels. + NaN pixels are excluded. The minimum and maximum are identical to + the ``nanmin`` and ``nanmax`` reductions over the segment pixels. + The flux is accumulated sequentially in float64 in raster order, + which is what ``np.cumsum(values, dtype=np.float64)[-1]`` computes. Parameters ---------- @@ -772,9 +776,10 @@ def deblend_source_extrema(const data_t[:, ::1] data, Returns ------- - source_min, source_max : 1D float64 `~numpy.ndarray` - The minimum and maximum data value of each source segment. Both - are NaN for a segment without any finite pixel. + source_min, source_max, source_sum : 1D float64 `~numpy.ndarray` + The minimum, maximum, and flux of each source segment. The + minimum and maximum are NaN, and the flux is 0, for a segment + without any finite pixel. """ cdef Py_ssize_t n_src = labels.shape[0] cdef Py_ssize_t img_nx = data.shape[1] @@ -782,16 +787,19 @@ def deblend_source_extrema(const data_t[:, ::1] data, smin_arr = np.empty(n_src, dtype=np.float64) smax_arr = np.empty(n_src, dtype=np.float64) + ssum_arr = np.empty(n_src, dtype=np.float64) cdef double[::1] smin_mv = smin_arr cdef double[::1] smax_mv = smax_arr + cdef double[::1] ssum_mv = ssum_arr with nogil: for isrc in range(n_src): - _source_extrema(&data[0, 0], &segm_data[0, 0], img_nx, - labels[isrc], y0[isrc], y1[isrc], x0[isrc], - x1[isrc], &smin_mv[isrc], &smax_mv[isrc]) + _source_stats(&data[0, 0], &segm_data[0, 0], img_nx, + labels[isrc], y0[isrc], y1[isrc], x0[isrc], + x1[isrc], &smin_mv[isrc], &smax_mv[isrc], + &ssum_mv[isrc]) - return smin_arr, smax_arr + return smin_arr, smax_arr, ssum_arr cdef Py_ssize_t _source_markers(const data_t* data, const segm_t* segm, @@ -847,7 +855,9 @@ def deblend_markers_chunk(const data_t[:, ::1] data, const long long[::1] y1, const long long[::1] x0, const long long[::1] x1, - const double[:, ::1] thresholds, *, + const double[:, ::1] thresholds, + int[::1] packed, + const Py_ssize_t[::1] starts, *, int n_pixels, int connectivity, int max_markers, saddle_limits=None): """ @@ -876,6 +886,18 @@ def deblend_markers_chunk(const data_t[:, ::1] data, The multithreshold levels of each source, with shape ``(n_sources, n_levels)`` and ascending along the second axis. + packed : 1D int32 `~numpy.ndarray` + The buffer that receives the marker image of every source. The + region of source ``i`` starts at ``starts[i]`` and holds its + ``(y1 - y0) * (x1 - x0)`` cutout pixels in raster order. Each + region is zeroed, and then the markers are written for the + sources that split into two or more markers (and no more than + ``max_markers`` when it is not negative). The other regions are + left at zero. + + starts : 1D intp `~numpy.ndarray` + The start index of each source's region in ``packed``. + n_pixels : int The minimum number of connected pixels an above-threshold component must have to be considered a source. @@ -885,9 +907,9 @@ def deblend_markers_chunk(const data_t[:, ::1] data, max_markers : int The number of markers above which the marker image of a source - is not built. Its marker count is still returned, so that the - caller can retry the source with other levels. A negative value - disables the limit. + is not kept (its region is left at zero). Its marker count is + still returned, so that the caller can retry the source with + other levels. A negative value disables the limit. saddle_limits : 1D float64 `~numpy.ndarray` or `None`, optional If given, the markers are selected with the saddle contrast @@ -899,28 +921,31 @@ def deblend_markers_chunk(const data_t[:, ::1] data, Returns ------- - markers_list : list of 2D int `~numpy.ndarray` or `None` - The cutout marker image of each source that splits into two or - more markers (and no more than ``max_markers``), otherwise - `None`. - n_markers : 1D intp `~numpy.ndarray` The number of markers found for each source. """ cdef Py_ssize_t n_src = labels.shape[0] cdef Py_ssize_t img_nx = data.shape[1] cdef int n_levels = thresholds.shape[1] - cdef Py_ssize_t isrc, n_tot, max_ntot, ny_c, nx_c + cdef Py_ssize_t isrc, n_tot, max_ntot, ny_c, nx_c, start, p cdef Py_ssize_t n_markers cdef bint use_saddle = saddle_limits is not None if thresholds.shape[0] != n_src: msg = 'thresholds must have one row per source' raise ValueError(msg) + if (y0.shape[0] != n_src or y1.shape[0] != n_src + or x0.shape[0] != n_src or x1.shape[0] != n_src + or starts.shape[0] != n_src): + msg = 'every per-source array must have one entry per source' + raise ValueError(msg) max_ntot = 1 for isrc in range(n_src): n_tot = (y1[isrc] - y0[isrc]) * (x1[isrc] - x0[isrc]) + if starts[isrc] + n_tot > packed.shape[0]: + msg = 'packed is too small for the source regions' + raise ValueError(msg) if n_tot > max_ntot: max_ntot = n_tot @@ -932,7 +957,6 @@ def deblend_markers_chunk(const data_t[:, ::1] data, node_of_root_arr = np.zeros(max_ntot, dtype=np.int32) order_arr = np.empty(max_ntot, dtype=np.int32) flood_arr = np.empty(max_ntot, dtype=np.int32) - markers_arr = np.zeros(max_ntot, dtype=np.int32) n_markers_arr = np.zeros(n_src, dtype=np.intp) cdef int[::1] q_mv = q_arr @@ -943,7 +967,6 @@ def deblend_markers_chunk(const data_t[:, ::1] data, cdef int[::1] node_of_root_mv = node_of_root_arr cdef int[::1] order_mv = order_arr cdef int[::1] flood_mv = flood_arr - cdef int[::1] markers_mv = markers_arr cdef Py_ssize_t[::1] n_markers_mv = n_markers_arr # The saddle criterion inputs, with workspaces used only by it @@ -965,14 +988,20 @@ def deblend_markers_chunk(const data_t[:, ::1] data, saddle.posimg = &posimg_mv[0] saddle.fsum = &fsum_mv[0] - markers_list = [] - for isrc in range(n_src): - ny_c = y1[isrc] - y0[isrc] - nx_c = x1[isrc] - x0[isrc] - if use_saddle: - saddle.limit = saddle_limits_mv[isrc] - saddle.thresholds = &thresholds[isrc, 0] - with nogil: + # The whole chunk runs without the GIL, so that the threads of a + # multithreaded deblend do not contend for it once per source. + cdef bint failed = False + with nogil: + for isrc in range(n_src): + ny_c = y1[isrc] - y0[isrc] + nx_c = x1[isrc] - x0[isrc] + n_tot = ny_c * nx_c + start = starts[isrc] + if use_saddle: + saddle.limit = saddle_limits_mv[isrc] + saddle.thresholds = &thresholds[isrc, 0] + for p in range(n_tot): + packed[start + p] = 0 n_markers = _source_markers( &data[0, 0], &segm_data[0, 0], img_nx, labels[isrc], y0[isrc], y1[isrc], x0[isrc], x1[isrc], n_pixels, @@ -980,18 +1009,16 @@ def deblend_markers_chunk(const data_t[:, ::1] data, &saddle, &q_mv[0], &parent_mv[0], &size_mv[0], &added_mv[0], &stamp_mv[0], &node_of_root_mv[0], &order_mv[0], - &flood_mv[0], &markers_mv[0]) - if n_markers < 0: - raise MemoryError - n_markers_mv[isrc] = n_markers - if n_markers >= 2 and (max_markers < 0 - or n_markers <= max_markers): - n_tot = ny_c * nx_c - quantized = q_arr[:n_tot].reshape(ny_c, nx_c) - markers = markers_arr[:n_tot].reshape(ny_c, nx_c) - markers_list.append(np.where(quantized > 0, markers, - np.int32(0))) - else: - markers_list.append(None) + &flood_mv[0], &packed[start]) + if n_markers < 0: + failed = True + break + if n_markers < 2 or (max_markers >= 0 + and n_markers > max_markers): + for p in range(n_tot): + packed[start + p] = 0 + n_markers_mv[isrc] = n_markers + if failed: + raise MemoryError - return markers_list, n_markers_arr + return n_markers_arr diff --git a/photutils/segmentation/_deblend_reference.py b/photutils/segmentation/_deblend_reference.py index 718947369..7d0841aaf 100644 --- a/photutils/segmentation/_deblend_reference.py +++ b/photutils/segmentation/_deblend_reference.py @@ -20,7 +20,7 @@ from photutils.segmentation._deblend_watershed import deblend_watershed from photutils.segmentation.core import _get_labels from photutils.segmentation.deblend import _MAX_MARKERS, _create_relabel_map -from photutils.utils._stats import nanmax, nanmin, nansum +from photutils.utils._stats import nanmax, nanmin def _detect_sources_deblend(data, threshold, n_pixels, *, footprint, @@ -153,7 +153,15 @@ def __init__(self, data, segment_data, label, deblend_params): data_values = data[self.segment_mask] self.source_min = nanmin(data_values) self.source_max = nanmax(data_values) - self.source_sum = nansum(data_values) + # The flux is accumulated sequentially in float64 in raster + # order, exactly as the compiled stats kernel does. np.cumsum + # is a sequential accumulation, unlike the pairwise np.nansum + finite_values = data_values[~np.isnan(data_values)] + if finite_values.size > 0: + self.source_sum = float(np.cumsum(finite_values, + dtype=np.float64)[-1]) + else: + self.source_sum = 0.0 self.warnings = {} @cached_property diff --git a/photutils/segmentation/_deblend_watershed.pyx b/photutils/segmentation/_deblend_watershed.pyx index f07e62cff..c30c6e60c 100644 --- a/photutils/segmentation/_deblend_watershed.pyx +++ b/photutils/segmentation/_deblend_watershed.pyx @@ -30,7 +30,8 @@ import numpy as np from libc.math cimport INFINITY, isnan from libc.stdlib cimport free, malloc -__all__ = ['deblend_source_contrast', 'deblend_watershed'] +__all__ = ['deblend_contrast_chunk', 'deblend_watershed', + 'write_deblended_labels'] ctypedef fused data_t: float @@ -229,7 +230,7 @@ def deblend_watershed(image, markers, mask, connectivity): This entry point is exported only for the pure-Python reference implementation in ``_deblend_reference`` and the cross-implementation tests. The production contrast loop calls the - watershed core directly through ``deblend_source_contrast``. + watershed core directly through ``deblend_contrast_chunk``. Parameters ---------- @@ -452,27 +453,35 @@ cdef int _contrast_core(const double* posimg, const double* negimg, return status -def deblend_source_contrast(const data_t[:, ::1] data, - const segm_t[:, ::1] segm_data, - long long label, Py_ssize_t y0, - Py_ssize_t y1, Py_ssize_t x0, - Py_ssize_t x1, markers, *, - int connectivity, double contrast, - double source_sum, double source_min, - bint apply_contrast=True): +def deblend_contrast_chunk(const data_t[:, ::1] data, + const segm_t[:, ::1] segm_data, + const long long[::1] labels, + const long long[::1] y0, + const long long[::1] y1, + const long long[::1] x0, + const long long[::1] x1, + int[::1] packed, + const Py_ssize_t[::1] starts, + const Py_ssize_t[::1] n_markers, *, + int connectivity, double contrast, + const double[::1] source_sum, + const double[::1] source_min, + bint apply_contrast): """ - Apply the watershed contrast loop to one source's markers. + Apply the watershed contrast loop to the markers of a chunk of + sources in place. - The watershed flooding, the basin flux measurements, the - below-contrast marker removal (single or batched), and the final - consecutive relabeling all run in compiled code that releases - the GIL, producing results identical to the per-step NumPy - implementation in ``_SingleSourceDeblender``. + For every source with two or more markers, the flooding, the basin + flux measurements, the below-contrast marker removal, and the final + consecutive relabeling run in compiled code that releases the GIL, + reusing one workspace sized to the largest cutout in the chunk. The + results are identical to the per-step NumPy implementation in + ``_SingleSourceDeblender``. Parameters ---------- data : 2D float `~numpy.ndarray` - The full data array. NaN pixels within the segment are flooded + The full data array. NaN pixels within a segment are flooded after all the finite pixels, so they are assigned to a neighboring basin. They contribute NaN to the flux of that basin, as in the NumPy implementation. @@ -480,15 +489,24 @@ def deblend_source_contrast(const data_t[:, ::1] data, segm_data : 2D int `~numpy.ndarray` The full segmentation array. - label : int - The label of the source segment. + labels : 1D int64 `~numpy.ndarray` + The label of each source in the chunk. - y0, y1, x0, x1 : int - The bounding-box slice bounds of the source. + y0, y1, x0, x1 : 1D int64 `~numpy.ndarray` + The bounding-box slice bounds of each source. - markers : 2D int `~numpy.ndarray` - The marker image cutout, with markers labeled from 1. It is - not modified. + packed : 1D int32 `~numpy.ndarray` + The packed marker buffer written by ``deblend_markers_chunk``. + On output, the region of every source that deblends holds its + final labels, consecutive from 1, and every other region is + zero. + + starts : 1D intp `~numpy.ndarray` + The start index of each source's region in ``packed``. + + n_markers : 1D intp `~numpy.ndarray` + The number of markers of each source. Sources with fewer than + two markers are skipped. connectivity : {8, 4} The pixel connectivity. @@ -497,14 +515,11 @@ def deblend_source_contrast(const data_t[:, ::1] data, The contrast criterion (the minimum fraction of the total source flux that a watershed basin must contain). - source_sum : float - The total flux of the source segment (NaN pixels excluded). + source_sum, source_min : 1D float64 `~numpy.ndarray` + The flux and the minimum data value of each source segment + (NaN pixels excluded). - source_min : float - The minimum data value of the source segment (NaN pixels - excluded). - - apply_contrast : bool, optional + apply_contrast : bool Whether to apply the contrast criterion. If `False`, a single watershed pass is run with no basin removal. This is used with the saddle contrast criterion, where the markers are already @@ -512,69 +527,188 @@ def deblend_source_contrast(const data_t[:, ::1] data, Returns ------- - output : 2D int `~numpy.ndarray` or `None` - The deblended cutout with consecutive labels starting from - 1, or `None` if only one basin remains after applying the - contrast criterion. + n_labels : 1D intp `~numpy.ndarray` + The number of final labels of each source. It is 0 for the + skipped sources and 1 for the sources whose basins were all but + one removed, and the packed regions of both are zero. Raises ------ ValueError - If the flooded basins do not cover the segment mask, which - happens when the detection and deblending connectivities - differ. + If the flooded basins of a source do not cover its segment, + which happens when the detection and deblending connectivities + differ, or if the per-source arrays do not have one entry per + source, or if ``packed`` is too small for the source regions. + + MemoryError + If a workspace allocation fails. """ - cdef Py_ssize_t ny_c = y1 - y0 - cdef Py_ssize_t nx_c = x1 - x0 + cdef Py_ssize_t n_src = labels.shape[0] cdef Py_ssize_t img_nx = data.shape[1] + cdef Py_ssize_t isrc, n_tot, max_ntot, ny_c, nx_c, start + cdef Py_ssize_t iy, ix, idx, p, n_max_labels + cdef Py_ssize_t failed = -1 + cdef int status = 0 + cdef bint conn8 = connectivity == 8 - posimg_arr = np.empty((ny_c, nx_c), dtype=np.float64) - negimg_arr = np.empty((ny_c, nx_c), dtype=np.float64) - mask_arr = np.empty((ny_c, nx_c), dtype=np.uint8) - output_arr = np.array(markers, dtype=np.int32, copy=True, order='C') + if (y0.shape[0] != n_src or y1.shape[0] != n_src + or x0.shape[0] != n_src or x1.shape[0] != n_src + or starts.shape[0] != n_src or n_markers.shape[0] != n_src + or source_sum.shape[0] != n_src + or source_min.shape[0] != n_src): + msg = 'every per-source array must have one entry per source' + raise ValueError(msg) - cdef double[:, ::1] posimg_mv = posimg_arr - cdef double[:, ::1] negimg_mv = negimg_arr - cdef unsigned char[:, ::1] mask_mv = mask_arr - cdef int[:, ::1] output_mv = output_arr - cdef double* posimg = &posimg_mv[0, 0] - cdef double* negimg = &negimg_mv[0, 0] - cdef unsigned char* mask = &mask_mv[0, 0] - cdef int* output = &output_mv[0, 0] + # The workspaces are sized to the largest cutout that is deblended + max_ntot = 1 + for isrc in range(n_src): + n_tot = (y1[isrc] - y0[isrc]) * (x1[isrc] - x0[isrc]) + if starts[isrc] + n_tot > packed.shape[0]: + msg = 'packed is too small for the source regions' + raise ValueError(msg) + if n_markers[isrc] >= 2 and n_tot > max_ntot: + max_ntot = n_tot + + posimg_arr = np.empty(max_ntot, dtype=np.float64) + negimg_arr = np.empty(max_ntot, dtype=np.float64) + mask_arr = np.empty(max_ntot, dtype=np.uint8) + output_arr = np.empty(max_ntot, dtype=np.int32) + n_labels_arr = np.zeros(n_src, dtype=np.intp) + cdef double[::1] posimg_mv = posimg_arr + cdef double[::1] negimg_mv = negimg_arr + cdef unsigned char[::1] mask_mv = mask_arr + cdef int[::1] output_mv = output_arr + cdef Py_ssize_t[::1] n_labels_mv = n_labels_arr + cdef double* posimg = &posimg_mv[0] + cdef double* negimg = &negimg_mv[0] + cdef unsigned char* mask = &mask_mv[0] + cdef int* output = &output_mv[0] cdef const data_t* data_ptr = &data[0, 0] cdef const segm_t* segm_ptr = &segm_data[0, 0] - cdef bint conn8 = connectivity == 8 - cdef Py_ssize_t iy, ix, idx, p - cdef Py_ssize_t n_max_labels = 0 - cdef int status - with nogil: - for iy in range(ny_c): - for ix in range(nx_c): - idx = (y0 + iy) * img_nx + x0 + ix - p = iy * nx_c + ix - posimg[p] = data_ptr[idx] - if isnan(posimg[p]): - # NaN pixels are flooded after all finite pixels - negimg[p] = INFINITY - else: - negimg[p] = -posimg[p] - mask[p] = segm_ptr[idx] == label - if output[p] > n_max_labels: - n_max_labels = output[p] - status = _contrast_core(posimg, negimg, mask, output, ny_c, - nx_c, conn8, contrast, source_sum, - source_min, n_max_labels, - apply_contrast) - - if status == -1: - raise MemoryError - if status == -2: - msg = (f'Deblending failed for source {int(label)!r}. ' + for isrc in range(n_src): + if n_markers[isrc] < 2: + continue + ny_c = y1[isrc] - y0[isrc] + nx_c = x1[isrc] - x0[isrc] + n_tot = ny_c * nx_c + start = starts[isrc] + n_max_labels = 0 + for iy in range(ny_c): + for ix in range(nx_c): + idx = (y0[isrc] + iy) * img_nx + x0[isrc] + ix + p = iy * nx_c + ix + posimg[p] = data_ptr[idx] + if isnan(posimg[p]): + # NaN pixels are flooded after all finite + # pixels + negimg[p] = INFINITY + else: + negimg[p] = -posimg[p] + mask[p] = segm_ptr[idx] == labels[isrc] + output[p] = packed[start + p] + if output[p] > n_max_labels: + n_max_labels = output[p] + status = _contrast_core(posimg, negimg, mask, output, ny_c, + nx_c, conn8, contrast, + source_sum[isrc], source_min[isrc], + n_max_labels, apply_contrast) + if status < 0: + failed = isrc + break + n_labels_mv[isrc] = status + if status >= 2: + for p in range(n_tot): + packed[start + p] = output[p] + else: + for p in range(n_tot): + packed[start + p] = 0 + + if failed >= 0: + if status == -1: + raise MemoryError + msg = (f'Deblending failed for source {int(labels[failed])!r}. ' 'Please ensure you used the same pixel connectivity ' 'in detect_sources and deblend_sources.') raise ValueError(msg) - if status == 1: # no deblending - return None - return output_arr + + return n_labels_arr + + +def write_deblended_labels(segm_t[:, ::1] segm_out, + const int[::1] packed, + const Py_ssize_t[::1] starts, + const long long[::1] y0, + const long long[::1] y1, + const long long[::1] x0, + const long long[::1] x1, + const Py_ssize_t[::1] n_labels, + const long long[::1] label_offsets): + """ + Write the deblended labels of a chunk into a segmentation array. + + For every source with two or more final labels, the nonzero + labels of its packed region are written into ``segm_out`` at its + bounding box, offset by ``label_offsets``. The other sources are + left untouched. + + Parameters + ---------- + segm_out : 2D int `~numpy.ndarray` + The segmentation array to write into. It must be a copy of the + input segmentation array. + + packed : 1D int32 `~numpy.ndarray` + The packed label buffer written by ``deblend_contrast_chunk``. + + starts : 1D intp `~numpy.ndarray` + The start index of each source's region in ``packed``. + + y0, y1, x0, x1 : 1D int64 `~numpy.ndarray` + The bounding-box slice bounds of each source. + + n_labels : 1D intp `~numpy.ndarray` + The number of final labels of each source. + + label_offsets : 1D int64 `~numpy.ndarray` + The value added to the labels of each source. + + Raises + ------ + ValueError + If the per-source arrays do not have one entry per source, or + if ``packed`` is too small for the source regions. + """ + cdef Py_ssize_t n_src = n_labels.shape[0] + cdef Py_ssize_t img_nx = segm_out.shape[1] + cdef Py_ssize_t isrc, ny_c, nx_c, iy, ix, idx, p, start + cdef int value + cdef segm_t* out_ptr = &segm_out[0, 0] + + if (y0.shape[0] != n_src or y1.shape[0] != n_src + or x0.shape[0] != n_src or x1.shape[0] != n_src + or starts.shape[0] != n_src + or label_offsets.shape[0] != n_src): + msg = 'every per-source array must have one entry per source' + raise ValueError(msg) + for isrc in range(n_src): + if (starts[isrc] + (y1[isrc] - y0[isrc]) * (x1[isrc] - x0[isrc]) + > packed.shape[0]): + msg = 'packed is too small for the source regions' + raise ValueError(msg) + + with nogil: + for isrc in range(n_src): + if n_labels[isrc] < 2: + continue + ny_c = y1[isrc] - y0[isrc] + nx_c = x1[isrc] - x0[isrc] + start = starts[isrc] + for iy in range(ny_c): + for ix in range(nx_c): + p = iy * nx_c + ix + value = packed[start + p] + if value != 0: + idx = (y0[isrc] + iy) * img_nx + x0[isrc] + ix + out_ptr[idx] = (value + label_offsets[isrc]) diff --git a/photutils/segmentation/deblend.py b/photutils/segmentation/deblend.py index e396a817c..b863f0801 100644 --- a/photutils/segmentation/deblend.py +++ b/photutils/segmentation/deblend.py @@ -12,14 +12,14 @@ from astropy.units import Quantity from photutils.segmentation._deblend_markers import (deblend_markers_chunk, - deblend_source_extrema) -from photutils.segmentation._deblend_watershed import deblend_source_contrast + deblend_source_stats) +from photutils.segmentation._deblend_watershed import (deblend_contrast_chunk, + write_deblended_labels) from photutils.segmentation.core import (SegmentationImage, _get_labels, _remap_deblend_label_map) from photutils.segmentation.flags import SEGMENTATION_FLAGS from photutils.segmentation.utils import _make_binary_structure from photutils.utils._deprecation import deprecated_renamed_argument -from photutils.utils._stats import nansum from photutils.utils.exceptions import DeblendWarning __all__ = ['deblend_sources'] @@ -83,6 +83,29 @@ class _DeblendParams: contrast_method: str = 'basin' +@dataclass +class _ChunkResult: + """ + The deblending results of a chunk of sources. + + The deblended labels of every source are stored in the packed + cutout layout of the compiled kernels. The region of source ``i`` + is ``packed[offsets[i]:offsets[i + 1]]``, its bounding-box cutout + in raster order. It holds consecutive labels from 1 if the source + deblended (``n_labels[i] >= 2``), and zeros otherwise. + """ + + n_labels: np.ndarray + packed: np.ndarray + offsets: np.ndarray + y0: np.ndarray + y1: np.ndarray + x0: np.ndarray + x1: np.ndarray + nonposmin: np.ndarray + n_markers_fallback: np.ndarray + + @deprecated_renamed_argument('segment_img', 'segmentation_image', '3.0', until='4.0') @deprecated_renamed_argument('npixels', 'n_pixels', '3.0', until='4.0') @@ -203,10 +226,10 @@ def deblend_sources(data, segmentation_image, n_pixels, *, labels=None, sources are divided into chunks and processed concurrently. The per-source results are independent and are assembled in the input-label order, so they are identical to the single-threaded - computation. The marker-building kernels release the Python - global interpreter lock (GIL), as do the watershed and most of - the array operations, so multithreading can significantly speed - up the deblending, especially for large sources. + computation. Each chunk is deblended by a few compiled calls + that release the Python global interpreter lock (GIL), so + multithreading speeds up the deblending of fields of many small + sources as well as of large sources. nproc : int, optional This keyword is deprecated and has no effect. It was the name of @@ -318,7 +341,6 @@ def deblend_sources(data, segmentation_image, n_pixels, *, labels=None, deblend_params = _DeblendParams(n_pixels, footprint, n_levels, contrast, mode, contrast_method) - segm_deblended = segmentation_image.data.copy() label_indices = segmentation_image.get_indices(labels) all_slices = [segmentation_image.slices[idx] for idx in label_indices] @@ -344,7 +366,7 @@ def deblend_sources(data, segmentation_image, n_pixels, *, labels=None, # their bounding-box area (the size of the cutouts the kernels # work on), so that the chunks carry similar amounts of work # regardless of the source ordering. The chunks are processed - # concurrently and the per-source results are scattered back + # concurrently and the per-source results are gathered back # into the input order, so they are identical to the # single-threaded computation. bbox_areas = np.array([(slc[0].stop - slc[0].start) @@ -360,39 +382,54 @@ def _run_chunk(indices): driver_segm, labels[indices], chunk_slices, deblend_params) - results = [None] * len(labels) with ThreadPoolExecutor(max_workers=n_chunks) as executor: - for indices, chunk_results in zip( - chunk_indices, executor.map(_run_chunk, chunk_indices), - strict=True): - for index, result in zip(indices, chunk_results, - strict=True): - results[index] = result + chunk_results = list(zip(chunk_indices, + executor.map(_run_chunk, + chunk_indices), + strict=True)) else: - results = _deblend_sources_chunk(data, segm_data, driver_data, - driver_segm, labels, - all_slices, deblend_params) - - deblend_label_map = {} - max_label = segmentation_image.max_label - nonposmin_labels = [] - n_markers_labels = [] - for label, source_slice, (source_deblended, warns) in zip( - labels, all_slices, results, strict=True): - if warns: - if 'nonposmin' in warns: - nonposmin_labels.append(label) - if 'n_markers' in warns: - n_markers_labels.append(label) - - if source_deblended is not None: - source_mask = source_deblended > 0 - new_segm = source_deblended[source_mask] # min label = 1 - segm_deblended[source_slice][source_mask] = ( - new_segm + max_label) - new_labels = _get_labels(new_segm) + max_label - deblend_label_map[label] = new_labels - max_label += len(new_labels) + chunk_results = [(np.arange(len(labels)), + _deblend_sources_chunk(data, segm_data, + driver_data, driver_segm, + labels, all_slices, + deblend_params))] + + # Gather the per-source counts and fallback flags in the input + # order. The deblended labels of each source follow its + # predecessors, so the label offsets are a cumulative sum of the + # counts, starting after the largest input label + n_labels = np.zeros(len(labels), dtype=np.intp) + nonposmin = np.zeros(len(labels), dtype=bool) + n_markers_fallback = np.zeros(len(labels), dtype=bool) + for indices, result in chunk_results: + n_labels[indices] = result.n_labels + nonposmin[indices] = result.nonposmin + n_markers_fallback[indices] = result.n_markers_fallback + deblended = n_labels >= 2 + counts = np.where(deblended, n_labels, 0) + label_offsets = np.zeros(len(labels), dtype=np.int64) + np.cumsum(counts[:-1], out=label_offsets[1:]) + label_offsets += segmentation_image.max_label + + segm_out = driver_segm.copy() + for indices, result in chunk_results: + write_deblended_labels(segm_out, result.packed, + result.offsets[:-1], result.y0, + result.y1, result.x0, result.x1, + result.n_labels, label_offsets[indices]) + segm_deblended = segm_out.astype(segm_data.dtype, copy=False) + + # The child labels carry the dtype of the output segmentation + # image, as the image itself does + deblend_label_map = { + int(label): (np.arange(1, count + 1) + offset).astype( + segm_deblended.dtype) + for label, count, offset in zip(labels[deblended], + counts[deblended], + label_offsets[deblended], + strict=True)} + nonposmin_labels = list(labels[nonposmin]) + n_markers_labels = list(labels[n_markers_fallback]) if nonposmin_labels or n_markers_labels: msg = ('The deblending mode of one or more source labels from the ' @@ -538,18 +575,24 @@ def _compute_thresholds(source_min, source_max, n_levels, mode): nonposmin) -def _deblend_sources_chunk(data, segm_data, driver_data, driver_segm, - labels, slices, deblend_params): +def _deblend_sources_chunk(data, segm_data, # noqa: ARG001 + driver_data, driver_segm, labels, slices, + deblend_params): """ Deblend a chunk of labeled sources. + The signature is shared with the pure-Python chunk function of the + cross-implementation tests, which needs ``segm_data`` even though + the compiled kernels read only ``driver_segm``. + The data extrema of every source in the chunk are computed by a compiled kernel, the multithreshold levels and the mode fallbacks are computed for all sources at once in NumPy, the level quantization and the marker construction run in the compiled chunk driver (which releases the GIL and reuses one workspace across the chunk), and the watershed and contrast steps then run for the - sources that split into two or more markers. + sources that split into two or more markers, in one compiled call + over the chunk. Parameters ---------- @@ -574,9 +617,8 @@ def _deblend_sources_chunk(data, segm_data, driver_data, driver_segm, Returns ------- - results : list of (2D `~numpy.ndarray` or `None`, dict) - For each source, the deblended cutout (`None` if the source - did not deblend) and the mode-fallback warnings dictionary. + result : `_ChunkResult` + The per-source deblended labels and mode-fallback flags. """ labels = np.asarray(labels, dtype=np.int64) y0 = np.array([slc[0].start for slc in slices], dtype=np.int64) @@ -589,28 +631,28 @@ def _deblend_sources_chunk(data, segm_data, driver_data, driver_segm, chunk_kwargs = {'n_pixels': int(deblend_params.n_pixels), 'connectivity': connectivity} - source_min, source_max = deblend_source_extrema( + source_min, source_max, source_sum = deblend_source_stats( driver_data, driver_segm, labels, y0, y1, x0, x1) # Constant (or all-NaN) sources do not deblend. The multithreshold # levels of the other sources are computed in the data dtype, as # the pure-Python reference implementation does active = np.flatnonzero(source_min < source_max) - markers_list = [None] * len(labels) nonposmin = np.zeros(len(labels), dtype=bool) n_markers_fallback = np.zeros(len(labels), dtype=bool) + n_markers = np.zeros(len(labels), dtype=np.intp) + + # One packed buffer holds the bounding-box cutout of every source + # in the chunk, back to back. The kernels write the markers and + # then the deblended labels into it, so no per-source arrays are + # created in Python + sizes = (y1 - y0) * (x1 - x0) + offsets = np.zeros(len(labels) + 1, dtype=np.intp) + np.cumsum(sizes, out=offsets[1:]) + packed = np.zeros(offsets[-1], dtype=np.int32) + starts = offsets[:-1] - # The saddle criterion selects the markers against the total source - # fluxes, so those are computed up front. The basin criterion needs - # them only for the sources that split. use_saddle = deblend_params.contrast_method == 'saddle' - source_sums = {} - if use_saddle: - for index in active: - slc = slices[index] - values = data[slc][segm_data[slc] == labels[index]] - source_sums[index] = float(nansum(values)) - if active.size > 0: values_dtype = data.dtype.newbyteorder('=') smin = source_min[active].astype(values_dtype) @@ -620,9 +662,7 @@ def _deblend_sources_chunk(data, segm_data, driver_data, driver_segm, nonposmin[active] = fallback saddle_limits = None if use_saddle: - saddle_limits = np.array([deblend_params.contrast - * source_sums[index] - for index in active]) + saddle_limits = deblend_params.contrast * source_sum[active] # Sources with too many markers are only retried with linearly # spaced levels (below), so only then may the kernel skip # building their markers. With the saddle criterion the markers @@ -630,64 +670,35 @@ def _deblend_sources_chunk(data, segm_data, driver_data, driver_segm, # watershed, so no fallback is needed. can_retry = mode != 'linear' and not use_saddle max_markers = _MAX_MARKERS if can_retry else -1 - markers, n_markers = deblend_markers_chunk( + n_markers[active] = deblend_markers_chunk( driver_data, driver_segm, labels[active], y0[active], - y1[active], x0[active], x1[active], thresholds, - max_markers=max_markers, saddle_limits=saddle_limits, - **chunk_kwargs) - for index, source_markers in zip(active, markers, strict=True): - markers_list[index] = source_markers + y1[active], x0[active], x1[active], thresholds, packed, + starts[active], max_markers=max_markers, + saddle_limits=saddle_limits, **chunk_kwargs) # Too many markers make the watershed step very slow, so such # sources are deblended again with linearly spaced levels if can_retry: - retry = np.flatnonzero(n_markers > _MAX_MARKERS) + retry = np.flatnonzero(n_markers[active] > _MAX_MARKERS) if retry.size > 0: thresholds, _ = _compute_thresholds( smin[retry], smax[retry], n_levels, 'linear') retry = active[retry] - markers, _ = deblend_markers_chunk( + n_markers[retry] = deblend_markers_chunk( driver_data, driver_segm, labels[retry], y0[retry], - y1[retry], x0[retry], x1[retry], thresholds, - max_markers=-1, **chunk_kwargs) - for index, source_markers in zip(retry, markers, - strict=True): - markers_list[index] = source_markers + y1[retry], x0[retry], x1[retry], thresholds, packed, + starts[retry], max_markers=-1, **chunk_kwargs) n_markers_fallback[retry] = True - results = [] - for index, (label, slc, markers) in enumerate( - zip(labels, slices, markers_list, strict=True)): - warns = {} - if nonposmin[index]: - warns['nonposmin'] = 'non-positive minimum' - if n_markers_fallback[index]: - warns['n_markers'] = 'too many markers' - if markers is None: - results.append((None, warns)) - continue - - # The total source flux is computed with the same reduction as - # the per-source Python path (its rounding depends on the - # summation order) and the source minimum comes from the - # compiled extrema. Both are passed to the compiled contrast - # loop. With the saddle criterion, the markers are already - # contrast-selected, so the basin removal is disabled. - if use_saddle: - source_sum = source_sums[index] - else: - values = data[slc][segm_data[slc] == label] - source_sum = float(nansum(values)) - source_deblended = deblend_source_contrast( - driver_data, driver_segm, int(label), int(y0[index]), - int(y1[index]), int(x0[index]), int(x1[index]), markers, - connectivity=connectivity, - contrast=float(deblend_params.contrast), - source_sum=source_sum, source_min=float(source_min[index]), - apply_contrast=not use_saddle) - results.append((source_deblended, warns)) - - return results + n_labels = deblend_contrast_chunk( + driver_data, driver_segm, labels, y0, y1, x0, x1, packed, starts, + n_markers, connectivity=connectivity, + contrast=float(deblend_params.contrast), source_sum=source_sum, + source_min=source_min, apply_contrast=not use_saddle) + + return _ChunkResult(n_labels=n_labels, packed=packed, offsets=offsets, + y0=y0, y1=y1, x0=x0, x1=x1, nonposmin=nonposmin, + n_markers_fallback=n_markers_fallback) def _make_flags_map(deblend_label_map, nonposmin_labels, n_markers_labels, diff --git a/photutils/segmentation/tests/test_deblend.py b/photutils/segmentation/tests/test_deblend.py index 48efe6568..8066c3798 100644 --- a/photutils/segmentation/tests/test_deblend.py +++ b/photutils/segmentation/tests/test_deblend.py @@ -17,10 +17,13 @@ from photutils.segmentation import SegmentationImage from photutils.segmentation import deblend as deblend_module from photutils.segmentation import deblend_sources, detect_sources -from photutils.segmentation._deblend_markers import deblend_markers_chunk +from photutils.segmentation._deblend_markers import (deblend_markers_chunk, + deblend_source_stats) from photutils.segmentation._deblend_reference import _SingleSourceDeblender -from photutils.segmentation._deblend_watershed import deblend_watershed -from photutils.segmentation.deblend import (_compute_thresholds, +from photutils.segmentation._deblend_watershed import (deblend_contrast_chunk, + deblend_watershed, + write_deblended_labels) +from photutils.segmentation.deblend import (_ChunkResult, _compute_thresholds, _create_relabel_map, _DeblendParams) from photutils.segmentation.flags import SEGMENTATION_FLAGS @@ -220,6 +223,8 @@ def test_deblend_sources_norelabel(self, mode): assert result.n_labels == 2 assert_equal(result.labels, [2, 3]) assert_equal(result.parent_to_deblended_labels, {1: [2, 3]}) + assert (result.parent_to_deblended_labels[1].dtype + == result.data.dtype) assert len(result.slices) <= result.max_label assert len(result.slices) == result.n_labels assert_allclose(np.nonzero(self.segm), np.nonzero(result)) @@ -673,16 +678,35 @@ def python_deblend_chunk(data, segm_data, driver_data, # noqa: ARG001 Returns ------- - results : list of (2D `~numpy.ndarray` or `None`, dict) - The deblended cutout and warnings for each source. + result : `_ChunkResult` + The per-source deblended labels in the packed layout of the + compiled chunk driver. """ - results = [] - for label, slc in zip(labels, slices, strict=True): + y0 = np.array([slc[0].start for slc in slices], dtype=np.int64) + y1 = np.array([slc[0].stop for slc in slices], dtype=np.int64) + x0 = np.array([slc[1].start for slc in slices], dtype=np.int64) + x1 = np.array([slc[1].stop for slc in slices], dtype=np.int64) + sizes = (y1 - y0) * (x1 - x0) + offsets = np.zeros(len(labels) + 1, dtype=np.intp) + np.cumsum(sizes, out=offsets[1:]) + packed = np.zeros(offsets[-1], dtype=np.int32) + n_labels = np.zeros(len(labels), dtype=np.intp) + nonposmin = np.zeros(len(labels), dtype=bool) + n_markers_fallback = np.zeros(len(labels), dtype=bool) + for index, (label, slc) in enumerate(zip(labels, slices, + strict=True)): deblender = _SingleSourceDeblender(data[slc], segm_data[slc], label, deblend_params) - results.append((deblender.deblend_source(), - deblender.warnings)) - return results + deblended = deblender.deblend_source() + nonposmin[index] = 'nonposmin' in deblender.warnings + n_markers_fallback[index] = 'n_markers' in deblender.warnings + if deblended is not None: + n_labels[index] = deblended.max() + packed[offsets[index]:offsets[index + 1]] = deblended.ravel() + return _ChunkResult(n_labels=n_labels, packed=packed, + offsets=offsets, y0=y0, y1=y1, x0=x0, x1=x1, + nonposmin=nonposmin, + n_markers_fallback=n_markers_fallback) @pytest.mark.parametrize('contrast_method', ['basin', 'saddle']) @@ -943,6 +967,265 @@ def test_compute_thresholds_matches_reference(dtype, mode, n_levels): assert nonposmin[i] == ('nonposmin' in deblender.warnings) +@pytest.mark.parametrize('dtype', ['float64', 'float32', 'int32']) +def test_source_stats_matches_reference(dtype): + """ + Test that the compiled per-source minimum, maximum, and flux are + identical to the reference implementation, with NaN pixels + excluded, and that the flux agrees with np.nansum. + """ + data, segm = make_multipeak_source() + data = data.astype(dtype) + if dtype != 'int32': + data[45:48, 40:43] = np.nan + driver_data = np.ascontiguousarray(data, dtype=np.float64) + labels = np.asarray(segm.labels, dtype=np.int64) + slc = segm.slices[0] + y0 = np.array([slc[0].start]) + y1 = np.array([slc[0].stop]) + x0 = np.array([slc[1].start]) + x1 = np.array([slc[1].stop]) + smin, smax, ssum = deblend_source_stats(driver_data, segm.data, + labels, y0, y1, x0, x1) + + params = _DeblendParams(5, np.ones((3, 3)), 32, 0.001, 'linear') + deblender = _SingleSourceDeblender(data[slc], segm.data[slc], 1, + params) + assert smin[0] == deblender.source_min + assert smax[0] == deblender.source_max + assert ssum[0] == deblender.source_sum + values = data[segm.data == 1] + assert_allclose(ssum[0], np.nansum(values, dtype=np.float64), + rtol=1e-12) + + +def make_packed_pair(): + """ + Return a two-source scene in the packed chunk layout. + + The first source is a blended pair that splits into two markers + and the second is a single Gaussian that does not split. + + Returns + ------- + data : 2D `~numpy.ndarray` + The image. + + segm : `~photutils.segmentation.SegmentationImage` + The segmentation image. + + bounds : tuple of 1D int64 `~numpy.ndarray` + The ``(labels, y0, y1, x0, x1)`` arrays of the two sources. + + offsets : 1D intp `~numpy.ndarray` + The packed region offsets, with one more entry than sources. + """ + y, x = np.mgrid[0:61, 0:141] + data = (Gaussian2D(100, 30, 30, 5, 5)(x, y) + + Gaussian2D(100, 45, 30, 5, 5)(x, y) + + Gaussian2D(50, 110, 30, 5, 5)(x, y)) + segm = detect_sources(data, 10, 5) + assert segm.n_labels == 2 + labels = np.asarray(segm.labels, dtype=np.int64) + y0 = np.array([slc[0].start for slc in segm.slices]) + y1 = np.array([slc[0].stop for slc in segm.slices]) + x0 = np.array([slc[1].start for slc in segm.slices]) + x1 = np.array([slc[1].stop for slc in segm.slices]) + sizes = (y1 - y0) * (x1 - x0) + offsets = np.concatenate(([0], np.cumsum(sizes))).astype(np.intp) + return data, segm, (labels, y0, y1, x0, x1), offsets + + +def run_packed_markers(data, segm, bounds, offsets, *, max_markers=-1): + """ + Build the markers of a packed two-source scene. + + Returns + ------- + packed : 1D int32 `~numpy.ndarray` + The packed marker buffer. + + n_markers : 1D intp `~numpy.ndarray` + The marker count of each source. + + stats : tuple of 1D float64 `~numpy.ndarray` + The ``(source_min, source_max, source_sum)`` arrays. + """ + labels, y0, y1, x0, x1 = bounds + stats = deblend_source_stats(data, segm.data, labels, y0, y1, x0, x1) + thresholds, _ = _compute_thresholds(stats[0], stats[1], 32, + 'exponential') + packed = np.zeros(offsets[-1], dtype=np.int32) + n_markers = deblend_markers_chunk(data, segm.data, labels, y0, y1, x0, + x1, thresholds, packed, offsets[:-1], + n_pixels=5, connectivity=8, + max_markers=max_markers) + return packed, n_markers, stats + + +@pytest.mark.parametrize('contrast', [0.001, 0.6]) +def test_contrast_chunk_packed_buffer(contrast): + """ + Test that the chunk contrast kernel relabels the packed region of a + source that deblends with consecutive labels covering its segment, + zeros the region of a source that collapses to one basin, and skips + the sources with fewer than two markers. + """ + data, segm, bounds, offsets = make_packed_pair() + labels, y0, y1, x0, x1 = bounds + packed, n_markers, (smin, _, ssum) = run_packed_markers( + data, segm, bounds, offsets) + assert_equal(n_markers, [2, 0]) + packed[offsets[1]:offsets[2]] = 7 # skipped regions are untouched + + n_labels = deblend_contrast_chunk(data, segm.data, labels, y0, y1, x0, + x1, packed, offsets[:-1], n_markers, + connectivity=8, contrast=contrast, + source_sum=ssum, source_min=smin, + apply_contrast=True) + region0 = packed[offsets[0]:offsets[1]] + segment0 = (segm.data[y0[0]:y1[0], x0[0]:x1[0]] == labels[0]).ravel() + if contrast == 0.6: + # Both basins hold about half the flux, so the fainter one is + # removed and the source is left with a single basin + assert_equal(n_labels, [1, 0]) + assert not region0.any() + else: + assert_equal(n_labels, [2, 0]) + assert_equal(np.unique(region0), [0, 1, 2]) + assert_equal(region0 > 0, segment0) + assert_equal(packed[offsets[1]:offsets[2]], 7) + + +@pytest.mark.parametrize('dtype', [np.int32, np.int64]) +def test_write_deblended_labels(dtype): + """ + Test that the write-out kernel adds the per-source offset to the + nonzero packed labels of the sources with two or more labels and + leaves every other pixel and source untouched. + """ + segm_out = np.full((4, 6), 9, dtype=dtype) + y0 = np.array([0, 2], dtype=np.int64) + y1 = np.array([2, 4], dtype=np.int64) + x0 = np.array([1, 3], dtype=np.int64) + x1 = np.array([4, 6], dtype=np.int64) + region0 = np.array([[1, 0, 2], + [1, 2, 0]], dtype=np.int32) + region1 = np.array([[1, 1, 1], + [0, 0, 1]], dtype=np.int32) + packed = np.concatenate((region0.ravel(), region1.ravel())) + starts = np.array([0, region0.size], dtype=np.intp) + n_labels = np.array([2, 1], dtype=np.intp) + label_offsets = np.array([10, 20], dtype=np.int64) + + write_deblended_labels(segm_out, packed, starts, y0, y1, x0, x1, + n_labels, label_offsets) + expected = np.full((4, 6), 9, dtype=dtype) + expected[0, 1] = 11 + expected[0, 3] = 12 + expected[1, 1] = 11 + expected[1, 2] = 12 + assert_equal(segm_out, expected) + assert segm_out.dtype == dtype + + +def test_chunk_kernels_validate_inputs(): + """ + Test that the chunk kernels reject per-source arrays with the wrong + length and a packed buffer that is too small for the regions. + """ + data, segm, bounds, offsets = make_packed_pair() + labels, y0, y1, x0, x1 = bounds + packed, n_markers, (smin, smax, ssum) = run_packed_markers( + data, segm, bounds, offsets) + thresholds, _ = _compute_thresholds(smin, smax, 32, 'exponential') + starts = offsets[:-1] + short = starts[:1] + small = packed[:-1] + n_labels = np.array([2, 0], dtype=np.intp) + label_offsets = np.zeros(2, dtype=np.int64) + segm_out = segm.data.copy() + + match_len = 'one entry per source' + match_size = 'packed is too small' + with pytest.raises(ValueError, match=match_len): + deblend_markers_chunk(data, segm.data, labels, y0, y1, x0, x1, + thresholds, packed, short, n_pixels=5, + connectivity=8, max_markers=-1) + with pytest.raises(ValueError, match=match_size): + deblend_markers_chunk(data, segm.data, labels, y0, y1, x0, x1, + thresholds, small, starts, n_pixels=5, + connectivity=8, max_markers=-1) + with pytest.raises(ValueError, match=match_len): + deblend_contrast_chunk(data, segm.data, labels, y0, y1, x0, x1, + packed, short, n_markers, connectivity=8, + contrast=0.001, source_sum=ssum, + source_min=smin, apply_contrast=True) + with pytest.raises(ValueError, match=match_size): + deblend_contrast_chunk(data, segm.data, labels, y0, y1, x0, x1, + small, starts, n_markers, connectivity=8, + contrast=0.001, source_sum=ssum, + source_min=smin, apply_contrast=True) + with pytest.raises(ValueError, match=match_len): + write_deblended_labels(segm_out, packed, short, y0, y1, x0, x1, + n_labels, label_offsets) + with pytest.raises(ValueError, match=match_size): + write_deblended_labels(segm_out, small, starts, y0, y1, x0, x1, + n_labels, label_offsets) + + +def test_source_stats_all_nan(): + """ + Test that a segment without any finite pixel has NaN extrema and + zero flux in both the compiled stats kernel and the reference + implementation, and that deblend_sources leaves it unchanged. + """ + data = np.full((3, 4), np.nan) + segment = np.ones((3, 4), dtype=np.int32) + labels = np.array([1], dtype=np.int64) + smin, smax, ssum = deblend_source_stats(data, segment, labels, + np.array([0]), np.array([3]), + np.array([0]), np.array([4])) + assert np.isnan(smin[0]) + assert np.isnan(smax[0]) + assert ssum[0] == 0.0 + + params = _DeblendParams(1, np.ones((3, 3)), 32, 0.001, 'linear') + with warnings.catch_warnings(): + warnings.simplefilter('ignore', RuntimeWarning) + deblender = _SingleSourceDeblender(data, segment, 1, params) + assert np.isnan(deblender.source_min) + assert np.isnan(deblender.source_max) + assert deblender.source_sum == 0.0 + + segm = SegmentationImage(segment) + result = deblend_sources(data, segm, 1) + assert_equal(result.data, segment) + assert result.parent_to_deblended_labels == {} + + +def test_markers_chunk_packed_buffer(): + """ + Test that the marker kernel writes each source's markers into its + packed region, leaves the regions of sources that do not split or + that exceed max_markers at zero, and reports the marker counts. + """ + data, segm, bounds, offsets = make_packed_pair() + packed, n_markers, _ = run_packed_markers(data, segm, bounds, offsets) + assert_equal(n_markers, [2, 0]) + region0 = packed[offsets[0]:offsets[1]] + region1 = packed[offsets[1]:offsets[2]] + assert_equal(np.unique(region0), [0, 1, 2]) + assert not region1.any() + + # A limit below the marker count leaves the region zero but still + # reports the count + packed, n_markers, _ = run_packed_markers(data, segm, bounds, offsets, + max_markers=1) + assert_equal(n_markers, [2, 0]) + assert not packed.any() + + @pytest.mark.parametrize('dtype', ['float64', 'float32']) @pytest.mark.parametrize('connectivity', [8, 4]) @pytest.mark.parametrize( @@ -1008,12 +1291,16 @@ def test_saddle_markers_match_reference(dtype, connectivity, scene, thresholds_2d = np.ascontiguousarray(thresholds[None, :], dtype=np.float64) limit = deblender.contrast * float(deblender.source_sum) - markers_list, _ = deblend_markers_chunk( + packed = np.zeros(cutout.size, dtype=np.int32) + starts = np.zeros(1, dtype=np.intp) + n_markers = deblend_markers_chunk( np.ascontiguousarray(data), segm.data, np.array([1], dtype=np.int64), y0, y1, x0, x1, thresholds_2d, - n_pixels=5, connectivity=connectivity, max_markers=-1, - saddle_limits=np.array([limit], dtype=np.float64)) - result = markers_list[0] + packed, starts, n_pixels=5, connectivity=connectivity, + max_markers=-1, saddle_limits=np.array([limit], dtype=np.float64)) + result = None + if n_markers[0] >= 2: + result = packed.reshape(cutout.shape) if expected is None or result is None: assert expected is None