def main(): import argparse parser = argparse.ArgumentParser() parser.add_argument('--plots', action='store_true') parser.add_argument('--brick', help='Brick name to run') parser.add_argument( '--input-dir', default='/global/projecta/projectdirs/cosmo/work/legacysurvey/dr7') #/global/cscratch1/sd/desiproc/dr7out') parser.add_argument('--survey-dir', default='/global/cscratch1/sd/dstn/dr7-depthcut') parser.add_argument('--output-dir', default='/global/cscratch1/sd/dstn/bright') opt = parser.parse_args() plots = opt.plots ps = PlotSequence('bright') brickname = opt.brick insurvey = LegacySurveyData(opt.input_dir, cache_dir=opt.survey_dir) outsurvey = LegacySurveyData(opt.output_dir, output_dir=opt.output_dir) bfn = insurvey.find_file('blobmap', brick=brickname) print('Found blob map', bfn) blobs = fitsio.read(bfn) h, w = blobs.shape brick = insurvey.get_brick_by_name(brickname) brickwcs = wcs_for_brick(brick) radius = np.sqrt(2.) * 0.25 * 1.01 neighbors = insurvey.get_bricks_near(brick.ra, brick.dec, radius) print(len(neighbors), 'bricks nearby') def showbool(X): d = downsample_max(X, 8) h, w = X.shape plt.imshow(d, interpolation='nearest', origin='lower', vmin=0, vmax=1, extent=[0, w, 0, h], cmap='gray') brightblobs = set() for nb in neighbors: if nb.brickname == brickname: # ignore myself! continue print('Neighbor:', nb.brickname) mfn = insurvey.find_file('maskbits', brick=nb.brickname) if not os.path.exists(mfn): print('No maskbits file:', mfn) continue maskbits = fitsio.read(mfn) bright = ((maskbits & MASKBITS['BRIGHT']) > 0) print(np.sum(bright > 0), 'BRIGHT pixels set') primary = (maskbits & MASKBITS['NPRIMARY'] == 0) print(np.sum(primary), 'PRIMARY pixels set') edge = binary_dilation(primary, structure=np.ones((3, 3), bool)) edge = edge * np.logical_not(primary) brightedge = edge & bright if plots: plt.clf() showbool(bright) plt.title('bright: brick %s' % nb.brickname) ps.savefig() # plt.clf() # showbool(primary) # plt.title('PRIMARY, brick %s' % nb.brickname) # ps.savefig() # # plt.clf() # showbool(edge) # plt.title('boundary, brick %s' % nb.brickname) # ps.savefig() plt.clf() showbool(brightedge) plt.title('bright at edge, brick %s' % nb.brickname) ps.savefig() nwcs = wcs_for_brick(nb) yy, xx = np.nonzero(brightedge) print(len(yy), 'bright edge pixels') if len(yy) == 0: continue rr, dd = nwcs.pixelxy2radec(xx + 1, yy + 1) print('RA range', rr.min(), rr.max(), 'vs brick', brick.ra1, brick.ra2) print('Dec range', dd.min(), dd.max(), 'vs brick', brick.dec1, brick.dec2) # Find pixels that are within this brick's unique area I, = np.nonzero((rr >= brick.ra1) * (rr <= brick.ra2) * (dd >= brick.dec1) * (dd <= brick.dec2)) if plots: plt.clf() plt.plot( [brick.ra1, brick.ra1, brick.ra2, brick.ra2, brick.ra1], [brick.dec1, brick.dec2, brick.dec2, brick.dec1, brick.dec1], 'b-') plt.plot(rr, dd, 'k.') plt.plot(rr[I], dd[I], 'r.') plt.title('Bright pixels from %s' % nb.brickname) ps.savefig() if len(I) == 0: print('No edge pixels touch') #plt.plot(br,bd, 'b-') continue #print('Edge pixels touch!') #plt.plot(br,bd, 'r-', zorder=20) ok, x, y = brickwcs.radec2pixelxy(rr[I], dd[I]) x = np.round(x).astype(int) - 1 y = np.round(y).astype(int) - 1 print('Pixel ranges X', x.min(), x.max(), 'Y', y.min(), y.max()) assert (np.all((x >= 0) * (x < w) * (y >= 0) * (y < h))) print('Adding blobs:', np.unique(blobs[y, x])) brightblobs.update(blobs[y, x]) print('Blobs touching bright pixels:', brightblobs) print() brightblobs.discard(-1) if len(brightblobs) == 0: print('No neighboring bright blobs to update!') return print('Updating', len(brightblobs), 'blobs:', brightblobs) tmap = np.zeros(blobs.max() + 2, bool) for b in brightblobs: tmap[b + 1] = True touching = tmap[blobs + 1] if plots: plt.clf() showbool(touching) plt.title('Blobs touching bright, brick %s' % brickname) ps.savefig() mfn = insurvey.find_file('maskbits', brick=brickname) maskbits, hdr = fitsio.read(mfn, header=True) updated = maskbits | (MASKBITS['BRIGHT'] * touching) if np.all(maskbits == updated): print('No bits updated! (Bright stars were already masked)') return maskbits = updated if plots: plt.clf() showbool((maskbits & MASKBITS['BRIGHT']) > 0) plt.title('New maskbits map for BRIGHT, brick %s' % brickname) ps.savefig() with outsurvey.write_output('maskbits', brick=brickname) as out: out.fits.write(maskbits, hdr=hdr) tfn = insurvey.find_file('tractor', brick=brickname) phdr = fitsio.read_header(tfn, ext=0) hdr = fitsio.read_header(tfn, ext=1) T = fits_table(tfn) print('Read', len(T), 'sources') print('Bright:', Counter(T.brightstarinblob)) iby = np.clip(np.round(T.by), 0, h - 1).astype(int) ibx = np.clip(np.round(T.bx), 0, w - 1).astype(int) if plots: before = np.flatnonzero(T.brightstarinblob) T.brightstarinblob |= touching[iby, ibx] print('Bright:', Counter(T.brightstarinblob)) # yuck -- copy the TUNIT headers from input to output. units = [ hdr.get('TUNIT%i' % (i + 1), '') for i in range(len(T.get_columns())) ] if plots: plt.clf() showbool((maskbits & MASKBITS['BRIGHT']) > 0) ax = plt.axis() after = np.flatnonzero(T.brightstarinblob) plt.plot(T.bx[before], T.by[before], 'gx') plt.plot(T.bx[after], T.by[after], 'r.') plt.axis(ax) plt.title('sources with brightstarinblob, brick %s' % brickname) ps.savefig() with outsurvey.write_output('tractor', brick=brickname) as out: T.writeto(None, fits_object=out.fits, primheader=phdr, header=hdr, units=units)