Sep measure function reworked

btk/measure.py

@@ -62,6 +62,8 @@ def measure_function(batch, idx, **kwargs):
 import astropy.table
 import numpy as np
 import sep
+from astropy import units
+from astropy.coordinates import SkyCoord
 from skimage.feature import peak_local_max
 
 from btk.multiprocess import multiprocess
@@ -135,18 +137,115 @@ def basic_measure(
     return {"catalog": catalog}
 
 
-def sep_measure(
+def sep_multiband_measure(
     batch,
     idx,
     channels_last=False,
     surveys=None,
+    matching_threshold=1.0,
     sigma_noise=1.5,
     is_multiresolution=False,
     **kwargs,
 ):
-    """Return detection, segmentation and deblending information with SEP.
+    """Returns centers detected with source extractor by combining predictions in different bands.
+
+    For each band in the input image we run sep for detection and append new detections to a running
+    list of detected coordinates. In order to avoid repeating detections, we run a KD-Tree algroithm
+    to calculate the angular distance between each new coordinate and its closest neighbour. Then we
+    discard those new coordinates that were closer than matching_threshold to one of already
+    detected coordinates.
+
+    NOTE: If this function is used with the multiresolution feature,
+    measurements will be carried on the first survey.
+
+    Args:
+        batch (dict): Output of DrawBlendsGenerator object's `__next__` method.
+        idx (int): Index number of blend scene in the batch to preform
+            measurement on.
+        sigma_noise (float): Sigma threshold for detection against noise.
+        matching_threshold (float): Match centers of objects that are closer than
+            this threshold to a single prediction.
+
+    Returns:
+            dict containing catalog with entries corresponding to measured peaks.
+    """
+    channel_indx = 0 if not channels_last else -1
+    # multiresolution
+    if is_multiresolution:
+        if surveys is None:
+            raise ValueError("surveys are required in order to use the MR feature.")
+        survey_name = surveys[0].name
+        image = batch["blend_images"][survey_name][idx]
+        wcs = batch["wcs"][survey_name]
+
+    # single-survey
+    else:
+        image = batch["blend_images"][idx]
+        wcs = batch["wcs"]
+
+    # run source extractor on the first band
+    band_image = image[0] if channel_indx == 0 else image[:, :, 0]
+    bkg = sep.Background(band_image)
+    catalog = sep.extract(band_image, sigma_noise, err=bkg.globalrms, segmentation_map=False)
+
+    # convert predictions to arcseconds
+    ra_coordinates, dec_coordinates = wcs.pixel_to_world_values(catalog["x"], catalog["y"])
+    ra_coordinates *= 3600
+    dec_coordinates *= 3600
+
+    # iterate over remaining bands and match predictions using KdTree
+    for band in range(1, image.shape[channel_indx]):
+        # run source extractor
+        band_image = image[band] if channel_indx == 0 else image[:, :, band]
+        bkg = sep.Background(band_image)
+        catalog = sep.extract(band_image, sigma_noise, err=bkg.globalrms, segmentation_map=False)
+
+        # convert predictions to arcseconds
+        ra_detections, dec_detections = wcs.pixel_to_world_values(catalog["x"], catalog["y"])
+        ra_detections *= 3600
+        dec_detections *= 3600
+
+        # convert to sky coordinates
+        c1 = SkyCoord(ra=ra_detections * units.arcsec, dec=dec_detections * units.arcsec)
+        c2 = SkyCoord(ra=ra_coordinates * units.arcsec, dec=dec_coordinates * units.arcsec)
+
+        # merge new detections with the running list of coordinates
+        if len(c1) > 0 and len(c2) > 0:
+            # runs KD-tree to get distances to the closest neighbours
+            idx, distance2d, _ = c1.match_to_catalog_sky(c2)
+            distance2d = distance2d.arcsec
+
+            # add new predictions, masking those that are closer than threshold
+            ra_coordinates = np.concatenate(
+                [ra_coordinates, ra_detections[distance2d > matching_threshold]]
+            )
+            dec_coordinates = np.concatenate(
+                [dec_coordinates, dec_detections[distance2d > matching_threshold]]
+            )
+        else:
+            ra_coordinates = np.concatenate([ra_coordinates, ra_detections])
+            dec_coordinates = np.concatenate([dec_coordinates, dec_detections])
+
+    # Wrap in the astropy table
+    t = astropy.table.Table()
+    t["ra"] = ra_coordinates
+    t["dec"] = dec_coordinates
+
+    return {"catalog": t}
+
+
+def sep_singleband_measure(
+    batch,
+    idx,
+    meas_band_num=3,
+    channels_last=False,
+    surveys=None,
+    sigma_noise=1.5,
+    is_multiresolution=False,
+    **kwargs,
+):
+    """Return detection, segmentation and deblending information running SEP on a single band.
 
-    For each potentially multi-band image, an average over the bands is taken before measurement.
     NOTE: If this function is used with the multiresolution feature,
     measurements will be carried on the first survey, and deblended images
     or segmentations will not be returned.
@@ -155,6 +254,7 @@ def sep_measure(
         batch (dict): Output of DrawBlendsGenerator object's `__next__` method.
         idx (int): Index number of blend scene in the batch to preform
             measurement on.
+        meas_band_num (int) – Indicates the index of band to use fo the measurement
         sigma_noise (float): Sigma threshold for detection against noise.
 
     Returns:
@@ -168,21 +268,22 @@ def sep_measure(
             raise ValueError("surveys are required in order to use the MR feature.")
         survey_name = surveys[0].name
         image = batch["blend_images"][survey_name][idx]
-        avg_image = np.mean(image, axis=channel_indx)
         wcs = batch["wcs"][survey_name]
 
     # single-survey
     else:
         image = batch["blend_images"][idx]
-        avg_image = np.mean(image, axis=channel_indx)
         wcs = batch["wcs"]
 
-    stamp_size = avg_image.shape[0]
-    bkg = sep.Background(avg_image)
+    # run source extractor
+    band_image = image[meas_band_num] if channel_indx == 0 else image[:, :, meas_band_num]
+    stamp_size = band_image.shape[0]
+    bkg = sep.Background(band_image)
     catalog, segmentation = sep.extract(
-        avg_image, sigma_noise, err=bkg.globalrms, segmentation_map=True
+        band_image, sigma_noise, err=bkg.globalrms, segmentation_map=True
     )
 
+    # reshape segmentation map
     n_objects = len(catalog)
     segmentation_exp = np.zeros((n_objects, stamp_size, stamp_size), dtype=bool)
     deblended_images = np.zeros((n_objects, *image.shape), dtype=image.dtype)
@@ -195,6 +296,7 @@ def sep_measure(
         seg_i_reshaped = np.moveaxis(seg_i_reshaped, 0, np.argmin(image.shape))
         deblended_images[i] = image * seg_i_reshaped
 
+    # wrap results in astropy table
     t = astropy.table.Table()
     t["ra"], t["dec"] = wcs.pixel_to_world_values(catalog["x"], catalog["y"])
     t["ra"] *= 3600
@@ -430,4 +532,8 @@ def __next__(self):
         return blend_output, measure_results
 
 
-available_measure_functions = {"basic": basic_measure, "sep": sep_measure}
+available_measure_functions = {
+    "basic": basic_measure,
+    "sep_singleband_measure": sep_singleband_measure,
+    "sep_multiband_measure": sep_multiband_measure,
+}

notebooks/01a-cosmos_tutorial.ipynb

@@ -411,7 +411,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.10"
+   "version": "3.8.12"
   },
   "toc": {
    "base_numbering": 1,

notebooks/01b-scarlet-measure.ipynb

@@ -206,7 +206,7 @@
    "source": [
     "measure_kwargs=[{\"sigma_noise\": 2.0}]\n",
     "meas_generator = btk.measure.MeasureGenerator(\n",
-    "    [btk.measure.sep_measure,scarlet_measure], draw_blend_generator, measure_kwargs=measure_kwargs\n",
+    "    [btk.measure.sep_singleband_measure,scarlet_measure], draw_blend_generator, measure_kwargs=measure_kwargs\n",
     ")\n",
     "metrics_generator = btk.metrics.MetricsGenerator(\n",
     "    meas_generator,\n",
@@ -397,9 +397,9 @@
  "metadata": {
   "celltoolbar": "Aucun(e)",
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "btk38",
    "language": "python",
-   "name": "python3"
+   "name": "btk38"
   },
   "language_info": {
    "codemirror_mode": {
@@ -411,7 +411,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.5"
+   "version": "3.8.13"
   },
   "toc": {
    "base_numbering": 1,

notebooks/tutorial.ipynb

@@ -801,13 +801,13 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "f61a8392",
+   "id": "a5b32bf1-1204-41de-a70f-8977831695bf",
    "metadata": {},
    "outputs": [],
    "source": [
     "import sep\n",
-    "def sep_measure(batch, idx, channels_last=False, sigma_noise=1.5, **kwargs):\n",
-    "    \"\"\"Return detection, segmentation and deblending information with SEP.\n",
+    "def sep_singleband_measure(batch, idx, meas_band_num=2, channels_last=False, sigma_noise=1.5, **kwargs):\n",
+    "    \"\"\"Return detection, segmentation and deblending information running SEP on a single band.\n",
     "\n",
     "    NOTE: This is a simplified version of the actual function that does not \n",
     "    support multi-resolution. \n",
@@ -816,27 +816,27 @@
     "        batch (dict): Output of DrawBlendsGenerator object's `__next__` method.\n",
     "        idx (int): Index number of blend scene in the batch to preform\n",
     "            measurement on.\n",
+    "        meas_band_num (int) – Indicates the index of band to use fo the measurement\n",
     "        sigma_noise (float): Sigma threshold for detection against noise.\n",
     "\n",
     "    Returns:\n",
     "        dict with the centers of sources detected by SEP detection algorithm.\n",
     "    \"\"\"\n",
     "    channel_indx = 0 if not channels_last else -1\n",
-    "\n",
-    "\n",
     "    image = batch[\"blend_images\"][idx]\n",
-    "    coadd = np.mean(image, axis=channel_indx)\n",
+    "    # select band of an image\n",
+    "    band_image = image[meas_band_num] if channel_indx == 0 else image[:, :, meas_band_num]\n",
     "    wcs = batch[\"wcs\"]\n",
     "\n",
-    "    stamp_size = coadd.shape[0]\n",
-    "    bkg = sep.Background(coadd)\n",
+    "    # run source extractor\n",
+    "    stamp_size = band_image.shape[0]\n",
+    "    bkg = sep.Background(band_image)\n",
     "    catalog, segmentation = sep.extract(\n",
-    "        coadd, sigma_noise, err=bkg.globalrms, segmentation_map=True\n",
+    "        band_image, sigma_noise, err=bkg.globalrms, segmentation_map=True\n",
     "    )\n",
     "\n",
+    "    # reshape segmentation map\n",
     "    n_objects = len(catalog)\n",
-    "\n",
-    "    # organizing returned images into numpy arrays\n",
     "    segmentation_exp = np.zeros((n_objects, stamp_size, stamp_size), dtype=bool)\n",
     "    deblended_images = np.zeros((n_objects, *image.shape), dtype=image.dtype)\n",
     "    for i in range(n_objects):\n",
@@ -848,9 +848,11 @@
     "        seg_i_reshaped = np.moveaxis(seg_i_reshaped, 0, np.argmin(image.shape))\n",
     "        deblended_images[i] = image * seg_i_reshaped\n",
     "\n",
-    "    # translate from pixel to sky coordinates. \n",
+    "    # wrap results in astropy table\n",
     "    t = astropy.table.Table()\n",
     "    t[\"ra\"], t[\"dec\"] = wcs.pixel_to_world_values(catalog[\"x\"], catalog[\"y\"])\n",
+    "    t[\"ra\"] *= 3600\n",
+    "    t[\"dec\"] *= 3600\n",
     "\n",
     "    return {\n",
     "        \"catalog\": t,\n",
@@ -1017,7 +1019,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.13"
+   "version": "3.8.12"
   },
   "toc": {
    "base_numbering": 1,