Merge 5e7dd7c into c8179c6

py/legacypipe/oneblob.py

@@ -504,14 +504,41 @@ def compute_segmentation_map(self):
         Iseg, = np.nonzero((self.refmap[iy, ix] & IN_BLOB['CLUSTER']) == 0)
         # Zero out the S/N in CLUSTER mask
         maxsn[(self.refmap & IN_BLOB['CLUSTER']) > 0] = 0.
-
-        # (also zero out the satmap)
+        # (also zero out the satmap in the CLUSTER mask)
         saturated_pix[(self.refmap & IN_BLOB['CLUSTER']) > 0] = False
 
-        Ibright = _argsort_by_brightness([self.srcs[i] for i in Iseg], self.bands)
-        rank = np.empty(len(Iseg), int)
-        rank[Ibright] = np.arange(len(Iseg), dtype=int)
-        rankmap = dict([(Iseg[i],r) for r,i in enumerate(Ibright)])
+        Irank = Iseg
+        # Start by setting the segmentation map for reference sources
+        # that are forced to be point sources, to the size of the
+        # point source model.
+        ref_radius = 25
+        refpsf = np.array([getattr(self.srcs[i], 'forced_point_source', False)
+                           for i in Iseg], dtype=bool)
+        Irefpsf = Iseg[refpsf]
+        refclaimed = None
+        if len(Irefpsf):
+            debug('Setting segmentation map for', len(Irefpsf), 'forced-point-source refs')
+            Ibr = _argsort_by_brightness([self.srcs[i] for i in Irefpsf], self.bands)
+            _set_kingdoms(segmap, ref_radius, Irefpsf[Ibr], ix, iy)
+            del Ibr
+            refclaimed = (segmap != -1)
+            # Now we remove these sources from the list of sources that we need
+            # to find segmentations for.  They remain in 'Irank', which are the ones
+            # that compete for rank.
+            Iseg = Iseg[np.logical_not(refpsf)]
+        del refpsf
+
+        # Iseg are the indices in self.srcs of sources to segment
+        Ibright = _argsort_by_brightness([self.srcs[i] for i in Irank], self.bands)
+        rank = np.empty(len(Irank), int)
+        rank[Ibright] = np.arange(len(Irank), dtype=int)
+        del Ibright
+        # rankmap goes from source index to rank
+        # (unlike 'rank', which is for sources in Iranks)
+        rankmap = dict(zip(Irank, rank))
+        print('Irank:', Irank)
+        print('rank:' ,rank)
+        print('rankmap:', rankmap)
 
         todo = set(Iseg)
         mx = int(np.ceil(maxsn.max()))
@@ -530,61 +557,67 @@ def compute_segmentation_map(self):
             debug('S/N', thresh, ':', len(todo), 'sources to find still')
             if len(todo) == 0:
                 break
-            # We previously filled saturated pixels with a max value,
-            # so this is maybe not necessary?
-            #hot = np.logical_or(maxsn >= thresh, saturated_pix)
             hot = (maxsn >= thresh)
             hot = binary_fill_holes(hot)
             blobs,_ = label(hot)
-            srcblobs = blobs[iy[Iseg], ix[Iseg]]
+            srcblobs = blobs[iy[Irank], ix[Irank]]
             done = set()
-
+            # We build up a map of blob -> (ranks of sources in blob)
             blobranks = {}
+            blobsrcs = {}
             for i,(b,r) in enumerate(zip(srcblobs, rank)):
                 if not b in blobranks:
                     blobranks[b] = []
+                    blobsrcs[b] = []
                 blobranks[b].append(r)
+                blobsrcs[b].append(i)
 
             for t in todo:
                 bl = blobs[iy[t], ix[t]]
                 if bl == 0:
-                    # ??
+                    # source not in a blob...
                     done.add(t)
                     continue
+                # Is this source the brightest in this blob?
                 if rankmap[t] == min(blobranks[bl]):
-                    #print('Source', t, 'has rank', rank[t], 'vs blob ranks', blobranks[bl])
+                    # Claim all the territory in my thresholded blob.
+                    # Other (fainter) sources may still claim sub-regions at higher
+                    # threshold levels.
                     segmap[blobs == bl] = t
                     #print('Source', t, 'is isolated at S/N', thresh)
                     done.add(t)
+                # Is this source outside the regions claimed by brighter (ref) sources?
+                else:
+                    ok = True
+                    for r,isrc in zip(blobranks[bl], blobsrcs[bl]):
+                        if not isrc in Irefpsf:
+                            # regular (non-ref-PSF) source; its claim holds
+                            ok = False
+                            break
+                        r2 = (ix[t] - ix[isrc])**2 + (iy[t] - iy[isrc])**2
+                        if r2 <= ref_radius**2:
+                            # within ref source's claimed region; its claim holds
+                            ok = False
+                            break
+                    if ok:
+                        # If we got here, we share a blob with a
+                        # brighter ref source, but are far enough away
+                        # from it.  We claim the blob, except for what
+                        # the ref source has already claimed.
+                        segmap[(blobs == bl) * (refclaimed == False)] = t
+                        done.add(t)
+
             todo.difference_update(done)
             del hot
         del maxsn, saturated_pix
 
-        # ensure that each source owns a tiny radius around its center in the segmentation map.
-        # If there is more than one source in that radius, each pixel gets assigned to its nearest source.
-        # record the current distance to nearest source
-        kingdom = np.empty(segmap.shape, np.uint8)
-        kingdom[:,:,] = 255
-        H,W = segmap.shape
-        xcoords = np.arange(W)
-        ycoords = np.arange(H)
-        for i in Ibright:
-            radius = 5
-            x,y = ix[i], iy[i]
-            yslc = slice(max(0, y-radius), min(H, y+radius+1))
-            xslc = slice(max(0, x-radius), min(W, x+radius+1))
-            slc = (yslc, xslc)
-            # Radius to nearest earlier source
-            oldr = kingdom[slc]
-            # Radius to new source
-            newr = np.hypot(xcoords[np.newaxis, xslc] - x, ycoords[yslc, np.newaxis] - y)
-            assert(newr.shape == oldr.shape)
-            newr = (newr + 0.5).astype(np.uint8)
-            # Pixels that are within range and closer to this source than any other.
-            owned = (newr <= radius) * (newr < oldr)
-            segmap[slc][owned] = i
-            kingdom[slc][owned] = newr[owned]
-        del kingdom, xcoords, ycoords
+        # ensure that each source owns a tiny radius around its center
+        # in the segmentation map.  If there is more than one source
+        # in that radius, each pixel gets assigned to its nearest
+        # source.
+        radius = 5
+        Ibright = _argsort_by_brightness([self.srcs[i] for i in Iseg], self.bands)
+        _set_kingdoms(segmap, radius, Iseg[Ibright], ix, iy)
 
         self.segmap = segmap
 
@@ -1882,6 +1915,39 @@ def create_tims(self, timargs):
             tims.append(tim)
         return tims
 
+def _set_kingdoms(segmap, radius, I, ix, iy):
+    '''
+    radius: int
+    ix,iy: int arrays
+    I: indices into ix,iy that will be placed into 'segmap'
+    '''
+    # ensure that each source owns a tiny radius around its center
+    # in the segmentation map.  If there is more than one source
+    # in that radius, each pixel gets assigned to its nearest
+    # source.
+    # 'kingdom' records the current distance to nearest source
+    assert(radius < 255)
+    kingdom = np.empty(segmap.shape, np.uint8)
+    kingdom[:,:,] = 255
+    H,W = segmap.shape
+    xcoords = np.arange(W)
+    ycoords = np.arange(H)
+    for i in I:
+        x,y = ix[i], iy[i]
+        yslc = slice(max(0, y-radius), min(H, y+radius+1))
+        xslc = slice(max(0, x-radius), min(W, x+radius+1))
+        slc = (yslc, xslc)
+        # Radius to nearest earlier source
+        oldr = kingdom[slc]
+        # Radius to new source
+        newr = np.hypot(xcoords[np.newaxis, xslc] - x, ycoords[yslc, np.newaxis] - y)
+        assert(newr.shape == oldr.shape)
+        newr = (newr + 0.5).astype(np.uint8)
+        # Pixels that are within range and closer to this source than any other.
+        owned = (newr <= radius) * (newr < oldr)
+        segmap[slc][owned] = i
+        kingdom[slc][owned] = newr[owned]
+
 def _convert_ellipses(src):
     if isinstance(src, (DevGalaxy, ExpGalaxy, SersicGalaxy)):
         src.shape = src.shape.toEllipseE()

py/legacypipe/runbrick.py

@@ -2032,7 +2032,7 @@ def _add_bit_description(header, BITS, bits, bnpat, bitpat, bitmapname):
         header.add_record(
             dict(name=bnpat % short, value=BITS[key],
                  comment='%s: %s' % (bitmapname, comm)))
-    revmap = dict([(bit,name) for name,bit in MASKBITS.items()])
+    revmap = dict([(bit,name) for name,bit in BITS.items()])
     nicemap = dict([(k,c) for k,short,c in bits])
     for bit in range(16):
         bitval = 1<<bit