"""
desisim.scripts.quicktransients
===============================
"""
import os
import numpy as np
import healpy as hp
from datetime import datetime
from desisim.templates import BGS, ELG, LRG
from desisim.transients import transients
from desitarget.mock.mockmaker import BGSMaker, ELGMaker, LRGMaker # requires desitarget<=3.2.0
from desitarget.cuts import isBGS_colors, isELG_colors, isLRG_colors
from desisim.simexp import reference_conditions
from desisim.transients import transients
from desisim.scripts.quickspectra import sim_spectra
from desispec.io import read_spectra, write_spectra
from desispec.coaddition import coadd_cameras
from desiutil.log import get_logger, DEBUG
from astropy.table import Table, hstack, vstack
import argparse
[docs]
def _set_wave(wavemin=None, wavemax=None, dw=0.8):
"""Set default wavelength grid for simulations.
Parameters
----------
wavemin : float or None
Minimum wavelength, in Angstroms.
wavemax : float or None
Maximum wavelength, in Angstroms.
dw : float
Bin size.
Returns
-------
wave : ndarray
Grid of wavelength values.
"""
from desimodel.io import load_throughput
if wavemin is None:
wavemin = load_throughput('b').wavemin - 10.0
if wavemax is None:
wavemax = load_throughput('z').wavemax + 10.0
return np.arange(round(wavemin, 1), wavemax, dw)
[docs]
def write_templates(filename, flux, wave, target, truth, objtruth):
"""Write galaxy templates to a FITS file.
Parameters
----------
filename : str
Path to output file.
flux : ndarray
Array of flux data for template spectra.
wave : ndarray
Array of wavelengths.
target : Table
Target information.
truth : Table
Template simulation truth.
objtruth : Table
Object-specific truth data.
"""
import astropy.units as u
from astropy.io import fits
hx = fits.HDUList()
# Write the wavelength table.
hdu_wave = fits.PrimaryHDU(wave)
hdu_wave.header['EXTNAME'] = 'WAVE'
hdu_wave.header['BUNIT'] = 'Angstrom'
hdu_wave.header['AIRORVAC'] = ('vac', 'Vacuum wavelengths')
hx.append(hdu_wave)
# Write the flux table.
fluxunits = 1e-17 * u.erg / (u.s * u.cm**2 * u.Angstrom)
hdu_flux = fits.ImageHDU(flux)
hdu_flux.header['EXTNAME'] = 'FLUX'
hdu_flux.header['BUNIT'] = str(fluxunits)
hx.append(hdu_flux)
# Write targets table.
hdu_targets = fits.table_to_hdu(target)
hdu_targets.header['EXTNAME'] = 'TARGETS'
hx.append(hdu_targets)
# Write truth table.
hdu_truth = fits.table_to_hdu(truth)
hdu_truth.header['EXTNAME'] = 'TRUTH'
hx.append(hdu_truth)
# Write objtruth table.
hdu_objtruth = fits.table_to_hdu(objtruth)
hdu_objtruth.header['EXTNAME'] = 'OBJTRUTH'
hx.append(hdu_objtruth)
hx.writeto(filename, overwrite=True)
[docs]
def parse(options=None):
"""Parse command line options.
"""
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
description='Fast galaxy + transient simulator')
# Set up observing conditions.
cond = parser.add_argument_group('Observing conditions')
for c, v in reference_conditions['DARK'].items():
cond.add_argument('--{}'.format(c.lower()), type=float,
default=v)
# Set up simulation.
sims = parser.add_argument_group('Simulation settings')
sims.add_argument('--nside', type=int, default=64,
help='HEALPix NSIDE')
sims.add_argument('--nspec', type=int, default=100,
help='Spectra per healpixel')
sims.add_argument('--nsim', type=int, default=10,
help='Number of simulations')
sims.add_argument('--seed', type=int, default=None,
help='RNG seed')
sims.add_argument('--host', choices=['bgs','elg','lrg'], default='bgs',
help='Host galaxy type')
sims.add_argument('--outdir', default='',
help='Absolute path to simulation output')
sims.add_argument('--coadd_cameras', action='store_true',
help='If true, coadd cameras in generated spectra')
# Set up transient model as a simulation setting. None==no transient.
tran = parser.add_argument_group('Transient settings')
tran.add_argument('--magrange', nargs=2, type=float, default=[0.,2.5],
help='Transient-host mag range; e.g., 0 5')
tran.add_argument('--epochrange', nargs=2, type=float, default=[-10.,10.],
help='Epoch range w.r.t. t0 in days; e.g., -10 10')
modelnames = []
mdict = transients.get_type_dict()
for t, models in mdict.items():
for m in models:
modelnames.append(m)
tran.add_argument('--transient', default=None,
help='None or one of these models: {}'.format(transients.get_type_dict()))
if options is None:
args = parser.parse_args()
else:
args = parser.parse_args(options)
return args
def main(args=None):
log = get_logger()
if isinstance(args, (list, tuple, type(None))):
args = parse(args)
thedate = datetime.now().strftime('%Y-%m-%d')
# Generate transient model if one is specified.
trans_model = None
if args.transient is not None:
trans_model = transients.get_model(args.transient)
# Generate list of HEALPix pixels to randomly sample the mocks.
rng = np.random.RandomState(args.seed)
nside = args.nside
healpixels = _get_healpixels_in_footprint(nside=args.nside)
npix = np.minimum(10*args.nsim, len(healpixels))
pixels = rng.choice(healpixels, size=npix, replace=False)
ipix = iter(pixels)
# Set up the template generator.
fluxratio_range = 10**(-np.sort(args.magrange)[::-1] / 2.5)
epoch_range = np.sort(args.epochrange)
if args.host == 'bgs':
maker = BGSMaker(seed=args.seed)
tmpl_maker = BGS
elif args.host == 'elg':
maker = ELGMaker(seed=args.seed)
tmpl_maker = ELG
elif args.host == 'lrg':
maker = LRGMaker(seed=args.seed)
tmpl_maker = LRG
else:
raise ValueError('Unusable host type {}'.format(args.host))
maker.template_maker = tmpl_maker(transient=trans_model,
tr_fluxratio=fluxratio_range,
tr_epoch=epoch_range)
for j in range(args.nsim):
# Loop until finding a non-empty healpixel with mock galaxies.
tdata = []
while len(tdata) == 0:
pixel = next(ipix)
tdata = maker.read(healpixels=pixel, nside=args.nside)
# Generate spectral templates and write them to truth files.
# Keep producing templates until we have enough to pass brightness cuts.
wave = None
flux, targ, truth, objtr = [], [], [], []
ntosim = np.min([args.nspec, len(tdata['RA'])])
ngood = 0
while ngood < args.nspec:
idx = rng.choice(len(tdata['RA']), ntosim)
tflux, twave, ttarg, ttruth, tobj = maker.make_spectra(tdata, indx=idx)
g, r, z, w1, w2 = [ttruth['FLUX_{}'.format(_)] for _ in ['G','R','Z','W1','W2']]
rfib = ttarg['FIBERFLUX_R']
# print(g, r, z, w1, w2, rfib)
# Apply color cuts.
is_bright = isBGS_colors(rfib, g, r, z, w1, w2, targtype='bright')
is_faint = isBGS_colors(rfib, g, r, z, w1, w2, targtype='faint')
is_wise = isBGS_colors(rfib, g, r, z, w1, w2, targtype='wise')
keep = np.logical_or.reduce([is_bright, is_faint, is_wise])
_ngood = np.count_nonzero(keep)
if _ngood > 0:
ngood += _ngood
flux.append(tflux[keep, :])
targ.append(ttarg[keep])
truth.append(ttruth[keep])
objtr.append(tobj[keep])
wave = maker.wave
flux = np.vstack(flux)[:args.nspec, :]
targ = vstack(targ)[:args.nspec]
truth = vstack(truth)[:args.nspec]
objtr = vstack(objtr)[:args.nspec]
# Set up and verify the TARGETID across all truth tables.
n = len(truth)
new_id = 10000*pixel + 100*j + np.arange(1, n+1)
targ['OBJID'][:] = new_id
truth['TARGETID'][:] = new_id
objtr['TARGETID'][:] = new_id
assert(len(truth) == args.nspec)
assert(np.all(targ['OBJID'] == truth['TARGETID']))
assert(len(targ) == len(np.unique(targ['OBJID'])))
assert(len(truth) == len(np.unique(truth['TARGETID'])))
assert(len(objtr) == len(np.unique(objtr['TARGETID'])))
truthfile = os.path.join(args.outdir,
'{}_{}_{:04d}s_{:03d}_truth.fits'.format(args.host, thedate, int(args.exptime), j+1))
write_templates(truthfile, flux, wave, targ, truth, objtr)
log.info('Wrote {}'.format(truthfile))
# Get observing conditions and generate spectra.
obs = dict(AIRMASS=args.airmass, EXPTIME=args.exptime,
MOONALT=args.moonalt, MOONFRAC=args.moonfrac,
MOONSEP=args.moonsep, SEEING=args.seeing)
fcols = dict(BRICKID=targ['BRICKID'],
BRICK_OBJID=targ['OBJID'],
FLUX_G=targ['FLUX_G'],
FLUX_R=targ['FLUX_R'],
FLUX_Z=targ['FLUX_Z'],
FLUX_W1=targ['FLUX_W1'],
FLUX_W2=targ['FLUX_W2'],
FLUX_IVAR_G=targ['FLUX_IVAR_G'],
FLUX_IVAR_R=targ['FLUX_IVAR_R'],
FLUX_IVAR_Z=targ['FLUX_IVAR_Z'],
FLUX_IVAR_W1=targ['FLUX_IVAR_W1'],
FLUX_IVAR_W2=targ['FLUX_IVAR_W2'],
FIBERFLUX_G=targ['FIBERFLUX_G'],
FIBERFLUX_R=targ['FIBERFLUX_R'],
FIBERFLUX_Z=targ['FIBERFLUX_Z'],
FIBERTOTFLUX_G=targ['FIBERTOTFLUX_G'],
FIBERTOTFLUX_R=targ['FIBERTOTFLUX_R'],
FIBERTOTFLUX_Z=targ['FIBERTOTFLUX_Z'],
MW_TRANSMISSION_G=targ['MW_TRANSMISSION_G'],
MW_TRANSMISSION_R=targ['MW_TRANSMISSION_R'],
MW_TRANSMISSION_Z=targ['MW_TRANSMISSION_Z'],
EBV=targ['EBV'])
specfile = os.path.join(args.outdir,
'{}_{}_{:04d}s_{:03d}_spect.fits'.format(args.host, thedate, int(args.exptime), j+1))
# redshifts = truth['TRUEZ'] if args.host=='bgs' else None
redshifts = None
sim_spectra(wave, flux, args.host, specfile,
sourcetype=args.host,
obsconditions=obs, meta=obs, fibermap_columns=fcols,
targetid=truth['TARGETID'], redshift=redshifts,
ra=targ['RA'], dec=targ['DEC'],
seed=args.seed, expid=j)
if args.coadd_cameras:
coaddfile = specfile.replace('spect', 'coadd')
spectra = read_spectra(specfile)
spectra = coadd_cameras(spectra)
write_spectra(coaddfile, spectra)
log.info('Wrote {}'.format(coaddfile))