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()
