def boxcar_extraction_from_filenames(image_filename, psf_filename, fibers=None, width=7): """ Fast boxcar extraction of spectra from a preprocessed image and a trace set Args: image_filename : input preprocessed fits filename psf_filename : input PSF fits filename Optional: fibers : 1D np.array of int (default is all fibers, the first fiber is always = 0) width : extraction boxcar width, default is 7 Returns: flux : 2D np.array of shape (nfibers,n0=image.shape[0]), sum of pixel values per row of length=width per fiber ivar : 2D np.array of shape (nfibers,n0), ivar[f,j] = 1/( sum_[j,b:e] (1/image.ivar) ), ivar=0 if at least 1 pixel in the row has image.ivar=0 or image.mask!=0 wave : 2D np.array of shape (nfibers,n0), determined from the traces """ tset = read_xytraceset(psf_filename) image = read_image(image_filename) qframe = qproc_boxcar_extraction(xytraceset, image, fibers=fibers, width=width) return qframe.flux, qframe.ivar, qframe.wave
def test_preproc_script(self): io.write_raw(self.rawfile, self.rawimage, self.header, camera='b0') io.write_raw(self.rawfile, self.rawimage, self.header, camera='b1') args = ['--infile', self.rawfile, '--cameras', 'b1', '--pixfile', self.pixfile] if os.path.exists(self.pixfile): os.remove(self.pixfile) desispec.scripts.preproc.main(args) img = io.read_image(self.pixfile) self.assertEqual(img.pix.shape, (2*self.ny, 2*self.nx))
def test_preproc_script(self): io.write_raw(self.rawfile, self.rawimage, self.header, primary_header = self.primary_header, camera='b0') io.write_raw(self.rawfile, self.rawimage, self.header, primary_header = self.primary_header, camera='b1') args = ['--infile', self.rawfile, '--cameras', 'b0', '--outfile', self.pixfile] if os.path.exists(self.pixfile): os.remove(self.pixfile) desispec.scripts.preproc.main(args) img = io.read_image(self.pixfile) self.assertEqual(img.pix.shape, (2*self.ny, 2*self.nx))
def main(args) : log= get_logger() log.info("degxx={} degxy={} degyx={} degyy={}".format(args.degxx,args.degxy,args.degyx,args.degyy)) # read preprocessed image image=read_image(args.image) log.info("read image {}".format(args.image)) if image.mask is not None : image.ivar *= (image.mask==0) tset = fit_trace_shifts(image=image,args=args) if args.outpsf is not None : write_traces_in_psf(args.psf,args.outpsf,tset) log.info("wrote modified PSF in %s"%args.outpsf)
def boxcar_extraction_from_filenames(image_filename,psf_filename,fibers=None, width=7) : """ Fast boxcar extraction of spectra from a preprocessed image and a trace set Args: image_filename : input preprocessed fits filename psf_filename : input PSF fits filename Optional: fibers : 1D np.array of int (default is all fibers, the first fiber is always = 0) width : extraction boxcar width, default is 7 Returns: flux : 2D np.array of shape (nfibers,n0=image.shape[0]), sum of pixel values per row of length=width per fiber ivar : 2D np.array of shape (nfibers,n0), ivar[f,j] = 1/( sum_[j,b:e] (1/image.ivar) ), ivar=0 if at least 1 pixel in the row has image.ivar=0 or image.mask!=0 wave : 2D np.array of shape (nfibers,n0), determined from the traces """ tset = read_xytraceset(psf_filename) image = read_image(image_filename) qframe = qproc_boxcar_extraction(xytraceset,image,fibers=fibers,width=width) return qframe.flux, qframe.ivar, qframe.wave
def main(args): global psfs log = get_logger() log.info("starting") # read preprocessed image image = read_image(args.image) log.info("read image {}".format(args.image)) if image.mask is not None: image.ivar *= (image.mask == 0) xcoef = None ycoef = None psf = None wavemin = None wavemax = None nfibers = None lines = None psf, xcoef, ycoef, wavemin, wavemax = read_psf_and_traces(args.psf) nfibers = xcoef.shape[0] log.info( "read PSF trace with xcoef.shape = {} , ycoef.shape = {} , and wavelength range {}:{}" .format(xcoef.shape, ycoef.shape, int(wavemin), int(wavemax))) if args.lines is not None: log.info("We will fit the image using the psf model and lines") # read lines lines = np.loadtxt(args.lines, usecols=[0]) ok = (lines > wavemin) & (lines < wavemax) log.info( "read {} lines in {}, with {} of them in traces wavelength range". format(len(lines), args.lines, np.sum(ok))) lines = lines[ok] else: log.info( "We will do an internal calibration of trace coordinates without using the psf shape in a first step" ) internal_wavelength_calib = (not args.continuum) external_wavelength_calib = args.sky | (args.spectrum is not None) if args.auto: log.debug("read flavor of input image {}".format(args.image)) hdus = pyfits.open(args.image) if "FLAVOR" not in hdus[0].header: log.error( "no FLAVOR keyword in image header, cannot run with --auto option" ) raise KeyError( "no FLAVOR keyword in image header, cannot run with --auto option" ) flavor = hdus[0].header["FLAVOR"].strip().lower() hdus.close() log.info("Input is a '{}' image".format(flavor)) if flavor == "flat": internal_wavelength_calib = False external_wavelength_calib = False elif flavor == "arc": internal_wavelength_calib = True external_wavelength_calib = False else: internal_wavelength_calib = True external_wavelength_calib = True log.info("wavelength calib, internal={}, external={}".format( internal_wavelength_calib, external_wavelength_calib)) spectrum_filename = args.spectrum if external_wavelength_calib and spectrum_filename is None: srch_file = "data/spec-sky.dat" if not resource_exists('desispec', srch_file): log.error("Cannot find sky spectrum file {:s}".format(srch_file)) raise RuntimeError( "Cannot find sky spectrum file {:s}".format(srch_file)) else: spectrum_filename = resource_filename('desispec', srch_file) log.info("Use external calibration from cross-correlation with {}". format(spectrum_filename)) if args.nfibers is not None: nfibers = args.nfibers # FOR DEBUGGING fibers = np.arange(nfibers) if lines is not None: # use a forward modeling of the image # it's slower and works only for individual lines # it's in principle more accurate # but gives systematic residuals for complex spectra like the sky x, y, dx, ex, dy, ey, fiber_xy, wave_xy = compute_dx_dy_using_psf( psf, image, fibers, lines) x_for_dx = x y_for_dx = y fiber_for_dx = fiber_xy wave_for_dx = wave_xy x_for_dy = x y_for_dy = y fiber_for_dy = fiber_xy wave_for_dy = wave_xy else: # internal calibration method that does not use the psf # nor a prior set of lines. this method is much faster # measure x shifts x_for_dx, y_for_dx, dx, ex, fiber_for_dx, wave_for_dx = compute_dx_from_cross_dispersion_profiles( xcoef, ycoef, wavemin, wavemax, image=image, fibers=fibers, width=7, deg=args.degxy) if internal_wavelength_calib: # measure y shifts x_for_dy, y_for_dy, dy, ey, fiber_for_dy, wave_for_dy = compute_dy_using_boxcar_extraction( xcoef, ycoef, wavemin, wavemax, image=image, fibers=fibers, width=7) else: # duplicate dx results with zero shift to avoid write special case code below x_for_dy = x_for_dx.copy() y_for_dy = y_for_dx.copy() dy = np.zeros(dx.shape) ey = 1.e-6 * np.ones(ex.shape) fiber_for_dy = fiber_for_dx.copy() wave_for_dy = wave_for_dx.copy() degxx = args.degxx degxy = args.degxy degyx = args.degyx degyy = args.degyy while (True): # loop because polynomial degrees could be reduced log.info( "polynomial fit of measured offsets with degx=(%d,%d) degy=(%d,%d)" % (degxx, degxy, degyx, degyy)) try: dx_coeff, dx_coeff_covariance, dx_errorfloor, dx_mod, dx_mask = polynomial_fit( z=dx, ez=ex, xx=x_for_dx, yy=y_for_dx, degx=degxx, degy=degxy) dy_coeff, dy_coeff_covariance, dy_errorfloor, dy_mod, dy_mask = polynomial_fit( z=dy, ez=ey, xx=x_for_dy, yy=y_for_dy, degx=degyx, degy=degyy) log.info("dx dy error floor = %4.3f %4.3f pixels" % (dx_errorfloor, dy_errorfloor)) log.info("check fit uncertainties are ok on edge of CCD") merr = 0. for fiber in [0, nfibers - 1]: for rw in [-1, 1]: tx = legval(rw, xcoef[fiber]) ty = legval(rw, ycoef[fiber]) m = monomials(tx, ty, degxx, degxy) tdx = np.inner(dx_coeff, m) tsx = np.sqrt(np.inner(m, dx_coeff_covariance.dot(m))) m = monomials(tx, ty, degyx, degyy) tdy = np.inner(dy_coeff, m) tsy = np.sqrt(np.inner(m, dy_coeff_covariance.dot(m))) merr = max(merr, tdx) merr = max(merr, tdy) #log.info("fiber=%d wave=%dA x=%d y=%d dx=%4.3f+-%4.3f dy=%4.3f+-%4.3f"%(fiber,int(wave),int(x),int(y),dx,sx,dy,sy)) log.info("max edge shift error = %4.3f pixels" % merr) if degxx == 0 and degxy == 0 and degyx == 0 and degyy == 0: break except (LinAlgError, ValueError): log.warning( "polynomial fit failed with degx=(%d,%d) degy=(%d,%d)" % (degxx, degxy, degyx, degyy)) if degxx == 0 and degxy == 0 and degyx == 0 and degyy == 0: log.error( "polynomial degrees are already 0. we can fit the offsets") raise RuntimeError( "polynomial degrees are already 0. we can fit the offsets") merr = 100000. # this will lower the pol. degree. if merr > args.max_error: if merr != 100000.: log.warning( "max edge shift error = %4.3f pixels is too large, reducing degrees" % merr) if degxy > 0 and degyy > 0 and degxy > degxx and degyy > degyx: # first along wavelength if degxy > 0: degxy -= 1 if degyy > 0: degyy -= 1 else: # then along fiber if degxx > 0: degxx -= 1 if degyx > 0: degyx -= 1 else: # error is ok, so we quit the loop break # write this for debugging if args.outoffsets: file = open(args.outoffsets, "w") file.write( "# axis wave fiber x y delta error polval (axis 0=y axis1=x)\n") for e in range(dy.size): file.write("0 %f %d %f %f %f %f %f\n" % (wave_for_dy[e], fiber_for_dy[e], x_for_dy[e], y_for_dy[e], dy[e], ey[e], dy_mod[e])) for e in range(dx.size): file.write("1 %f %d %f %f %f %f %f\n" % (wave_for_dx[e], fiber_for_dx[e], x_for_dx[e], y_for_dx[e], dx[e], ex[e], dx_mod[e])) file.close() log.info("wrote offsets in ASCII file %s" % args.outoffsets) # print central shift mx = np.median(x_for_dx) my = np.median(y_for_dx) m = monomials(mx, my, degxx, degxy) mdx = np.inner(dx_coeff, m) mex = np.sqrt(np.inner(m, dx_coeff_covariance.dot(m))) mx = np.median(x_for_dy) my = np.median(y_for_dy) m = monomials(mx, my, degyx, degyy) mdy = np.inner(dy_coeff, m) mey = np.sqrt(np.inner(m, dy_coeff_covariance.dot(m))) log.info("central shifts dx = %4.3f +- %4.3f dy = %4.3f +- %4.3f " % (mdx, mex, mdy, mey)) # for each fiber, apply offsets and recompute legendre polynomial log.info("for each fiber, apply offsets and recompute legendre polynomial") xcoef, ycoef = recompute_legendre_coefficients(xcoef=xcoef, ycoef=ycoef, wavemin=wavemin, wavemax=wavemax, degxx=degxx, degxy=degxy, degyx=degyx, degyy=degyy, dx_coeff=dx_coeff, dy_coeff=dy_coeff) # use an input spectrum as an external calibration of wavelength if spectrum_filename: log.info( "write and reread PSF to be sure predetermined shifts were propagated" ) write_traces_in_psf(args.psf, args.outpsf, xcoef, ycoef, wavemin, wavemax) psf, xcoef, ycoef, wavemin, wavemax = read_psf_and_traces(args.outpsf) ycoef = shift_ycoef_using_external_spectrum( psf=psf, xcoef=xcoef, ycoef=ycoef, wavemin=wavemin, wavemax=wavemax, image=image, fibers=fibers, spectrum_filename=spectrum_filename, degyy=args.degyy, width=7) write_traces_in_psf(args.psf, args.outpsf, xcoef, ycoef, wavemin, wavemax) log.info("wrote modified PSF in %s" % args.outpsf) else: write_traces_in_psf(args.psf, args.outpsf, xcoef, ycoef, wavemin, wavemax) log.info("wrote modified PSF in %s" % args.outpsf)
def main_mpi(args, comm=None): log = get_logger() psf_file = args.psf input_file = args.input # these parameters are interpreted as the *global* spec range, # to be divided among processes. specmin = args.specmin nspec = args.nspec #- Load input files and broadcast # FIXME: after we have fixed the serialization # of the PSF, read and broadcast here, to reduce # disk contention. img = None if comm is None: img = io.read_image(input_file) else: if comm.rank == 0: img = io.read_image(input_file) img = comm.bcast(img, root=0) psf = load_psf(psf_file) # get spectral range if nspec is None: nspec = psf.nspec specmax = specmin + nspec camera = img.meta['CAMERA'].lower() #- b0, r1, .. z9 spectrograph = int(camera[1]) fibermin = spectrograph * psf.nspec + specmin if args.fibermap is not None: fibermap = io.read_fibermap(args.fibermap) fibermap = fibermap[fibermin:fibermin + nspec] fibers = fibermap['FIBER'] else: fibermap = None fibers = np.arange(fibermin, fibermin + nspec, dtype='i4') #- Get wavelength grid from options if args.wavelength is not None: wstart, wstop, dw = [float(tmp) for tmp in args.wavelength.split(',')] else: wstart = np.ceil(psf.wmin_all) wstop = np.floor(psf.wmax_all) dw = 0.5 wave = np.arange(wstart, wstop + dw / 2.0, dw) nwave = len(wave) #- Confirm that this PSF covers these wavelengths for these spectra psf_wavemin = np.max(psf.wavelength(list(range(specmin, specmax)), y=0)) psf_wavemax = np.min( psf.wavelength(list(range(specmin, specmax)), y=psf.npix_y - 1)) if psf_wavemin > wstart: raise ValueError( 'Start wavelength {:.2f} < min wavelength {:.2f} for these fibers'. format(wstart, psf_wavemin)) if psf_wavemax < wstop: raise ValueError( 'Stop wavelength {:.2f} > max wavelength {:.2f} for these fibers'. format(wstop, psf_wavemax)) # Now we divide our spectra into bundles bundlesize = args.bundlesize checkbundles = set() checkbundles.update( np.floor_divide(np.arange(specmin, specmax), bundlesize * np.ones(nspec)).astype(int)) bundles = sorted(checkbundles) nbundle = len(bundles) bspecmin = {} bnspec = {} for b in bundles: if specmin > b * bundlesize: bspecmin[b] = specmin else: bspecmin[b] = b * bundlesize if (b + 1) * bundlesize > specmax: bnspec[b] = specmax - bspecmin[b] else: bnspec[b] = bundlesize # Now we assign bundles to processes nproc = 1 rank = 0 if comm is not None: nproc = comm.size rank = comm.rank mynbundle = int(nbundle // nproc) myfirstbundle = 0 leftover = nbundle % nproc if rank < leftover: mynbundle += 1 myfirstbundle = rank * mynbundle else: myfirstbundle = ((mynbundle + 1) * leftover) + (mynbundle * (rank - leftover)) if rank == 0: #- Print parameters log.info("extract: input = {}".format(input_file)) log.info("extract: psf = {}".format(psf_file)) log.info("extract: specmin = {}".format(specmin)) log.info("extract: nspec = {}".format(nspec)) log.info("extract: wavelength = {},{},{}".format(wstart, wstop, dw)) log.info("extract: nwavestep = {}".format(args.nwavestep)) log.info("extract: regularize = {}".format(args.regularize)) # get the root output file outpat = re.compile(r'(.*)\.fits') outmat = outpat.match(args.output) if outmat is None: raise RuntimeError( "extraction output file should have .fits extension") outroot = outmat.group(1) outdir = os.path.normpath(os.path.dirname(outroot)) if rank == 0: if not os.path.isdir(outdir): os.makedirs(outdir) if comm is not None: comm.barrier() failcount = 0 for b in range(myfirstbundle, myfirstbundle + mynbundle): outbundle = "{}_{:02d}.fits".format(outroot, b) outmodel = "{}_model_{:02d}.fits".format(outroot, b) log.info('extract: Rank {} starting {} spectra {}:{} at {}'.format( rank, os.path.basename(input_file), bspecmin[b], bspecmin[b] + bnspec[b], time.asctime(), )) sys.stdout.flush() #- The actual extraction try: results = ex2d(img.pix, img.ivar * (img.mask == 0), psf, bspecmin[b], bnspec[b], wave, regularize=args.regularize, ndecorr=True, bundlesize=bundlesize, wavesize=args.nwavestep, verbose=args.verbose, full_output=True) flux = results['flux'] ivar = results['ivar'] Rdata = results['resolution_data'] chi2pix = results['chi2pix'] mask = np.zeros(flux.shape, dtype=np.uint32) mask[results['pixmask_fraction'] > 0.5] |= specmask.SOMEBADPIX mask[results['pixmask_fraction'] == 1.0] |= specmask.ALLBADPIX mask[chi2pix > 100.0] |= specmask.BAD2DFIT #- Augment input image header for output img.meta['NSPEC'] = (nspec, 'Number of spectra') img.meta['WAVEMIN'] = (wstart, 'First wavelength [Angstroms]') img.meta['WAVEMAX'] = (wstop, 'Last wavelength [Angstroms]') img.meta['WAVESTEP'] = (dw, 'Wavelength step size [Angstroms]') img.meta['SPECTER'] = (specter.__version__, 'https://github.com/desihub/specter') img.meta['IN_PSF'] = (_trim(psf_file), 'Input spectral PSF') img.meta['IN_IMG'] = (_trim(input_file), 'Input image') if fibermap is not None: bfibermap = fibermap[bspecmin[b] - specmin:bspecmin[b] + bnspec[b] - specmin] else: bfibermap = None bfibers = fibers[bspecmin[b] - specmin:bspecmin[b] + bnspec[b] - specmin] frame = Frame(wave, flux, ivar, mask=mask, resolution_data=Rdata, fibers=bfibers, meta=img.meta, fibermap=bfibermap, chi2pix=chi2pix) #- Write output io.write_frame(outbundle, frame, units='photon/bin') if args.model is not None: from astropy.io import fits fits.writeto(outmodel, results['modelimage'], header=frame.meta) log.info('extract: Done {} spectra {}:{} at {}'.format( os.path.basename(input_file), bspecmin[b], bspecmin[b] + bnspec[b], time.asctime())) sys.stdout.flush() except: # Log the error and increment the number of failures log.error( "extract: FAILED bundle {}, spectrum range {}:{}".format( b, bspecmin[b], bspecmin[b] + bnspec[b])) exc_type, exc_value, exc_traceback = sys.exc_info() lines = traceback.format_exception(exc_type, exc_value, exc_traceback) log.error(''.join(lines)) failcount += 1 sys.stdout.flush() if comm is not None: failcount = comm.allreduce(failcount) if failcount > 0: # all processes throw raise RuntimeError("some extraction bundles failed") if rank == 0: mergeopts = ['--output', args.output, '--force', '--delete'] mergeopts.extend( ["{}_{:02d}.fits".format(outroot, b) for b in bundles]) mergeargs = mergebundles.parse(mergeopts) mergebundles.main(mergeargs) if args.model is not None: model = None for b in bundles: outmodel = "{}_model_{:02d}.fits".format(outroot, b) if model is None: model = fits.getdata(outmodel) else: #- TODO: test and warn if models overlap for pixels with #- non-zero values model += fits.getdata(outmodel) os.remove(outmodel) fits.writeto(args.model, model)
def main(args): psf_file = args.psf input_file = args.input specmin = args.specmin nspec = args.nspec #- Load input files psf = load_psf(psf_file) img = io.read_image(input_file) if nspec is None: nspec = psf.nspec specmax = specmin + nspec camera = img.meta['CAMERA'].lower() #- b0, r1, .. z9 spectrograph = int(camera[1]) fibermin = spectrograph * psf.nspec + specmin print('Starting {} spectra {}:{} at {}'.format( os.path.basename(input_file), specmin, specmin + nspec, time.asctime())) if args.fibermap is not None: fibermap = io.read_fibermap(args.fibermap) fibermap = fibermap[fibermin:fibermin + nspec] fibers = fibermap['FIBER'] else: fibermap = None fibers = np.arange(fibermin, fibermin + nspec, dtype='i4') #- Get wavelength grid from options if args.wavelength is not None: wstart, wstop, dw = [float(tmp) for tmp in args.wavelength.split(',')] else: wstart = np.ceil(psf.wmin_all) wstop = np.floor(psf.wmax_all) dw = 0.5 wave = np.arange(wstart, wstop + dw / 2.0, dw) nwave = len(wave) bundlesize = args.bundlesize #- Confirm that this PSF covers these wavelengths for these spectra psf_wavemin = np.max(psf.wavelength(list(range(specmin, specmax)), y=0)) psf_wavemax = np.min( psf.wavelength(list(range(specmin, specmax)), y=psf.npix_y - 1)) if psf_wavemin > wstart: raise ValueError( 'Start wavelength {:.2f} < min wavelength {:.2f} for these fibers'. format(wstart, psf_wavemin)) if psf_wavemax < wstop: raise ValueError( 'Stop wavelength {:.2f} > max wavelength {:.2f} for these fibers'. format(wstop, psf_wavemax)) #- Print parameters print("""\ #--- Extraction Parameters --- input: {input} psf: {psf} output: {output} wavelength: {wstart} - {wstop} AA steps {dw} specmin: {specmin} nspec: {nspec} regularize: {regularize} #-----------------------------\ """.format(input=input_file, psf=psf_file, output=args.output, wstart=wstart, wstop=wstop, dw=dw, specmin=specmin, nspec=nspec, regularize=args.regularize)) #- The actual extraction results = ex2d(img.pix, img.ivar * (img.mask == 0), psf, specmin, nspec, wave, regularize=args.regularize, ndecorr=True, bundlesize=bundlesize, wavesize=args.nwavestep, verbose=args.verbose, full_output=True) flux = results['flux'] ivar = results['ivar'] Rdata = results['resolution_data'] chi2pix = results['chi2pix'] mask = np.zeros(flux.shape, dtype=np.uint32) mask[results['pixmask_fraction'] > 0.5] |= specmask.SOMEBADPIX mask[results['pixmask_fraction'] == 1.0] |= specmask.ALLBADPIX mask[chi2pix > 100.0] |= specmask.BAD2DFIT #- Augment input image header for output img.meta['NSPEC'] = (nspec, 'Number of spectra') img.meta['WAVEMIN'] = (wstart, 'First wavelength [Angstroms]') img.meta['WAVEMAX'] = (wstop, 'Last wavelength [Angstroms]') img.meta['WAVESTEP'] = (dw, 'Wavelength step size [Angstroms]') img.meta['SPECTER'] = (specter.__version__, 'https://github.com/desihub/specter') img.meta['IN_PSF'] = (_trim(psf_file), 'Input spectral PSF') img.meta['IN_IMG'] = (_trim(input_file), 'Input image') frame = Frame(wave, flux, ivar, mask=mask, resolution_data=Rdata, fibers=fibers, meta=img.meta, fibermap=fibermap, chi2pix=chi2pix) #- Write output io.write_frame(args.output, frame, units='photon/bin') if args.model is not None: from astropy.io import fits fits.writeto(args.model, results['modelimage'], header=frame.meta, clobber=True) print('Done {} spectra {}:{} at {}'.format(os.path.basename(input_file), specmin, specmin + nspec, time.asctime()))
def main(args): if args.mpi: from mpi4py import MPI comm = MPI.COMM_WORLD return main_mpi(args, comm) psf_file = args.psf input_file = args.input specmin = args.specmin nspec = args.nspec #- Load input files psf = load_psf(psf_file) img = io.read_image(input_file) if nspec is None: nspec = psf.nspec specmax = specmin + nspec if args.fibermap_index is not None : fibermin = args.fibermap_index else : camera = img.meta['CAMERA'].lower() #- b0, r1, .. z9 spectrograph = int(camera[1]) fibermin = spectrograph * 500 + specmin print('Starting {} spectra {}:{} at {}'.format(os.path.basename(input_file), specmin, specmin+nspec, time.asctime())) if args.fibermap is not None: fibermap = io.read_fibermap(args.fibermap) else: try: fibermap = io.read_fibermap(args.input) except (AttributeError, IOError, KeyError): fibermap = None #- Trim fibermap to matching fiber range and create fibers array if fibermap: ii = np.in1d(fibermap['FIBER'], np.arange(fibermin, fibermin+nspec)) fibermap = fibermap[ii] fibers = fibermap['FIBER'] else: fibers = np.arange(fibermin, fibermin+nspec, dtype='i4') #- Get wavelength grid from options if args.wavelength is not None: wstart, wstop, dw = [float(tmp) for tmp in args.wavelength.split(',')] else: wstart = np.ceil(psf.wmin_all) wstop = np.floor(psf.wmax_all) dw = 0.7 if args.heliocentric_correction : heliocentric_correction_factor = heliocentric_correction_multiplicative_factor(img.meta) wstart /= heliocentric_correction_factor wstop /= heliocentric_correction_factor dw /= heliocentric_correction_factor else : heliocentric_correction_factor = 1. wave = np.arange(wstart, wstop+dw/2.0, dw) nwave = len(wave) bundlesize = args.bundlesize #- Confirm that this PSF covers these wavelengths for these spectra psf_wavemin = np.max(psf.wavelength(list(range(specmin, specmax)), y=0)) psf_wavemax = np.min(psf.wavelength(list(range(specmin, specmax)), y=psf.npix_y-1)) if psf_wavemin-5 > wstart: raise ValueError('Start wavelength {:.2f} < min wavelength {:.2f} for these fibers'.format(wstart, psf_wavemin)) if psf_wavemax+5 < wstop: raise ValueError('Stop wavelength {:.2f} > max wavelength {:.2f} for these fibers'.format(wstop, psf_wavemax)) #- Print parameters print("""\ #--- Extraction Parameters --- input: {input} psf: {psf} output: {output} wavelength: {wstart} - {wstop} AA steps {dw} specmin: {specmin} nspec: {nspec} regularize: {regularize} #-----------------------------\ """.format(input=input_file, psf=psf_file, output=args.output, wstart=wstart, wstop=wstop, dw=dw, specmin=specmin, nspec=nspec, regularize=args.regularize)) #- The actual extraction results = ex2d(img.pix, img.ivar*(img.mask==0), psf, specmin, nspec, wave, regularize=args.regularize, ndecorr=args.decorrelate_fibers, bundlesize=bundlesize, wavesize=args.nwavestep, verbose=args.verbose, full_output=True, nsubbundles=args.nsubbundles,psferr=args.psferr) flux = results['flux'] ivar = results['ivar'] Rdata = results['resolution_data'] chi2pix = results['chi2pix'] mask = np.zeros(flux.shape, dtype=np.uint32) mask[results['pixmask_fraction']>0.5] |= specmask.SOMEBADPIX mask[results['pixmask_fraction']==1.0] |= specmask.ALLBADPIX mask[chi2pix>100.0] |= specmask.BAD2DFIT if heliocentric_correction_factor != 1 : #- Apply heliocentric correction factor to the wavelength #- without touching the spectra, that is the whole point wave *= heliocentric_correction_factor wstart *= heliocentric_correction_factor wstop *= heliocentric_correction_factor dw *= heliocentric_correction_factor img.meta['HELIOCOR'] = heliocentric_correction_factor #- Augment input image header for output img.meta['NSPEC'] = (nspec, 'Number of spectra') img.meta['WAVEMIN'] = (wstart, 'First wavelength [Angstroms]') img.meta['WAVEMAX'] = (wstop, 'Last wavelength [Angstroms]') img.meta['WAVESTEP']= (dw, 'Wavelength step size [Angstroms]') img.meta['SPECTER'] = (specter.__version__, 'https://github.com/desihub/specter') img.meta['IN_PSF'] = (_trim(psf_file), 'Input spectral PSF') img.meta['IN_IMG'] = (_trim(input_file), 'Input image') frame = Frame(wave, flux, ivar, mask=mask, resolution_data=Rdata, fibers=fibers, meta=img.meta, fibermap=fibermap, chi2pix=chi2pix) #- Add unit # In specter.extract.ex2d one has flux /= dwave # to convert the measured total number of electrons per # wavelength node to an electron 'density' frame.meta['BUNIT'] = 'count/Angstrom' #- Add scores to frame if not args.no_scores : compute_and_append_frame_scores(frame,suffix="RAW") #- Write output io.write_frame(args.output, frame) if args.model is not None: from astropy.io import fits fits.writeto(args.model, results['modelimage'], header=frame.meta, overwrite=True) print('Done {} spectra {}:{} at {}'.format(os.path.basename(input_file), specmin, specmin+nspec, time.asctime()))
def main_mpi(args, comm=None, timing=None): freeze_iers() nproc = 1 rank = 0 if comm is not None: nproc = comm.size rank = comm.rank mark_start = time.time() log = get_logger() psf_file = args.psf input_file = args.input # these parameters are interpreted as the *global* spec range, # to be divided among processes. specmin = args.specmin nspec = args.nspec #- Load input files and broadcast # FIXME: after we have fixed the serialization # of the PSF, read and broadcast here, to reduce # disk contention. img = None if rank == 0: img = io.read_image(input_file) if comm is not None: img = comm.bcast(img, root=0) psf = load_psf(psf_file) mark_read_input = time.time() # get spectral range if nspec is None: nspec = psf.nspec if args.fibermap is not None: fibermap = io.read_fibermap(args.fibermap) else: try: fibermap = io.read_fibermap(args.input) except (AttributeError, IOError, KeyError): fibermap = None if fibermap is not None: fibermap = fibermap[specmin:specmin + nspec] if nspec > len(fibermap): log.warning( "nspec {} > len(fibermap) {}; reducing nspec to {}".format( nspec, len(fibermap), len(fibermap))) nspec = len(fibermap) fibers = fibermap['FIBER'] else: fibers = np.arange(specmin, specmin + nspec) specmax = specmin + nspec #- Get wavelength grid from options if args.wavelength is not None: raw_wstart, raw_wstop, raw_dw = [ float(tmp) for tmp in args.wavelength.split(',') ] else: raw_wstart = np.ceil(psf.wmin_all) raw_wstop = np.floor(psf.wmax_all) raw_dw = 0.7 raw_wave = np.arange(raw_wstart, raw_wstop + raw_dw / 2.0, raw_dw) nwave = len(raw_wave) bundlesize = args.bundlesize if args.barycentric_correction: if ('RA' in img.meta) or ('TARGTRA' in img.meta): barycentric_correction_factor = \ barycentric_correction_multiplicative_factor(img.meta) #- Early commissioning has RA/TARGTRA in fibermap but not HDU 0 elif fibermap is not None and \ (('RA' in fibermap.meta) or ('TARGTRA' in fibermap.meta)): barycentric_correction_factor = \ barycentric_correction_multiplicative_factor(fibermap.meta) else: msg = 'Barycentric corr requires (TARGT)RA in HDU 0 or fibermap' log.critical(msg) raise KeyError(msg) else: barycentric_correction_factor = 1. # Explictly define the correct wavelength values to avoid confusion of reference frame # If correction applied, otherwise divide by 1 and use the same raw values wstart = raw_wstart / barycentric_correction_factor wstop = raw_wstop / barycentric_correction_factor dw = raw_dw / barycentric_correction_factor wave = raw_wave / barycentric_correction_factor #- Confirm that this PSF covers these wavelengths for these spectra psf_wavemin = np.max(psf.wavelength(list(range(specmin, specmax)), y=-0.5)) psf_wavemax = np.min( psf.wavelength(list(range(specmin, specmax)), y=psf.npix_y - 0.5)) if psf_wavemin - 5 > wstart: raise ValueError( 'Start wavelength {:.2f} < min wavelength {:.2f} for these fibers'. format(wstart, psf_wavemin)) if psf_wavemax + 5 < wstop: raise ValueError( 'Stop wavelength {:.2f} > max wavelength {:.2f} for these fibers'. format(wstop, psf_wavemax)) if rank == 0: #- Print parameters log.info("extract: input = {}".format(input_file)) log.info("extract: psf = {}".format(psf_file)) log.info("extract: specmin = {}".format(specmin)) log.info("extract: nspec = {}".format(nspec)) log.info("extract: wavelength = {},{},{}".format(wstart, wstop, dw)) log.info("extract: nwavestep = {}".format(args.nwavestep)) log.info("extract: regularize = {}".format(args.regularize)) if barycentric_correction_factor != 1.: img.meta['HELIOCOR'] = barycentric_correction_factor #- Augment input image header for output img.meta['NSPEC'] = (nspec, 'Number of spectra') img.meta['WAVEMIN'] = (raw_wstart, 'First wavelength [Angstroms]') img.meta['WAVEMAX'] = (raw_wstop, 'Last wavelength [Angstroms]') img.meta['WAVESTEP'] = (raw_dw, 'Wavelength step size [Angstroms]') img.meta['SPECTER'] = (specter.__version__, 'https://github.com/desihub/specter') img.meta['IN_PSF'] = (io.shorten_filename(psf_file), 'Input spectral PSF') img.meta['IN_IMG'] = io.shorten_filename(input_file) depend.add_dependencies(img.meta) #- Check if input PSF was itself a traceshifted version of another PSF orig_psf = None if rank == 0: try: psfhdr = fits.getheader(psf_file, 'PSF') orig_psf = psfhdr['IN_PSF'] except KeyError: #- could happen due to PSF format not having "PSF" extension, #- or due to PSF header not having 'IN_PSF' keyword. Either is OK pass if comm is not None: orig_psf = comm.bcast(orig_psf, root=0) if orig_psf is not None: img.meta['ORIG_PSF'] = orig_psf #- If not using MPI, use a single call to each of these and then end this function call # Otherwise, continue on to splitting things up for the different ranks if comm is None: _extract_and_save(img, psf, specmin, nspec, specmin, wave, raw_wave, fibers, fibermap, args.output, args.model, bundlesize, args, log) #- This is it if we aren't running MPI, so return return #else: # # Continue to the MPI section, which could go under this else statment # # But to save on indentation we'll just pass on to the rest of the function # # since the alternative has already returned # pass # Now we divide our spectra into bundles checkbundles = set() checkbundles.update( np.floor_divide(np.arange(specmin, specmax), bundlesize * np.ones(nspec)).astype(int)) bundles = sorted(checkbundles) nbundle = len(bundles) bspecmin = {} bnspec = {} for b in bundles: if specmin > b * bundlesize: bspecmin[b] = specmin else: bspecmin[b] = b * bundlesize if (b + 1) * bundlesize > specmax: bnspec[b] = specmax - bspecmin[b] else: bnspec[b] = bundlesize # Now we assign bundles to processes mynbundle = int(nbundle // nproc) myfirstbundle = 0 leftover = nbundle % nproc if rank < leftover: mynbundle += 1 myfirstbundle = rank * mynbundle else: myfirstbundle = ((mynbundle + 1) * leftover) + (mynbundle * (rank - leftover)) # get the root output file outpat = re.compile(r'(.*)\.fits') outmat = outpat.match(args.output) if outmat is None: raise RuntimeError( "extraction output file should have .fits extension") outroot = outmat.group(1) outdir = os.path.normpath(os.path.dirname(outroot)) if rank == 0: if not os.path.isdir(outdir): os.makedirs(outdir) if comm is not None: comm.barrier() mark_preparation = time.time() time_total_extraction = 0.0 time_total_write_output = 0.0 failcount = 0 for b in range(myfirstbundle, myfirstbundle + mynbundle): mark_iteration_start = time.time() outbundle = "{}_{:02d}.fits".format(outroot, b) outmodel = "{}_model_{:02d}.fits".format(outroot, b) log.info('extract: Rank {} extracting {} spectra {}:{} at {}'.format( rank, os.path.basename(input_file), bspecmin[b], bspecmin[b] + bnspec[b], time.asctime(), )) sys.stdout.flush() #- The actual extraction try: mark_extraction = _extract_and_save(img, psf, bspecmin[b], bnspec[b], specmin, wave, raw_wave, fibers, fibermap, outbundle, outmodel, bundlesize, args, log) mark_write_output = time.time() time_total_extraction += mark_extraction - mark_iteration_start time_total_write_output += mark_write_output - mark_extraction except: # Log the error and increment the number of failures log.error( "extract: FAILED bundle {}, spectrum range {}:{}".format( b, bspecmin[b], bspecmin[b] + bnspec[b])) exc_type, exc_value, exc_traceback = sys.exc_info() lines = traceback.format_exception(exc_type, exc_value, exc_traceback) log.error(''.join(lines)) failcount += 1 sys.stdout.flush() if comm is not None: failcount = comm.allreduce(failcount) if failcount > 0: # all processes throw raise RuntimeError("some extraction bundles failed") time_merge = None if rank == 0: mark_merge_start = time.time() mergeopts = ['--output', args.output, '--force', '--delete'] mergeopts.extend( ["{}_{:02d}.fits".format(outroot, b) for b in bundles]) mergeargs = mergebundles.parse(mergeopts) mergebundles.main(mergeargs) if args.model is not None: model = None for b in bundles: outmodel = "{}_model_{:02d}.fits".format(outroot, b) if model is None: model = fits.getdata(outmodel) else: #- TODO: test and warn if models overlap for pixels with #- non-zero values model += fits.getdata(outmodel) os.remove(outmodel) fits.writeto(args.model, model) mark_merge_end = time.time() time_merge = mark_merge_end - mark_merge_start # Resolve difference timer data if type(timing) is dict: timing["read_input"] = mark_read_input - mark_start timing["preparation"] = mark_preparation - mark_read_input timing["total_extraction"] = time_total_extraction timing["total_write_output"] = time_total_write_output timing["merge"] = time_merge
def main(args) : global psfs log=get_logger() log.info("starting") # read preprocessed image image=read_image(args.image) log.info("read image {}".format(args.image)) if image.mask is not None : image.ivar *= (image.mask==0) xytraceset = read_xytraceset(args.psf) wavemin = xytraceset.wavemin wavemax = xytraceset.wavemax xcoef = xytraceset.x_vs_wave_traceset._coeff ycoef = xytraceset.y_vs_wave_traceset._coeff nfibers=xcoef.shape[0] log.info("read PSF trace with xcoef.shape = {} , ycoef.shape = {} , and wavelength range {}:{}".format(xcoef.shape,ycoef.shape,int(wavemin),int(wavemax))) lines=None if args.lines is not None : log.info("We will fit the image using the psf model and lines") # read lines lines=np.loadtxt(args.lines,usecols=[0]) ok=(lines>wavemin)&(lines<wavemax) log.info("read {} lines in {}, with {} of them in traces wavelength range".format(len(lines),args.lines,np.sum(ok))) lines=lines[ok] else : log.info("We will do an internal calibration of trace coordinates without using the psf shape in a first step") internal_wavelength_calib = (not args.continuum) external_wavelength_calib = args.sky | ( args.spectrum is not None ) if args.auto : log.debug("read flavor of input image {}".format(args.image)) hdus = pyfits.open(args.image) if "FLAVOR" not in hdus[0].header : log.error("no FLAVOR keyword in image header, cannot run with --auto option") raise KeyError("no FLAVOR keyword in image header, cannot run with --auto option") flavor = hdus[0].header["FLAVOR"].strip().lower() hdus.close() log.info("Input is a '{}' image".format(flavor)) if flavor == "flat" : internal_wavelength_calib = False external_wavelength_calib = False elif flavor == "arc" : internal_wavelength_calib = True external_wavelength_calib = False else : internal_wavelength_calib = True external_wavelength_calib = True log.info("wavelength calib, internal={}, external={}".format(internal_wavelength_calib,external_wavelength_calib)) spectrum_filename = args.spectrum if external_wavelength_calib and spectrum_filename is None : srch_file = "data/spec-sky.dat" if not resource_exists('desispec', srch_file): log.error("Cannot find sky spectrum file {:s}".format(srch_file)) raise RuntimeError("Cannot find sky spectrum file {:s}".format(srch_file)) else : spectrum_filename=resource_filename('desispec', srch_file) log.info("Use external calibration from cross-correlation with {}".format(spectrum_filename)) if args.nfibers is not None : nfibers = args.nfibers # FOR DEBUGGING fibers=np.arange(nfibers) if lines is not None : # use a forward modeling of the image # it's slower and works only for individual lines # it's in principle more accurate # but gives systematic residuals for complex spectra like the sky psf = read_specter_psf(args.psf) x,y,dx,ex,dy,ey,fiber_xy,wave_xy=compute_dx_dy_using_psf(psf,image,fibers,lines) x_for_dx=x y_for_dx=y fiber_for_dx=fiber_xy wave_for_dx=wave_xy x_for_dy=x y_for_dy=y fiber_for_dy=fiber_xy wave_for_dy=wave_xy else : # internal calibration method that does not use the psf # nor a prior set of lines. this method is much faster # measure x shifts x_for_dx,y_for_dx,dx,ex,fiber_for_dx,wave_for_dx = compute_dx_from_cross_dispersion_profiles(xcoef,ycoef,wavemin,wavemax, image=image, fibers=fibers, width=args.width, deg=args.degxy,image_rebin=args.ccd_rows_rebin) if internal_wavelength_calib : # measure y shifts x_for_dy,y_for_dy,dy,ey,fiber_for_dy,wave_for_dy = compute_dy_using_boxcar_extraction(xytraceset, image=image, fibers=fibers, width=args.width) mdy = np.median(dy) log.info("Subtract median(dy)={}".format(mdy)) dy -= mdy # remove median, because this is an internal calibration else : # duplicate dx results with zero shift to avoid write special case code below x_for_dy = x_for_dx.copy() y_for_dy = y_for_dx.copy() dy = np.zeros(dx.shape) ey = 1.e-6*np.ones(ex.shape) fiber_for_dy = fiber_for_dx.copy() wave_for_dy = wave_for_dx.copy() degxx=args.degxx degxy=args.degxy degyx=args.degyx degyy=args.degyy while(True) : # loop because polynomial degrees could be reduced log.info("polynomial fit of measured offsets with degx=(%d,%d) degy=(%d,%d)"%(degxx,degxy,degyx,degyy)) try : dx_coeff,dx_coeff_covariance,dx_errorfloor,dx_mod,dx_mask=polynomial_fit(z=dx,ez=ex,xx=x_for_dx,yy=y_for_dx,degx=degxx,degy=degxy) dy_coeff,dy_coeff_covariance,dy_errorfloor,dy_mod,dy_mask=polynomial_fit(z=dy,ez=ey,xx=x_for_dy,yy=y_for_dy,degx=degyx,degy=degyy) log.info("dx dy error floor = %4.3f %4.3f pixels"%(dx_errorfloor,dy_errorfloor)) log.info("check fit uncertainties are ok on edge of CCD") merr=0. for fiber in [0,nfibers-1] : for rw in [-1,1] : tx = legval(rw,xcoef[fiber]) ty = legval(rw,ycoef[fiber]) m=monomials(tx,ty,degxx,degxy) tdx=np.inner(dx_coeff,m) tsx=np.sqrt(np.inner(m,dx_coeff_covariance.dot(m))) m=monomials(tx,ty,degyx,degyy) tdy=np.inner(dy_coeff,m) tsy=np.sqrt(np.inner(m,dy_coeff_covariance.dot(m))) merr=max(merr,tsx) merr=max(merr,tsy) log.info("max edge shift error = %4.3f pixels"%merr) if degxx==0 and degxy==0 and degyx==0 and degyy==0 : break except ( LinAlgError , ValueError ) : log.warning("polynomial fit failed with degx=(%d,%d) degy=(%d,%d)"%(degxx,degxy,degyx,degyy)) if degxx==0 and degxy==0 and degyx==0 and degyy==0 : log.error("polynomial degrees are already 0. we can fit the offsets") raise RuntimeError("polynomial degrees are already 0. we can fit the offsets") merr = 100000. # this will lower the pol. degree. if merr > args.max_error : if merr != 100000. : log.warning("max edge shift error = %4.3f pixels is too large, reducing degrees"%merr) if degxy>0 and degyy>0 and degxy>degxx and degyy>degyx : # first along wavelength if degxy>0 : degxy-=1 if degyy>0 : degyy-=1 else : # then along fiber if degxx>0 : degxx-=1 if degyx>0 : degyx-=1 else : # error is ok, so we quit the loop break # write this for debugging if args.outoffsets : file=open(args.outoffsets,"w") file.write("# axis wave fiber x y delta error polval (axis 0=y axis1=x)\n") for e in range(dy.size) : file.write("0 %f %d %f %f %f %f %f\n"%(wave_for_dy[e],fiber_for_dy[e],x_for_dy[e],y_for_dy[e],dy[e],ey[e],dy_mod[e])) for e in range(dx.size) : file.write("1 %f %d %f %f %f %f %f\n"%(wave_for_dx[e],fiber_for_dx[e],x_for_dx[e],y_for_dx[e],dx[e],ex[e],dx_mod[e])) file.close() log.info("wrote offsets in ASCII file %s"%args.outoffsets) # print central shift mx=np.median(x_for_dx) my=np.median(y_for_dx) m=monomials(mx,my,degxx,degxy) mdx=np.inner(dx_coeff,m) mex=np.sqrt(np.inner(m,dx_coeff_covariance.dot(m))) mx=np.median(x_for_dy) my=np.median(y_for_dy) m=monomials(mx,my,degyx,degyy) mdy=np.inner(dy_coeff,m) mey=np.sqrt(np.inner(m,dy_coeff_covariance.dot(m))) log.info("central shifts dx = %4.3f +- %4.3f dy = %4.3f +- %4.3f "%(mdx,mex,mdy,mey)) # for each fiber, apply offsets and recompute legendre polynomial log.info("for each fiber, apply offsets and recompute legendre polynomial") # compute x y to record max deviations u = np.linspace(-1,1,5) x0 = np.zeros((xcoef.shape[0],u.size)) y0 = np.zeros((ycoef.shape[0],u.size)) for f in range(xcoef.shape[0]) : x0[f]=legval(u,xcoef[f]) y0[f]=legval(u,ycoef[f]) xcoef,ycoef = recompute_legendre_coefficients(xcoef=xcoef,ycoef=ycoef,wavemin=wavemin,wavemax=wavemax,degxx=degxx,degxy=degxy,degyx=degyx,degyy=degyy,dx_coeff=dx_coeff,dy_coeff=dy_coeff) # use an input spectrum as an external calibration of wavelength if spectrum_filename : log.info("write and reread PSF to be sure predetermined shifts were propagated") write_traces_in_psf(args.psf,args.outpsf,xcoef,ycoef,wavemin,wavemax) #psf,xcoef,ycoef,wavemin,wavemax = read_psf_and_traces(args.outpsf) xytraceset = read_xytraceset(args.outpsf) wavemin = xytraceset.wavemin wavemax = xytraceset.wavemax xcoef = xytraceset.x_vs_wave_traceset._coeff ycoef = xytraceset.y_vs_wave_traceset._coeff psf = read_specter_psf(args.outpsf) ycoef=shift_ycoef_using_external_spectrum(psf=psf,xytraceset=xytraceset, image=image,fibers=fibers,spectrum_filename=spectrum_filename,degyy=args.degyy,width=7) x = np.zeros((xcoef.shape[0],u.size)) y = np.zeros((ycoef.shape[0],u.size)) for f in range(xcoef.shape[0]) : x[f]=legval(u,xcoef[f]) y[f]=legval(u,ycoef[f]) dx = x-x0 dy = y-y0 header_keywords = {} header_keywords["MEANDX"]=np.mean(dx) header_keywords["MINDX"]=np.min(dx) header_keywords["MAXDX"]=np.max(dx) header_keywords["MEANDY"]=np.mean(dy) header_keywords["MINDY"]=np.min(dy) header_keywords["MAXDY"]=np.max(dy) write_traces_in_psf(args.psf,args.outpsf,xcoef,ycoef,wavemin,wavemax,header_keywords=header_keywords) log.info("wrote modified PSF in %s"%args.outpsf) else : x = np.zeros((xcoef.shape[0],u.size)) y = np.zeros((ycoef.shape[0],u.size)) for f in range(xcoef.shape[0]) : x[f]=legval(u,xcoef[f]) y[f]=legval(u,ycoef[f]) dx = x-x0 dy = y-y0 header_keywords = {} header_keywords["MEANDX"]=np.mean(dx) header_keywords["MINDX"]=np.min(dx) header_keywords["MAXDX"]=np.max(dx) header_keywords["MEANDY"]=np.mean(dy) header_keywords["MINDY"]=np.min(dy) header_keywords["MAXDY"]=np.max(dy) write_traces_in_psf(args.psf,args.outpsf,xcoef,ycoef,wavemin,wavemax,header_keywords=header_keywords) log.info("wrote modified PSF in %s"%args.outpsf)
def main(args=None): if args is None: args = parse() elif isinstance(args, (list, tuple)): args = parse(args) t0 = time.time() log = get_logger() # guess if it is a preprocessed or a raw image hdulist = fits.open(args.image) is_input_preprocessed = ("IMAGE" in hdulist) & ("IVAR" in hdulist) primary_header = hdulist[0].header hdulist.close() if is_input_preprocessed: image = read_image(args.image) else: if args.camera is None: print( "ERROR: Need to specify camera to open a raw fits image (with all cameras in different fits HDUs)" ) print( "Try adding the option '--camera xx', with xx in {brz}{0-9}, like r7, or type 'desi_qproc --help' for more options" ) sys.exit(12) image = read_raw(args.image, args.camera, fill_header=[ 1, ]) if args.auto: log.debug("AUTOMATIC MODE") try: night = image.meta['NIGHT'] if not 'EXPID' in image.meta: if 'EXPNUM' in image.meta: log.warning('using EXPNUM {} for EXPID'.format( image.meta['EXPNUM'])) image.meta['EXPID'] = image.meta['EXPNUM'] expid = image.meta['EXPID'] except KeyError as e: log.error( "Need at least NIGHT and EXPID (or EXPNUM) to run in auto mode. Retry without the --auto option." ) log.error(str(e)) sys.exit(12) indir = os.path.dirname(args.image) if args.fibermap is None: filename = '{}/fibermap-{:08d}.fits'.format(indir, expid) if os.path.isfile(filename): log.debug("auto-mode: found a fibermap, {}, using it!".format( filename)) args.fibermap = filename if args.output_preproc is None: if not is_input_preprocessed: args.output_preproc = '{}/preproc-{}-{:08d}.fits'.format( args.auto_output_dir, args.camera.lower(), expid) log.debug("auto-mode: will write preproc in " + args.output_preproc) else: log.debug( "auto-mode: will not write preproc because input is a preprocessed image" ) if args.auto_output_dir != '.': if not os.path.isdir(args.auto_output_dir): log.debug("auto-mode: creating directory " + args.auto_output_dir) os.makedirs(args.auto_output_dir) if args.output_preproc is not None: write_image(args.output_preproc, image) cfinder = None if args.psf is None: if cfinder is None: cfinder = CalibFinder([image.meta, primary_header]) args.psf = cfinder.findfile("PSF") log.info(" Using PSF {}".format(args.psf)) tset = read_xytraceset(args.psf) # add fibermap if args.fibermap: if os.path.isfile(args.fibermap): fibermap = read_fibermap(args.fibermap) else: log.error("no fibermap file {}".format(args.fibermap)) fibermap = None else: fibermap = None if "OBSTYPE" in image.meta: obstype = image.meta["OBSTYPE"].upper() image.meta["OBSTYPE"] = obstype # make sure it's upper case qframe = None else: log.warning("No OBSTYPE keyword, trying to guess ...") qframe = qproc_boxcar_extraction(tset, image, width=args.width, fibermap=fibermap) obstype = check_qframe_flavor( qframe, input_flavor=image.meta["FLAVOR"]).upper() image.meta["OBSTYPE"] = obstype log.info("OBSTYPE = '{}'".format(obstype)) if args.auto: # now set the things to do if obstype == "SKY" or obstype == "TWILIGHT" or obstype == "SCIENCE": args.shift_psf = True args.output_psf = '{}/psf-{}-{:08d}.fits'.format( args.auto_output_dir, args.camera, expid) args.output_rawframe = '{}/qframe-{}-{:08d}.fits'.format( args.auto_output_dir, args.camera, expid) args.apply_fiberflat = True args.skysub = True args.output_skyframe = '{}/qsky-{}-{:08d}.fits'.format( args.auto_output_dir, args.camera, expid) args.fluxcalib = True args.outframe = '{}/qcframe-{}-{:08d}.fits'.format( args.auto_output_dir, args.camera, expid) elif obstype == "ARC" or obstype == "TESTARC": args.shift_psf = True args.output_psf = '{}/psf-{}-{:08d}.fits'.format( args.auto_output_dir, args.camera, expid) args.output_rawframe = '{}/qframe-{}-{:08d}.fits'.format( args.auto_output_dir, args.camera, expid) args.compute_lsf_sigma = True elif obstype == "FLAT" or obstype == "TESTFLAT": args.shift_psf = True args.output_psf = '{}/psf-{}-{:08d}.fits'.format( args.auto_output_dir, args.camera, expid) args.output_rawframe = '{}/qframe-{}-{:08d}.fits'.format( args.auto_output_dir, args.camera, expid) args.compute_fiberflat = '{}/qfiberflat-{}-{:08d}.fits'.format( args.auto_output_dir, args.camera, expid) if args.shift_psf: # using the trace shift script if args.auto: options = option_list({ "psf": args.psf, "image": "dummy", "outpsf": "dummy", "continuum": ((obstype == "FLAT") | (obstype == "TESTFLAT")), "sky": ((obstype == "SCIENCE") | (obstype == "SKY")) }) else: options = option_list({ "psf": args.psf, "image": "dummy", "outpsf": "dummy" }) tmp_args = trace_shifts_script.parse(options=options) tset = trace_shifts_script.fit_trace_shifts(image=image, args=tmp_args) qframe = qproc_boxcar_extraction(tset, image, width=args.width, fibermap=fibermap) if tset.meta is not None: # add traceshift info in the qframe, this will be saved in the qframe header if qframe.meta is None: qframe.meta = dict() for k in tset.meta.keys(): qframe.meta[k] = tset.meta[k] if args.output_rawframe is not None: write_qframe(args.output_rawframe, qframe) log.info("wrote raw extracted frame in {}".format( args.output_rawframe)) if args.compute_lsf_sigma: tset = process_arc(qframe, tset, linelist=None, npoly=2, nbins=2) if args.output_psf is not None: for k in qframe.meta: if k not in tset.meta: tset.meta[k] = qframe.meta[k] write_xytraceset(args.output_psf, tset) if args.compute_fiberflat is not None: fiberflat = qproc_compute_fiberflat(qframe) #write_qframe(args.compute_fiberflat,qflat) write_fiberflat(args.compute_fiberflat, fiberflat, header=qframe.meta) log.info("wrote fiberflat in {}".format(args.compute_fiberflat)) if args.apply_fiberflat or args.input_fiberflat: if args.input_fiberflat is None: if cfinder is None: cfinder = CalibFinder([image.meta, primary_header]) try: args.input_fiberflat = cfinder.findfile("FIBERFLAT") except KeyError as e: log.error("no FIBERFLAT for this spectro config") sys.exit(12) log.info("applying fiber flat {}".format(args.input_fiberflat)) flat = read_fiberflat(args.input_fiberflat) qproc_apply_fiberflat(qframe, flat) if args.skysub: log.info("sky subtraction") if args.output_skyframe is not None: skyflux = qproc_sky_subtraction(qframe, return_skymodel=True) sqframe = QFrame(qframe.wave, skyflux, np.ones(skyflux.shape)) write_qframe(args.output_skyframe, sqframe) log.info("wrote sky model in {}".format(args.output_skyframe)) else: qproc_sky_subtraction(qframe) if args.fluxcalib: if cfinder is None: cfinder = CalibFinder([image.meta, primary_header]) # check for flux calib if cfinder.haskey("FLUXCALIB"): fluxcalib_filename = cfinder.findfile("FLUXCALIB") fluxcalib = read_average_flux_calibration(fluxcalib_filename) log.info("read average calib in {}".format(fluxcalib_filename)) seeing = qframe.meta["SEEING"] airmass = qframe.meta["AIRMASS"] exptime = qframe.meta["EXPTIME"] exposure_calib = fluxcalib.value(seeing=seeing, airmass=airmass) for q in range(qframe.nspec): fiber_calib = np.interp(qframe.wave[q], fluxcalib.wave, exposure_calib) * exptime inv_calib = (fiber_calib > 0) / (fiber_calib + (fiber_calib == 0)) qframe.flux[q] *= inv_calib qframe.ivar[q] *= fiber_calib**2 * (fiber_calib > 0) # add keyword in header giving the calibration factor applied at a reference wavelength band = qframe.meta["CAMERA"].upper()[0] if band == "B": refwave = 4500 elif band == "R": refwave = 6500 else: refwave = 8500 calvalue = np.interp(refwave, fluxcalib.wave, exposure_calib) * exptime qframe.meta["CALWAVE"] = refwave qframe.meta["CALVALUE"] = calvalue else: log.error( "Cannot calibrate fluxes because no FLUXCALIB keywork in calibration files" ) fibers = parse_fibers(args.fibers) if fibers is None: fibers = qframe.flux.shape[0] else: ii = np.arange(qframe.fibers.size)[np.in1d(qframe.fibers, fibers)] if ii.size == 0: log.error("no such fibers in frame,") log.error("fibers are in range [{}:{}]".format( qframe.fibers[0], qframe.fibers[-1] + 1)) sys.exit(12) qframe = qframe[ii] if args.outframe is not None: write_qframe(args.outframe, qframe) log.info("wrote {}".format(args.outframe)) t1 = time.time() log.info("all done in {:3.1f} sec".format(t1 - t0)) if args.plot: log.info("plotting {} spectra".format(qframe.wave.shape[0])) import matplotlib.pyplot as plt fig = plt.figure() for i in range(qframe.wave.shape[0]): j = (qframe.ivar[i] > 0) plt.plot(qframe.wave[i, j], qframe.flux[i, j]) plt.grid() plt.xlabel("wavelength") plt.ylabel("flux") plt.show()
def main_mpi(args, comm=None): psf_file = args.psf input_file = args.input # these parameters are interpreted as the *global* spec range, # to be divided among processes. specmin = args.specmin nspec = args.nspec #- Load input files and broadcast # FIXME: after we have fixed the serialization # of the PSF, read and broadcast here, to reduce # disk contention. img = None if comm is None: img = io.read_image(input_file) else: if comm.rank == 0: img = io.read_image(input_file) img = comm.bcast(img, root=0) psf = load_psf(psf_file) # get spectral range if nspec is None: nspec = psf.nspec specmax = specmin + nspec camera = img.meta['CAMERA'].lower() #- b0, r1, .. z9 spectrograph = int(camera[1]) fibermin = spectrograph * psf.nspec + specmin if args.fibermap is not None: fibermap = io.read_fibermap(args.fibermap) fibermap = fibermap[fibermin:fibermin + nspec] fibers = fibermap['FIBER'] else: fibermap = None fibers = np.arange(fibermin, fibermin + nspec, dtype='i4') #- Get wavelength grid from options if args.wavelength is not None: wstart, wstop, dw = map(float, args.wavelength.split(',')) else: wstart = np.ceil(psf.wmin_all) wstop = np.floor(psf.wmax_all) dw = 0.5 wave = np.arange(wstart, wstop + dw / 2.0, dw) nwave = len(wave) #- Confirm that this PSF covers these wavelengths for these spectra psf_wavemin = np.max(psf.wavelength(range(specmin, specmax), y=0)) psf_wavemax = np.min( psf.wavelength(range(specmin, specmax), y=psf.npix_y - 1)) if psf_wavemin > wstart: raise ValueError, 'Start wavelength {:.2f} < min wavelength {:.2f} for these fibers'.format( wstart, psf_wavemin) if psf_wavemax < wstop: raise ValueError, 'Stop wavelength {:.2f} > max wavelength {:.2f} for these fibers'.format( wstop, psf_wavemax) # Now we divide our spectra into bundles bundlesize = args.bundlesize checkbundles = set() checkbundles.update( np.floor_divide(np.arange(specmin, specmax), bundlesize * np.ones(nspec)).astype(int)) bundles = sorted(list(checkbundles)) nbundle = len(bundles) bspecmin = {} bnspec = {} for b in bundles: if specmin > b * bundlesize: bspecmin[b] = specmin else: bspecmin[b] = b * bundlesize if (b + 1) * bundlesize > specmax: bnspec[b] = specmax - bspecmin[b] else: bnspec[b] = bundlesize # Now we assign bundles to processes nproc = 1 rank = 0 if comm is not None: nproc = comm.size rank = comm.rank mynbundle = int(nbundle / nproc) myfirstbundle = 0 leftover = nbundle % nproc if rank < leftover: mynbundle += 1 myfirstbundle = rank * mynbundle else: myfirstbundle = ((mynbundle + 1) * leftover) + (mynbundle * (rank - leftover)) if rank == 0: #- Print parameters print "extract: input = {}".format(input_file) print "extract: psf = {}".format(psf_file) print "extract: specmin = {}".format(specmin) print "extract: nspec = {}".format(nspec) print "extract: wavelength = {},{},{}".format(wstart, wstop, dw) print "extract: nwavestep = {}".format(args.nwavestep) print "extract: regularize = {}".format(args.regularize) # get the root output file outpat = re.compile(r'(.*)\.fits') outmat = outpat.match(args.output) if outmat is None: raise RuntimeError( "extraction output file should have .fits extension") outroot = outmat.group(1) outdir = os.path.dirname(outroot) if rank == 0: if not os.path.isdir(outdir): os.makedirs(outdir) if comm is not None: comm.barrier() failcount = 0 for b in range(myfirstbundle, myfirstbundle + mynbundle): outbundle = "{}_{:02d}.fits".format(outroot, b) print('extract: Starting {} spectra {}:{} at {}'.format( os.path.basename(input_file), bspecmin[b], bspecmin[b] + bnspec[b], time.asctime())) #- The actual extraction try: flux, ivar, Rdata = ex2d(img.pix, img.ivar * (img.mask == 0), psf, bspecmin[b], bnspec[b], wave, regularize=args.regularize, ndecorr=True, bundlesize=bundlesize, wavesize=args.nwavestep, verbose=args.verbose) #- Augment input image header for output img.meta['NSPEC'] = (nspec, 'Number of spectra') img.meta['WAVEMIN'] = (wstart, 'First wavelength [Angstroms]') img.meta['WAVEMAX'] = (wstop, 'Last wavelength [Angstroms]') img.meta['WAVESTEP'] = (dw, 'Wavelength step size [Angstroms]') img.meta['SPECTER'] = (specter.__version__, 'https://github.com/desihub/specter') img.meta['IN_PSF'] = (_trim(psf_file), 'Input spectral PSF') img.meta['IN_IMG'] = (_trim(input_file), 'Input image') bfibermap = fibermap[bspecmin[b] - specmin:bspecmin[b] + bnspec[b] - specmin] bfibers = fibers[bspecmin[b] - specmin:bspecmin[b] + bnspec[b] - specmin] frame = Frame(wave, flux, ivar, resolution_data=Rdata, fibers=bfibers, meta=img.meta, fibermap=bfibermap) #- Write output io.write_frame(outbundle, frame) print('extract: Done {} spectra {}:{} at {}'.format( os.path.basename(input_file), bspecmin[b], bspecmin[b] + bnspec[b], time.asctime())) except: failcount += 1 if comm is not None: failcount = comm.allreduce(failcount) if failcount > 0: # all processes throw raise RuntimeError("some extraction bundles failed") if rank == 0: opts = ['--output', args.output, '--force', '--delete'] opts.extend(["{}_{:02d}.fits".format(outroot, b) for b in bundles]) args = merge.parse(opts) merge.main(args)
def main(args): psf_file = args.psf input_file = args.input specmin = args.specmin nspec = args.nspec #- Load input files psf = load_psf(psf_file) img = io.read_image(input_file) if nspec is None: nspec = psf.nspec specmax = specmin + nspec camera = img.meta['CAMERA'].lower() #- b0, r1, .. z9 spectrograph = int(camera[1]) fibermin = spectrograph * psf.nspec + specmin print('Starting {} spectra {}:{} at {}'.format( os.path.basename(input_file), specmin, specmin + nspec, time.asctime())) if args.fibermap is not None: fibermap = io.read_fibermap(args.fibermap) fibermap = fibermap[fibermin:fibermin + nspec] fibers = fibermap['FIBER'] else: fibermap = None fibers = np.arange(fibermin, fibermin + nspec, dtype='i4') #- Get wavelength grid from options if args.wavelength is not None: wstart, wstop, dw = map(float, args.wavelength.split(',')) else: wstart = np.ceil(psf.wmin_all) wstop = np.floor(psf.wmax_all) dw = 0.5 wave = np.arange(wstart, wstop + dw / 2.0, dw) nwave = len(wave) bundlesize = args.bundlesize #- Confirm that this PSF covers these wavelengths for these spectra psf_wavemin = np.max(psf.wavelength(range(specmin, specmax), y=0)) psf_wavemax = np.min( psf.wavelength(range(specmin, specmax), y=psf.npix_y - 1)) if psf_wavemin > wstart: raise ValueError, 'Start wavelength {:.2f} < min wavelength {:.2f} for these fibers'.format( wstart, psf_wavemin) if psf_wavemax < wstop: raise ValueError, 'Stop wavelength {:.2f} > max wavelength {:.2f} for these fibers'.format( wstop, psf_wavemax) #- Print parameters print """\ #--- Extraction Parameters --- input: {input} psf: {psf} output: {output} wavelength: {wstart} - {wstop} AA steps {dw} specmin: {specmin} nspec: {nspec} regularize: {regularize} #-----------------------------\ """.format(input=input_file, psf=psf_file, output=args.output, wstart=wstart, wstop=wstop, dw=dw, specmin=specmin, nspec=nspec, regularize=args.regularize) #- The actual extraction flux, ivar, Rdata = ex2d(img.pix, img.ivar * (img.mask == 0), psf, specmin, nspec, wave, regularize=args.regularize, ndecorr=True, bundlesize=bundlesize, wavesize=args.nwavestep, verbose=args.verbose) #- Augment input image header for output img.meta['NSPEC'] = (nspec, 'Number of spectra') img.meta['WAVEMIN'] = (wstart, 'First wavelength [Angstroms]') img.meta['WAVEMAX'] = (wstop, 'Last wavelength [Angstroms]') img.meta['WAVESTEP'] = (dw, 'Wavelength step size [Angstroms]') img.meta['SPECTER'] = (specter.__version__, 'https://github.com/desihub/specter') img.meta['IN_PSF'] = (_trim(psf_file), 'Input spectral PSF') img.meta['IN_IMG'] = (_trim(input_file), 'Input image') frame = Frame(wave, flux, ivar, resolution_data=Rdata, fibers=fibers, meta=img.meta, fibermap=fibermap) #- Write output io.write_frame(args.output, frame) print('Done {} spectra {}:{} at {}'.format(os.path.basename(input_file), specmin, specmin + nspec, time.asctime()))
def main_mpi(args, comm=None): psf_file = args.psf input_file = args.input # these parameters are interpreted as the *global* spec range, # to be divided among processes. specmin = args.specmin nspec = args.nspec #- Load input files and broadcast # FIXME: after we have fixed the serialization # of the PSF, read and broadcast here, to reduce # disk contention. img = None if comm is None: img = io.read_image(input_file) else: if comm.rank == 0: img = io.read_image(input_file) img = comm.bcast(img, root=0) psf = load_psf(psf_file) # get spectral range if nspec is None: nspec = psf.nspec specmax = specmin + nspec camera = img.meta['CAMERA'].lower() #- b0, r1, .. z9 spectrograph = int(camera[1]) fibermin = spectrograph * psf.nspec + specmin if args.fibermap is not None: fibermap = io.read_fibermap(args.fibermap) fibermap = fibermap[fibermin:fibermin+nspec] fibers = fibermap['FIBER'] else: fibermap = None fibers = np.arange(fibermin, fibermin+nspec, dtype='i4') #- Get wavelength grid from options if args.wavelength is not None: wstart, wstop, dw = map(float, args.wavelength.split(',')) else: wstart = np.ceil(psf.wmin_all) wstop = np.floor(psf.wmax_all) dw = 0.5 wave = np.arange(wstart, wstop+dw/2.0, dw) nwave = len(wave) #- Confirm that this PSF covers these wavelengths for these spectra psf_wavemin = np.max(psf.wavelength(range(specmin, specmax), y=0)) psf_wavemax = np.min(psf.wavelength(range(specmin, specmax), y=psf.npix_y-1)) if psf_wavemin > wstart: raise ValueError, 'Start wavelength {:.2f} < min wavelength {:.2f} for these fibers'.format(wstart, psf_wavemin) if psf_wavemax < wstop: raise ValueError, 'Stop wavelength {:.2f} > max wavelength {:.2f} for these fibers'.format(wstop, psf_wavemax) # Now we divide our spectra into bundles bundlesize = args.bundlesize checkbundles = set() checkbundles.update(np.floor_divide(np.arange(specmin, specmax), bundlesize*np.ones(nspec)).astype(int)) bundles = sorted(list(checkbundles)) nbundle = len(bundles) bspecmin = {} bnspec = {} for b in bundles: if specmin > b * bundlesize: bspecmin[b] = specmin else: bspecmin[b] = b * bundlesize if (b+1) * bundlesize > specmax: bnspec[b] = specmax - bspecmin[b] else: bnspec[b] = bundlesize # Now we assign bundles to processes nproc = 1 rank = 0 if comm is not None: nproc = comm.size rank = comm.rank mynbundle = int(nbundle // nproc) myfirstbundle = 0 leftover = nbundle % nproc if rank < leftover: mynbundle += 1 myfirstbundle = rank * mynbundle else: myfirstbundle = ((mynbundle + 1) * leftover) + (mynbundle * (rank - leftover)) if rank == 0: #- Print parameters print "extract: input = {}".format(input_file) print "extract: psf = {}".format(psf_file) print "extract: specmin = {}".format(specmin) print "extract: nspec = {}".format(nspec) print "extract: wavelength = {},{},{}".format(wstart, wstop, dw) print "extract: nwavestep = {}".format(args.nwavestep) print "extract: regularize = {}".format(args.regularize) # get the root output file outpat = re.compile(r'(.*)\.fits') outmat = outpat.match(args.output) if outmat is None: raise RuntimeError("extraction output file should have .fits extension") outroot = outmat.group(1) outdir = os.path.normpath(os.path.dirname(outroot)) if rank == 0: if not os.path.isdir(outdir): os.makedirs(outdir) if comm is not None: comm.barrier() failcount = 0 for b in range(myfirstbundle, myfirstbundle+mynbundle): outbundle = "{}_{:02d}.fits".format(outroot, b) outmodel = "{}_model_{:02d}.fits".format(outroot, b) print('extract: Rank {} starting {} spectra {}:{} at {}'.format( rank, os.path.basename(input_file), bspecmin[b], bspecmin[b]+bnspec[b], time.asctime(), ) ) #- The actual extraction try: results = ex2d(img.pix, img.ivar*(img.mask==0), psf, bspecmin[b], bnspec[b], wave, regularize=args.regularize, ndecorr=True, bundlesize=bundlesize, wavesize=args.nwavestep, verbose=args.verbose, full_output=True) flux = results['flux'] ivar = results['ivar'] Rdata = results['resolution_data'] chi2pix = results['chi2pix'] mask = np.zeros(flux.shape, dtype=np.uint32) mask[results['pixmask_fraction']>0.5] |= specmask.SOMEBADPIX mask[results['pixmask_fraction']==1.0] |= specmask.ALLBADPIX mask[chi2pix>100.0] |= specmask.BAD2DFIT #- Augment input image header for output img.meta['NSPEC'] = (nspec, 'Number of spectra') img.meta['WAVEMIN'] = (wstart, 'First wavelength [Angstroms]') img.meta['WAVEMAX'] = (wstop, 'Last wavelength [Angstroms]') img.meta['WAVESTEP']= (dw, 'Wavelength step size [Angstroms]') img.meta['SPECTER'] = (specter.__version__, 'https://github.com/desihub/specter') img.meta['IN_PSF'] = (_trim(psf_file), 'Input spectral PSF') img.meta['IN_IMG'] = (_trim(input_file), 'Input image') if fibermap is not None: bfibermap = fibermap[bspecmin[b]-specmin:bspecmin[b]+bnspec[b]-specmin] else: bfibermap = None bfibers = fibers[bspecmin[b]-specmin:bspecmin[b]+bnspec[b]-specmin] frame = Frame(wave, flux, ivar, mask=mask, resolution_data=Rdata, fibers=bfibers, meta=img.meta, fibermap=bfibermap, chi2pix=chi2pix) #- Write output io.write_frame(outbundle, frame) if args.model is not None: from astropy.io import fits fits.writeto(outmodel, results['modelimage'], header=frame.meta) print('extract: Done {} spectra {}:{} at {}'.format(os.path.basename(input_file), bspecmin[b], bspecmin[b]+bnspec[b], time.asctime())) except: failcount += 1 if comm is not None: failcount = comm.allreduce(failcount) if failcount > 0: # all processes throw raise RuntimeError("some extraction bundles failed") if rank == 0: mergeopts = [ '--output', args.output, '--force', '--delete' ] mergeopts.extend([ "{}_{:02d}.fits".format(outroot, b) for b in bundles ]) mergeargs = mergebundles.parse(mergeopts) mergebundles.main(mergeargs) if args.model is not None: model = None for b in bundles: outmodel = "{}_model_{:02d}.fits".format(outroot, b) if model is None: model = fits.getdata(outmodel) else: #- TODO: test and warn if models overlap for pixels with #- non-zero values model += fits.getdata(outmodel) os.remove(outmodel) fits.writeto(args.model, model)
def main(args): psf_file = args.psf input_file = args.input specmin = args.specmin nspec = args.nspec #- Load input files psf = load_psf(psf_file) img = io.read_image(input_file) if nspec is None: nspec = psf.nspec specmax = specmin + nspec camera = img.meta['CAMERA'].lower() #- b0, r1, .. z9 spectrograph = int(camera[1]) fibermin = spectrograph * psf.nspec + specmin print('Starting {} spectra {}:{} at {}'.format(os.path.basename(input_file), specmin, specmin+nspec, time.asctime())) if args.fibermap is not None: fibermap = io.read_fibermap(args.fibermap) fibermap = fibermap[fibermin:fibermin+nspec] fibers = fibermap['FIBER'] else: fibermap = None fibers = np.arange(fibermin, fibermin+nspec, dtype='i4') #- Get wavelength grid from options if args.wavelength is not None: wstart, wstop, dw = map(float, args.wavelength.split(',')) else: wstart = np.ceil(psf.wmin_all) wstop = np.floor(psf.wmax_all) dw = 0.5 wave = np.arange(wstart, wstop+dw/2.0, dw) nwave = len(wave) bundlesize = args.bundlesize #- Confirm that this PSF covers these wavelengths for these spectra psf_wavemin = np.max(psf.wavelength(range(specmin, specmax), y=0)) psf_wavemax = np.min(psf.wavelength(range(specmin, specmax), y=psf.npix_y-1)) if psf_wavemin > wstart: raise ValueError, 'Start wavelength {:.2f} < min wavelength {:.2f} for these fibers'.format(wstart, psf_wavemin) if psf_wavemax < wstop: raise ValueError, 'Stop wavelength {:.2f} > max wavelength {:.2f} for these fibers'.format(wstop, psf_wavemax) #- Print parameters print """\ #--- Extraction Parameters --- input: {input} psf: {psf} output: {output} wavelength: {wstart} - {wstop} AA steps {dw} specmin: {specmin} nspec: {nspec} regularize: {regularize} #-----------------------------\ """.format(input=input_file, psf=psf_file, output=args.output, wstart=wstart, wstop=wstop, dw=dw, specmin=specmin, nspec=nspec, regularize=args.regularize) #- The actual extraction results = ex2d(img.pix, img.ivar*(img.mask==0), psf, specmin, nspec, wave, regularize=args.regularize, ndecorr=True, bundlesize=bundlesize, wavesize=args.nwavestep, verbose=args.verbose, full_output=True) flux = results['flux'] ivar = results['ivar'] Rdata = results['resolution_data'] chi2pix = results['chi2pix'] mask = np.zeros(flux.shape, dtype=np.uint32) mask[results['pixmask_fraction']>0.5] |= specmask.SOMEBADPIX mask[results['pixmask_fraction']==1.0] |= specmask.ALLBADPIX mask[chi2pix>100.0] |= specmask.BAD2DFIT #- Augment input image header for output img.meta['NSPEC'] = (nspec, 'Number of spectra') img.meta['WAVEMIN'] = (wstart, 'First wavelength [Angstroms]') img.meta['WAVEMAX'] = (wstop, 'Last wavelength [Angstroms]') img.meta['WAVESTEP']= (dw, 'Wavelength step size [Angstroms]') img.meta['SPECTER'] = (specter.__version__, 'https://github.com/desihub/specter') img.meta['IN_PSF'] = (_trim(psf_file), 'Input spectral PSF') img.meta['IN_IMG'] = (_trim(input_file), 'Input image') frame = Frame(wave, flux, ivar, mask=mask, resolution_data=Rdata, fibers=fibers, meta=img.meta, fibermap=fibermap, chi2pix=chi2pix) #- Write output io.write_frame(args.output, frame) if args.model is not None: from astropy.io import fits fits.writeto(args.model, results['modelimage'], header=frame.meta, clobber=True) print('Done {} spectra {}:{} at {}'.format(os.path.basename(input_file), specmin, specmin+nspec, time.asctime()))
def main(args=None): log = get_logger() if args is None: args = parse() first_preproc = None images = [] ivars = [] masks = [] log.info("read inputs ...") for filename in args.infile: log.info(" read {}".format(filename)) tmp = read_image(filename) if first_preproc is None: first_preproc = tmp images.append(tmp.pix.ravel()) # make sure we don't include data with ivar=0 mask = tmp.mask + (tmp.ivar == 0).astype(int) masks.append(mask.ravel()) ivars.append((tmp.ivar * (tmp.mask == 0)).ravel()) images = np.array(images) masks = np.array(masks) ivars = np.array(ivars) smask = np.sum(masks, axis=0) log.info("compute masked median image ...") medimage = masked_median(images, masks) log.info("use masked median image to discard outlier ...") good = ((images - medimage)**2 * (ivars > 0) / ((medimage > 0) * (medimage * args.fracerr)**2 + 1. / (ivars + (ivars == 0)))) < args.nsig**3 ivars *= good.astype(float) if args.weighted: log.info("compute weighted mean ...") sw = np.sum(ivars, axis=0) swf = np.sum(ivars * images, axis=0) meanimage = swf / (sw + (sw == 0)) meanivar = sw else: log.info("compute unweighted mean ...") s1 = np.sum((ivars > 0).astype(int), axis=0) sf = np.sum((ivars > 0) * images, axis=0) meanimage = sf / (s1 + (s1 == 0)) meanvar = np.sum( (ivars > 0) / (ivars + (ivars == 0))) / (s1 + (s1 == 0))**2 meanivar = (meanvar > 0) / (meanvar + (meanvar == 0)) log.info("write nimages.fits ...") fitsio.write("nimages.fits", s1.reshape(first_preproc.pix.shape), clobber=True) log.info("compute mask ...") meanmask = masks[0] for mask in masks[1:]: meanmask &= mask log.info("write image ...") preproc = first_preproc shape = preproc.pix.shape preproc.pix = meanimage.reshape(shape) preproc.ivar = meanivar.reshape(shape) preproc.mask = meanmask.reshape(shape) write_image(args.outfile, preproc) log.info("wrote {}".format(args.outfile))