SSI depth maps from detection fraction

txpipe/auxiliary_maps.py

@@ -1,7 +1,7 @@
 from .maps import TXBaseMaps, map_config_options
 import numpy as np
 from .base_stage import PipelineStage
-from .mapping import make_dask_shear_maps, make_dask_flag_maps, make_dask_bright_object_map, make_dask_depth_map
+from .mapping import make_dask_shear_maps, make_dask_flag_maps, make_dask_bright_object_map, make_dask_depth_map, make_dask_depth_map_det_prob
 from .data_types import MapsFile, HDFFile, ShearCatalog
 from .utils import choose_pixelization, import_dask
 from .maps import map_config_options
@@ -268,8 +268,9 @@ class TXAuxiliarySSIMaps(TXBaseMaps):
     name = "TXAuxiliarySSIMaps"
     dask_parallel = True
     inputs = [
-        ("injection_catalog", HDFFile),  # for injection locations
-        ("matched_ssi_photometry_catalog", HDFFile),
+        ("matched_ssi_photometry_catalog", HDFFile), # injected objhects that were detected
+        ("injection_catalog", HDFFile),  # injection locations
+        ("ssi_detection_catalog", HDFFile), # detection info on each injection
     ]
     outputs = [
         ("aux_ssi_maps", MapsFile),
@@ -283,32 +284,46 @@ class TXAuxiliarySSIMaps(TXBaseMaps):
         "depth_band": "i",  # Make depth maps for this band
         "snr_threshold": 10.0,  # The S/N value to generate maps for (e.g. 5 for 5-sigma depth)
         "snr_delta": 1.0,  # The range threshold +/- delta is used for finding objects at the boundary
+        "det_prob_threshold": 0.8, #detection probability threshold for SSI depth (i.e. 0.9 for magnitude at which 90% of brighter objects are detected)
+        "mag_delta": 0.01,  # Size of the magnitude bins used to determine detection probability depth
+        "min_depth": 18, # Min magnitude used in detection probability depth
+        "max_depth": 26, # Max magnitude used in detection probability depth
+        "smooth_det_frac": True, # Apply savgol filtering to frac det vs magnitude for each pixel
+        "smooth_window":0.5 # Size of smoothing window in magnitudes
     }
 
     def run(self):
         # Import dask and alias it as 'da'
         _, da = import_dask()
 
-
         # Retrieve configuration parameters
         block_size = self.config["block_size"]
         if block_size == 0:
             block_size = "auto"
         band = self.config["depth_band"]
 
-        # Open the input photometry catalog file.
-        # We can't use a "with" statement because we need to keep the file open
+        # Open the input catalog files
+        # We can't use "with" statements because we need to keep the file open
         # while we're using dask.
-        f = self.open_input("matched_ssi_photometry_catalog", wrapper=True)
+        f_matched = self.open_input("matched_ssi_photometry_catalog", wrapper=True)
+        f_inj = self.open_input("injection_catalog", wrapper=True)
+        f_det = self.open_input("ssi_detection_catalog", wrapper=True)
 
-        # Load photometry data into dask arrays.
+        # Load matched catalog data into dask arrays.
         # This is lazy in dask, so we're not actually loading the data here.
-        ra = da.from_array(f.file["photometry/ra"], block_size)
+        ra = da.from_array(f_matched.file["photometry/ra"], block_size)
         block_size = ra.chunksize
-        dec = da.from_array(f.file["photometry/dec"], block_size)
-        snr = da.from_array(f.file[f"photometry/snr_{band}"], block_size)
-        mag_meas = da.from_array(f.file[f"photometry/mag_{band}"], block_size)
-        mag_true = da.from_array(f.file[f"photometry/inj_mag_{band}"], block_size)
+        dec = da.from_array(f_matched.file["photometry/dec"], block_size)
+        snr = da.from_array(f_matched.file[f"photometry/snr_{band}"], block_size)
+        mag_meas = da.from_array(f_matched.file[f"photometry/mag_{band}"], block_size)
+        mag_true = da.from_array(f_matched.file[f"photometry/inj_mag_{band}"], block_size)
+
+        # Load detection catalog data into dask arrays.
+        # This is lazy in dask, so we're not actually loading the data here.
+        ra_inj = da.from_array(f_inj.file["photometry/ra"], block_size)
+        dec_inj = da.from_array(f_inj.file["photometry/dec"], block_size)
+        inj_mag = da.from_array(f_inj.file[f"photometry/inj_mag_{band}"], block_size)
+        det = da.from_array(f_det.file[f"photometry/detected"], block_size)
 
         # Choose the pixelization scheme based on the configuration.
         # Might need to review this to make sure we use the same scheme everywhere
@@ -321,17 +336,27 @@ def run(self):
         # Create depth maps using dask and measured magnitudes
         pix2, count_map, depth_map, depth_var = make_dask_depth_map(
             ra, dec, mag_meas, snr, self.config["snr_threshold"], self.config["snr_delta"], pixel_scheme)
-        maps["depth/depth_meas"] = (pix2, depth_map[pix2])
-        maps["depth/depth_meas_count"] = (pix2, count_map[pix2])
-        maps["depth/depth_meas_var"] = (pix2, depth_var[pix2])
-
+        maps["depth_meas/depth"] = (pix2, depth_map[pix2])
+        maps["depth_meas/depth_count"] = (pix2, count_map[pix2])
+        maps["depth_meas/depth_var"] = (pix2, depth_var[pix2])
 
-        # Create depth maps using daskand true magnitudes
+        # Create depth maps using dask and true magnitudes
         pix2, count_map, depth_map, depth_var = make_dask_depth_map(
             ra, dec, mag_true, snr, self.config["snr_threshold"], self.config["snr_delta"], pixel_scheme)
-        maps["depth/depth_true"] = (pix2, depth_map[pix2])
-        maps["depth/depth_true_count"] = (pix2, count_map[pix2])
-        maps["depth/depth_true_var"] = (pix2, depth_var[pix2])
+        maps["depth_true/depth"] = (pix2, depth_map[pix2])
+        maps["depth_true/depth_count"] = (pix2, count_map[pix2])
+        maps["depth_true/depth_var"] = (pix2, depth_var[pix2])
+
+        # Create depth maps using injection catalog
+        # depth is defined at given detection probability
+        pix2, det_count_map, inj_count_map, depth_map, frac_stack, mag_edges = make_dask_depth_map_det_prob(
+            ra_inj, dec_inj, inj_mag, det, self.config["det_prob_threshold"], self.config["mag_delta"], 
+            self.config["min_depth"], self.config["max_depth"], pixel_scheme, 
+            self.config["smooth_det_frac"], self.config["smooth_window"])
+        maps["depth_det_prob/depth"] = (pix2, depth_map[pix2])
+        maps["depth_det_prob/depth_det_count"] = (pix2, det_count_map[pix2])
+        maps["depth_det_prob/depth_inj_count"] = (pix2, inj_count_map[pix2])
+        maps["depth_det_prob/frac_stack"] = (pix2, frac_stack[:,pix2])
 
         maps, = da.compute(maps)
 
@@ -346,6 +371,10 @@ def run(self):
         metadata["depth_band"] = band
         metadata["depth_snr_threshold"] = self.config["snr_threshold"]
         metadata["depth_snr_delta"] = self.config["snr_delta"]
+        metadata["mag_delta"] = self.config["mag_delta"]
+        metadata["min_depth"] = self.config["min_depth"]
+        metadata["max_depth"] = self.config["max_depth"]
+        metadata["mag_edges"] = mag_edges
         metadata.update(pixel_scheme.metadata)
 
         # Write the output maps to the output file

txpipe/data_types/types.py

@@ -261,11 +261,16 @@ def write_map(self, map_name, pixel, value, metadata):
             self.file.create_group("maps")
         if not "pixelization" in metadata:
             raise ValueError("Map metadata should include pixelization")
-        if not pixel.shape == value.shape:
-            raise ValueError(
-                f"Map pixels and values should be same shape "
-                f"but are {pixel.shape} vs {value.shape}"
-            )
+        pixerr = ValueError(
+                    f"Map pixels and values should be same shape "
+                    f"but are {pixel.shape} vs {value.shape}"
+                )
+        if value.ndim == 2: #value.ndim == 2 allows for 2d stacked maps
+            if not pixel.shape[0] == value.shape[1]:
+                raise pixerr
+        else:
+            if not pixel.shape == value.shape:
+                raise pixerr
 
         if not 'maps' in self.file.keys():
             self.file.create_group('maps')

txpipe/map_plots.py

@@ -205,6 +205,7 @@ class TXMapPlotsSSI(TXMapPlots):
     outputs = [
         ("depth_ssi_meas_map", PNGFile),
         ("depth_ssi_true_map", PNGFile),
+        ("depth_det_prob_map", PNGFile),
     ]
 
     def run(self):
@@ -233,18 +234,26 @@ def run(self):
     def aux_ssi_plots(self):
         import matplotlib.pyplot as plt
         if self.get_input("aux_ssi_maps") == "none":
+            # Make empty plots if no data available, so that the
+            # pipeline thinks it is complete.
             with self.open_output("depth_ssi_meas_map", wrapper=True) as f:
                 pass
             with self.open_output("depth_ssi_true_map", wrapper=True) as f:
                 pass
+            with self.open_output("depth_det_prob_map", wrapper=True) as f:
+                pass
             return
 
         m = self.open_input("aux_ssi_maps", wrapper=True)
 
         # Depth plots (measured magnitude)
         with self.open_output("depth_ssi_meas_map", wrapper=True, figsize=(5, 5)) as fig:
-            m.plot("depth/depth_meas", view=self.config["projection"], rot180=self.config["rot180"])
+            m.plot("depth_meas/depth", view=self.config["projection"], rot180=self.config["rot180"])
 
         # Depth plots (true magnitude)
         with self.open_output("depth_ssi_true_map", wrapper=True, figsize=(5, 5)) as fig:
-            m.plot("depth/depth_true", view=self.config["projection"], rot180=self.config["rot180"])
+            m.plot("depth_true/depth", view=self.config["projection"], rot180=self.config["rot180"])
+
+        # Depth plots (true magnitude)
+        with self.open_output("depth_det_prob_map", wrapper=True, figsize=(5, 5)) as fig:
+            m.plot("depth_det_prob/depth", view=self.config["projection"], rot180=self.config["rot180"])

txpipe/mapping/__init__.py

@@ -1,2 +1,2 @@
-from .dr1 import make_dask_bright_object_map, make_dask_depth_map
+from .dr1 import make_dask_bright_object_map, make_dask_depth_map, make_dask_depth_map_det_prob
 from .basic_maps import make_dask_flag_maps, make_dask_shear_maps, make_dask_lens_maps

txpipe/mapping/dr1.py

@@ -86,4 +86,99 @@ def make_dask_depth_map(ra, dec, mag, snr, threshold, delta, pixel_scheme):
     depth_var = depth2_map - depth_map**2
 
     pix = da.unique(pix)
-    return pix, count_map, depth_map, depth_var
+    return pix, count_map, depth_map, depth_var
+
+def make_dask_depth_map_det_prob(
+    ra, dec, mag, det, det_prob_threshold, mag_delta, min_depth, max_depth, pixel_scheme, 
+    smooth_det_frac=False, smooth_window=0.5,
+):
+    """
+    Generate a depth map using Dask, by finding the mean magnitude of
+    objects with a signal-to-noise ratio close to a given threshold.
+
+    Parameters
+    ----------
+    ra : dask.array
+        Right Ascension coordinates in degrees.
+    dec : dask.array
+        Declination coordinates in degrees.
+    mag : dask.array
+        Magnitudes of the objects, in band of user's choice
+    det : dask.array
+        dask array of boolean detection parameter
+    det_prob_threshold : float
+        Detection probability threshold for SSI depth 
+        (i.e. 0.9 to get magnitude at which 90% of brighter objects are detected)
+    pixel_scheme : PixelScheme
+        An object that provides pixelization scheme with methods `npix` and `ang2pix`.
+    smooth_det_frac: bool
+        if True apply a savgol filtering to the individual detection frac vs magnitude cut signals
+    smooth_window: float
+        if smooth_det_frac==True, this is the window size of the filter (in magnitudes)
+
+    Returns
+    -------
+    tuple
+        A tuple containing:
+        - pix (dask.array): Unique pixel indices.
+        - det_count_map (dask.array): Count of detected objects per pixel.
+        - inj_count_map (dask.array): Count of injected objects per pixel.
+        - depth_map (dask.array): Mean depth per pixel.
+        - frac_stack (dask.array): stack of maps. fraction of detections brighter than mag_edges
+        - mag_edges (dask.array): grid of magnitudes at which frac_stack was evaluated
+    """
+    from scipy.signal import savgol_filter
+    _, da = import_dask()
+    npix = pixel_scheme.npix
+    pix = pixel_scheme.ang2pix(ra, dec)
+
+    det_count_map = da.bincount(pix, weights=det, minlength=npix)
+    inj_count_map = da.bincount(pix, minlength=npix)
+
+    # Make array of magnitude bins
+    mag_edges = da.arange(min_depth, max_depth, mag_delta)
+    n_depth_bins = len(mag_edges)
+
+    # loop over mag bins
+    # TODO: add option to compute fraction *at* each magnitude, rather than below
+    frac_list = []
+    for mag_thresh in mag_edges:
+        above_thresh = mag < mag_thresh
+        ntot = da.bincount(pix, weights=above_thresh, minlength=npix)
+        ndet = da.bincount(pix, weights=above_thresh * det, minlength=npix)
+        frac_det = da.where(ntot != 0, ndet / ntot, np.nan)
+        frac_list.append(frac_det)
+    frac_stack = da.stack(frac_list)
+
+    # Optional smoothing of the stacked detection fractions
+    if smooth_det_frac:
+        window_length = int(n_depth_bins*smooth_window/(max_depth-min_depth)) #converting to units of depth bin
+        poly_order = 2 # TODO: could make config option
+
+        # Here extend the chunks of the dask array when applying a local filter to avoid boundary issues
+        frac_stack = frac_stack.map_overlap( 
+            lambda a: savgol_filter(a, window_length, poly_order, axis=0), 
+            depth=window_length // 2,   # Extend chunks by half window size
+            boundary="reflect",         # Reflect at edges to avoid NaNs
+            dtype=frac_stack.dtype, 
+            )
+
+    # In order for pixel to give a valid depth estimate it must have 
+    # (1) at least one mag_thresh with a computed frac_det above the threshold
+    # (2) at least one mag_thresh with a computed frac_det below the threshold
+    has_high = (frac_stack > det_prob_threshold).any(axis=0)
+    has_low = (frac_stack < det_prob_threshold).any(axis=0)
+    valid_pix_mask = has_high & has_low
+
+    # find the first element smaller than the threshold
+    below_threshold = frac_stack < det_prob_threshold
+    masked = da.where(
+        below_threshold, da.arange(frac_stack.shape[0])[:, None], n_depth_bins - 1
+    )
+    thres_index = da.nanmin(masked, axis=0)
+
+    depth_map = mag_edges[thres_index]
+    depth_map[~valid_pix_mask] = np.nan
+
+    pix = da.unique(pix)
+    return pix, det_count_map, inj_count_map, depth_map, frac_stack, mag_edges