def get_current_results(): thisdir = os.path.split(os.path.abspath(__file__))[0] psffile = thisdir + '/../etc/psfnight-r0.fits' psf = specter.psf.load_psf(psffile) np.random.seed(0) ny, nx = psf.npix_y, psf.npix_x image = np.random.normal(loc=0, scale=1, size=(ny, nx)) imageivar = np.ones_like(image) #- Spectra and wavelengths to extract nspec = 10 wave = np.arange(psf.wmin_all, psf.wmin_all + 200, 1) #- Wake up code flux, ivar, R = ex2d(image, imageivar, psf, 0, nspec, wave) #- Now do it for real t0 = time.time() flux, ivar, R = ex2d(image, imageivar, psf, 0, nspec, wave, bundlesize=nspec) runtime = time.time() - t0 results = dict(flux=flux, ivar=ivar, R=R, runtime=runtime, nspec=nspec, wave=wave) return results
def run_pa(self, input_image, psf, specmin, nspec, wave, regularize=None, ndecorr=True, bundlesize=25, wavesize=50, outfile=None, fibers=None, fibermap=None): import specter from specter.extract import ex2d from desispec.frame import Frame as fr flux, ivar, Rdata = ex2d(input_image.pix, input_image.ivar * (input_image.mask == 0), psf, specmin, nspec, wave, regularize=regularize, ndecorr=ndecorr, bundlesize=bundlesize, wavesize=wavesize) #- Augment input image header for output input_image.meta['NSPEC'] = (nspec, 'Number of spectra') input_image.meta['WAVEMIN'] = (wave[0], 'First wavelength [Angstroms]') input_image.meta['WAVEMAX'] = (wave[-1], 'Last wavelength [Angstroms]') input_image.meta['WAVESTEP'] = (wave[1] - wave[0], 'Wavelength step size [Angstroms]') input_image.meta['SPECTER'] = (specter.__version__, 'https://github.com/desihub/specter') #input_image.meta['IN_PSF'] = (_trim(psf_file), 'Input spectral PSF') #input_image.meta['IN_IMG'] = (_trim(input_file), 'Input image') frame = fr(wave, flux, ivar, resolution_data=Rdata, fibers=fibers, meta=input_image.meta, fibermap=fibermap) if outfile is not None: #- writing to a frame file if needed. from desispec import io io.write_frame(outfile, frame) log.info("wrote frame output file %s" % outfile) return frame
def wrap_ex2d(icpu, qin, qout): p = psutil.Process(os.getpid()) p.cpu_affinity([ icpu, ]) while True: orig_cpu = get_cpu() i, specmin, nspec, wave = qin.get() results = ex2d(image, imageivar, psf, specmin, nspec, wave) # A = np.random.uniform(size=(500,500)) # results = np.linalg.svd(A.T.dot(A)) # results = 1 # time.sleep(np.random.uniform(0,1)) final_cpu = get_cpu() qout.put((i, results, icpu, orig_cpu, final_cpu))
def run_pa(self,input_image,psf,specmin,nspec,wave,regularize=None,ndecorr=True,bundlesize=25,wavesize=50, outfile=None,fibers=None,fibermap=None): import specter from specter.extract import ex2d from desispec.frame import Frame as fr flux,ivar,Rdata=ex2d(input_image.pix,input_image.ivar*(input_image.mask==0),psf,specmin,nspec,wave,regularize=regularize,ndecorr=ndecorr,bundlesize=bundlesize,wavesize=wavesize) #- Augment input image header for output input_image.meta['NSPEC'] = (nspec, 'Number of spectra') input_image.meta['WAVEMIN'] = (wave[0], 'First wavelength [Angstroms]') input_image.meta['WAVEMAX'] = (wave[-1], 'Last wavelength [Angstroms]') input_image.meta['WAVESTEP']= (wave[1]-wave[0], 'Wavelength step size [Angstroms]') input_image.meta['SPECTER'] = (specter.__version__, 'https://github.com/desihub/specter') #input_image.meta['IN_PSF'] = (_trim(psf_file), 'Input spectral PSF') #input_image.meta['IN_IMG'] = (_trim(input_file), 'Input image') frame = fr(wave, flux, ivar, resolution_data=Rdata,fibers=fibers, meta=input_image.meta, fibermap=fibermap) if outfile is not None: #- writing to a frame file if needed. from desispec import io io.write_frame(outfile,frame) log.info("wrote frame output file %s"%outfile) return frame
opts.psf = thisdir + '/../etc/psfnight-r0.fits' assert os.path.exists(opts.psf) psf = specter.psf.load_psf(opts.psf) #- Create fake noisy image np.random.seed(opts.seed) ny, nx = psf.npix_y, psf.npix_x image = np.random.normal(loc=0, scale=1, size=(ny, nx)) imageivar = np.ones_like(image) #- Spectra and wavelengths to extract wave = np.arange(psf.wmin_all, psf.wmin_all + opts.numwave, 1) #- Wake up the code in case there is library loading overhead flux, ivar, R = ex2d(image, imageivar, psf, 0, 2, wave[0:10]) print('Running on {}/{}'.format(platform.node(), platform.processor())) if 'OMP_NUM_THREADS' not in os.environ: print('$OMP_NUM_THREADS not set') else: print('$OMP_NUM_THREADS={}'.format(os.getenv('OMP_NUM_THREADS'))) print('nspec rate') for n in nspec: t0 = time.time() flux, ivar, R = ex2d(image, imageivar, psf, 0, n, wave,
thisdir = os.path.split(knltest.__file__)[0] opts.psf = thisdir + '/../etc/psfnight-r0.fits' assert os.path.exists(opts.psf) psf = specter.psf.load_psf(opts.psf) #- Create fake noisy image ny, nx = psf.npix_y, psf.npix_x image = np.random.normal(loc=0, scale=1, size=(ny, nx)) imageivar = np.ones_like(image) #- Spectra and wavelengths to extract w = np.linspace(psf.wmin_all, psf.wmax_all, 2000) #- Wake up the code in case there is library loading overhead flux, ivar, R = ex2d(image, imageivar, psf, 0, 2, w[0:10]) #- Setup sub extractions extract_args = list() iarg = 0 for specmin in range(0, psf.nspec, opts.bundlesize): for i in range(0, len(w), opts.numwave): x = (iarg, specmin, opts.bundlesize, w[i:i + opts.numwave]) extract_args.append(x) iarg += 1 #- add comm.barrier() here to make sure everyone is woken up? comm.barrier() if comm.rank == 0: t2 = time.time() print('setup time {:.1f}'.format(t2 - t1))
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 wrap_ex2d(args): specmin, nspec, ww, xyrange = args xmin, xmax, ymin, ymax = xyrange return ex2d(image[ymin:ymax, xmin:xmax], image_ivar[ymin:ymax, xmin:xmax], psf, specmin, nspec, ww, xyrange)
#- Load point spread function model if opts.psf is None: thisdir = os.path.split(knltest.__file__)[0] opts.psf = thisdir + '/../etc/psfnight-r0.fits' assert os.path.exists(opts.psf) psf = specter.psf.load_psf(opts.psf) #- Create fake noisy image ny, nx = psf.npix_y, psf.npix_x image = np.random.normal(loc=0, scale=1, size=(ny, nx)) imageivar = np.ones_like(image) #- Spectra and wavelengths to extract w = np.arange(psf.wmin_all, psf.wmin_all + opts.numwave, 1) args = [image, imageivar, psf] kwargs = dict(specmin=0, nspec=opts.numspec, wavelengths=w) #- Wake up the code in case there is library loading overhead flux, ivar, R = ex2d(image, imageivar, psf, 0, 2, w[0:10]) print('Running on {}/{} with {} logical cores'.format( platform.node(), platform.processor(), multiprocessing.cpu_count())) print('{} spectra x {} wavelengths extracted'.format(opts.numspec, opts.numwave)) print('OMP_NUM_THREADS time') for n in ntest: os.environ['OMP_NUM_THREADS'] = str(n) t = knltest.timeit(ex2d, args, kwargs) print("{:3} {:5.1f}".format(n, t), flush=True)
def wrap_ex2d(x): i, specmin, nspec, wave = x return ex2d(image, imageivar, psf, specmin, nspec, wave)
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 _extract_and_save(img, psf, bspecmin, bnspec, specmin, wave, raw_wave, fibers, fibermap, outbundle, outmodel, bundlesize, args, log): ''' Performs the main extraction and saving of extracted frames found in the body of the main loop. Refactored to be callable by both MPI and non-MPI versions of the code. This should be viewed as a shorthand for the following commands. ''' results = ex2d(img.pix, img.ivar * (img.mask == 0), psf, bspecmin, bnspec, wave, regularize=args.regularize, ndecorr=args.decorrelate_fibers, bundlesize=bundlesize, wavesize=args.nwavestep, verbose=args.verbose, full_output=True, nsubbundles=args.nsubbundles) 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 fibermap is not None: bfibermap = fibermap[bspecmin - specmin:bspecmin + bnspec - specmin] else: bfibermap = None bfibers = fibers[bspecmin - specmin:bspecmin + bnspec - specmin] #- Save the raw wavelength, not the corrected one (if corrected) frame = Frame(raw_wave, flux, ivar, mask=mask, resolution_data=Rdata, fibers=bfibers, meta=img.meta, fibermap=bfibermap, 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'] = 'electron/Angstrom' #- Add scores to frame if not args.no_scores: compute_and_append_frame_scores(frame, suffix="RAW") mark_extraction = time.time() #- Write output io.write_frame(outbundle, frame) if args.model is not None: fits.writeto(outmodel, results['modelimage'], header=frame.meta) log.info('extract: Done {} spectra {}:{} at {}'.format( os.path.basename(args.input), bspecmin, bspecmin + bnspec, time.asctime())) sys.stdout.flush() return mark_extraction