<commit_before>"""
desisim.scripts.quickgen
========================

Quickgen quickly simulates pipeline outputs if given input files or brickname.
This can be run in two distinct modes:

* QUICKGEN MODE

  - must provide simspec and fibermap files via newexp-desi
  - Number of spectra to be simulated can be given as an argument for quickgen,
    but the number of spectra in the simspec file is taken by default
  - For this option, airmass and exposure time are keywords given to newexp-desi
  - The keywords provided in the examples are all required, additional keywords
    are provided below
  - Collect a set of templates to simulate as a new exposure::

        newexp-desi --nspec 500 --night 20150915 --expid 0 --flavor dark

  - newexp-desi keyword arguments::

        --flavor : arc/flat/dark/gray/bright/bgs/mws/elg/lrg/qso, type=str, default='dark'
        --tileid : tile id, type=int
        --expid : exposure id, type=int
        --exptime : exposure time in seconds, default for arc = 5s, flat = 10s, dark/elg/lrg/qso=1000s, bright/bgs/mws=300s, type=int
        --night : YEARMMDD, type=str
        --nspec : Number of spectra to simulate, type=int, default=5000
        --airmass : type=float, default=1.0
        --seed : random number seed, type=int
        --testslit : test slit simulation with fewer fibers, action="store_true"
        --arc-lines : alternate arc lines filename, type=str, default=None
        --flat-spectrum : alternate flat spectrum filename, type=str

  - Actually do the simulation::

        simdir=$DESI_SPECTRO_SIM/$PIXPROD/20150915
        quickgen --simspec $simdir/simspec-00000000.fits --fibermap $simdir/fibermap-00000000.fits

  - quickgen output (can also provide frame file only as keyword)::

        1. frame-{camera}-{expid}.fits : raw extracted photons with no calibration at all
        2. sky-{camera}-{expid}.fits : the sky model in photons
        3. fluxcalib-{camera}-{expid}.fits] : the flux calibration vectors
        4. cframe-{camera}-{expid}.fits : flux calibrated spectra

        These files are written to $simdir/{expid}

* QUICKBRICK MODE

  - must provide brickname
  - Simulates 100 spectra by default, unless nspec keyword is provided
  - For this option, airmass and exposure time are given directly to quickgen
  - Brickname is the only required keyword
  - Simulate bricks (can provide anything as brick name)::

        quickgen --brickname brick

  - quickbrick output::

        1. brick-{arm}-{brickname}.fits : flux calibrated spectra in bricks

        These files are written to the current working directory by default

* additional keyword arguments (descriptions below):

  - common to both modes: nspec, config, seed, moon-phase, moon-angle, moon-zenith
  - quickgen only: simspec, fibermap, nstart, spectrograph, frameonly
  - quickbrick only: brickname, objtype, airmass, exptime, outdir, verbose, outdir-truth,
    zrange-qso, zrange-elg, zrange-lrg, zrange-bgs, sne-rfluxratiorange,
    add-SNeIa
"""
from __future__ import absolute_import, division, print_function

import argparse
import astropy.units as u
from astropy.io import fits
from astropy.table import Table, Column, vstack
from time import asctime

import os
import os.path
import numpy as np
import scipy.special
import scipy.interpolate
import scipy.constants as const
import sys

import desispec
import desispec.io
import desisim
import desisim.io
import desisim.templates
import desiutil.io
from desispec.resolution import Resolution
from desispec.io import write_flux_calibration, write_fiberflat, read_fibermap, specprod_root, fitsheader, empty_fibermap, Brick
from desispec.interpolation import resample_flux
from desispec.frame import Frame
from desispec.fiberflat import FiberFlat
from desispec.sky import SkyModel
from desispec.fluxcalibration import FluxCalib
from desispec.log import get_logger, DEBUG, INFO
from desisim.obs import get_night
from desisim.targets import sample_objtype
from desisim.specsim import get_simulator
from desimodel.io import load_desiparams

def _add_truth(hdus, header, meta, trueflux, sflux, wave, channel):
    """Utility function for adding truth to an output FITS file."""
    hdus.append(
        fits.ImageHDU(trueflux[channel], name='_TRUEFLUX', header=header))
    if channel == 'b':
        swave = wave.astype(np.float32)
        hdus.append(fits.ImageHDU(swave, name='_SOURCEWAVE', header=header))
        hdus.append(fits.ImageHDU(sflux, name='_SOURCEFLUX', header=header))
        metatable = desiutil.io.encode_table(meta, encoding='ascii')
        metahdu = fits.convenience.table_to_hdu(meta)
        metahdu.header['EXTNAME'] = '_TRUTH'
        hdus.append(metahdu)

def parse(options=None):
    parser=argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    parser.add_argument('--simspec',type=str, help="input simspec file")
    parser.add_argument('--fibermap',type=str, help='input fibermap file')
    parser.add_argument('-n','--nspec',type=int,default=100,help='number of spectra to be simulated, starting from first')
    parser.add_argument('--nstart', type=int, default=0,help='starting spectra # for simulation 0-4999')
    parser.add_argument('--spectrograph',type=int, default=None,help='Spectrograph no. 0-9')
    parser.add_argument('--config', type=str, default='desi', help='specsim configuration')
    parser.add_argument('-s','--seed', type=int, default=0,  help="random seed")
    # Only produce uncalibrated output
    parser.add_argument('--frameonly', action="store_true", help="only output frame files")

    # Moon options if bright or gray time
    parser.add_argument('--moon-phase', type=float,  help='moon phase (0=full, 1=new)', default=None, metavar='')
    parser.add_argument('--moon-angle', type=float,  help='separation angle to the moon (0-180 deg)', default=None, metavar='')
    parser.add_argument('--moon-zenith', type=float,  help='zenith angle of the moon (0-90 deg)', default=None, metavar='')

    # Brick options
    parser.add_argument('-b','--brickname', type=str, help='unique output brickname suffix', metavar='')
    parser.add_argument('--objtype', type=str,  help='ELG, LRG, QSO, BGS, MWS, WD, DARK_MIX, or BRIGHT_MIX', default='DARK_MIX', metavar='')
    parser.add_argument('-a','--airmass', type=float,  help='airmass', default=1.25, metavar='') # Science Req. Doc L3.3.2
    parser.add_argument('-e','--exptime', type=float,  help='exposure time (s)', default=None,metavar='')
    parser.add_argument('-o','--outdir', type=str,  help='output directory', default='.', metavar='')
    parser.add_argument('-v','--verbose', action='store_true', help='toggle on verbose output')
    parser.add_argument('--outdir-truth', type=str,  help='optional alternative output directory for truth files', metavar='')

    # Object type specific options
    parser.add_argument('--zrange-qso', type=float, default=(0.5, 4.0), nargs=2, metavar='',
                        help='minimum and maximum redshift range for QSO')
    parser.add_argument('--zrange-elg', type=float, default=(0.6, 1.6), nargs=2, metavar='',
                        help='minimum and maximum redshift range for ELG')
    parser.add_argument('--zrange-lrg', type=float, default=(0.5, 1.1), nargs=2, metavar='',
                        help='minimum and maximum redshift range for LRG')
    parser.add_argument('--zrange-bgs', type=float, default=(0.01, 0.4), nargs=2, metavar='',
                        help='minimum and maximum redshift range for BGS')
    parser.add_argument('--rmagrange-bgs', type=float, default=(15.0, 19.5), nargs=2, metavar='',
                        help='Minimum and maximum BGS r-band (AB) magnitude range')
    parser.add_argument('--sne-rfluxratiorange', type=float, default=(0.1, 1.0), nargs=2, metavar='',
                        help='r-band flux ratio of the SNeIa spectrum relative to the galaxy')
    parser.add_argument('--add-SNeIa', action='store_true', help='include SNeIa spectra')

    if options is None:
        args = parser.parse_args()
    else:
        args = parser.parse_args(options)

    log = get_logger()
    if args.simspec:
        args.objtype = None
        if args.fibermap is None:
            msg = 'If simspec set, must also set fibermap'
            log.error(msg)
            raise ValueError(msg)
        hdr = fits.getheader(args.simspec)
        night = str(hdr['NIGHT'])
        expid = int(hdr['EXPID'])
        args.simspec = desisim.io.findfile('simspec', night, expid)
        args.fibermap = desispec.io.findfile('fibermap', night, expid)
<commit_msg>use the input fibermap filename<commit_after>"""
desisim.scripts.quickgen
========================

Quickgen quickly simulates pipeline outputs if given input files or brickname.
This can be run in two distinct modes:

* QUICKGEN MODE

  - must provide simspec and fibermap files via newexp-desi
  - Number of spectra to be simulated can be given as an argument for quickgen,
    but the number of spectra in the simspec file is taken by default
  - For this option, airmass and exposure time are keywords given to newexp-desi
  - The keywords provided in the examples are all required, additional keywords
    are provided below
  - Collect a set of templates to simulate as a new exposure::

        newexp-desi --nspec 500 --night 20150915 --expid 0 --flavor dark

  - newexp-desi keyword arguments::

        --flavor : arc/flat/dark/gray/bright/bgs/mws/elg/lrg/qso, type=str, default='dark'
        --tileid : tile id, type=int
        --expid : exposure id, type=int
        --exptime : exposure time in seconds, default for arc = 5s, flat = 10s, dark/elg/lrg/qso=1000s, bright/bgs/mws=300s, type=int
        --night : YEARMMDD, type=str
        --nspec : Number of spectra to simulate, type=int, default=5000
        --airmass : type=float, default=1.0
        --seed : random number seed, type=int
        --testslit : test slit simulation with fewer fibers, action="store_true"
        --arc-lines : alternate arc lines filename, type=str, default=None
        --flat-spectrum : alternate flat spectrum filename, type=str

  - Actually do the simulation::

        simdir=$DESI_SPECTRO_SIM/$PIXPROD/20150915
        quickgen --simspec $simdir/simspec-00000000.fits --fibermap $simdir/fibermap-00000000.fits

  - quickgen output (can also provide frame file only as keyword)::

        1. frame-{camera}-{expid}.fits : raw extracted photons with no calibration at all
        2. sky-{camera}-{expid}.fits : the sky model in photons
        3. fluxcalib-{camera}-{expid}.fits] : the flux calibration vectors
        4. cframe-{camera}-{expid}.fits : flux calibrated spectra

        These files are written to $simdir/{expid}

* QUICKBRICK MODE

  - must provide brickname
  - Simulates 100 spectra by default, unless nspec keyword is provided
  - For this option, airmass and exposure time are given directly to quickgen
  - Brickname is the only required keyword
  - Simulate bricks (can provide anything as brick name)::

        quickgen --brickname brick

  - quickbrick output::

        1. brick-{arm}-{brickname}.fits : flux calibrated spectra in bricks

        These files are written to the current working directory by default

* additional keyword arguments (descriptions below):

  - common to both modes: nspec, config, seed, moon-phase, moon-angle, moon-zenith
  - quickgen only: simspec, fibermap, nstart, spectrograph, frameonly
  - quickbrick only: brickname, objtype, airmass, exptime, outdir, verbose, outdir-truth,
    zrange-qso, zrange-elg, zrange-lrg, zrange-bgs, sne-rfluxratiorange,
    add-SNeIa
"""
from __future__ import absolute_import, division, print_function

import argparse
import astropy.units as u
from astropy.io import fits
from astropy.table import Table, Column, vstack
from time import asctime

import os
import os.path
import numpy as np
import scipy.special
import scipy.interpolate
import scipy.constants as const
import sys

import desispec
import desispec.io
import desisim
import desisim.io
import desisim.templates
import desiutil.io
from desispec.resolution import Resolution
from desispec.io import write_flux_calibration, write_fiberflat, read_fibermap, specprod_root, fitsheader, empty_fibermap, Brick
from desispec.interpolation import resample_flux
from desispec.frame import Frame
from desispec.fiberflat import FiberFlat
from desispec.sky import SkyModel
from desispec.fluxcalibration import FluxCalib
from desispec.log import get_logger, DEBUG, INFO
from desisim.obs import get_night
from desisim.targets import sample_objtype
from desisim.specsim import get_simulator
from desimodel.io import load_desiparams

def _add_truth(hdus, header, meta, trueflux, sflux, wave, channel):
    """Utility function for adding truth to an output FITS file."""
    hdus.append(
        fits.ImageHDU(trueflux[channel], name='_TRUEFLUX', header=header))
    if channel == 'b':
        swave = wave.astype(np.float32)
        hdus.append(fits.ImageHDU(swave, name='_SOURCEWAVE', header=header))
        hdus.append(fits.ImageHDU(sflux, name='_SOURCEFLUX', header=header))
        metatable = desiutil.io.encode_table(meta, encoding='ascii')
        metahdu = fits.convenience.table_to_hdu(meta)
        metahdu.header['EXTNAME'] = '_TRUTH'
        hdus.append(metahdu)

def parse(options=None):
    parser=argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    parser.add_argument('--simspec',type=str, help="input simspec file")
    parser.add_argument('--fibermap',type=str, help='input fibermap file')
    parser.add_argument('-n','--nspec',type=int,default=100,help='number of spectra to be simulated, starting from first')
    parser.add_argument('--nstart', type=int, default=0,help='starting spectra # for simulation 0-4999')
    parser.add_argument('--spectrograph',type=int, default=None,help='Spectrograph no. 0-9')
    parser.add_argument('--config', type=str, default='desi', help='specsim configuration')
    parser.add_argument('-s','--seed', type=int, default=0,  help="random seed")
    # Only produce uncalibrated output
    parser.add_argument('--frameonly', action="store_true", help="only output frame files")

    # Moon options if bright or gray time
    parser.add_argument('--moon-phase', type=float,  help='moon phase (0=full, 1=new)', default=None, metavar='')
    parser.add_argument('--moon-angle', type=float,  help='separation angle to the moon (0-180 deg)', default=None, metavar='')
    parser.add_argument('--moon-zenith', type=float,  help='zenith angle of the moon (0-90 deg)', default=None, metavar='')

    # Brick options
    parser.add_argument('-b','--brickname', type=str, help='unique output brickname suffix', metavar='')
    parser.add_argument('--objtype', type=str,  help='ELG, LRG, QSO, BGS, MWS, WD, DARK_MIX, or BRIGHT_MIX', default='DARK_MIX', metavar='')
    parser.add_argument('-a','--airmass', type=float,  help='airmass', default=1.25, metavar='') # Science Req. Doc L3.3.2
    parser.add_argument('-e','--exptime', type=float,  help='exposure time (s)', default=None,metavar='')
    parser.add_argument('-o','--outdir', type=str,  help='output directory', default='.', metavar='')
    parser.add_argument('-v','--verbose', action='store_true', help='toggle on verbose output')
    parser.add_argument('--outdir-truth', type=str,  help='optional alternative output directory for truth files', metavar='')

    # Object type specific options
    parser.add_argument('--zrange-qso', type=float, default=(0.5, 4.0), nargs=2, metavar='',
                        help='minimum and maximum redshift range for QSO')
    parser.add_argument('--zrange-elg', type=float, default=(0.6, 1.6), nargs=2, metavar='',
                        help='minimum and maximum redshift range for ELG')
    parser.add_argument('--zrange-lrg', type=float, default=(0.5, 1.1), nargs=2, metavar='',
                        help='minimum and maximum redshift range for LRG')
    parser.add_argument('--zrange-bgs', type=float, default=(0.01, 0.4), nargs=2, metavar='',
                        help='minimum and maximum redshift range for BGS')
    parser.add_argument('--rmagrange-bgs', type=float, default=(15.0, 19.5), nargs=2, metavar='',
                        help='Minimum and maximum BGS r-band (AB) magnitude range')
    parser.add_argument('--sne-rfluxratiorange', type=float, default=(0.1, 1.0), nargs=2, metavar='',
                        help='r-band flux ratio of the SNeIa spectrum relative to the galaxy')
    parser.add_argument('--add-SNeIa', action='store_true', help='include SNeIa spectra')

    if options is None:
        args = parser.parse_args()
    else:
        args = parser.parse_args(options)

    log = get_logger()
    if args.simspec:
        args.objtype = None
        if args.fibermap is None:
            msg = 'If simspec set, must also set fibermap'
            log.error(msg)
            raise ValueError(msg)
        #hdr = fits.getheader(args.simspec)
        #night = str(hdr['NIGHT'])
        #expid = int(hdr['EXPID'])
        #args.simspec = desisim.io.findfile('simspec', night, expid)
        #args.fibermap = desispec.io.findfile('fibermap', night, expid)

{'fraction_non_alphanumeric': 0.12133015522353352, 'fraction_numerical': 0.014185221789430847, 'mean_word_length': 3.665581773799837, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 4, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4601585', 'n_tokens_mistral': 5606, 'n_tokens_neox': 5283, 'n_words': 1532}