In [1]:
import os
import numpy as np
import matplotlib.pyplot as plt
from astropy.io import fits
from astropy.table import Table, vstack
import seaborn as sns
from desispec import io
from desispec.log import get_logger
from desispec.io.util import write_bintable, makepath
%matplotlib inline
sns.set(style='ticks', font_scale=1.5, palette='Set2')
col = sns.color_palette()
In [2]:
# Set the paths.
desidir = os.getenv('DESI_ROOT')
simsdir = os.path.join(desidir, 'spectro', 'sim', 'desi2')
brickdir = os.path.join(simsdir, 'bricks')
log = get_logger()
In [3]:
seed = 678245
brickroot = 'sim01'
zenith = 180
nspec = 10
# Mapping from limiting r-band magnitude and median redshift is
# from cosmos_zmedian_vs_rmag.fits from Risa.
exptime = np.array([5.0, 8.0, 10.0, 15.0, 20.0, 30.0, 40.0, 60.0])*60.0
rmag = [(20.0, 20.0), (20.5, 20.5), (21.0, 21.0), (21.5, 21.5), (22.0, 22.0), (22.5, 22.5)]
redshift = [0.2, 0.25, 0.25, 0.28, 0.33, 0.35]
nexp = len(exptime)
nrmag = np.shape(rmag)[0]
nbrick = nexp * nrmag
nobj = nbrick*nspec
rand = np.random.RandomState(seed)
brickseed = rand.randint(2**32, size=nbrick)
In [4]:
def make_bricks(objtype=BGS):
'''Generate the bricks.'''
brickcount = 0
for ie in range(nexp):
for ir in range(nrmag):
brickname = '{}-{:03d}'.format(brickroot, brickcount)
brightoptions = [
'--brickname', '{}'.format(brickname),
'--nbrick', '{}'.format(1),
'--nspec', '{}'.format(nspec),
'--outdir', '{}'.format(simsdir),
'--brickdir', '{}'.format(brickdir),
'--seed', '{}'.format(brickseed[brickcount]),
'--objtype', objtype]
simoptions = [
'--exptime-range', '{}'.format(exptime[ie]), '{}'.format(exptime[ie]),
'--rmagrange-bgs', '{}'.format(rmag[ir][0]), '{}'.format(rmag[ir][1]),
'--zrange-bgs', '{}'.format(redshift[ir]), '{}'.format(redshift[ir]),
'--moon-zenith-range', '{}'.format(zenith), '{}'.format(zenith)]
brightargs = brightsims.parse(np.hstack((brightoptions, simoptions)))
brightargs.verbose = True
brightsims.main(brightargs)
brickcount += 1
In [5]:
def parse_results():
'''Parse the results.'''
for ib in range(nbrick):
inputfile = os.path.join(simsdir, brickroot+'-{:03d}-input.fits'.format(ib))
#log.info('Reading {}'.format(inputfile))
cat1 = Table(fits.getdata(inputfile, 1))
if ib == 0:
cat = cat1
else:
cat = vstack((cat, cat1))
# Build a results table.
resultfile = makepath(os.path.join(simsdir, '{}-results.fits'.format(brickroot)))
resultcols = [
('BRICKNAME', 'S20'),
('EXPTIME', 'f4'),
('AIRMASS', 'f4'),
('SNR_B', 'f4'),
('SNR_R', 'f4'),
('SNR_Z', 'f4'),
('TARGETID', 'i8'),
('RMAG', 'f4'),
('D4000', 'f4'),
('EWHBETA', 'f4'),
('ZTRUE', 'f4'),
('Z', 'f4'),
('ZERR', 'f4'),
('ZWARNING', 'f4')]
result = Table(np.zeros(nobj, dtype=resultcols))
avgresultfile = makepath(os.path.join(simsdir, '{}-avgresults.fits'.format(brickroot)))
avgresultcols = [
('BRICKNAME', 'S20'),
('EXPTIME', 'f4'),
('AIRMASS', 'f4'),
('SNR_B', 'f4'),
('SNR_R', 'f4'),
('SNR_Z', 'f4'),
('RMAG', 'f4')]
avgresult = Table(np.zeros(nbrick, dtype=avgresultcols))
for ib in range(nbrick):
# Copy over some data.
thisbrick = cat['BRICKNAME'][ib]
result['BRICKNAME'][nspec*ib:nspec*(ib+1)] = thisbrick
result['EXPTIME'][nspec*ib:nspec*(ib+1)] = cat['EXPTIME'][ib]
result['AIRMASS'][nspec*ib:nspec*(ib+1)] = cat['AIRMASS'][ib]
avgresult['BRICKNAME'][ib] = thisbrick
avgresult['EXPTIME'][ib] = cat['EXPTIME'][ib]
avgresult['AIRMASS'][ib] = cat['AIRMASS'][ib]
# Read the truth file of the first channel to get the metadata.
truthfile = os.path.join(brickdir, thisbrick, 'truth-brick-{}-{}.fits'.format('b', thisbrick))
log.info('Reading {}'.format(truthfile))
truth = io.Brick(truthfile).hdu_list[4].data
result['TARGETID'][nspec*ib:nspec*(ib+1)] = truth['TARGETID']
result['RMAG'][nspec*ib:nspec*(ib+1)] = 22.5-2.5*np.log10(truth['DECAM_FLUX'][:, 2])
result['D4000'][nspec*ib:nspec*(ib+1)] = truth['D4000']
result['EWHBETA'][nspec*ib:nspec*(ib+1)] = truth['EWHBETA']
result['ZTRUE'][nspec*ib:nspec*(ib+1)] = truth['TRUEZ']
avgresult['RMAG'][ib] = np.mean(22.5-2.5*np.log10(truth['DECAM_FLUX'][:, 2]))
# Finally read the zbest file.
zbestfile = os.path.join(brickdir, thisbrick, 'zbest-{}.fits'.format(thisbrick))
if os.path.isfile(zbestfile):
log.info('Reading {}'.format(zbestfile))
zbest = read_zbest(zbestfile)
# There's gotta be a better way than looping here!
for ii in range(nspec):
this = np.where(zbest.targetid[ii] == result['TARGETID'])[0]
result['Z'][this] = zbest.z[ii]
result['ZERR'][this] = zbest.zerr[ii]
result['ZWARNING'][this] = zbest.zwarn[ii]
# Finally, read the spectra and truth tables, one per channel.
for channel in ('b','r','z'):
brickfile = os.path.join(brickdir, thisbrick, 'brick-{}-{}.fits'.format(channel, thisbrick))
log.info('Reading {}'.format(brickfile))
brick = io.Brick(brickfile)
wave = brick.get_wavelength_grid()
for iobj in range(nspec):
flux = brick.hdu_list[0].data[iobj,:]
ivar = brick.hdu_list[1].data[iobj,:]
these = np.where((wave>np.mean(wave)-50)*(wave<np.mean(wave)+50)*(flux>0))[0]
#print(np.mean(wave[these]))
result['SNR_'+channel.upper()][nspec*ib+iobj] = \
np.median(np.sqrt(flux[these]*ivar[these]))
avgresult['SNR_'+channel.upper()][ib] = np.mean(
result['SNR_'+channel.upper()][nspec*ib:nspec*(ib+1)])
log.info('Writing {}'.format(resultfile))
write_bintable(resultfile, result, extname='RESULTS', clobber=True)
log.info('Writing {}'.format(avgresultfile))
write_bintable(avgresultfile, avgresult, extname='AVGRESULTS', clobber=True)
In [ ]:
# Build the bricks.
make_bricks()
In [8]:
# Parse the results.
parse_results()
In [6]:
# Make some plots.
avgres = fits.getdata(os.path.join(simsdir, 'sim01-avgresults.fits'), 1)
avgres.columns
np.unique(avgres['RMAG'])
Out[6]:
In [7]:
fig, ax = plt.subplots(1, 1, figsize = (10, 7))
for ii, rm in enumerate(np.unique(rmag)):
these = np.where(avgres['RMAG'] == rm)[0]
ax.plot(avgres['EXPTIME'][these]/60, avgres['SNR_R'][these],
marker='s', markersize=10, color=col[ii], label='r<{:.1f}'.format(rm))
ax.set_xlabel('Exposure Time (m)', fontsize=20)
ax.set_ylabel('Median S/N (pixel$^{-1}$ around 6700 $\AA$)', fontsize=20)
ax.legend(loc='upper left')
ax.set_xlim(4, 61)
Out[7]:
In [15]:
from scipy.interpolate import interp1d
for ii, rm in enumerate(np.unique(rmag)):
these = np.where(avgres['RMAG'] == rm)[0]
ff = interp1d(avgres['SNR_R'][these], avgres['EXPTIME'][these],
fill_value=0, bounds_error=False)
print(rm, ff(1.0)/60, ff(10)/60)
In [ ]: