def _write_hpx_fits(pixlist): """from files that touch a pixel, write out objects in each pixel""" pixnum, files = pixlist # ADM only proceed if some files touch a pixel. if len(files) > 0: # ADM track if it's our first time through the files loop. first = True # ADM Read in files that touch a pixel. for file in files: filename = os.path.join(fitsdir, file) objs = fitsio.read(filename) # ADM only retain objects in the correct pixel. pix = radec2pix(nside, objs["RA"], objs["DEC"]) if first: done = objs[pix == pixnum] first = False else: done = np.hstack([done, objs[pix == pixnum]]) # ADM construct the name of the output file. outfilename = 'healpix-{:05d}.fits'.format(pixnum) outfile = os.path.join(hpxdir, outfilename) # ADM write out the file. hdr = fitsio.FITSHDR() hdr['HPXNSIDE'] = nside hdr['HPXNEST'] = True fitsio.write(outfile, done, extname='URATHPX', header=hdr) return
def _write_urat_fits(infile): """read an input name for a csv file and write it to FITS""" outbase = os.path.basename(infile) outfilename = "{}.fits".format(outbase.split(".")[0]) outfile = os.path.join(fitsdir, outfilename) # ADM astropy understands without specifying format='csv'. fitstable = ascii.read(infile) # ADM map the ascii-read csv to typical DESI quantities. nobjs = len(fitstable) done = np.zeros(nobjs, dtype=uratdatamodel.dtype) # ADM have to do this one-by-one, given the format. done["RA"] = fitstable['col1'] / 1000. / 3600. done["DEC"] = fitstable['col2'] / 1000. / 3600. - 90. done["PMRA"] = fitstable['col16'] / 10. done["PMDEC"] = fitstable['col17'] / 10. done["PM_ERROR"] = fitstable['col18'] / 10. done["APASS_G_MAG"] = fitstable['col36'] / 1000. done["APASS_R_MAG"] = fitstable['col37'] / 1000. done["APASS_I_MAG"] = fitstable['col38'] / 1000. done["APASS_G_MAG_ERROR"] = fitstable['col41'] / 1000. done["APASS_R_MAG_ERROR"] = fitstable['col42'] / 1000. done["APASS_I_MAG_ERROR"] = fitstable['col43'] / 1000. done["URAT_ID"] = fitstable['col46'] fitsio.write(outfile, done, extname='URATFITS') # ADM return the HEALPixels that this file touches. pix = set(radec2pix(nside, done["RA"], done["DEC"])) return [pix, os.path.basename(outfile)]
def _write_gaia_fits(infile): """read an input name for a csv file and write it to FITS""" outbase = os.path.basename(infile) outfilename = "{}.fits".format(outbase.split(".")[0]) outfile = os.path.join(fitsdir, outfilename) fitstable = ascii.read(infile, format='csv') # ADM need to convert 5-string values to boolean. cols = np.array(fitstable.dtype.names) boolcols = cols[np.hstack(fitstable.dtype.descr)[1::2] == '<U5'] for col in boolcols: fitstable[col] = fitstable[col] == 'true' # ADM only write out the columns we need for targeting. nobjs = len(fitstable) done = np.zeros(nobjs, dtype=ingaiadatamodel.dtype) for col in done.dtype.names: if col == 'REF_CAT': done[col] = 'G2' else: done[col] = fitstable[col.lower()] fitsio.write(outfile, done, extname='GAIAFITS') # ADM return the HEALPixels that this file touches. pix = set(radec2pix(nside, fitstable["ra"], fitstable["dec"])) return [pix, os.path.basename(outfile)]
def join_zcat_fastspec(survey, faprgrm, specred='everest'): """ Join the zpix* redshift catalog with the fastspec catalog for a given 'survey' and 'faprgrm' Parameters ---------- survey : str sv1|sv2|sv3|main faprgrm : str Fiber Assignment Program: dark|bright specred : str DESI Spectral Reduction Pipeline. Default = 'everest' """ # Directory Paths for zcatalogs and fastspec in NERSC zcat_dir = f'/global/cfs/cdirs/desi/spectro/redux/{specred}/zcatalog' fast_dir = f'/global/cfs/cdirs/desi/spectro/fastspecfit/{specred}/catalogs' # Filenames for zcatalog and fastspec for the given survey and faprgrm zcat_file = f'{zcat_dir}/zpix-{survey}-{faprgrm}.fits' fastspec_file = f'{fast_dir}/fastspec-{specred}-{survey}-{faprgrm}.fits' # Reading in the fits files hzcat = fits.open(zcat_file) hspec = fits.open(fastspec_file) # Reading the tables from the hdus zcat = Table(hzcat['ZCATALOG'].data) tspec = Table(hspec['FASTSPEC'].data) mspec = Table(hspec['METADATA'].data) # zcat does not have HPXPIXEL column in everest - it will be added in fuji. # Adding the HPXPIXEL column here for consistency hpx = radec2pix(nside=64, ra=zcat['TARGET_RA'].data, dec=zcat['TARGET_DEC']) hpx_col = Column(hpx, name='HPXPIXEL') zcat.add_column(hpx_col, 1) # Selecting only some columns from zcatalog zcat = zcat['TARGETID', 'HPXPIXEL', 'Z', 'ZERR', 'ZWARN', 'SPECTYPE', \ 'COADD_FIBERSTATUS', 'TARGET_RA', 'TARGET_DEC'] # Selecting important columns from the metadeta mspec_sel = mspec['TARGETID', 'SURVEY', 'FAPRGRM', 'HPXPIXEL'] # Joining the two fastspec tables to create a single table spec = join(mspec_sel, tspec, keys='TARGETID') # Left-joining the zcatalog and fastspec table res_table = join(zcat, spec, keys=['TARGETID', 'HPXPIXEL'], join_type='left') # Should add 'SURVEY' and 'FAPRGM' if we start from summary catalogs return (res_table)
def trim_to_footprint(self, ra, dec): ''' Given RA and Dec values, return a np.ndarray of which pixels are in the DESI footprint. In principle `desimodel.footprint` should return appropriate systematic weight given ra and dec. I dont' think this is the case at the moment. ''' assert len(ra) == len(dec) pixweight = load_pixweight(256) healpix = footprint.radec2pix(256, ra, dec) # ra dec of the targets weights = pixweight[healpix] # weights=1 in footprint return weights
def find_zbest_files(fibermap_data): from desimodel.footprint import radec2pix # Init zbest_files = [] # Search for zbest files with healpy ra_targ = fibermap_data['TARGET_RA'].data dec_targ = fibermap_data['TARGET_DEC'].data # Getting some NAN in RA/DEC good = np.isfinite(ra_targ) & np.isfinite(dec_targ) pixels = radec2pix(64, ra_targ[good], dec_targ[good]) uni_pixels = np.unique(pixels) for uni_pix in uni_pixels: zbest_files.append(desispec.io.findfile('zbest', groupname=uni_pix, nside=64)) # Return return zbest_files
def target_density(cat, nside=128): """Determine the target density by grouping targets in healpix pixels. The code below was code shamelessly taken from desiutil.plot.plot_sky_binned (by D. Kirkby). Args: cat: Table with columns RA and DEC Optional: nside: healpix nside, integer power of 2 """ npix = hp.nside2npix(nside) bin_area = hp.nside2pixarea(nside, degrees=True) pixels = radec2pix(nside, cat['RA'], cat['DEC']) #pixels = hp.ang2pix(nside, np.radians(90 - cat['DEC']), # np.radians(cat['RA']), nest=False) counts = np.bincount(pixels, weights=None, minlength=npix) dens = counts[np.flatnonzero(counts)] / bin_area return dens
def tycho_fits_to_healpix(): """Convert files in $TYCHO_DIR/fits to files in $TYCHO_DIR/healpix. Returns ------- Nothing But the archived Tycho FITS files in $TYCHO_DIR/fits are rearranged by HEALPixel in the directory $TYCHO_DIR/healpix. The HEALPixel sense is nested with nside=get_tycho_nside(), and each file in $TYCHO_DIR/healpix is called healpix-xxxxx.fits, where xxxxx corresponds to the HEALPixel number. Notes ----- - The environment variable $TYCHO_DIR must be set. """ # ADM the resolution at which the Tycho HEALPix files are stored. nside = get_tycho_nside() npix = hp.nside2npix(nside) # ADM check that the TYCHO_DIR is set. tychodir = get_tycho_dir() # ADM construct the directories for reading/writing files. fitsdir = os.path.join(tychodir, "fits") tychofn = os.path.join(fitsdir, "tycho2.kd.fits") hpxdir = os.path.join(tychodir, "healpix") # ADM make sure the output directory is empty. if os.path.exists(hpxdir): if len(os.listdir(hpxdir)) > 0: msg = "{} must be empty to make Tycho HEALPix files!".format(hpxdir) log.critical(msg) raise ValueError(msg) # ADM make the output directory, if needed. else: log.info("Making Tycho directory for storing HEALPix files") os.makedirs(hpxdir) # ADM read in the Tycho file and assing Tycho objects to HEALPixels. objs, allhdr = fitsio.read(tychofn, header=True, upper=True) pix = radec2pix(nside, objs["RA"], objs["DEC"]) # ADM loop through the pixels and write out the files. for pixnum in range(npix): # ADM construct the name of the output file. outfilename = io.hpx_filename(pixnum) outfile = os.path.join(hpxdir, outfilename) # ADM update the header with new information. hdr = dict(allhdr).copy() hdr["HPXNSIDE"] = nside hdr["HPXNEST"] = True hdr["HPXDATE"] = datetime.utcnow().isoformat(timespec='seconds') # ADM determine which objects are in this pixel and write out. done = objs[pix == pixnum] fitsio.write(outfile, done, extname="TYCHOHPX", header=hdr) log.info('Wrote Tycho HEALPix files...t={:.1f}s'.format(time()-start)) return
def integration_test(night=None, nspec=5, clobber=False): """Run an integration test from raw data simulations through redshifts. Args: night (str, optional): YEARMMDD, defaults to current night nspec (int, optional): number of spectra to include clobber (bool, optional): rerun steps even if outputs already exist Raises: RuntimeError if any script fails """ log = get_logger() #- YEARMMDD string, rolls over at noon not midnight if night is None: night = time.strftime('%Y%m%d', time.localtime(time.time() - 12 * 3600)) #- check for required environment variables check_env() #- parameter dictionary that will later be used for formatting commands params = dict(night=night, nspec=nspec) #----- #- Input fibermaps, spectra, and pixel-level raw data # raw_dict = {0: 'flat', 1: 'arc', 2: 'dark'} programs = ('flat', 'arc', 'dark') channels = ('b', 'r', 'z') cameras = ('b0', 'r0', 'z0') # for expid, program in raw_dict.items(): for expid, program in enumerate(programs): cmd = "newexp-random --program {program} --nspec {nspec} --night {night} --expid {expid}".format( expid=expid, program=program, **params) fibermap = io.findfile('fibermap', night, expid) simspec = '{}/simspec-{:08d}.fits'.format(os.path.dirname(fibermap), expid) inputs = [] outputs = [fibermap, simspec] if runcmd(cmd, inputs=inputs, outputs=outputs, clobber=clobber) != 0: raise RuntimeError( 'pixsim newexp failed for {} exposure {}'.format( program, expid)) cmd = "pixsim --preproc --nspec {nspec} --night {night} --expid {expid}".format( expid=expid, **params) inputs = [fibermap, simspec] outputs = list() outputs.append(fibermap.replace('fibermap-', 'simpix-')) for camera in cameras: pixfile = io.findfile('pix', night, expid, camera) outputs.append(pixfile) # outputs.append(os.path.join(os.path.dirname(pixfile), os.path.basename(pixfile).replace('pix-', 'simpix-'))) if runcmd(cmd, inputs=inputs, outputs=outputs, clobber=clobber) != 0: raise RuntimeError('pixsim failed for {} exposure {}'.format( program, expid)) #----- #- Extract waverange = dict(b="3570,5940,1.0", r="5630,7740,1.0", z="7440,9830,1.0") for expid, program in enumerate(programs): for ic, channel in enumerate(channels): pixfile = io.findfile('pix', night, expid, cameras[ic]) fiberfile = io.findfile('fibermap', night, expid) psffile = '{}/data/specpsf/psf-{}.fits'.format( os.getenv('DESIMODEL'), channel) framefile = io.findfile('frame', night, expid, cameras[ic]) # cmd = "exspec -i {pix} -p {psf} --specmin 0 --nspec {nspec} -w {wave} -o {frame}".format( # pix=pixfile, psf=psffile, wave=waverange[channel], frame=framefile, **params) cmd = "desi_extract_spectra -i {pix} -p {psf} -f {fibermap} --specmin 0 --nspec {nspec} -o {frame}".format( pix=pixfile, psf=psffile, frame=framefile, fibermap=fiberfile, **params) inputs = [pixfile, psffile, fiberfile] outputs = [ framefile, ] if runcmd(cmd, inputs=inputs, outputs=outputs, clobber=clobber) != 0: raise RuntimeError('extraction failed for {} expid {}'.format( cameras[ic], expid)) #----- #- Fiber flat expid = 0 for ic, channel in enumerate(channels): framefile = io.findfile('frame', night, expid, cameras[ic]) fiberflat = io.findfile('fiberflat', night, expid, cameras[ic]) fibermap = io.findfile('fibermap', night, expid) # for QA qafile = io.findfile('qa_calib', night, expid, cameras[ic]) qafig = io.findfile('qa_flat_fig', night, expid, cameras[ic]) cmd = "desi_compute_fiberflat --infile {frame} --outfile {fiberflat} --qafile {qafile} --qafig {qafig}".format( frame=framefile, fiberflat=fiberflat, qafile=qafile, qafig=qafig, **params) inputs = [ framefile, fibermap, ] outputs = [ fiberflat, qafile, qafig, ] if runcmd(cmd, inputs=inputs, outputs=outputs, clobber=clobber) != 0: raise RuntimeError('fiberflat failed for ' + cameras[ic]) #----- #- Sky model flat_expid = 0 expid = 2 for ic, channel in enumerate(channels): framefile = io.findfile('frame', night, expid, cameras[ic]) fibermap = io.findfile('fibermap', night, expid) fiberflat = io.findfile('fiberflat', night, flat_expid, cameras[ic]) skyfile = io.findfile('sky', night, expid, cameras[ic]) qafile = io.findfile('qa_data', night, expid, cameras[ic]) qafig = io.findfile('qa_sky_fig', night, expid, cameras[ic]) cmd = "desi_compute_sky --infile {frame} --fiberflat {fiberflat} --outfile {sky} --qafile {qafile} --qafig {qafig}".format( frame=framefile, fiberflat=fiberflat, sky=skyfile, qafile=qafile, qafig=qafig, **params) inputs = [framefile, fibermap, fiberflat] outputs = [ skyfile, qafile, qafig, ] if runcmd(cmd, inputs=inputs, outputs=outputs, clobber=clobber) != 0: raise RuntimeError('sky model failed for ' + cameras[ic]) #----- #- Fit standard stars if 'STD_TEMPLATES' in os.environ: std_templates = os.getenv('STD_TEMPLATES') else: std_templates = os.getenv( 'DESI_ROOT' ) + '/spectro/templates/star_templates/v1.1/star_templates_v1.1.fits' stdstarfile = io.findfile('stdstars', night, expid, spectrograph=0) flats = list() frames = list() skymodels = list() for ic, channel in enumerate(channels): frames.append(io.findfile('frame', night, expid, cameras[ic])) flats.append(io.findfile('fiberflat', night, flat_expid, cameras[ic])) skymodels.append(io.findfile('sky', night, expid, cameras[ic])) frames = ' '.join(frames) flats = ' '.join(flats) skymodels = ' '.join(skymodels) cmd = """desi_fit_stdstars \ --frames {frames} \ --fiberflats {flats} \ --skymodels {skymodels} \ --starmodels {std_templates} \ -o {stdstars}""".format(frames=frames, flats=flats, skymodels=skymodels, std_templates=std_templates, stdstars=stdstarfile) inputs = [fibermap, std_templates] outputs = [ stdstarfile, ] if runcmd(cmd, inputs=inputs, outputs=outputs, clobber=clobber) != 0: raise RuntimeError('fitting stdstars failed') #----- #- Flux calibration for ic, channel in enumerate(channels): framefile = io.findfile('frame', night, expid, cameras[ic]) fibermap = io.findfile('fibermap', night, expid) fiberflat = io.findfile('fiberflat', night, flat_expid, cameras[ic]) skyfile = io.findfile('sky', night, expid, cameras[ic]) calibfile = io.findfile('calib', night, expid, cameras[ic]) qafile = io.findfile('qa_data', night, expid, cameras[ic]) qafig = io.findfile('qa_flux_fig', night, expid, cameras[ic]) #- Compute flux calibration vector cmd = """desi_compute_fluxcalibration \ --infile {frame} --fiberflat {fiberflat} --sky {sky} \ --models {stdstars} --outfile {calib} --qafile {qafile} --qafig {qafig}""".format( frame=framefile, fiberflat=fiberflat, sky=skyfile, stdstars=stdstarfile, calib=calibfile, qafile=qafile, qafig=qafig) inputs = [framefile, fibermap, fiberflat, skyfile, stdstarfile] outputs = [calibfile, qafile, qafig] if runcmd(cmd, inputs=inputs, outputs=outputs, clobber=clobber) != 0: raise RuntimeError('flux calibration failed for ' + cameras[ic]) #- Apply the flux calibration to write a cframe file cframefile = io.findfile('cframe', night, expid, cameras[ic]) cmd = """desi_process_exposure \ --infile {frame} --fiberflat {fiberflat} --sky {sky} --calib {calib} \ --outfile {cframe}""".format(frame=framefile, fibermap=fibermap, fiberflat=fiberflat, sky=skyfile, calib=calibfile, cframe=cframefile) inputs = [framefile, fiberflat, skyfile, calibfile] outputs = [ cframefile, ] if runcmd(cmd, inputs=inputs, outputs=outputs, clobber=clobber) != 0: raise RuntimeError('combining calibration steps failed for ' + cameras[ic]) #----- #- Collate QA # Collate data QA program2flavor = dict(arc='arc', flat='flat') for program in ('dark', 'gray', 'bright', 'elg', 'lrg', 'qso', 'bgs', 'mws'): program2flavor[program] = 'science' expid = 2 qafile = io.findfile('qa_data_exp', night, expid) if clobber or not os.path.exists(qafile): flavor = program2flavor[programs[expid]] qaexp_data = QA_Exposure(expid, night, flavor) # Removes camera files io.write_qa_exposure(os.path.splitext(qafile)[0], qaexp_data) if not os.path.exists(qafile): raise RuntimeError( 'FAILED data QA_Exposure({},{}, ...) -> {}'.format( expid, night, qafile)) # Collate calib QA calib_expid = [0, 1] for expid in calib_expid: qafile = io.findfile('qa_calib_exp', night, expid) if clobber or not os.path.exists(qafile): qaexp_calib = QA_Exposure(expid, night, programs[expid]) io.write_qa_exposure(os.path.splitext(qafile)[0], qaexp_calib) if not os.path.exists(qafile): raise RuntimeError( 'FAILED calib QA_Exposure({},{}, ...) -> {}'.format( expid, night, qafile)) #----- #- Regroup cframe -> spectra expid = 2 inputs = list() for camera in cameras: inputs.append(io.findfile('cframe', night, expid, camera)) outputs = list() fibermap = io.read_fibermap(io.findfile('fibermap', night, expid)) from desimodel.footprint import radec2pix nside = 64 pixels = np.unique( radec2pix(nside, fibermap['RA_TARGET'], fibermap['DEC_TARGET'])) for pix in pixels: outputs.append(io.findfile('spectra', groupname=pix)) cmd = "desi_group_spectra" if runcmd(cmd, inputs=inputs, outputs=outputs, clobber=clobber) != 0: raise RuntimeError('spectra regrouping failed') #----- #- Redshifts! for pix in pixels: specfile = io.findfile('spectra', groupname=pix) zbestfile = io.findfile('zbest', groupname=pix) inputs = [ specfile, ] outputs = [ zbestfile, ] cmd = "rrdesi {} --zbest {}".format(specfile, zbestfile) if runcmd(cmd, inputs=inputs, outputs=outputs, clobber=clobber) != 0: raise RuntimeError('rrdesi failed for healpixel {}'.format(pix)) # # Load redshifts into database # options = get_options( '--clobber', '--filename', 'dailytest.db', os.path.join(os.environ['DESI_SPECTRO_REDUX'], os.environ['SPECPROD'])) postgresql = setup_db(options) load_zcat(options.datapath) # ztruth QA # qafile = io.findfile('qa_ztruth', night) # qafig = io.findfile('qa_ztruth_fig', night) # cmd = "desi_qa_zfind --night {night} --qafile {qafile} --qafig {qafig} --verbose".format( # night=night, qafile=qafile, qafig=qafig) # inputs = [] # outputs = [qafile, qafig] # if runcmd(cmd, inputs=inputs, outputs=outputs, clobber=clobber) != 0: # raise RuntimeError('redshift QA failed for night '+night) #----- #- Did it work? #- (this combination of fibermap, simspec, and zbest is a pain) simdir = os.path.dirname(io.findfile('fibermap', night=night, expid=expid)) simspec = '{}/simspec-{:08d}.fits'.format(simdir, expid) try: siminfo = fits.getdata(simspec, 'TRUTH') except KeyError: siminfo = fits.getdata(simspec, 'METADATA') print() print("--------------------------------------------------") print("Pixel True z -> Class z zwarn") # print("3338p190 SKY 0.00000 -> QSO 1.60853 12 - ok") for pix in pixels: zbest = io.read_zbest(io.findfile('zbest', groupname=pix)) for i in range(len(zbest.z)): objtype = zbest.spectype[i] z, zwarn = zbest.z[i], zbest.zwarn[i] j = np.where(fibermap['TARGETID'] == zbest.targetid[i])[0][0] truetype = siminfo['OBJTYPE'][j] oiiflux = siminfo['OIIFLUX'][j] truez = siminfo['REDSHIFT'][j] dv = 3e5 * (z - truez) / (1 + truez) if truetype == 'SKY' and zwarn > 0: status = 'ok' elif truetype == 'ELG' and zwarn > 0 and oiiflux < 8e-17: status = 'ok ([OII] flux {:.2g})'.format(oiiflux) elif zwarn == 0: if truetype == 'LRG' and objtype == 'GALAXY' and abs(dv) < 150: status = 'ok' elif truetype == 'ELG' and objtype == 'GALAXY': if abs(dv) < 150: status = 'ok' elif oiiflux < 8e-17: status = 'ok ([OII] flux {:.2g})'.format(oiiflux) else: status = 'OOPS ([OII] flux {:.2g})'.format(oiiflux) elif truetype == 'QSO' and objtype == 'QSO' and abs(dv) < 750: status = 'ok' elif truetype in ('STD', 'FSTD') and objtype == 'STAR': status = 'ok' else: status = 'OOPS' else: status = 'OOPS' print('{0:<8d} {1:4s} {2:8.5f} -> {3:6s} {4:8.5f} {5:4d} - {6}'. format(pix, truetype, truez, objtype, z, zwarn, status)) print("--------------------------------------------------")
def integration_test(night=None, nspec=5, clobber=False): """Run an integration test from raw data simulations through redshifts Args: night (str, optional): YEARMMDD, defaults to current night nspec (int, optional): number of spectra to include clobber (bool, optional): rerun steps even if outputs already exist Raises: RuntimeError if any script fails """ import argparse parser = argparse.ArgumentParser(usage="{prog} [options]") # parser.add_argument("-i", "--input", type=str, help="input data") # parser.add_argument("-o", "--output", type=str, help="output data") parser.add_argument("--skip-psf", action="store_true", help="Skip PSF fitting step") args = parser.parse_args() log = logging.get_logger() # YEARMMDD string, rolls over at noon not midnight if night is None: night = "20160726" # check for required environment variables check_env() # simulate inputs sim(night, nspec=nspec, clobber=clobber) # raw and production locations rawdir = os.path.abspath(io.rawdata_root()) proddir = os.path.abspath(io.specprod_root()) # create production if clobber and os.path.isdir(proddir): shutil.rmtree(proddir) dbfile = io.get_pipe_database() if not os.path.exists(dbfile): com = "desi_pipe create --db-sqlite" log.info('Running {}'.format(com)) sp.check_call(com, shell=True) else: log.info("Using pre-existing production database {}".format(dbfile)) # Modify options file to restrict the spectral range optpath = os.path.join(proddir, "run", "options.yaml") opts = pipe.prod.yaml_read(optpath) opts['extract']['specmin'] = 0 opts['extract']['nspec'] = nspec opts['psf']['specmin'] = 0 opts['psf']['nspec'] = nspec opts['traceshift']['nfibers'] = nspec pipe.prod.yaml_write(optpath, opts) if args.skip_psf: #- Copy desimodel psf into this production instead of fitting psf import shutil for channel in ['b', 'r', 'z']: refpsf = '{}/data/specpsf/psf-{}.fits'.format( os.getenv('DESIMODEL'), channel) nightpsf = io.findfile('psfnight', night, camera=channel + '0') shutil.copy(refpsf, nightpsf) for expid in [0, 1, 2]: exppsf = io.findfile('psf', night, expid, camera=channel + '0') shutil.copy(refpsf, exppsf) #- Resync database to current state dbpath = io.get_pipe_database() db = pipe.load_db(dbpath, mode="w") db.sync(night) # Run the pipeline tasks in order from desispec.pipeline.tasks.base import default_task_chain for tasktype in default_task_chain: #- if we skip psf/psfnight/traceshift, update state prior to extractions if tasktype == 'traceshift' and args.skip_psf: db.getready() run_pipeline_step(tasktype) # #----- # #- Did it work? # #- (this combination of fibermap, simspec, and zbest is a pain) expid = 2 fmfile = io.findfile('fibermap', night=night, expid=expid) fibermap = io.read_fibermap(fmfile) simdir = os.path.dirname(fmfile) simspec = '{}/simspec-{:08d}.fits'.format(simdir, expid) siminfo = fits.getdata(simspec, 'TRUTH') try: elginfo = fits.getdata(simspec, 'TRUTH_ELG') except: elginfo = None from desimodel.footprint import radec2pix nside = 64 pixels = np.unique( radec2pix(nside, fibermap['TARGET_RA'], fibermap['TARGET_DEC'])) num_missing = 0 for pix in pixels: zfile = io.findfile('zbest', groupname=pix) if not os.path.exists(zfile): log.error('Missing {}'.format(zfile)) num_missing += 1 if num_missing > 0: log.critical('{} zbest files missing'.format(num_missing)) sys.exit(1) print() print("--------------------------------------------------") print("Pixel True z -> Class z zwarn") # print("3338p190 SKY 0.00000 -> QSO 1.60853 12 - ok") for pix in pixels: zfile = io.findfile('zbest', groupname=pix) if not os.path.exists(zfile): log.error('Missing {}'.format(zfile)) continue zfx = fits.open(zfile, memmap=False) zbest = zfx['ZBEST'].data for i in range(len(zbest['Z'])): objtype = zbest['SPECTYPE'][i] z, zwarn = zbest['Z'][i], zbest['ZWARN'][i] j = np.where(fibermap['TARGETID'] == zbest['TARGETID'][i])[0][0] truetype = siminfo['OBJTYPE'][j] oiiflux = 0.0 if truetype == 'ELG': k = np.where(elginfo['TARGETID'] == zbest['TARGETID'][i])[0][0] oiiflux = elginfo['OIIFLUX'][k] truez = siminfo['REDSHIFT'][j] dv = C_LIGHT * (z - truez) / (1 + truez) status = None if truetype == 'SKY' and zwarn > 0: status = 'ok' elif truetype == 'ELG' and zwarn > 0 and oiiflux < 8e-17: status = 'ok ([OII] flux {:.2g})'.format(oiiflux) elif zwarn == 0: if truetype == 'LRG' and objtype == 'GALAXY' and abs(dv) < 150: status = 'ok' elif truetype == 'ELG' and objtype == 'GALAXY': if abs(dv) < 150: status = 'ok' elif oiiflux < 8e-17: status = 'ok ([OII] flux {:.2g})'.format(oiiflux) else: status = 'OOPS ([OII] flux {:.2g})'.format(oiiflux) elif truetype == 'QSO' and objtype == 'QSO' and abs(dv) < 750: status = 'ok' elif truetype in ('STD', 'FSTD') and objtype == 'STAR': status = 'ok' else: status = 'OOPS' else: status = 'OOPS' print('{0:<8d} {1:4s} {2:8.5f} -> {3:5s} {4:8.5f} {5:4d} - {6}'. format(pix, truetype, truez, objtype, z, zwarn, status)) print("--------------------------------------------------")
def simulate_one_healpix(ifilename, args, model, obsconditions, decam_and_wise_filters, footprint_healpix_weight, footprint_healpix_nside, seed, bal=None): log = get_logger() # set seed now # we need a seed per healpix because # the spectra simulator REQUIRES a seed np.random.seed(seed) # read the header of the tranmission file to find the healpix pixel number, nside # and if we are lucky the scheme. # if this fails, try to guess it from the filename (for backward compatibility) healpix = -1 nside = -1 hpxnest = True hdulist = pyfits.open(ifilename) if "METADATA" in hdulist: head = hdulist["METADATA"].header for k in ["HPXPIXEL", "PIXNUM"]: if k in head: healpix = int(head[k]) log.info("healpix={}={}".format(k, healpix)) break for k in ["HPXNSIDE", "NSIDE"]: if k in head: nside = int(head[k]) log.info("nside={}={}".format(k, nside)) break for k in ["HPXNEST", "NESTED", "SCHEME"]: if k in head: if k == "SCHEME": hpxnest = (head[k] == "NEST") else: hpxnest = bool(head[k]) log.info("hpxnest from {} = {}".format(k, hpxnest)) break if healpix >= 0 and nside < 0: log.error("Read healpix in header but not nside.") raise ValueError("Read healpix in header but not nside.") if healpix < 0: vals = os.path.basename(ifilename).split(".")[0].split("-") if len(vals) < 3: log.error("Cannot guess nside and healpix from filename {}".format( ifilename)) raise ValueError( "Cannot guess nside and healpix from filename {}".format( ifilename)) try: healpix = int(vals[-1]) nside = int(vals[-2]) except ValueError: raise ValueError( "Cannot guess nside and healpix from filename {}".format( ifilename)) log.warning( "Guessed healpix and nside from filename, assuming the healpix scheme is 'NESTED'" ) zbest_filename = None if args.outfile: ofilename = args.outfile else: ofilename = os.path.join( args.outdir, "{}/{}/spectra-{}-{}.fits".format(healpix // 100, healpix, nside, healpix)) pixdir = os.path.dirname(ofilename) if args.zbest: zbest_filename = os.path.join( pixdir, "zbest-{}-{}.fits".format(nside, healpix)) if not args.overwrite: # check whether output exists or not if args.zbest: if os.path.isfile(ofilename) and os.path.isfile(zbest_filename): log.info("skip existing {} and {}".format( ofilename, zbest_filename)) return else: # only test spectra file if os.path.isfile(ofilename): log.info("skip existing {}".format(ofilename)) return log.info("Read skewers in {}, random seed = {}".format(ifilename, seed)) ##ALMA: It reads only the skewers only if there are no DLAs or if they are added randomly. if (not args.dla or args.dla == 'random'): trans_wave, transmission, metadata = read_lya_skewers(ifilename) ok = np.where((metadata['Z'] >= args.zmin) & (metadata['Z'] <= args.zmax))[0] transmission = transmission[ok] metadata = metadata[:][ok] ##ALMA:Added to read dla_info elif (args.dla == 'file'): log.info("Read DLA information in {}".format(ifilename)) trans_wave, transmission, metadata, dla_info = read_lya_skewers( ifilename, dla_='TRUE') ok = np.where((metadata['Z'] >= args.zmin) & (metadata['Z'] <= args.zmax))[0] transmission = transmission[ok] metadata = metadata[:][ok] else: log.error( 'Not a valid option to add DLAs. Valid options are "random" or "file"' ) sys.exit(1) if args.dla: dla_NHI, dla_z, dla_id = [], [], [] dla_filename = os.path.join(pixdir, "dla-{}-{}.fits".format(nside, healpix)) if args.desi_footprint: footprint_healpix = footprint.radec2pix(footprint_healpix_nside, metadata["RA"], metadata["DEC"]) selection = np.where( footprint_healpix_weight[footprint_healpix] > 0.99)[0] log.info("Select QSOs in DESI footprint {} -> {}".format( transmission.shape[0], selection.size)) if selection.size == 0: log.warning("No intersection with DESI footprint") return transmission = transmission[selection] metadata = metadata[:][selection] nqso = transmission.shape[0] if args.downsampling is not None: if args.downsampling <= 0 or args.downsampling > 1: log.error( "Down sampling fraction={} must be between 0 and 1".format( args.downsampling)) raise ValueError( "Down sampling fraction={} must be between 0 and 1".format( args.downsampling)) indices = np.where(np.random.uniform(size=nqso) < args.downsampling)[0] if indices.size == 0: log.warning( "Down sampling from {} to 0 (by chance I presume)".format( nqso)) return transmission = transmission[indices] metadata = metadata[:][indices] nqso = transmission.shape[0] ##ALMA:added to set transmission to 1 for z>zqso, this can be removed when transmission is corrected. for ii in range(len(metadata)): transmission[ii][trans_wave > 1215.67 * (metadata[ii]['Z'] + 1)] = 1.0 if (args.dla == 'file'): log.info('Adding DLAs from transmision file') min_trans_wave = np.min(trans_wave / 1215.67 - 1) for ii in range(len(metadata)): if min_trans_wave < metadata[ii]['Z']: idd = metadata['MOCKID'][ii] dlas = dla_info[dla_info['MOCKID'] == idd] dlass = [] for i in range(len(dlas)): ##Adding only dlas between zqso and 1.95, check again for the next version of London mocks... if (dlas[i]['Z_DLA'] < metadata[ii]['Z']) and (dlas[i]['Z_DLA'] > 1.95): dlass.append( dict(z=dlas[i]['Z_DLA'] + dlas[i]['DZ_DLA'], N=dlas[i]['N_HI_DLA'])) if len(dlass) > 0: dla_model = dla_spec(trans_wave, dlass) transmission[ii] = dla_model * transmission[ii] dla_z += [idla['z'] for idla in dlass] dla_NHI += [idla['N'] for idla in dlass] dla_id += [idd] * len(dlass) elif (args.dla == 'random'): log.info('Adding DLAs randomly') min_trans_wave = np.min(trans_wave / 1215.67 - 1) for ii in range(len(metadata)): if min_trans_wave < metadata[ii]['Z']: idd = metadata['MOCKID'][ii] dlass, dla_model = insert_dlas(trans_wave, metadata[ii]['Z']) if len(dlass) > 0: transmission[ii] = dla_model * transmission[ii] dla_z += [idla['z'] for idla in dlass] dla_NHI += [idla['N'] for idla in dlass] dla_id += [idd] * len(dlass) if args.dla: if len(dla_id) > 0: dla_meta = Table() dla_meta['NHI'] = dla_NHI dla_meta['z'] = dla_z dla_meta['ID'] = dla_id if args.nmax is not None: if args.nmax < nqso: log.info( "Limit number of QSOs from {} to nmax={} (random subsample)". format(nqso, args.nmax)) # take a random subsample indices = (np.random.uniform(size=args.nmax) * nqso).astype(int) transmission = transmission[indices] metadata = metadata[:][indices] nqso = args.nmax if args.dla: dla_meta = dla_meta[:][dla_meta['ID'] == metadata['MOCKID']] if args.target_selection or args.mags: wanted_min_wave = 3329. # needed to compute magnitudes for decam2014-r (one could have trimmed the transmission file ...) wanted_max_wave = 55501. # needed to compute magnitudes for wise2010-W2 if trans_wave[0] > wanted_min_wave: log.info( "Increase wavelength range from {}:{} to {}:{} to compute magnitudes" .format(int(trans_wave[0]), int(trans_wave[-1]), int(wanted_min_wave), int(trans_wave[-1]))) # pad with zeros at short wavelength because we assume transmission = 0 # and we don't need any wavelength resolution here new_trans_wave = np.append([wanted_min_wave, trans_wave[0] - 0.01], trans_wave) new_transmission = np.zeros( (transmission.shape[0], new_trans_wave.size)) new_transmission[:, 2:] = transmission trans_wave = new_trans_wave transmission = new_transmission if trans_wave[-1] < wanted_max_wave: log.info( "Increase wavelength range from {}:{} to {}:{} to compute magnitudes" .format(int(trans_wave[0]), int(trans_wave[-1]), int(trans_wave[0]), int(wanted_max_wave))) # pad with ones at long wavelength because we assume transmission = 1 coarse_dwave = 2. # we don't care about resolution, we just need a decent QSO spectrum, there is no IGM transmission in this range n = int((wanted_max_wave - trans_wave[-1]) / coarse_dwave) + 1 new_trans_wave = np.append( trans_wave, np.linspace(trans_wave[-1] + coarse_dwave, trans_wave[-1] + coarse_dwave * (n + 1), n)) new_transmission = np.ones( (transmission.shape[0], new_trans_wave.size)) new_transmission[:, :trans_wave.size] = transmission trans_wave = new_trans_wave transmission = new_transmission log.info("Simulate {} QSOs".format(nqso)) tmp_qso_flux, tmp_qso_wave, meta = model.make_templates( nmodel=nqso, redshift=metadata['Z'], lyaforest=False, nocolorcuts=True, noresample=True, seed=seed) log.info("Resample to transmission wavelength grid") # because we don't want to alter the transmission field with resampling here qso_flux = np.zeros((tmp_qso_flux.shape[0], trans_wave.size)) for q in range(tmp_qso_flux.shape[0]): qso_flux[q] = np.interp(trans_wave, tmp_qso_wave, tmp_qso_flux[q]) tmp_qso_flux = qso_flux tmp_qso_wave = trans_wave ##To add BALs to be checked by Luz and Jaime if (args.balprob): if (args.balprob <= 1. and args.balprob > 0): log.info("Adding BALs with probability {}".format(args.balprob)) tmp_qso_flux, meta_bal = bal.insert_bals(tmp_qso_wave, tmp_qso_flux, metadata['Z'], balprob=args.balprob, seed=seed) else: log.error("Probability to add BALs is not between 0 and 1") sys.exit(1) log.info("Apply lya") tmp_qso_flux = apply_lya_transmission(tmp_qso_wave, tmp_qso_flux, trans_wave, transmission) if args.metals is not None: lstMetals = '' for m in args.metals: lstMetals += m + ', ' log.info("Apply metals: {}".format(lstMetals[:-2])) tmp_qso_flux = apply_metals_transmission(tmp_qso_wave, tmp_qso_flux, trans_wave, transmission, args.metals) bbflux = None if args.target_selection or args.mags: bands = ['FLUX_G', 'FLUX_R', 'FLUX_Z', 'FLUX_W1', 'FLUX_W2'] bbflux = dict() # need to recompute the magnitudes to account for lya transmission log.info("Compute QSO magnitudes") maggies = decam_and_wise_filters.get_ab_maggies( 1e-17 * tmp_qso_flux, tmp_qso_wave) for band, filt in zip(bands, [ 'decam2014-g', 'decam2014-r', 'decam2014-z', 'wise2010-W1', 'wise2010-W2' ]): bbflux[band] = np.ma.getdata(1e9 * maggies[filt]) # nanomaggies if args.target_selection: log.info("Apply target selection") isqso = isQSO_colors(gflux=bbflux['FLUX_G'], rflux=bbflux['FLUX_R'], zflux=bbflux['FLUX_Z'], w1flux=bbflux['FLUX_W1'], w2flux=bbflux['FLUX_W2']) log.info("Target selection: {}/{} QSOs selected".format( np.sum(isqso), nqso)) selection = np.where(isqso)[0] if selection.size == 0: return tmp_qso_flux = tmp_qso_flux[selection] metadata = metadata[:][selection] meta = meta[:][selection] for band in bands: bbflux[band] = bbflux[band][selection] nqso = selection.size log.info("Resample to a linear wavelength grid (needed by DESI sim.)") # we need a linear grid. for this resampling we take care of integrating in bins # we do not do a simple interpolation qso_wave = np.linspace(args.wmin, args.wmax, int((args.wmax - args.wmin) / args.dwave) + 1) qso_flux = np.zeros((tmp_qso_flux.shape[0], qso_wave.size)) for q in range(tmp_qso_flux.shape[0]): qso_flux[q] = resample_flux(qso_wave, tmp_qso_wave, tmp_qso_flux[q]) log.info("Simulate DESI observation and write output file") pixdir = os.path.dirname(ofilename) if len(pixdir) > 0: if not os.path.isdir(pixdir): log.info("Creating dir {}".format(pixdir)) os.makedirs(pixdir) if "MOCKID" in metadata.dtype.names: #log.warning("Using MOCKID as TARGETID") targetid = np.array(metadata["MOCKID"]).astype(int) elif "ID" in metadata.dtype.names: log.warning("Using ID as TARGETID") targetid = np.array(metadata["ID"]).astype(int) else: log.warning("No TARGETID") targetid = None meta = {"HPXNSIDE": nside, "HPXPIXEL": healpix, "HPXNEST": hpxnest} if args.target_selection or args.mags: # today we write mags because that's what is in the fibermap mags = np.zeros((qso_flux.shape[0], 5)) for i, band in enumerate(bands): jj = (bbflux[band] > 0) mags[jj, i] = 22.5 - 2.5 * np.log10(bbflux[band][jj]) # AB magnitudes fibermap_columns = {"MAG": mags} else: fibermap_columns = None sim_spectra(qso_wave, qso_flux, args.program, obsconditions=obsconditions, spectra_filename=ofilename, sourcetype="qso", skyerr=args.skyerr, ra=metadata["RA"], dec=metadata["DEC"], targetid=targetid, meta=meta, seed=seed, fibermap_columns=fibermap_columns) if args.zbest: log.info("Read fibermap") fibermap = read_fibermap(ofilename) log.info("Writing a zbest file {}".format(zbest_filename)) columns = [('CHI2', 'f8'), ('COEFF', 'f8', (4, )), ('Z', 'f8'), ('ZERR', 'f8'), ('ZWARN', 'i8'), ('SPECTYPE', (str, 96)), ('SUBTYPE', (str, 16)), ('TARGETID', 'i8'), ('DELTACHI2', 'f8'), ('BRICKNAME', (str, 8))] zbest = Table(np.zeros(nqso, dtype=columns)) zbest["CHI2"][:] = 0. zbest["Z"] = metadata['Z'] zbest["ZERR"][:] = 0. zbest["ZWARN"][:] = 0 zbest["SPECTYPE"][:] = "QSO" zbest["SUBTYPE"][:] = "" zbest["TARGETID"] = fibermap["TARGETID"] zbest["DELTACHI2"][:] = 25. hzbest = pyfits.convenience.table_to_hdu(zbest) hzbest.name = "ZBEST" hfmap = pyfits.convenience.table_to_hdu(fibermap) hfmap.name = "FIBERMAP" hdulist = pyfits.HDUList([pyfits.PrimaryHDU(), hzbest, hfmap]) hdulist.writeto(zbest_filename, clobber=True) hdulist.close() # see if this helps with memory issue if args.dla: #This will change according to discussion log.info("Updating the spectra file to add DLA metadata {}".format( ofilename)) hdudla = pyfits.table_to_hdu(dla_meta) hdudla.name = "DLA_META" hdul = pyfits.open(ofilename, mode='update') hdul.append(hdudla) hdul.flush() hdul.close()
def main(args): import os.path import sys import yaml import numpy as np import pdb import matplotlib matplotlib.use('agg') import desispec.io from desiutil.log import get_logger from desisim.spec_qa import redshifts as dsqa_z from desiutil.io import yamlify import desiutil.depend from desimodel.footprint import radec2pix log = get_logger() if args.load_simz_table is not None: from astropy.table import Table log.info("Loading simz info from {:s}".format(args.load_simz_table)) simz_tab = Table.read(args.load_simz_table) else: # Grab list of fibermap files fibermap_files = [] zbest_files = [] nights = desispec.io.get_nights() for night in nights: for exposure in desispec.io.get_exposures(night, raw=True, rawdata_dir=args.rawdir): # Ignore exposures with no fibermap, assuming they are calibration data. fibermap_path = desispec.io.findfile(filetype='fibermap', night=night, expid=exposure, rawdata_dir=args.rawdir) if not os.path.exists(fibermap_path): log.debug('Skipping exposure %08d with no fibermap.' % exposure) continue ''' # Search for zbest files fibermap_data = desispec.io.read_fibermap(fibermap_path) flavor = fibermap_data.meta['FLAVOR'] if flavor.lower() in ('arc', 'flat', 'bias'): log.debug('Skipping calibration {} exposure {:08d}'.format(flavor, exposure)) continue brick_names = set(fibermap_data['BRICKNAME']) import pdb; pdb.set_trace() for brick in brick_names: zbest_path=desispec.io.findfile('zbest', groupname=brick, specprod_dir=args.reduxdir) if os.path.exists(zbest_path): log.debug('Found {}'.format(os.path.basename(zbest_path))) zbest_files.append(zbest_path) else: log.warn('Missing {}'.format(os.path.basename(zbest_path))) #pdb.set_trace() ''' # Load data fibermap_data = desispec.io.read_fibermap(fibermap_path) # Skip calib if fibermap_data['OBJTYPE'][0] in ['FLAT', 'ARC', 'BIAS']: continue elif fibermap_data['OBJTYPE'][0] in ['SKY', 'STD', 'SCIENCE']: pass else: pdb.set_trace() # Append fibermap file fibermap_files.append(fibermap_path) # Search for zbest files with healpy ra_targ = fibermap_data['RA_TARGET'].data dec_targ = fibermap_data['DEC_TARGET'].data # Getting some NAN in RA/DEC good = np.isfinite(ra_targ) & np.isfinite(dec_targ) pixels = radec2pix(64, ra_targ[good], dec_targ[good]) uni_pixels = np.unique(pixels) for uni_pix in uni_pixels: zbest_files.append( desispec.io.findfile('zbest', groupname=uni_pix, nside=64)) # Cut down zbest_files to unique ones zbest_files = list(set(zbest_files)) if len(zbest_files) == 0: log.fatal('No zbest files found') sys.exit(1) # Write? Table simz_tab = dsqa_z.load_z(fibermap_files, zbest_files) if args.write_simz_table is not None: simz_tab.write(args.write_simz_table, overwrite=True) # Meta data meta = dict( DESISIM=desiutil.depend.getdep(simz_tab.meta, 'desisim'), SPECTER=desiutil.depend.getdep(simz_tab.meta, 'specter'), SPECPROD=os.getenv('SPECPROD', 'unknown'), PIXPROD=os.getenv('PIXPROD', 'unknown'), ) # Run stats log.info("Running stats..") summ_dict = dsqa_z.summ_stats(simz_tab) if args.qafile is not None: log.info("Generating yaml file: {:s}".format(args.qafile)) # yamlify # Write yaml with open(args.qafile, 'w') as outfile: outfile.write(yaml.dump(yamlify(meta), default_flow_style=False)) outfile.write( yaml.dump(yamlify(summ_dict), default_flow_style=False)) if args.qafig_root is not None: log.info("Generating QA files") # Summary for dz of all types outfile = args.qafig_root + '_dzsumm.png' #dsqa_z.dz_summ(simz_tab, outfile=outfile) # Summary of individual types #outfile = args.qafig_root+'_summ_fig.png' #dsqa_z.summ_fig(simz_tab, summ_dict, meta, outfile=outfile) for objtype in ['BGS', 'MWS', 'ELG', 'LRG', 'QSO_T', 'QSO_L']: outfile = args.qafig_root + '_zfind_{:s}.png'.format(objtype) dsqa_z.obj_fig(simz_tab, objtype, summ_dict, outfile=outfile)
def simulate_one_healpix(ifilename, args, model, obsconditions, decam_and_wise_filters, footprint_healpix_weight, footprint_healpix_nside): log = get_logger() healpix = 0 nside = 0 vals = os.path.basename(ifilename).split(".")[0].split("-") if len(vals) < 3: log.error("Cannot guess nside and healpix from filename {}".format( ifilename)) raise ValueError( "Cannot guess nside and healpix from filename {}".format( ifilename)) try: healpix = int(vals[-1]) nside = int(vals[-2]) except ValueError: raise ValueError( "Cannot guess nside and healpix from filename {}".format( ifilename)) zbest_filename = None if args.outfile: ofilename = args.outfile else: ofilename = os.path.join( args.outdir, "{}/{}/spectra-{}-{}.fits".format(healpix // 100, healpix, nside, healpix)) pixdir = os.path.dirname(ofilename) if not args.overwrite: # check whether output exists or not if args.zbest: zbest_filename = os.path.join( pixdir, "zbest-{}-{}.fits".format(nside, healpix)) if os.path.isfile(ofilename) and os.path.isfile(zbest_filename): log.info("skip existing {} and {}".format( ofilename, zbest_filename)) return else: # only test spectra file if os.path.isfile(ofilename): log.info("skip existing {}".format(ofilename)) return log.info("Read skewers in {}".format(ifilename)) trans_wave, transmission, metadata = read_lya_skewers(ifilename) ok = np.where((metadata['Z'] >= args.zmin) & (metadata['Z'] <= args.zmax))[0] transmission = transmission[ok] metadata = metadata[:][ok] # set seed now in case we are downsampling np.random.seed(args.seed) # create quasars if args.desi_footprint: footprint_healpix = footprint.radec2pix(footprint_healpix_nside, metadata["RA"], metadata["DEC"]) selection = np.where( footprint_healpix_weight[footprint_healpix] > 0.99)[0] log.info("Select QSOs in DESI footprint {} -> {}".format( transmission.shape[0], selection.size)) if selection.size == 0: log.warning("No intersection with DESI footprint") return transmission = transmission[selection] metadata = metadata[:][selection] nqso = transmission.shape[0] if args.downsampling is not None: if args.downsampling <= 0 or args.downsampling > 1: log.error( "Down sampling fraction={} must be between 0 and 1".format( args.downsampling)) raise ValueError( "Down sampling fraction={} must be between 0 and 1".format( args.downsampling)) indices = np.where(np.random.uniform(size=nqso) < args.downsampling)[0] if indices.size == 0: log.warning( "Down sampling from {} to 0 (by chance I presume)".format( nqso)) return transmission = transmission[indices] metadata = metadata[:][indices] nqso = transmission.shape[0] if args.nmax is not None: if args.nmax < nqso: log.info( "Limit number of QSOs from {} to nmax={} (random subsample)". format(nqso, args.nmax)) # take a random subsample indices = (np.random.uniform(size=args.nmax) * nqso).astype(int) transmission = transmission[indices] metadata = metadata[:][indices] nqso = args.nmax log.info("Simulate {} QSOs".format(nqso)) tmp_qso_flux, tmp_qso_wave, meta = model.make_templates( nmodel=nqso, redshift=metadata['Z'], seed=args.seed, lyaforest=False, nocolorcuts=True, noresample=True) log.info("Resample to transmission wavelength grid") # because we don't want to alter the transmission field with resampling here qso_flux = np.zeros((tmp_qso_flux.shape[0], trans_wave.size)) for q in range(tmp_qso_flux.shape[0]): qso_flux[q] = np.interp(trans_wave, tmp_qso_wave, tmp_qso_flux[q]) tmp_qso_flux = qso_flux tmp_qso_wave = trans_wave log.info("Apply lya") tmp_qso_flux = apply_lya_transmission(tmp_qso_wave, tmp_qso_flux, trans_wave, transmission) if args.target_selection: log.info("Compute QSO magnitudes for target selection") maggies = decam_and_wise_filters.get_ab_maggies(1e-17 * tmp_qso_flux, tmp_qso_wave.copy(), mask_invalid=True) for band, filt in zip( ('FLUX_G', 'FLUX_R', 'FLUX_Z', 'FLUX_W1', 'FLUX_W2'), ('decam2014-g', 'decam2014-r', 'decam2014-z', 'wise2010-W1', 'wise2010-W2')): meta[band] = np.ma.getdata(1e9 * maggies[filt]) # nanomaggies isqso = isQSO_colors(gflux=meta['FLUX_G'], rflux=meta['FLUX_R'], zflux=meta['FLUX_Z'], w1flux=meta['FLUX_W1'], w2flux=meta['FLUX_W2']) log.info("Target selection: {}/{} QSOs selected".format( np.sum(isqso), nqso)) selection = np.where(isqso)[0] if selection.size == 0: return tmp_qso_flux = tmp_qso_flux[selection] metadata = metadata[:][selection] meta = meta[:][selection] nqso = selection.size log.info("Resample to a linear wavelength grid (needed by DESI sim.)") # we need a linear grid. for this resampling we take care of integrating in bins # we do not do a simple interpolation qso_wave = np.linspace(args.wmin, args.wmax, int((args.wmax - args.wmin) / args.dwave) + 1) qso_flux = np.zeros((tmp_qso_flux.shape[0], qso_wave.size)) for q in range(tmp_qso_flux.shape[0]): qso_flux[q] = resample_flux(qso_wave, tmp_qso_wave, tmp_qso_flux[q]) log.info("Simulate DESI observation and write output file") pixdir = os.path.dirname(ofilename) if not os.path.isdir(pixdir): log.info("Creating dir {}".format(pixdir)) os.makedirs(pixdir) if "MOCKID" in metadata.dtype.names: #log.warning("Using MOCKID as TARGETID") targetid = np.array(metadata["MOCKID"]).astype(int) elif "ID" in metadata.dtype.names: log.warning("Using ID as TARGETID") targetid = np.array(metadata["ID"]).astype(int) else: log.warning("No TARGETID") targetid = None sim_spectra(qso_wave, qso_flux, args.program, obsconditions=obsconditions, spectra_filename=ofilename, seed=args.seed, sourcetype="qso", skyerr=args.skyerr, ra=metadata["RA"], dec=metadata["DEC"], targetid=targetid) if args.zbest: log.info("Read fibermap") fibermap = read_fibermap(ofilename) log.info("Writing a zbest file {}".format(zbest_filename)) columns = [('CHI2', 'f8'), ('COEFF', 'f8', (4, )), ('Z', 'f8'), ('ZERR', 'f8'), ('ZWARN', 'i8'), ('SPECTYPE', (str, 96)), ('SUBTYPE', (str, 16)), ('TARGETID', 'i8'), ('DELTACHI2', 'f8'), ('BRICKNAME', (str, 8))] zbest = Table(np.zeros(nqso, dtype=columns)) zbest["CHI2"][:] = 0. zbest["Z"] = metadata['Z'] zbest["ZERR"][:] = 0. zbest["ZWARN"][:] = 0 zbest["SPECTYPE"][:] = "QSO" zbest["SUBTYPE"][:] = "" zbest["TARGETID"] = fibermap["TARGETID"] zbest["DELTACHI2"][:] = 25. hzbest = pyfits.convenience.table_to_hdu(zbest) hzbest.name = "ZBEST" hfmap = pyfits.convenience.table_to_hdu(fibermap) hfmap.name = "FIBERMAP" hdulist = pyfits.HDUList([pyfits.PrimaryHDU(), hzbest, hfmap]) hdulist.writeto(zbest_filename, clobber=True) hdulist.close() # see if this helps with memory issue
def integration_test(night=None, nspec=5, clobber=False): """Run an integration test from raw data simulations through redshifts. Args: night (str, optional): YEARMMDD, defaults to current night nspec (int, optional): number of spectra to include clobber (bool, optional): rerun steps even if outputs already exist Raises: RuntimeError if any script fails """ log = get_logger() #- YEARMMDD string, rolls over at noon not midnight if night is None: night = time.strftime('%Y%m%d', time.localtime(time.time()-12*3600)) #- check for required environment variables check_env() #- parameter dictionary that will later be used for formatting commands params = dict(night=night, nspec=nspec) #----- #- Input fibermaps, spectra, and pixel-level raw data # raw_dict = {0: 'flat', 1: 'arc', 2: 'dark'} programs = ('flat', 'arc', 'dark') channels = ('b', 'r', 'z') cameras = ('b0', 'r0', 'z0') # for expid, program in raw_dict.items(): for expid, program in enumerate(programs): cmd = "newexp-random --program {program} --nspec {nspec} --night {night} --expid {expid}".format( expid=expid, program=program, **params) fibermap = io.findfile('fibermap', night, expid) simspec = '{}/simspec-{:08d}.fits'.format(os.path.dirname(fibermap), expid) inputs = [] outputs = [fibermap, simspec] if runcmd(cmd, inputs=inputs, outputs=outputs, clobber=clobber) != 0: raise RuntimeError('pixsim newexp failed for {} exposure {}'.format(program, expid)) cmd = "pixsim --nspec {nspec} --night {night} --expid {expid}".format(expid=expid, **params) inputs = [fibermap, simspec] outputs = [fibermap.replace('fibermap-', 'simpix-'), ] if runcmd(cmd, inputs=inputs, outputs=outputs, clobber=clobber) != 0: raise RuntimeError('pixsim failed for {} exposure {}'.format(program, expid)) #----- #- Preproc for expid, program in enumerate(programs): rawfile = io.findfile('desi', night, expid) outdir = os.path.dirname(io.findfile('preproc', night, expid, 'b0')) cmd = "desi_preproc --infile {} --outdir {}".format(rawfile, outdir) inputs = [rawfile,] outputs = list() for camera in cameras: outputs.append(io.findfile('preproc', night, expid, camera)) if runcmd(cmd, inputs=inputs, outputs=outputs, clobber=clobber) != 0: raise RuntimeError('preproc failed for expid {}'.format(expid)) #----- #- Extract waverange = dict(b="3570,5940,1.0", r="5630,7740,1.0", z="7440,9830,1.0") for expid, program in enumerate(programs): for ic, channel in enumerate(channels): pixfile = io.findfile('preproc', night, expid, cameras[ic]) fiberfile = io.findfile('fibermap', night, expid) psffile = '{}/data/specpsf/psf-{}.fits'.format(os.getenv('DESIMODEL'), channel) framefile = io.findfile('frame', night, expid, cameras[ic]) # cmd = "exspec -i {pix} -p {psf} --specmin 0 --nspec {nspec} -w {wave} -o {frame}".format( # pix=pixfile, psf=psffile, wave=waverange[channel], frame=framefile, **params) cmd = "desi_extract_spectra -i {pix} -p {psf} -f {fibermap} --specmin 0 --nspec {nspec} -o {frame}".format( pix=pixfile, psf=psffile, frame=framefile, fibermap=fiberfile, **params) inputs = [pixfile, psffile, fiberfile] outputs = [framefile,] if runcmd(cmd, inputs=inputs, outputs=outputs, clobber=clobber) != 0: raise RuntimeError('extraction failed for {} expid {}'.format(cameras[ic], expid)) #----- #- Fiber flat expid = 0 for ic, channel in enumerate(channels): framefile = io.findfile('frame', night, expid, cameras[ic]) fiberflat = io.findfile('fiberflat', night, expid, cameras[ic]) fibermap = io.findfile('fibermap', night, expid) # for QA qafile = io.findfile('qa_calib', night, expid, cameras[ic]) qafig = io.findfile('qa_flat_fig', night, expid, cameras[ic]) cmd = "desi_compute_fiberflat --infile {frame} --outfile {fiberflat} --qafile {qafile} --qafig {qafig}".format( frame=framefile, fiberflat=fiberflat, qafile=qafile, qafig=qafig, **params) inputs = [framefile,fibermap,] outputs = [fiberflat,qafile,qafig,] if runcmd(cmd, inputs=inputs, outputs=outputs, clobber=clobber) != 0: raise RuntimeError('fiberflat failed for '+cameras[ic]) #----- #- Sky model flat_expid = 0 expid = 2 for ic, channel in enumerate(channels): framefile = io.findfile('frame', night, expid, cameras[ic]) fibermap = io.findfile('fibermap', night, expid) fiberflat = io.findfile('fiberflat', night, flat_expid, cameras[ic]) skyfile = io.findfile('sky', night, expid, cameras[ic]) qafile = io.findfile('qa_data', night, expid, cameras[ic]) qafig = io.findfile('qa_sky_fig', night, expid, cameras[ic]) cmd="desi_compute_sky --infile {frame} --fiberflat {fiberflat} --outfile {sky} --qafile {qafile} --qafig {qafig}".format( frame=framefile, fiberflat=fiberflat, sky=skyfile, qafile=qafile, qafig=qafig, **params) inputs = [framefile, fibermap, fiberflat] outputs = [skyfile, qafile, qafig,] if runcmd(cmd, inputs=inputs, outputs=outputs, clobber=clobber) != 0: raise RuntimeError('sky model failed for '+cameras[ic]) #----- #- Fit standard stars if 'STD_TEMPLATES' in os.environ: std_templates = os.getenv('STD_TEMPLATES') else: std_templates = os.getenv('DESI_ROOT')+'/spectro/templates/star_templates/v1.1/star_templates_v1.1.fits' stdstarfile = io.findfile('stdstars', night, expid, spectrograph=0) flats = list() frames = list() skymodels = list() for ic, channel in enumerate(channels): frames.append( io.findfile('frame', night, expid, cameras[ic]) ) flats.append( io.findfile('fiberflat', night, flat_expid, cameras[ic]) ) skymodels.append( io.findfile('sky', night, expid, cameras[ic]) ) frames = ' '.join(frames) flats = ' '.join(flats) skymodels = ' '.join(skymodels) cmd = """desi_fit_stdstars \ --frames {frames} \ --fiberflats {flats} \ --skymodels {skymodels} \ --starmodels {std_templates} \ -o {stdstars}""".format( frames=frames, flats=flats, skymodels=skymodels, std_templates=std_templates, stdstars=stdstarfile) inputs = [fibermap, std_templates] outputs = [stdstarfile,] if runcmd(cmd, inputs=inputs, outputs=outputs, clobber=clobber) != 0: raise RuntimeError('fitting stdstars failed') #----- #- Flux calibration for ic, channel in enumerate(channels): framefile = io.findfile('frame', night, expid, cameras[ic]) fibermap = io.findfile('fibermap', night, expid) fiberflat = io.findfile('fiberflat', night, flat_expid, cameras[ic]) skyfile = io.findfile('sky', night, expid, cameras[ic]) calibfile = io.findfile('calib', night, expid, cameras[ic]) qafile = io.findfile('qa_data', night, expid, cameras[ic]) qafig = io.findfile('qa_flux_fig', night, expid, cameras[ic]) #- Compute flux calibration vector cmd = """desi_compute_fluxcalibration \ --infile {frame} --fiberflat {fiberflat} --sky {sky} \ --models {stdstars} --outfile {calib} --qafile {qafile} --qafig {qafig}""".format( frame=framefile, fiberflat=fiberflat, sky=skyfile, stdstars=stdstarfile, calib=calibfile, qafile=qafile, qafig=qafig ) inputs = [framefile, fibermap, fiberflat, skyfile, stdstarfile] outputs = [calibfile, qafile, qafig] if runcmd(cmd, inputs=inputs, outputs=outputs, clobber=clobber) != 0: raise RuntimeError('flux calibration failed for '+cameras[ic]) #- Apply the flux calibration to write a cframe file cframefile = io.findfile('cframe', night, expid, cameras[ic]) cmd = """desi_process_exposure \ --infile {frame} --fiberflat {fiberflat} --sky {sky} --calib {calib} \ --outfile {cframe}""".format(frame=framefile, fibermap=fibermap, fiberflat=fiberflat, sky=skyfile, calib=calibfile, cframe=cframefile) inputs = [framefile, fiberflat, skyfile, calibfile] outputs = [cframefile, ] if runcmd(cmd, inputs=inputs, outputs=outputs, clobber=clobber) != 0: raise RuntimeError('combining calibration steps failed for '+cameras[ic]) #----- #- Collate QA # Collate data QA program2flavor = dict(arc='arc', flat='flat') for program in ('dark', 'gray', 'bright', 'elg', 'lrg', 'qso', 'bgs', 'mws'): program2flavor[program] = 'science' expid = 2 qafile = io.findfile('qa_data_exp', night, expid) if clobber or not os.path.exists(qafile): flavor = program2flavor[programs[expid]] qaexp_data = QA_Exposure(expid, night, flavor) # Removes camera files io.write_qa_exposure(os.path.splitext(qafile)[0], qaexp_data) if not os.path.exists(qafile): raise RuntimeError('FAILED data QA_Exposure({},{}, ...) -> {}'.format(expid, night, qafile)) # Collate calib QA calib_expid = [0,1] for expid in calib_expid: qafile = io.findfile('qa_calib_exp', night, expid) if clobber or not os.path.exists(qafile): qaexp_calib = QA_Exposure(expid, night, programs[expid]) io.write_qa_exposure(os.path.splitext(qafile)[0], qaexp_calib) if not os.path.exists(qafile): raise RuntimeError('FAILED calib QA_Exposure({},{}, ...) -> {}'.format(expid, night, qafile)) #----- #- Regroup cframe -> spectra expid = 2 inputs = list() for camera in cameras: inputs.append( io.findfile('cframe', night, expid, camera) ) outputs = list() fibermap = io.read_fibermap(io.findfile('fibermap', night, expid)) from desimodel.footprint import radec2pix nside=64 pixels = np.unique(radec2pix(nside, fibermap['TARGET_RA'], fibermap['TARGET_DEC'])) for pix in pixels: outputs.append( io.findfile('spectra', groupname=pix) ) cmd = "desi_group_spectra" if runcmd(cmd, inputs=inputs, outputs=outputs, clobber=clobber) != 0: raise RuntimeError('spectra regrouping failed') #----- #- Redshifts! for pix in pixels: specfile = io.findfile('spectra', groupname=pix) zbestfile = io.findfile('zbest', groupname=pix) inputs = [specfile, ] outputs = [zbestfile, ] cmd = "rrdesi {} --zbest {}".format(specfile, zbestfile) if runcmd(cmd, inputs=inputs, outputs=outputs, clobber=clobber) != 0: raise RuntimeError('rrdesi failed for healpixel {}'.format(pix)) # # Load redshifts into database # options = get_options('--overwrite', '--filename', 'dailytest.db', os.path.join(os.environ['DESI_SPECTRO_REDUX'], os.environ['SPECPROD'])) postgresql = setup_db(options) load_zbest(options.datapath) # ztruth QA # qafile = io.findfile('qa_ztruth', night) # qafig = io.findfile('qa_ztruth_fig', night) # cmd = "desi_qa_zfind --night {night} --qafile {qafile} --qafig {qafig} --verbose".format( # night=night, qafile=qafile, qafig=qafig) # inputs = [] # outputs = [qafile, qafig] # if runcmd(cmd, inputs=inputs, outputs=outputs, clobber=clobber) != 0: # raise RuntimeError('redshift QA failed for night '+night) #----- #- Did it work? #- (this combination of fibermap, simspec, and zbest is a pain) simdir = os.path.dirname(io.findfile('fibermap', night=night, expid=expid)) simspec = '{}/simspec-{:08d}.fits'.format(simdir, expid) siminfo = fits.getdata(simspec, 'TRUTH') try: elginfo = fits.getdata(simspec, 'TRUTH_ELG') except: elginfo = None print() print("--------------------------------------------------") print("Pixel True z -> Class z zwarn") # print("3338p190 SKY 0.00000 -> QSO 1.60853 12 - ok") for pix in pixels: zbest = fits.getdata(io.findfile('zbest', groupname=pix)) for i in range(len(zbest)): objtype = zbest['SPECTYPE'][i] z, zwarn = zbest['Z'][i], zbest['ZWARN'][i] j = np.where(fibermap['TARGETID'] == zbest['TARGETID'][i])[0][0] truetype = siminfo['OBJTYPE'][j] oiiflux = 0.0 if truetype == 'ELG': k = np.where(elginfo['TARGETID'] == zbest['TARGETID'][i])[0][0] oiiflux = elginfo['OIIFLUX'][k] truez = siminfo['REDSHIFT'][j] dv = C_LIGHT*(z-truez)/(1+truez) if truetype == 'SKY' and zwarn > 0: status = 'ok' elif truetype == 'ELG' and zwarn > 0 and oiiflux < 8e-17: status = 'ok ([OII] flux {:.2g})'.format(oiiflux) elif zwarn == 0: if truetype == 'LRG' and objtype == 'GALAXY' and abs(dv) < 150: status = 'ok' elif truetype == 'ELG' and objtype == 'GALAXY': if abs(dv) < 150: status = 'ok' elif oiiflux < 8e-17: status = 'ok ([OII] flux {:.2g})'.format(oiiflux) else: status = 'OOPS ([OII] flux {:.2g})'.format(oiiflux) elif truetype == 'QSO' and objtype == 'QSO' and abs(dv) < 750: status = 'ok' elif truetype in ('STD', 'FSTD') and objtype == 'STAR': status = 'ok' else: status = 'OOPS' else: status = 'OOPS' print('{0:<8d} {1:4s} {2:8.5f} -> {3:6s} {4:8.5f} {5:4d} - {6}'.format( pix, truetype, truez, objtype, z, zwarn, status)) print("--------------------------------------------------")
def simulate_one_healpix(ifilename, args, model, obsconditions, decam_and_wise_filters, footprint_healpix_weight, footprint_healpix_nside, seed): log = get_logger() # set seed now # we need a seed per healpix because # the spectra simulator REQUIRES a seed np.random.seed(seed) healpix = 0 nside = 0 vals = os.path.basename(ifilename).split(".")[0].split("-") if len(vals) < 3: log.error("Cannot guess nside and healpix from filename {}".format( ifilename)) raise ValueError( "Cannot guess nside and healpix from filename {}".format( ifilename)) try: healpix = int(vals[-1]) nside = int(vals[-2]) except ValueError: raise ValueError( "Cannot guess nside and healpix from filename {}".format( ifilename)) zbest_filename = None if args.outfile: ofilename = args.outfile else: ofilename = os.path.join( args.outdir, "{}/{}/spectra-{}-{}.fits".format(healpix // 100, healpix, nside, healpix)) pixdir = os.path.dirname(ofilename) if args.zbest: zbest_filename = os.path.join( pixdir, "zbest-{}-{}.fits".format(nside, healpix)) if not args.overwrite: # check whether output exists or not if args.zbest: if os.path.isfile(ofilename) and os.path.isfile(zbest_filename): log.info("skip existing {} and {}".format( ofilename, zbest_filename)) return else: # only test spectra file if os.path.isfile(ofilename): log.info("skip existing {}".format(ofilename)) return log.info("Read skewers in {}, random seed = {}".format(ifilename, seed)) trans_wave, transmission, metadata = read_lya_skewers(ifilename) ok = np.where((metadata['Z'] >= args.zmin) & (metadata['Z'] <= args.zmax))[0] transmission = transmission[ok] metadata = metadata[:][ok] # create quasars if args.desi_footprint: footprint_healpix = footprint.radec2pix(footprint_healpix_nside, metadata["RA"], metadata["DEC"]) selection = np.where( footprint_healpix_weight[footprint_healpix] > 0.99)[0] log.info("Select QSOs in DESI footprint {} -> {}".format( transmission.shape[0], selection.size)) if selection.size == 0: log.warning("No intersection with DESI footprint") return transmission = transmission[selection] metadata = metadata[:][selection] nqso = transmission.shape[0] if args.downsampling is not None: if args.downsampling <= 0 or args.downsampling > 1: log.error( "Down sampling fraction={} must be between 0 and 1".format( args.downsampling)) raise ValueError( "Down sampling fraction={} must be between 0 and 1".format( args.downsampling)) indices = np.where(np.random.uniform(size=nqso) < args.downsampling)[0] if indices.size == 0: log.warning( "Down sampling from {} to 0 (by chance I presume)".format( nqso)) return transmission = transmission[indices] metadata = metadata[:][indices] nqso = transmission.shape[0] if args.nmax is not None: if args.nmax < nqso: log.info( "Limit number of QSOs from {} to nmax={} (random subsample)". format(nqso, args.nmax)) # take a random subsample indices = (np.random.uniform(size=args.nmax) * nqso).astype(int) transmission = transmission[indices] metadata = metadata[:][indices] nqso = args.nmax if args.target_selection or args.mags: wanted_min_wave = 3329. # needed to compute magnitudes for decam2014-r (one could have trimmed the transmission file ...) wanted_max_wave = 55501. # needed to compute magnitudes for wise2010-W2 if trans_wave[0] > wanted_min_wave: log.info( "Increase wavelength range from {}:{} to {}:{} to compute magnitudes" .format(int(trans_wave[0]), int(trans_wave[-1]), int(wanted_min_wave), int(trans_wave[-1]))) # pad with zeros at short wavelength because we assume transmission = 0 # and we don't need any wavelength resolution here new_trans_wave = np.append([wanted_min_wave, trans_wave[0] - 0.01], trans_wave) new_transmission = np.zeros( (transmission.shape[0], new_trans_wave.size)) new_transmission[:, 2:] = transmission trans_wave = new_trans_wave transmission = new_transmission if trans_wave[-1] < wanted_max_wave: log.info( "Increase wavelength range from {}:{} to {}:{} to compute magnitudes" .format(int(trans_wave[0]), int(trans_wave[-1]), int(trans_wave[0]), int(wanted_max_wave))) # pad with ones at long wavelength because we assume transmission = 1 coarse_dwave = 2. # we don't care about resolution, we just need a decent QSO spectrum, there is no IGM transmission in this range n = int((wanted_max_wave - trans_wave[-1]) / coarse_dwave) + 1 new_trans_wave = np.append( trans_wave, np.linspace(trans_wave[-1] + coarse_dwave, trans_wave[-1] + coarse_dwave * (n + 1), n)) new_transmission = np.ones( (transmission.shape[0], new_trans_wave.size)) new_transmission[:, :trans_wave.size] = transmission trans_wave = new_trans_wave transmission = new_transmission log.info("Simulate {} QSOs".format(nqso)) tmp_qso_flux, tmp_qso_wave, meta = model.make_templates( nmodel=nqso, redshift=metadata['Z'], lyaforest=False, nocolorcuts=True, noresample=True, seed=seed) log.info("Resample to transmission wavelength grid") # because we don't want to alter the transmission field with resampling here qso_flux = np.zeros((tmp_qso_flux.shape[0], trans_wave.size)) for q in range(tmp_qso_flux.shape[0]): qso_flux[q] = np.interp(trans_wave, tmp_qso_wave, tmp_qso_flux[q]) tmp_qso_flux = qso_flux tmp_qso_wave = trans_wave log.info("Apply lya") tmp_qso_flux = apply_lya_transmission(tmp_qso_wave, tmp_qso_flux, trans_wave, transmission) bbflux = None if args.target_selection or args.mags: bands = ['FLUX_G', 'FLUX_R', 'FLUX_Z', 'FLUX_W1', 'FLUX_W2'] bbflux = dict() # need to recompute the magnitudes to account for lya transmission log.info("Compute QSO magnitudes") maggies = decam_and_wise_filters.get_ab_maggies( 1e-17 * tmp_qso_flux, tmp_qso_wave) for band, filt in zip(bands, [ 'decam2014-g', 'decam2014-r', 'decam2014-z', 'wise2010-W1', 'wise2010-W2' ]): bbflux[band] = np.ma.getdata(1e9 * maggies[filt]) # nanomaggies if args.target_selection: log.info("Apply target selection") isqso = isQSO_colors(gflux=bbflux['FLUX_G'], rflux=bbflux['FLUX_R'], zflux=bbflux['FLUX_Z'], w1flux=bbflux['FLUX_W1'], w2flux=bbflux['FLUX_W2']) log.info("Target selection: {}/{} QSOs selected".format( np.sum(isqso), nqso)) selection = np.where(isqso)[0] if selection.size == 0: return tmp_qso_flux = tmp_qso_flux[selection] metadata = metadata[:][selection] meta = meta[:][selection] for band in bands: bbflux[band] = bbflux[band][selection] nqso = selection.size log.info("Resample to a linear wavelength grid (needed by DESI sim.)") # we need a linear grid. for this resampling we take care of integrating in bins # we do not do a simple interpolation qso_wave = np.linspace(args.wmin, args.wmax, int((args.wmax - args.wmin) / args.dwave) + 1) qso_flux = np.zeros((tmp_qso_flux.shape[0], qso_wave.size)) for q in range(tmp_qso_flux.shape[0]): qso_flux[q] = resample_flux(qso_wave, tmp_qso_wave, tmp_qso_flux[q]) log.info("Simulate DESI observation and write output file") pixdir = os.path.dirname(ofilename) if len(pixdir) > 0: if not os.path.isdir(pixdir): log.info("Creating dir {}".format(pixdir)) os.makedirs(pixdir) if "MOCKID" in metadata.dtype.names: #log.warning("Using MOCKID as TARGETID") targetid = np.array(metadata["MOCKID"]).astype(int) elif "ID" in metadata.dtype.names: log.warning("Using ID as TARGETID") targetid = np.array(metadata["ID"]).astype(int) else: log.warning("No TARGETID") targetid = None log.warning("Assuming the healpix scheme is 'NESTED'") meta = {"HPXNSIDE": nside, "HPXPIXEL": healpix, "HPXNEST": True} if args.target_selection or args.mags: # today we write mags because that's what is in the fibermap mags = np.zeros((qso_flux.shape[0], 5)) for i, band in enumerate(bands): jj = (bbflux[band] > 0) mags[jj, i] = 22.5 - 2.5 * np.log10(bbflux[band][jj]) # AB magnitudes fibermap_columns = {"MAG": mags} else: fibermap_columns = None sim_spectra(qso_wave, qso_flux, args.program, obsconditions=obsconditions, spectra_filename=ofilename, sourcetype="qso", skyerr=args.skyerr, ra=metadata["RA"], dec=metadata["DEC"], targetid=targetid, meta=meta, seed=seed, fibermap_columns=fibermap_columns) if args.zbest: log.info("Read fibermap") fibermap = read_fibermap(ofilename) log.info("Writing a zbest file {}".format(zbest_filename)) columns = [('CHI2', 'f8'), ('COEFF', 'f8', (4, )), ('Z', 'f8'), ('ZERR', 'f8'), ('ZWARN', 'i8'), ('SPECTYPE', (str, 96)), ('SUBTYPE', (str, 16)), ('TARGETID', 'i8'), ('DELTACHI2', 'f8'), ('BRICKNAME', (str, 8))] zbest = Table(np.zeros(nqso, dtype=columns)) zbest["CHI2"][:] = 0. zbest["Z"] = metadata['Z'] zbest["ZERR"][:] = 0. zbest["ZWARN"][:] = 0 zbest["SPECTYPE"][:] = "QSO" zbest["SUBTYPE"][:] = "" zbest["TARGETID"] = fibermap["TARGETID"] zbest["DELTACHI2"][:] = 25. hzbest = pyfits.convenience.table_to_hdu(zbest) hzbest.name = "ZBEST" hfmap = pyfits.convenience.table_to_hdu(fibermap) hfmap.name = "FIBERMAP" hdulist = pyfits.HDUList([pyfits.PrimaryHDU(), hzbest, hfmap]) hdulist.writeto(zbest_filename, clobber=True) hdulist.close() # see if this helps with memory issue
def simulate_one_healpix(ifilename,args,model,obsconditions,decam_and_wise_filters, bassmzls_and_wise_filters,footprint_healpix_weight, footprint_healpix_nside, bal=None,sfdmap=None,eboss=None) : log = get_logger() # open filename and extract basic HEALPix information pixel, nside, hpxnest = get_healpix_info(ifilename) # using global seed (could be None) get seed for this particular pixel global_seed = args.seed seed = get_pixel_seed(pixel, nside, global_seed) # use this seed to generate future random numbers np.random.seed(seed) # get output file (we will write there spectra for this HEALPix pixel) ofilename = get_spectra_filename(args,nside,pixel) # get directory name (we will also write there zbest file) pixdir = os.path.dirname(ofilename) # get filename for truth file truth_filename = get_truth_filename(args,pixdir,nside,pixel) # get filename for zbest file zbest_filename = get_zbest_filename(args,pixdir,nside,pixel) if not args.overwrite : # check whether output exists or not if args.zbest : if os.path.isfile(ofilename) and os.path.isfile(zbest_filename) : log.info("skip existing {} and {}".format(ofilename,zbest_filename)) return else : # only test spectra file if os.path.isfile(ofilename) : log.info("skip existing {}".format(ofilename)) return # create sub-directories if required if len(pixdir)>0 : if not os.path.isdir(pixdir) : log.info("Creating dir {}".format(pixdir)) os.makedirs(pixdir) log.info("Read skewers in {}, random seed = {}".format(ifilename,seed)) # Read transmission from files. It might include DLA information, and it # might add metal transmission as well (from the HDU file). log.info("Read transmission file {}".format(ifilename)) trans_wave, transmission, metadata, dla_info = read_lya_skewers(ifilename,read_dlas=(args.dla=='file'),add_metals=args.metals_from_file,add_lyb=args.add_LYB) ### Add Finger-of-God, before generate the continua log.info("Add FOG to redshift with sigma {} to quasar redshift".format(args.sigma_kms_fog)) DZ_FOG = args.sigma_kms_fog/c*(1.+metadata['Z'])*np.random.normal(0,1,metadata['Z'].size) metadata['Z'] += DZ_FOG ### Select quasar within a given redshift range w = (metadata['Z']>=args.zmin) & (metadata['Z']<=args.zmax) transmission = transmission[w] metadata = metadata[:][w] DZ_FOG = DZ_FOG[w] # option to make for BOSS+eBOSS if not eboss is None: if args.downsampling or args.desi_footprint: raise ValueError("eboss option can not be run with " +"desi_footprint or downsampling") # Get the redshift distribution from SDSS selection = sdss_subsample_redshift(metadata["RA"],metadata["DEC"],metadata['Z'],eboss['redshift']) log.info("Select QSOs in BOSS+eBOSS redshift distribution {} -> {}".format(metadata['Z'].size,selection.sum())) if selection.sum()==0: log.warning("No intersection with BOSS+eBOSS redshift distribution") return transmission = transmission[selection] metadata = metadata[:][selection] DZ_FOG = DZ_FOG[selection] # figure out the density of all quasars N_highz = metadata['Z'].size # area of healpix pixel, in degrees area_deg2 = healpy.pixelfunc.nside2pixarea(nside,degrees=True) input_highz_dens_deg2 = N_highz/area_deg2 selection = sdss_subsample(metadata["RA"], metadata["DEC"], input_highz_dens_deg2,eboss['footprint']) log.info("Select QSOs in BOSS+eBOSS footprint {} -> {}".format(transmission.shape[0],selection.size)) if selection.size == 0 : log.warning("No intersection with BOSS+eBOSS footprint") return transmission = transmission[selection] metadata = metadata[:][selection] DZ_FOG = DZ_FOG[selection] if args.desi_footprint : footprint_healpix = footprint.radec2pix(footprint_healpix_nside, metadata["RA"], metadata["DEC"]) selection = np.where(footprint_healpix_weight[footprint_healpix]>0.99)[0] log.info("Select QSOs in DESI footprint {} -> {}".format(transmission.shape[0],selection.size)) if selection.size == 0 : log.warning("No intersection with DESI footprint") return transmission = transmission[selection] metadata = metadata[:][selection] DZ_FOG = DZ_FOG[selection] nqso=transmission.shape[0] if args.downsampling is not None : if args.downsampling <= 0 or args.downsampling > 1 : log.error("Down sampling fraction={} must be between 0 and 1".format(args.downsampling)) raise ValueError("Down sampling fraction={} must be between 0 and 1".format(args.downsampling)) indices = np.where(np.random.uniform(size=nqso)<args.downsampling)[0] if indices.size == 0 : log.warning("Down sampling from {} to 0 (by chance I presume)".format(nqso)) return transmission = transmission[indices] metadata = metadata[:][indices] DZ_FOG = DZ_FOG[indices] nqso = transmission.shape[0] if args.nmax is not None : if args.nmax < nqso : log.info("Limit number of QSOs from {} to nmax={} (random subsample)".format(nqso,args.nmax)) # take a random subsample indices = np.random.choice(np.arange(nqso),args.nmax,replace=False) ##Use random.choice instead of random.uniform (rarely but it does cause a duplication of qsos) transmission = transmission[indices] metadata = metadata[:][indices] DZ_FOG = DZ_FOG[indices] nqso = args.nmax # In previous versions of the London mocks we needed to enforce F=1 for # z > z_qso here, but this is not needed anymore. Moreover, now we also # have metal absorption that implies F < 1 for z > z_qso #for ii in range(len(metadata)): # transmission[ii][trans_wave>lambda_RF_LYA*(metadata[ii]['Z']+1)]=1.0 # if requested, add DLA to the transmission skewers if args.dla is not None : # if adding random DLAs, we will need a new random generator if args.dla=='random': log.info('Adding DLAs randomly') random_state_just_for_dlas = np.random.RandomState(seed) elif args.dla=='file': log.info('Adding DLAs from transmission file') else: log.error("Wrong option for args.dla: "+args.dla) sys.exit(1) # if adding DLAs, the information will be printed here dla_filename=os.path.join(pixdir,"dla-{}-{}.fits".format(nside,pixel)) dla_NHI, dla_z, dla_qid,dla_id = [], [], [],[] # identify minimum Lya redshift in transmission files min_lya_z = np.min(trans_wave/lambda_RF_LYA - 1) # loop over quasars in pixel for ii in range(len(metadata)): # quasars with z < min_z will not have any DLA in spectrum if min_lya_z>metadata['Z'][ii]: continue # quasar ID idd=metadata['MOCKID'][ii] dlas=[] if args.dla=='file': for dla in dla_info[dla_info['MOCKID']==idd]: # Adding only DLAs with z < zqso if dla['Z_DLA_RSD']>=metadata['Z'][ii]: continue dlas.append(dict(z=dla['Z_DLA_RSD'],N=dla['N_HI_DLA'],dlaid=dla['DLAID'])) transmission_dla = dla_spec(trans_wave,dlas) elif args.dla=='random': dlas, transmission_dla = insert_dlas(trans_wave, metadata['Z'][ii], rstate=random_state_just_for_dlas) for idla in dlas: idla['dlaid']+=idd*1000 #Added to have unique DLA ids. Same format as DLAs from file. # multiply transmissions and store information for the DLA file if len(dlas)>0: transmission[ii] = transmission_dla * transmission[ii] dla_z += [idla['z'] for idla in dlas] dla_NHI += [idla['N'] for idla in dlas] dla_id += [idla['dlaid'] for idla in dlas] dla_qid += [idd]*len(dlas) log.info('Added {} DLAs'.format(len(dla_id))) # write file with DLA information if len(dla_id)>0: dla_meta=Table() dla_meta['NHI'] = dla_NHI dla_meta['Z_DLA'] = dla_z #This is Z_DLA_RSD in transmision. dla_meta['TARGETID']=dla_qid dla_meta['DLAID'] = dla_id hdu_dla = pyfits.convenience.table_to_hdu(dla_meta) hdu_dla.name="DLA_META" del(dla_meta) log.info("DLA metadata to be saved in {}".format(truth_filename)) else: hdu_dla=pyfits.PrimaryHDU() hdu_dla.name="DLA_META" # if requested, extend transmission skewers to cover full spectrum if args.target_selection or args.bbflux : wanted_min_wave = 3329. # needed to compute magnitudes for decam2014-r (one could have trimmed the transmission file ...) wanted_max_wave = 55501. # needed to compute magnitudes for wise2010-W2 if trans_wave[0]>wanted_min_wave : log.info("Increase wavelength range from {}:{} to {}:{} to compute magnitudes".format(int(trans_wave[0]),int(trans_wave[-1]),int(wanted_min_wave),int(trans_wave[-1]))) # pad with ones at short wavelength, we assume F = 1 for z <~ 1.7 # we don't need any wavelength resolution here new_trans_wave = np.append([wanted_min_wave,trans_wave[0]-0.01],trans_wave) new_transmission = np.ones((transmission.shape[0],new_trans_wave.size)) new_transmission[:,2:] = transmission trans_wave = new_trans_wave transmission = new_transmission if trans_wave[-1]<wanted_max_wave : log.info("Increase wavelength range from {}:{} to {}:{} to compute magnitudes".format(int(trans_wave[0]),int(trans_wave[-1]),int(trans_wave[0]),int(wanted_max_wave))) # pad with ones at long wavelength because we assume F = 1 coarse_dwave = 2. # we don't care about resolution, we just need a decent QSO spectrum, there is no IGM transmission in this range n = int((wanted_max_wave-trans_wave[-1])/coarse_dwave)+1 new_trans_wave = np.append(trans_wave,np.linspace(trans_wave[-1]+coarse_dwave,trans_wave[-1]+coarse_dwave*(n+1),n)) new_transmission = np.ones((transmission.shape[0],new_trans_wave.size)) new_transmission[:,:trans_wave.size] = transmission trans_wave = new_trans_wave transmission = new_transmission # whether to use QSO or SIMQSO to generate quasar continua. Simulate # spectra in the north vs south separately because they're on different # photometric systems. south = np.where( is_south(metadata['DEC']) )[0] north = np.where( ~is_south(metadata['DEC']) )[0] meta, qsometa = empty_metatable(nqso, objtype='QSO', simqso=not args.no_simqso) if args.no_simqso: log.info("Simulate {} QSOs with QSO templates".format(nqso)) tmp_qso_flux = np.zeros([nqso, len(model.eigenwave)], dtype='f4') tmp_qso_wave = np.zeros_like(tmp_qso_flux) else: log.info("Simulate {} QSOs with SIMQSO templates".format(nqso)) tmp_qso_flux = np.zeros([nqso, len(model.basewave)], dtype='f4') tmp_qso_wave = model.basewave for these, issouth in zip( (north, south), (False, True) ): # number of quasars in these nt = len(these) if nt<=0: continue if not eboss is None: # for eBOSS, generate only quasars with r<22 magrange = (17.0, 21.3) _tmp_qso_flux, _tmp_qso_wave, _meta, _qsometa \ = model.make_templates(nmodel=nt, redshift=metadata['Z'][these], magrange=magrange, lyaforest=False, nocolorcuts=True, noresample=True, seed=seed, south=issouth) else: _tmp_qso_flux, _tmp_qso_wave, _meta, _qsometa \ = model.make_templates(nmodel=nt, redshift=metadata['Z'][these], lyaforest=False, nocolorcuts=True, noresample=True, seed=seed, south=issouth) _meta['TARGETID'] = metadata['MOCKID'][these] _qsometa['TARGETID'] = metadata['MOCKID'][these] meta[these] = _meta qsometa[these] = _qsometa tmp_qso_flux[these, :] = _tmp_qso_flux if args.no_simqso: tmp_qso_wave[these, :] = _tmp_qso_wave log.info("Resample to transmission wavelength grid") qso_flux=np.zeros((tmp_qso_flux.shape[0],trans_wave.size)) if args.no_simqso: for q in range(tmp_qso_flux.shape[0]) : qso_flux[q]=np.interp(trans_wave,tmp_qso_wave[q],tmp_qso_flux[q]) else: for q in range(tmp_qso_flux.shape[0]) : qso_flux[q]=np.interp(trans_wave,tmp_qso_wave,tmp_qso_flux[q]) tmp_qso_flux = qso_flux tmp_qso_wave = trans_wave # if requested, add BAL features to the quasar continua if args.balprob: if args.balprob<=1. and args.balprob >0: log.info("Adding BALs with probability {}".format(args.balprob)) # save current random state rnd_state = np.random.get_state() tmp_qso_flux,meta_bal=bal.insert_bals(tmp_qso_wave,tmp_qso_flux, metadata['Z'], balprob=args.balprob,seed=seed) # restore random state to get the same random numbers later # as when we don't insert BALs np.random.set_state(rnd_state) meta_bal['TARGETID'] = metadata['MOCKID'] w = meta_bal['TEMPLATEID']!=-1 meta_bal = meta_bal[:][w] hdu_bal=pyfits.convenience.table_to_hdu(meta_bal); hdu_bal.name="BAL_META" del meta_bal else: balstr=str(args.balprob) log.error("BAL probability is not between 0 and 1 : "+balstr) sys.exit(1) # Multiply quasar continua by transmitted flux fraction # (at this point transmission file might include Ly-beta, metals and DLAs) log.info("Apply transmitted flux fraction") if not args.no_transmission: tmp_qso_flux = apply_lya_transmission(tmp_qso_wave,tmp_qso_flux, trans_wave,transmission) # if requested, compute metal transmission on the fly # (if not included already from the transmission file) if args.metals is not None: if args.metals_from_file : log.error('you cannot add metals twice') raise ValueError('you cannot add metals twice') if args.no_transmission: log.error('you cannot add metals if asking for no-transmission') raise ValueError('can not add metals if using no-transmission') lstMetals = '' for m in args.metals: lstMetals += m+', ' log.info("Apply metals: {}".format(lstMetals[:-2])) tmp_qso_flux = apply_metals_transmission(tmp_qso_wave,tmp_qso_flux, trans_wave,transmission,args.metals) # if requested, compute magnitudes and apply target selection. Need to do # this calculation separately for QSOs in the north vs south. bbflux=None if args.target_selection or args.bbflux : bands=['FLUX_G','FLUX_R','FLUX_Z', 'FLUX_W1', 'FLUX_W2'] bbflux=dict() bbflux['SOUTH'] = is_south(metadata['DEC']) for band in bands: bbflux[band] = np.zeros(nqso) # need to recompute the magnitudes to account for lya transmission log.info("Compute QSO magnitudes") for these, filters in zip( (~bbflux['SOUTH'], bbflux['SOUTH']), (bassmzls_and_wise_filters, decam_and_wise_filters) ): if np.count_nonzero(these) > 0: maggies = filters.get_ab_maggies(1e-17 * tmp_qso_flux[these, :], tmp_qso_wave) for band, filt in zip( bands, maggies.colnames ): bbflux[band][these] = np.ma.getdata(1e9 * maggies[filt]) # nanomaggies if args.target_selection : log.info("Apply target selection") isqso = np.ones(nqso, dtype=bool) for these, issouth in zip( (~bbflux['SOUTH'], bbflux['SOUTH']), (False, True) ): if np.count_nonzero(these) > 0: # optical cuts only if using QSO vs SIMQSO isqso[these] &= isQSO_colors(gflux=bbflux['FLUX_G'][these], rflux=bbflux['FLUX_R'][these], zflux=bbflux['FLUX_Z'][these], w1flux=bbflux['FLUX_W1'][these], w2flux=bbflux['FLUX_W2'][these], south=issouth, optical=args.no_simqso) log.info("Target selection: {}/{} QSOs selected".format(np.sum(isqso),nqso)) selection=np.where(isqso)[0] if selection.size==0 : return tmp_qso_flux = tmp_qso_flux[selection] metadata = metadata[:][selection] meta = meta[:][selection] qsometa = qsometa[:][selection] DZ_FOG = DZ_FOG[selection] for band in bands : bbflux[band] = bbflux[band][selection] bbflux['SOUTH']=bbflux['SOUTH'][selection] nqso = selection.size log.info("Resample to a linear wavelength grid (needed by DESI sim.)") # careful integration of bins, not just a simple interpolation qso_wave=np.linspace(args.wmin,args.wmax,int((args.wmax-args.wmin)/args.dwave)+1) qso_flux=np.zeros((tmp_qso_flux.shape[0],qso_wave.size)) for q in range(tmp_qso_flux.shape[0]) : qso_flux[q]=resample_flux(qso_wave,tmp_qso_wave,tmp_qso_flux[q]) log.info("Simulate DESI observation and write output file") if "MOCKID" in metadata.dtype.names : #log.warning("Using MOCKID as TARGETID") targetid=np.array(metadata["MOCKID"]).astype(int) elif "ID" in metadata.dtype.names : log.warning("Using ID as TARGETID") targetid=np.array(metadata["ID"]).astype(int) else : log.warning("No TARGETID") targetid=None specmeta={"HPXNSIDE":nside,"HPXPIXEL":pixel, "HPXNEST":hpxnest} if args.target_selection or args.bbflux : fibermap_columns = dict( FLUX_G = bbflux['FLUX_G'], FLUX_R = bbflux['FLUX_R'], FLUX_Z = bbflux['FLUX_Z'], FLUX_W1 = bbflux['FLUX_W1'], FLUX_W2 = bbflux['FLUX_W2'], ) photsys = np.full(len(bbflux['FLUX_G']), 'N', dtype='S1') photsys[bbflux['SOUTH']] = b'S' fibermap_columns['PHOTSYS'] = photsys else : fibermap_columns=None # Attenuate the spectra for extinction if not sfdmap is None: Rv=3.1 #set by default indx=np.arange(metadata['RA'].size) extinction =Rv*ext_odonnell(qso_wave) EBV = sfdmap.ebv(metadata['RA'],metadata['DEC'], scaling=1.0) qso_flux *=10**( -0.4 * EBV[indx, np.newaxis] * extinction) if fibermap_columns is not None: fibermap_columns['EBV']=EBV EBV0=0.0 EBV_med=np.median(EBV) Ag = 3.303 * (EBV_med - EBV0) exptime_fact=np.power(10.0, (2.0 * Ag / 2.5)) obsconditions['EXPTIME']*=exptime_fact log.info("Dust extinction added") log.info('exposure time adjusted to {}'.format(obsconditions['EXPTIME'])) sim_spectra(qso_wave,qso_flux, args.program, obsconditions=obsconditions,spectra_filename=ofilename, sourcetype="qso", skyerr=args.skyerr,ra=metadata["RA"],dec=metadata["DEC"],targetid=targetid, meta=specmeta,seed=seed,fibermap_columns=fibermap_columns,use_poisson=False) # use Poisson = False to get reproducible results. ### Keep input redshift Z_spec = metadata['Z'].copy() Z_input = metadata['Z'].copy()-DZ_FOG ### Add a shift to the redshift, simulating the systematic imprecision of redrock DZ_sys_shift = args.shift_kms_los/c*(1.+Z_input) log.info('Added a shift of {} km/s to the redshift'.format(args.shift_kms_los)) meta['REDSHIFT'] += DZ_sys_shift metadata['Z'] += DZ_sys_shift ### Add a shift to the redshift, simulating the statistic imprecision of redrock if args.gamma_kms_zfit: log.info("Added zfit error with gamma {} to zbest".format(args.gamma_kms_zfit)) DZ_stat_shift = mod_cauchy(loc=0,scale=args.gamma_kms_zfit,size=nqso,cut=3000)/c*(1.+Z_input) meta['REDSHIFT'] += DZ_stat_shift metadata['Z'] += DZ_stat_shift ## Write the truth file, including metadata for DLAs and BALs log.info('Writing a truth file {}'.format(truth_filename)) meta.rename_column('REDSHIFT','Z') meta.add_column(Column(Z_spec,name='TRUEZ')) meta.add_column(Column(Z_input,name='Z_INPUT')) meta.add_column(Column(DZ_FOG,name='DZ_FOG')) meta.add_column(Column(DZ_sys_shift,name='DZ_SYS')) if args.gamma_kms_zfit: meta.add_column(Column(DZ_stat_shift,name='DZ_STAT')) if 'Z_noRSD' in metadata.dtype.names: meta.add_column(Column(metadata['Z_noRSD'],name='Z_NORSD')) else: log.info('Z_noRSD field not present in transmission file. Z_NORSD not saved to truth file') #Save global seed and pixel seed to primary header hdr=pyfits.Header() hdr['GSEED']=global_seed hdr['PIXSEED']=seed hdu = pyfits.convenience.table_to_hdu(meta) hdu.header['EXTNAME'] = 'TRUTH' hduqso=pyfits.convenience.table_to_hdu(qsometa) hduqso.header['EXTNAME'] = 'QSO_META' hdulist=pyfits.HDUList([pyfits.PrimaryHDU(header=hdr),hdu,hduqso]) if args.dla: hdulist.append(hdu_dla) if args.balprob: hdulist.append(hdu_bal) hdulist.writeto(truth_filename, overwrite=True) hdulist.close() if args.zbest : log.info("Read fibermap") fibermap = read_fibermap(ofilename) log.info("Writing a zbest file {}".format(zbest_filename)) columns = [ ('CHI2', 'f8'), ('COEFF', 'f8' , (4,)), ('Z', 'f8'), ('ZERR', 'f8'), ('ZWARN', 'i8'), ('SPECTYPE', (str,96)), ('SUBTYPE', (str,16)), ('TARGETID', 'i8'), ('DELTACHI2', 'f8'), ('BRICKNAME', (str,8))] zbest = Table(np.zeros(nqso, dtype=columns)) zbest['CHI2'][:] = 0. zbest['Z'][:] = metadata['Z'] zbest['ZERR'][:] = 0. zbest['ZWARN'][:] = 0 zbest['SPECTYPE'][:] = 'QSO' zbest['SUBTYPE'][:] = '' zbest['TARGETID'][:] = metadata['MOCKID'] zbest['DELTACHI2'][:] = 25. hzbest = pyfits.convenience.table_to_hdu(zbest); hzbest.name='ZBEST' hfmap = pyfits.convenience.table_to_hdu(fibermap); hfmap.name='FIBERMAP' hdulist =pyfits.HDUList([pyfits.PrimaryHDU(),hzbest,hfmap]) hdulist.writeto(zbest_filename, overwrite=True) hdulist.close() # see if this helps with memory issue
def integration_test(night=None, nspec=5, clobber=False): """Run an integration test from raw data simulations through redshifts Args: night (str, optional): YEARMMDD, defaults to current night nspec (int, optional): number of spectra to include clobber (bool, optional): rerun steps even if outputs already exist Raises: RuntimeError if any script fails """ log = logging.get_logger() log.setLevel(logging.DEBUG) # YEARMMDD string, rolls over at noon not midnight # TODO: fix usage of night to be something other than today if night is None: #night = time.strftime('%Y%m%d', time.localtime(time.time()-12*3600)) night = "20160726" # check for required environment variables check_env() # simulate inputs sim(night, nspec=nspec, clobber=clobber) # create production # FIXME: someday run PSF estimation too... ### com = "desi_pipe --spectrographs 0 --fakeboot --fakepsf" com = "desi_pipe --spectrographs 0 --fakeboot --fakepsf" sp.check_call(com, shell=True) # raw and production locations rawdir = os.path.abspath(io.rawdata_root()) proddir = os.path.abspath(io.specprod_root()) # Modify options file to restrict the spectral range optpath = os.path.join(proddir, "run", "options.yaml") opts = pipe.yaml_read(optpath) opts['extract']['specmin'] = 0 opts['extract']['nspec'] = nspec pipe.yaml_write(optpath, opts) # run the generated shell scripts # FIXME: someday run PSF estimation too... # print("Running bootcalib script...") # com = os.path.join(proddir, "run", "scripts", "bootcalib_all.sh") # sp.check_call(["bash", com]) # print("Running specex script...") # com = os.path.join(proddir, "run", "scripts", "specex_all.sh") # sp.check_call(["bash", com]) # print("Running psfcombine script...") # com = os.path.join(proddir, "run", "scripts", "psfcombine_all.sh") # sp.check_call(["bash", com]) com = os.path.join(proddir, "run", "scripts", "run_shell.sh") print("Running extraction through calibration: " + com) sp.check_call(["bash", com]) com = os.path.join(proddir, "run", "scripts", "spectra.sh") print("Running spectral regrouping: " + com) sp.check_call(["bash", com]) com = os.path.join(proddir, "run", "scripts", "redshift.sh") print("Running redshift script " + com) sp.check_call(["bash", com]) # #----- # #- Did it work? # #- (this combination of fibermap, simspec, and zbest is a pain) expid = 2 fmfile = io.findfile('fibermap', night=night, expid=expid) fibermap = io.read_fibermap(fmfile) simdir = os.path.dirname(fmfile) simspec = '{}/simspec-{:08d}.fits'.format(simdir, expid) siminfo = fits.getdata(simspec, 'TRUTH') from desimodel.footprint import radec2pix nside = 64 pixels = np.unique( radec2pix(nside, fibermap['RA_TARGET'], fibermap['DEC_TARGET'])) print() print("--------------------------------------------------") print("Pixel True z -> Class z zwarn") # print("3338p190 SKY 0.00000 -> QSO 1.60853 12 - ok") for pix in pixels: zbest = io.read_zbest(io.findfile('zbest', groupname=pix)) for i in range(len(zbest.z)): objtype = zbest.spectype[i] z, zwarn = zbest.z[i], zbest.zwarn[i] j = np.where(fibermap['TARGETID'] == zbest.targetid[i])[0][0] truetype = siminfo['OBJTYPE'][j] oiiflux = siminfo['OIIFLUX'][j] truez = siminfo['REDSHIFT'][j] dv = 3e5 * (z - truez) / (1 + truez) if truetype == 'SKY' and zwarn > 0: status = 'ok' elif truetype == 'ELG' and zwarn > 0 and oiiflux < 8e-17: status = 'ok ([OII] flux {:.2g})'.format(oiiflux) elif zwarn == 0: if truetype == 'LRG' and objtype == 'GALAXY' and abs(dv) < 150: status = 'ok' elif truetype == 'ELG' and objtype == 'GALAXY': if abs(dv) < 150: status = ok elif oiiflux < 8e-17: status = 'ok ([OII] flux {:.2g})'.format(oiiflux) else: status = 'OOPS ([OII] flux {:.2g})'.format(oiiflux) elif truetype == 'QSO' and objtype == 'QSO' and abs(dv) < 750: status = 'ok' elif truetype in ('STD', 'FSTD') and objtype == 'STAR': status = 'ok' else: status = 'OOPS' else: status = 'OOPS' print('{0:<8d} {1:4s} {2:8.5f} -> {3:5s} {4:8.5f} {5:4d} - {6}'. format(pix, truetype, truez, objtype, z, zwarn, status)) print("--------------------------------------------------")
def simulate_one_healpix(ifilename,args,model,obsconditions,decam_and_wise_filters, bassmzls_and_wise_filters,footprint_healpix_weight, footprint_healpix_nside, bal=None,sfdmap=None,eboss=None) : log = get_logger() # open filename and extract basic HEALPix information pixel, nside, hpxnest = get_healpix_info(ifilename) # using global seed (could be None) get seed for this particular pixel global_seed = args.seed seed = get_pixel_seed(pixel, nside, global_seed) # use this seed to generate future random numbers np.random.seed(seed) # get output file (we will write there spectra for this HEALPix pixel) ofilename = get_spectra_filename(args,nside,pixel) # get directory name (we will also write there zbest file) pixdir = os.path.dirname(ofilename) # get filename for truth file truth_filename = get_truth_filename(args,pixdir,nside,pixel) # get filename for zbest file zbest_filename = get_zbest_filename(args,pixdir,nside,pixel) if not args.overwrite : # check whether output exists or not if args.zbest : if os.path.isfile(ofilename) and os.path.isfile(zbest_filename) : log.info("skip existing {} and {}".format(ofilename,zbest_filename)) return else : # only test spectra file if os.path.isfile(ofilename) : log.info("skip existing {}".format(ofilename)) return # create sub-directories if required if len(pixdir)>0 : if not os.path.isdir(pixdir) : log.info("Creating dir {}".format(pixdir)) os.makedirs(pixdir) log.info("Read skewers in {}, random seed = {}".format(ifilename,seed)) # Read transmission from files. It might include DLA information, and it # might add metal transmission as well (from the HDU file). log.info("Read transmission file {}".format(ifilename)) trans_wave, transmission, metadata, dla_info = read_lya_skewers(ifilename,read_dlas=(args.dla=='file'),add_metals=args.metals_from_file) ### Add Finger-of-God, before generate the continua log.info("Add FOG to redshift with sigma {} to quasar redshift".format(args.sigma_kms_fog)) DZ_FOG = args.sigma_kms_fog/c*(1.+metadata['Z'])*np.random.normal(0,1,metadata['Z'].size) metadata['Z'] += DZ_FOG ### Select quasar within a given redshift range w = (metadata['Z']>=args.zmin) & (metadata['Z']<=args.zmax) transmission = transmission[w] metadata = metadata[:][w] DZ_FOG = DZ_FOG[w] # option to make for BOSS+eBOSS if not eboss is None: if args.downsampling or args.desi_footprint: raise ValueError("eboss option can not be run with " +"desi_footprint or downsampling") # Get the redshift distribution from SDSS selection = sdss_subsample_redshift(metadata["RA"],metadata["DEC"],metadata['Z'],eboss['redshift']) log.info("Select QSOs in BOSS+eBOSS redshift distribution {} -> {}".format(metadata['Z'].size,selection.sum())) if selection.sum()==0: log.warning("No intersection with BOSS+eBOSS redshift distribution") return transmission = transmission[selection] metadata = metadata[:][selection] DZ_FOG = DZ_FOG[selection] # figure out the density of all quasars N_highz = metadata['Z'].size # area of healpix pixel, in degrees area_deg2 = healpy.pixelfunc.nside2pixarea(nside,degrees=True) input_highz_dens_deg2 = N_highz/area_deg2 selection = sdss_subsample(metadata["RA"], metadata["DEC"], input_highz_dens_deg2,eboss['footprint']) log.info("Select QSOs in BOSS+eBOSS footprint {} -> {}".format(transmission.shape[0],selection.size)) if selection.size == 0 : log.warning("No intersection with BOSS+eBOSS footprint") return transmission = transmission[selection] metadata = metadata[:][selection] DZ_FOG = DZ_FOG[selection] if args.desi_footprint : footprint_healpix = footprint.radec2pix(footprint_healpix_nside, metadata["RA"], metadata["DEC"]) selection = np.where(footprint_healpix_weight[footprint_healpix]>0.99)[0] log.info("Select QSOs in DESI footprint {} -> {}".format(transmission.shape[0],selection.size)) if selection.size == 0 : log.warning("No intersection with DESI footprint") return transmission = transmission[selection] metadata = metadata[:][selection] DZ_FOG = DZ_FOG[selection] nqso=transmission.shape[0] if args.downsampling is not None : if args.downsampling <= 0 or args.downsampling > 1 : log.error("Down sampling fraction={} must be between 0 and 1".format(args.downsampling)) raise ValueError("Down sampling fraction={} must be between 0 and 1".format(args.downsampling)) indices = np.where(np.random.uniform(size=nqso)<args.downsampling)[0] if indices.size == 0 : log.warning("Down sampling from {} to 0 (by chance I presume)".format(nqso)) return transmission = transmission[indices] metadata = metadata[:][indices] DZ_FOG = DZ_FOG[indices] nqso = transmission.shape[0] if args.nmax is not None : if args.nmax < nqso : log.info("Limit number of QSOs from {} to nmax={} (random subsample)".format(nqso,args.nmax)) # take a random subsample indices = (np.random.uniform(size=args.nmax)*nqso).astype(int) transmission = transmission[indices] metadata = metadata[:][indices] DZ_FOG = DZ_FOG[indices] nqso = args.nmax # In previous versions of the London mocks we needed to enforce F=1 for # z > z_qso here, but this is not needed anymore. Moreover, now we also # have metal absorption that implies F < 1 for z > z_qso #for ii in range(len(metadata)): # transmission[ii][trans_wave>lambda_RF_LYA*(metadata[ii]['Z']+1)]=1.0 # if requested, add DLA to the transmission skewers if args.dla is not None : # if adding random DLAs, we will need a new random generator if args.dla=='random': log.info('Adding DLAs randomly') random_state_just_for_dlas = np.random.RandomState(seed) elif args.dla=='file': log.info('Adding DLAs from transmission file') else: log.error("Wrong option for args.dla: "+args.dla) sys.exit(1) # if adding DLAs, the information will be printed here dla_filename=os.path.join(pixdir,"dla-{}-{}.fits".format(nside,pixel)) dla_NHI, dla_z, dla_qid,dla_id = [], [], [],[] # identify minimum Lya redshift in transmission files min_lya_z = np.min(trans_wave/lambda_RF_LYA - 1) # loop over quasars in pixel for ii in range(len(metadata)): # quasars with z < min_z will not have any DLA in spectrum if min_lya_z>metadata['Z'][ii]: continue # quasar ID idd=metadata['MOCKID'][ii] dlas=[] if args.dla=='file': for dla in dla_info[dla_info['MOCKID']==idd]: # Adding only DLAs with z < zqso if dla['Z_DLA_RSD']>=metadata['Z'][ii]: continue dlas.append(dict(z=dla['Z_DLA_RSD'],N=dla['N_HI_DLA'],dlaid=dla['DLAID'])) transmission_dla = dla_spec(trans_wave,dlas) elif args.dla=='random': dlas, transmission_dla = insert_dlas(trans_wave, metadata['Z'][ii], rstate=random_state_just_for_dlas) for idla in dlas: idla['dlaid']+=idd*1000 #Added to have unique DLA ids. Same format as DLAs from file. # multiply transmissions and store information for the DLA file if len(dlas)>0: transmission[ii] = transmission_dla * transmission[ii] dla_z += [idla['z'] for idla in dlas] dla_NHI += [idla['N'] for idla in dlas] dla_id += [idla['dlaid'] for idla in dlas] dla_qid += [idd]*len(dlas) log.info('Added {} DLAs'.format(len(dla_id))) # write file with DLA information if len(dla_id)>0: dla_meta=Table() dla_meta['NHI'] = dla_NHI dla_meta['Z_DLA'] = dla_z #This is Z_DLA_RSD in transmision. dla_meta['TARGETID']=dla_qid dla_meta['DLAID'] = dla_id hdu_dla = pyfits.convenience.table_to_hdu(dla_meta) hdu_dla.name="DLA_META" del(dla_meta) log.info("DLA metadata to be saved in {}".format(truth_filename)) else: hdu_dla=pyfits.PrimaryHDU() hdu_dla.name="DLA_META" # if requested, extend transmission skewers to cover full spectrum if args.target_selection or args.bbflux : wanted_min_wave = 3329. # needed to compute magnitudes for decam2014-r (one could have trimmed the transmission file ...) wanted_max_wave = 55501. # needed to compute magnitudes for wise2010-W2 if trans_wave[0]>wanted_min_wave : log.info("Increase wavelength range from {}:{} to {}:{} to compute magnitudes".format(int(trans_wave[0]),int(trans_wave[-1]),int(wanted_min_wave),int(trans_wave[-1]))) # pad with ones at short wavelength, we assume F = 1 for z <~ 1.7 # we don't need any wavelength resolution here new_trans_wave = np.append([wanted_min_wave,trans_wave[0]-0.01],trans_wave) new_transmission = np.ones((transmission.shape[0],new_trans_wave.size)) new_transmission[:,2:] = transmission trans_wave = new_trans_wave transmission = new_transmission if trans_wave[-1]<wanted_max_wave : log.info("Increase wavelength range from {}:{} to {}:{} to compute magnitudes".format(int(trans_wave[0]),int(trans_wave[-1]),int(trans_wave[0]),int(wanted_max_wave))) # pad with ones at long wavelength because we assume F = 1 coarse_dwave = 2. # we don't care about resolution, we just need a decent QSO spectrum, there is no IGM transmission in this range n = int((wanted_max_wave-trans_wave[-1])/coarse_dwave)+1 new_trans_wave = np.append(trans_wave,np.linspace(trans_wave[-1]+coarse_dwave,trans_wave[-1]+coarse_dwave*(n+1),n)) new_transmission = np.ones((transmission.shape[0],new_trans_wave.size)) new_transmission[:,:trans_wave.size] = transmission trans_wave = new_trans_wave transmission = new_transmission # whether to use QSO or SIMQSO to generate quasar continua. Simulate # spectra in the north vs south separately because they're on different # photometric systems. south = np.where( is_south(metadata['DEC']) )[0] north = np.where( ~is_south(metadata['DEC']) )[0] meta, qsometa = empty_metatable(nqso, objtype='QSO', simqso=not args.no_simqso) if args.no_simqso: log.info("Simulate {} QSOs with QSO templates".format(nqso)) tmp_qso_flux = np.zeros([nqso, len(model.eigenwave)], dtype='f4') tmp_qso_wave = np.zeros_like(tmp_qso_flux) else: log.info("Simulate {} QSOs with SIMQSO templates".format(nqso)) tmp_qso_flux = np.zeros([nqso, len(model.basewave)], dtype='f4') tmp_qso_wave = model.basewave for these, issouth in zip( (north, south), (False, True) ): # number of quasars in these nt = len(these) if nt<=0: continue if not eboss is None: # for eBOSS, generate only quasars with r<22 magrange = (17.0, 21.3) _tmp_qso_flux, _tmp_qso_wave, _meta, _qsometa \ = model.make_templates(nmodel=nt, redshift=metadata['Z'][these], magrange=magrange, lyaforest=False, nocolorcuts=True, noresample=True, seed=seed, south=issouth) else: _tmp_qso_flux, _tmp_qso_wave, _meta, _qsometa \ = model.make_templates(nmodel=nt, redshift=metadata['Z'][these], lyaforest=False, nocolorcuts=True, noresample=True, seed=seed, south=issouth) _meta['TARGETID'] = metadata['MOCKID'][these] _qsometa['TARGETID'] = metadata['MOCKID'][these] meta[these] = _meta qsometa[these] = _qsometa tmp_qso_flux[these, :] = _tmp_qso_flux if args.no_simqso: tmp_qso_wave[these, :] = _tmp_qso_wave log.info("Resample to transmission wavelength grid") qso_flux=np.zeros((tmp_qso_flux.shape[0],trans_wave.size)) if args.no_simqso: for q in range(tmp_qso_flux.shape[0]) : qso_flux[q]=np.interp(trans_wave,tmp_qso_wave[q],tmp_qso_flux[q]) else: for q in range(tmp_qso_flux.shape[0]) : qso_flux[q]=np.interp(trans_wave,tmp_qso_wave,tmp_qso_flux[q]) tmp_qso_flux = qso_flux tmp_qso_wave = trans_wave # if requested, add BAL features to the quasar continua if args.balprob: if args.balprob<=1. and args.balprob >0: log.info("Adding BALs with probability {}".format(args.balprob)) # save current random state rnd_state = np.random.get_state() tmp_qso_flux,meta_bal=bal.insert_bals(tmp_qso_wave,tmp_qso_flux, metadata['Z'], balprob=args.balprob,seed=seed) # restore random state to get the same random numbers later # as when we don't insert BALs np.random.set_state(rnd_state) meta_bal['TARGETID'] = metadata['MOCKID'] w = meta_bal['TEMPLATEID']!=-1 meta_bal = meta_bal[:][w] hdu_bal=pyfits.convenience.table_to_hdu(meta_bal); hdu_bal.name="BAL_META" del meta_bal else: balstr=str(args.balprob) log.error("BAL probability is not between 0 and 1 : "+balstr) sys.exit(1) # Multiply quasar continua by transmitted flux fraction # (at this point transmission file might include Ly-beta, metals and DLAs) log.info("Apply transmitted flux fraction") if not args.no_transmission: tmp_qso_flux = apply_lya_transmission(tmp_qso_wave,tmp_qso_flux, trans_wave,transmission) # if requested, compute metal transmission on the fly # (if not included already from the transmission file) if args.metals is not None: if args.metals_from_file: log.error('you cannot add metals twice') raise ValueError('you cannot add metals twice') if args.no_transmission: log.error('you cannot add metals if asking for no-transmission') raise ValueError('can not add metals if using no-transmission') lstMetals = '' for m in args.metals: lstMetals += m+', ' log.info("Apply metals: {}".format(lstMetals[:-2])) tmp_qso_flux = apply_metals_transmission(tmp_qso_wave,tmp_qso_flux, trans_wave,transmission,args.metals) # if requested, compute magnitudes and apply target selection. Need to do # this calculation separately for QSOs in the north vs south. bbflux=None if args.target_selection or args.bbflux : bands=['FLUX_G','FLUX_R','FLUX_Z', 'FLUX_W1', 'FLUX_W2'] bbflux=dict() bbflux['SOUTH'] = is_south(metadata['DEC']) for band in bands: bbflux[band] = np.zeros(nqso) # need to recompute the magnitudes to account for lya transmission log.info("Compute QSO magnitudes") for these, filters in zip( (~bbflux['SOUTH'], bbflux['SOUTH']), (bassmzls_and_wise_filters, decam_and_wise_filters) ): if np.count_nonzero(these) > 0: maggies = filters.get_ab_maggies(1e-17 * tmp_qso_flux[these, :], tmp_qso_wave) for band, filt in zip( bands, maggies.colnames ): bbflux[band][these] = np.ma.getdata(1e9 * maggies[filt]) # nanomaggies if args.target_selection : log.info("Apply target selection") isqso = np.ones(nqso, dtype=bool) for these, issouth in zip( (~bbflux['SOUTH'], bbflux['SOUTH']), (False, True) ): if np.count_nonzero(these) > 0: # optical cuts only if using QSO vs SIMQSO isqso[these] &= isQSO_colors(gflux=bbflux['FLUX_G'][these], rflux=bbflux['FLUX_R'][these], zflux=bbflux['FLUX_Z'][these], w1flux=bbflux['FLUX_W1'][these], w2flux=bbflux['FLUX_W2'][these], south=issouth, optical=args.no_simqso) log.info("Target selection: {}/{} QSOs selected".format(np.sum(isqso),nqso)) selection=np.where(isqso)[0] if selection.size==0 : return tmp_qso_flux = tmp_qso_flux[selection] metadata = metadata[:][selection] meta = meta[:][selection] qsometa = qsometa[:][selection] DZ_FOG = DZ_FOG[selection] for band in bands : bbflux[band] = bbflux[band][selection] nqso = selection.size log.info("Resample to a linear wavelength grid (needed by DESI sim.)") # careful integration of bins, not just a simple interpolation qso_wave=np.linspace(args.wmin,args.wmax,int((args.wmax-args.wmin)/args.dwave)+1) qso_flux=np.zeros((tmp_qso_flux.shape[0],qso_wave.size)) for q in range(tmp_qso_flux.shape[0]) : qso_flux[q]=resample_flux(qso_wave,tmp_qso_wave,tmp_qso_flux[q]) log.info("Simulate DESI observation and write output file") if "MOCKID" in metadata.dtype.names : #log.warning("Using MOCKID as TARGETID") targetid=np.array(metadata["MOCKID"]).astype(int) elif "ID" in metadata.dtype.names : log.warning("Using ID as TARGETID") targetid=np.array(metadata["ID"]).astype(int) else : log.warning("No TARGETID") targetid=None specmeta={"HPXNSIDE":nside,"HPXPIXEL":pixel, "HPXNEST":hpxnest} if args.target_selection or args.bbflux : fibermap_columns = dict( FLUX_G = bbflux['FLUX_G'], FLUX_R = bbflux['FLUX_R'], FLUX_Z = bbflux['FLUX_Z'], FLUX_W1 = bbflux['FLUX_W1'], FLUX_W2 = bbflux['FLUX_W2'], ) photsys = np.full(len(bbflux['FLUX_G']), 'N', dtype='S1') photsys[bbflux['SOUTH']] = b'S' fibermap_columns['PHOTSYS'] = photsys else : fibermap_columns=None # Attenuate the spectra for extinction if not sfdmap is None: Rv=3.1 #set by default indx=np.arange(metadata['RA'].size) extinction =Rv*ext_odonnell(qso_wave) EBV = sfdmap.ebv(metadata['RA'],metadata['DEC'], scaling=1.0) qso_flux *=10**( -0.4 * EBV[indx, np.newaxis] * extinction) if fibermap_columns is not None: fibermap_columns['EBV']=EBV EBV0=0.0 EBV_med=np.median(EBV) Ag = 3.303 * (EBV_med - EBV0) exptime_fact=np.power(10.0, (2.0 * Ag / 2.5)) obsconditions['EXPTIME']*=exptime_fact log.info("Dust extinction added") log.info('exposure time adjusted to {}'.format(obsconditions['EXPTIME'])) sim_spectra(qso_wave,qso_flux, args.program, obsconditions=obsconditions,spectra_filename=ofilename, sourcetype="qso", skyerr=args.skyerr,ra=metadata["RA"],dec=metadata["DEC"],targetid=targetid, meta=specmeta,seed=seed,fibermap_columns=fibermap_columns,use_poisson=False) # use Poisson = False to get reproducible results. ### Keep input redshift Z_spec = metadata['Z'].copy() Z_input = metadata['Z'].copy()-DZ_FOG ### Add a shift to the redshift, simulating the systematic imprecision of redrock DZ_sys_shift = args.shift_kms_los/c*(1.+Z_input) log.info('Added a shift of {} km/s to the redshift'.format(args.shift_kms_los)) meta['REDSHIFT'] += DZ_sys_shift metadata['Z'] += DZ_sys_shift ### Add a shift to the redshift, simulating the statistic imprecision of redrock if args.gamma_kms_zfit: log.info("Added zfit error with gamma {} to zbest".format(args.gamma_kms_zfit)) DZ_stat_shift = mod_cauchy(loc=0,scale=args.gamma_kms_zfit,size=nqso,cut=3000)/c*(1.+Z_input) meta['REDSHIFT'] += DZ_stat_shift metadata['Z'] += DZ_stat_shift ## Write the truth file, including metadata for DLAs and BALs log.info('Writing a truth file {}'.format(truth_filename)) meta.rename_column('REDSHIFT','Z') meta.add_column(Column(Z_spec,name='TRUEZ')) meta.add_column(Column(Z_input,name='Z_INPUT')) meta.add_column(Column(DZ_FOG,name='DZ_FOG')) meta.add_column(Column(DZ_sys_shift,name='DZ_SYS')) if args.gamma_kms_zfit: meta.add_column(Column(DZ_stat_shift,name='DZ_STAT')) if 'Z_noRSD' in metadata.dtype.names: meta.add_column(Column(metadata['Z_noRSD'],name='Z_NORSD')) else: log.info('Z_noRSD field not present in transmission file. Z_NORSD not saved to truth file') hdu = pyfits.convenience.table_to_hdu(meta) hdu.header['EXTNAME'] = 'TRUTH' hduqso=pyfits.convenience.table_to_hdu(qsometa) hduqso.header['EXTNAME'] = 'QSO_META' hdulist=pyfits.HDUList([pyfits.PrimaryHDU(),hdu,hduqso]) if args.dla: hdulist.append(hdu_dla) if args.balprob: hdulist.append(hdu_bal) hdulist.writeto(truth_filename, overwrite=True) hdulist.close() if args.zbest : log.info("Read fibermap") fibermap = read_fibermap(ofilename) log.info("Writing a zbest file {}".format(zbest_filename)) columns = [ ('CHI2', 'f8'), ('COEFF', 'f8' , (4,)), ('Z', 'f8'), ('ZERR', 'f8'), ('ZWARN', 'i8'), ('SPECTYPE', (str,96)), ('SUBTYPE', (str,16)), ('TARGETID', 'i8'), ('DELTACHI2', 'f8'), ('BRICKNAME', (str,8))] zbest = Table(np.zeros(nqso, dtype=columns)) zbest['CHI2'][:] = 0. zbest['Z'][:] = metadata['Z'] zbest['ZERR'][:] = 0. zbest['ZWARN'][:] = 0 zbest['SPECTYPE'][:] = 'QSO' zbest['SUBTYPE'][:] = '' zbest['TARGETID'][:] = metadata['MOCKID'] zbest['DELTACHI2'][:] = 25. hzbest = pyfits.convenience.table_to_hdu(zbest); hzbest.name='ZBEST' hfmap = pyfits.convenience.table_to_hdu(fibermap); hfmap.name='FIBERMAP' hdulist =pyfits.HDUList([pyfits.PrimaryHDU(),hzbest,hfmap]) hdulist.writeto(zbest_filename, overwrite=True) hdulist.close() # see if this helps with memory issue
def read_gaia_file(filename, header=False, addobjid=False): """Read in a Gaia healpix file in the appropriate format for desitarget. Parameters ---------- filename : :class:`str` File name of a single Gaia "healpix-" file. header : :class:`bool`, optional, defaults to ``False`` If ``True`` then return (data, header) instead of just data. addobjid : :class:`bool`, optional, defaults to ``False`` Include, in the output, two additional columns. A column "GAIA_OBJID" that is the integer number of each row read from file and a column "GAIA_BRICKID" that is the integer number of the file itself. Returns ------- :class:`~numpy.ndarray` Gaia data translated to targeting format (upper-case etc.) with the columns corresponding to `desitarget.gaiamatch.gaiadatamodel` Notes ----- - A better location for this might be in `desitarget.io`? """ # ADM check for an epic fail on the the version of fitsio. check_fitsio_version() # ADM prepare to read in the Gaia data by reading in columns. fx = fitsio.FITS(filename, upper=True) fxcolnames = fx[1].get_colnames() hdr = fx[1].read_header() # ADM the default list of columns. readcolumns = list(ingaiadatamodel.dtype.names) # ADM read 'em in. outdata = fx[1].read(columns=readcolumns) # ADM change the data model to what we want for each column. outdata.dtype.names = gaiadatamodel.dtype.names # ADM the proper motion ERRORS need to be converted to IVARs. # ADM remember to leave 0 entries as 0. for col in ['PMRA_IVAR', 'PMDEC_IVAR', 'PARALLAX_IVAR']: w = np.where(outdata[col] != 0)[0] outdata[col][w] = 1. / (outdata[col][w]**2.) # ADM if requested, add an object identifier for each file row. if addobjid: newdt = outdata.dtype.descr for tup in ('GAIA_BRICKID', '>i4'), ('GAIA_OBJID', '>i4'): newdt.append(tup) nobjs = len(outdata) newoutdata = np.zeros(nobjs, dtype=newdt) for col in outdata.dtype.names: newoutdata[col] = outdata[col] newoutdata['GAIA_OBJID'] = np.arange(nobjs) nside = _get_gaia_nside() hpnum = radec2pix(nside, outdata["GAIA_RA"], outdata["GAIA_DEC"]) # ADM int should fail if HEALPix in the file aren't unique. newoutdata['GAIA_BRICKID'] = int(np.unique(hpnum)) outdata = newoutdata # ADM return data from the Gaia file, with the header if requested. if header: fx.close() return outdata, hdr else: fx.close() return outdata
def integration_test(night=None, nspec=5, clobber=False): """Run an integration test from raw data simulations through redshifts Args: night (str, optional): YEARMMDD, defaults to current night nspec (int, optional): number of spectra to include clobber (bool, optional): rerun steps even if outputs already exist Raises: RuntimeError if any script fails """ import argparse parser = argparse.ArgumentParser(usage = "{prog} [options]") # parser.add_argument("-i", "--input", type=str, help="input data") # parser.add_argument("-o", "--output", type=str, help="output data") parser.add_argument("--skip-psf", action="store_true", help="Skip PSF fitting step") args = parser.parse_args() log = logging.get_logger() # YEARMMDD string, rolls over at noon not midnight if night is None: night = "20160726" # check for required environment variables check_env() # simulate inputs sim(night, nspec=nspec, clobber=clobber) # raw and production locations rawdir = os.path.abspath(io.rawdata_root()) proddir = os.path.abspath(io.specprod_root()) # create production if clobber and os.path.isdir(proddir): shutil.rmtree(proddir) dbfile = io.get_pipe_database() if not os.path.exists(dbfile): com = "desi_pipe create --db-sqlite" log.info('Running {}'.format(com)) sp.check_call(com, shell=True) else: log.info("Using pre-existing production database {}".format(dbfile)) # Modify options file to restrict the spectral range optpath = os.path.join(proddir, "run", "options.yaml") opts = pipe.prod.yaml_read(optpath) opts['extract']['specmin'] = 0 opts['extract']['nspec'] = nspec opts['psf']['specmin'] = 0 opts['psf']['nspec'] = nspec opts['traceshift']['nfibers'] = nspec pipe.prod.yaml_write(optpath, opts) if args.skip_psf: #- Copy desimodel psf into this production instead of fitting psf import shutil for channel in ['b', 'r', 'z']: refpsf = '{}/data/specpsf/psf-{}.fits'.format( os.getenv('DESIMODEL'), channel) nightpsf = io.findfile('psfnight', night, camera=channel+'0') shutil.copy(refpsf, nightpsf) for expid in [0,1,2]: exppsf = io.findfile('psf', night, expid, camera=channel+'0') shutil.copy(refpsf, exppsf) #- Resync database to current state dbpath = io.get_pipe_database() db = pipe.load_db(dbpath, mode="w") db.sync(night) # Run the pipeline tasks in order from desispec.pipeline.tasks.base import default_task_chain for tasktype in default_task_chain: #- if we skip psf/psfnight/traceshift, update state prior to extractions if tasktype == 'traceshift' and args.skip_psf: db.getready() run_pipeline_step(tasktype) # #----- # #- Did it work? # #- (this combination of fibermap, simspec, and zbest is a pain) expid = 2 fmfile = io.findfile('fibermap', night=night, expid=expid) fibermap = io.read_fibermap(fmfile) simdir = os.path.dirname(fmfile) simspec = '{}/simspec-{:08d}.fits'.format(simdir, expid) siminfo = fits.getdata(simspec, 'TRUTH') try: elginfo = fits.getdata(simspec, 'TRUTH_ELG') except: elginfo = None from desimodel.footprint import radec2pix nside=64 pixels = np.unique(radec2pix(nside, fibermap['TARGET_RA'], fibermap['TARGET_DEC'])) num_missing = 0 for pix in pixels: zfile = io.findfile('zbest', groupname=pix) if not os.path.exists(zfile): log.error('Missing {}'.format(zfile)) num_missing += 1 if num_missing > 0: log.critical('{} zbest files missing'.format(num_missing)) sys.exit(1) print() print("--------------------------------------------------") print("Pixel True z -> Class z zwarn") # print("3338p190 SKY 0.00000 -> QSO 1.60853 12 - ok") for pix in pixels: zfile = io.findfile('zbest', groupname=pix) if not os.path.exists(zfile): log.error('Missing {}'.format(zfile)) continue zfx = fits.open(zfile, memmap=False) zbest = zfx['ZBEST'].data for i in range(len(zbest['Z'])): objtype = zbest['SPECTYPE'][i] z, zwarn = zbest['Z'][i], zbest['ZWARN'][i] j = np.where(fibermap['TARGETID'] == zbest['TARGETID'][i])[0][0] truetype = siminfo['OBJTYPE'][j] oiiflux = 0.0 if truetype == 'ELG': k = np.where(elginfo['TARGETID'] == zbest['TARGETID'][i])[0][0] oiiflux = elginfo['OIIFLUX'][k] truez = siminfo['REDSHIFT'][j] dv = C_LIGHT*(z-truez)/(1+truez) status = None if truetype == 'SKY' and zwarn > 0: status = 'ok' elif truetype == 'ELG' and zwarn > 0 and oiiflux < 8e-17: status = 'ok ([OII] flux {:.2g})'.format(oiiflux) elif zwarn == 0: if truetype == 'LRG' and objtype == 'GALAXY' and abs(dv) < 150: status = 'ok' elif truetype == 'ELG' and objtype == 'GALAXY': if abs(dv) < 150: status = 'ok' elif oiiflux < 8e-17: status = 'ok ([OII] flux {:.2g})'.format(oiiflux) else: status = 'OOPS ([OII] flux {:.2g})'.format(oiiflux) elif truetype == 'QSO' and objtype == 'QSO' and abs(dv) < 750: status = 'ok' elif truetype in ('STD', 'FSTD') and objtype == 'STAR': status = 'ok' else: status = 'OOPS' else: status = 'OOPS' print('{0:<8d} {1:4s} {2:8.5f} -> {3:5s} {4:8.5f} {5:4d} - {6}'.format( pix, truetype, truez, objtype, z, zwarn, status)) print("--------------------------------------------------")
def qa_targets_truth(output_dir, verbose=True): """Generate QA plots from the joined targets and truth catalogs. time select_mock_targets --output_dir debug --qa """ import fitsio from desiutil.log import get_logger, DEBUG from desitarget import desi_mask, bgs_mask, mws_mask, contam_mask if verbose: log = get_logger(DEBUG) else: log = get_logger() qadir = os.path.join(output_dir, 'qa') if os.path.exists(qadir): if os.listdir(qadir): log.warning('Output directory {} is not empty.'.format(qadir)) else: log.info('Creating directory {}'.format(qadir)) os.makedirs(qadir) log.info('Writing to output QA directory {}'.format(qadir)) # Read the catalogs. targfile = os.path.join(output_dir, 'targets.fits') truthfile = os.path.join(output_dir, 'truth.fits') skyfile = os.path.join(output_dir, 'sky.fits') stddarkfile = os.path.join(output_dir, 'standards-dark.fits') stdbrightfile = os.path.join(output_dir, 'standards-bright.fits') cat = list() for ff in (targfile, truthfile, skyfile, stddarkfile, stdbrightfile): if os.path.exists(ff): log.info('Reading {}'.format(ff)) cat.append(fitsio.read(ff, ext=1, upper=True)) else: log.warning('File {} not found.'.format(ff)) cat.append(None) targets, truth, sky, stddark, stdbright = [cc for cc in cat] # Do some sanity checking of the catalogs. nobj, nsky, ndark, nbright = [ len(cc) for cc in (targets, sky, stddark, stdbright) ] if nobj != len(truth): log.fatal( 'Mismatch in the number of objects in targets.fits (N={}) and truth.fits (N={})!' .format(nobj, len(truth))) raise ValueError # Assign healpixels to estimate the area covered by the catalog. nside = 256 npix = hp.nside2npix(nside) areaperpix = hp.nside2pixarea(nside, degrees=True) pix = radec2pix(nside, targets['RA'], targets['DEC']) counts = np.bincount(pix, weights=None, minlength=npix) area = np.sum(counts > 10) * areaperpix log.info('Approximate area spanned by catalog = {:.2f} deg2'.format(area)) binarea = 1.0 htmlfile = os.path.join(qadir, 'index.html') log.info('Building {}'.format(htmlfile)) html = open(htmlfile, 'w') html.write('<html><body>\n') html.write('<h1>QA directory: {}</h1>\n'.format(qadir)) html.write('<ul>\n') html.write('<li>Approximate total area = {:.1f} deg2</li>\n'.format(area)) html.write('<li>Science targets = {}</li>\n'.format(nobj)) html.write('<li>Sky targets = {}</li>\n'.format(nsky)) html.write('<li>Dark standards = {}</li>\n'.format(ndark)) html.write('<li>Bright standards = {}</li>\n'.format(nbright)) html.write('</ul>\n') # Desired target densities, including contaminants. html.write('<hr>\n') html.write('<hr>\n') html.write('<h2>Raw target densities (including contaminants)</h2>\n') targdens = {'ELG': 2400, 'LRG': 350, 'QSO': 260, 'SKY': 1400} if nobj > 0: html.write( '<h3>Science targets - ELG, LRG, QSO, BGS_ANY, MWS_ANY</h3>\n') for obj in ('ELG', 'LRG', 'QSO', 'BGS_ANY', 'MWS_ANY'): these = np.where( (targets['DESI_TARGET'] & desi_mask.mask(obj)) != 0)[0] if len(these) > 0: pngfile = qadensity(targets[these], obj, targdens, qadir=qadir, max_bin_area=binarea) _img2html(html, pngfile, log) html.write('<h3>Science targets - BGS_BRIGHT, BGS_FAINT</h3>\n') for obj in ('BGS_BRIGHT', 'BGS_FAINT'): these = np.where( (targets['BGS_TARGET'] & bgs_mask.mask(obj)) != 0)[0] if len(these) > 0: pngfile = qadensity(targets[these], obj, targdens, qadir=qadir, max_bin_area=binarea) _img2html(html, pngfile, log) html.write( '<h3>Science targets - MWS_MAIN, MWS_MAIN_VERY_FAINT, MWS_NEARBY, MWS_WD</h3>\n' ) for obj in ('MWS_MAIN', 'MWS_MAIN_VERY_FAINT', 'MWS_NEARBY', 'MWS_WD'): these = np.where( (targets['MWS_TARGET'] & mws_mask.mask(obj)) != 0)[0] if len(these) > 0: pngfile = qadensity(targets[these], obj, targdens, qadir=qadir, max_bin_area=binarea) _img2html(html, pngfile, log) html.write('<hr>\n') if nsky > 0: html.write('<h3>Sky targets</h3>\n') obj = 'SKY' these = np.where((sky['DESI_TARGET'] & desi_mask.mask(obj)) != 0)[0] if len(these) > 0: pngfile = qadensity(sky[these], obj, targdens, qadir=qadir, max_bin_area=binarea) _img2html(html, pngfile, log) html.write('<hr>\n') if ndark > 0 or nbright > 0: html.write('<h3>Standard Stars</h3>\n') for cat, nn, obj in zip((stddark, stdbright), (ndark, nbright), ('STD_FSTAR', 'STD_BRIGHT')): if nn > 0: these = np.where( (cat['DESI_TARGET'] & desi_mask.mask(obj)) != 0)[0] if len(these) > 0: pngfile = qadensity(cat[these], obj, targdens, qadir=qadir, max_bin_area=binarea) _img2html(html, pngfile, log) html.write('<hr>\n') html.write('<hr>\n') # Desired target densities, including contaminants. html.write('<h2>Detailed target densities</h2>\n') html.write('<h3>QSOs</h3>\n') pngfile = qa_qso(targets, truth, qadir=qadir) _img2html(html, pngfile, log) html.write('<hr>\n') html.write('<h3>ELGs</h3>\n') pngfile = qa_elg(targets, truth, qadir=qadir) _img2html(html, pngfile, log) html.write('<hr>\n') html.write('</html></body>\n') html.close()