Example #1
0
def main():

    seed = 123
    ngal = 100
    npergal = 50

    # Read the mock catalog.
    cat = h5py.File('galaxy_catalogue_0.hdf5')

    gr = cat['Data']['g_r'][:ngal]
    mag = cat['Data']['app_mag'][:ngal]
    z = cat['Data']['z_obs'][:ngal]
    zmin, zmax = np.min(z), np.max(z)

    # Read the basis templates and initialize the output metadata table.
    baseflux, basewave, basemeta = read_basis_templates('BGS')
    meta = empty_metatable(objtype='BGS', nmodel=ngal)

    bgs = BGS(minwave=3000, maxwave=11000.0)

    for igal in range(ngal):
        flux, wave, meta = bgs.make_templates(npergal, zrange=(zmin, zmax),
                                              nocolorcuts=True)

        pdb.set_trace()
Example #2
0
    def bgs(self, data, index=None, mockformat='durham_mxxl_hdf5'):
        """Generate spectra for BGS.

        Currently only the MXXL (durham_mxxl_hdf5) mock is supported.  DATA
        needs to have Z, SDSS_absmag_r01, SDSS_01gr, VDISP, and SEED, which are
        assigned in mock.io.read_durham_mxxl_hdf5.  See also BGSKDTree.bgs().

        """
        from desisim.io import empty_metatable

        objtype = 'BGS'
        if index is None:
            index = np.arange(len(data['Z']))

        input_meta = empty_metatable(nmodel=len(index), objtype=objtype)
        for inkey, datakey in zip(('SEED', 'MAG', 'REDSHIFT', 'VDISP'),
                                  ('SEED', 'MAG', 'Z', 'VDISP')):
            input_meta[inkey] = data[datakey][index]

        if mockformat.lower() == 'durham_mxxl_hdf5':
            alldata = np.vstack(
                (data['Z'][index], data['SDSS_absmag_r01'][index],
                 data['SDSS_01gr'][index])).T
            _, templateid = self.tree.query(objtype, alldata)
        else:
            raise ValueError('Unrecognized mockformat {}!'.format(mockformat))

        input_meta['TEMPLATEID'] = templateid
        flux, _, meta = self.bgs_templates.make_templates(
            input_meta=input_meta,
            nocolorcuts=True,
            novdisp=False,
            verbose=self.verbose)

        return flux, meta
Example #3
0
    def lrg(self, data, index=None, mockformat='gaussianfield'):
        """Generate spectra for the LRG sample.

        Currently only the GaussianField mock sample is supported.  DATA needs
        to have Z, GR, RZ, VDISP, and SEED, which are assigned in
        mock.io.read_gaussianfield.  See also TemplateKDTree.lrg().

        """
        objtype = 'LRG'
        if index is None:
            index = np.arange(len(data['Z']))
        nobj = len(index)

        input_meta = empty_metatable(nmodel=len(index), objtype=objtype)
        for inkey, datakey in zip(('SEED', 'MAG', 'REDSHIFT', 'VDISP'),
                                  ('SEED', 'MAG', 'Z', 'VDISP')):
            input_meta[inkey] = data[datakey][index]

        if mockformat.lower() == 'gaussianfield':
            # This is wrong: choose a template with equal probability.
            templateid = self.rand.choice(self.tree.lrg_meta['TEMPLATEID'], len(index))
        else:
            raise ValueError('Unrecognized mockformat {}!'.format(mockformat))

        input_meta['TEMPLATEID'] = templateid
        flux, _, meta = self.lrg_templates.make_templates(input_meta=input_meta,
                                                          nocolorcuts=True, novdisp=False,
                                                          verbose=self.verbose)

        return flux, meta
Example #4
0
    def mws(self, data, index=None, mockformat='galaxia'):
        """Generate spectra for the MWS_NEARBY and MWS_MAIN samples.

        """
        objtype = 'STAR'
        if index is None:
            index = np.arange(len(data['Z']))

        input_meta = empty_metatable(nmodel=len(index), objtype=objtype)
        for inkey, datakey in zip(('SEED', 'MAG', 'REDSHIFT', 'TEFF', 'LOGG', 'FEH'),
                                  ('SEED', 'MAG', 'Z', 'TEFF', 'LOGG', 'FEH')):
            input_meta[inkey] = data[datakey][index]

        if mockformat.lower() == '100pc':
            alldata = np.vstack((data['TEFF'][index],
                                 data['LOGG'][index],
                                 data['FEH'][index])).T
            _, templateid = self.tree.query(objtype, alldata)
            
        elif mockformat.lower() == 'galaxia':
            alldata = np.vstack((data['TEFF'][index],
                                 data['LOGG'][index],
                                 data['FEH'][index])).T
            _, templateid = self.tree.query(objtype, alldata)
        else:
            raise ValueError('Unrecognized mockformat {}!'.format(mockformat))

        input_meta['TEMPLATEID'] = templateid
        flux, _, meta = self.star_templates.make_templates(input_meta=input_meta,
                                                          verbose=self.verbose) # Note! No colorcuts.

        return flux, meta
Example #5
0
    def lrg(self, data, index=None, mockformat='gaussianfield'):
        """Generate magnitudes for the LRG sample.
        Currently only the GaussianField mock sample is supported.  DATA needs
        to have Z, GR, RZ, VDISP, and SEED, which are assigned in
        mock.io.read_gaussianfield.  
        """
        objtype = 'LRG'
        if index is None:
            index = np.arange(len(data['Z']))
        nobj = len(index)

        meta = empty_metatable(nmodel=len(index), objtype=objtype)
        for inkey, datakey in zip(('SEED', 'MAG', 'REDSHIFT', 'VDISP'),
                                  ('SEED', 'MAG', 'Z', 'VDISP')):
            meta[inkey] = data[datakey][index]

        if mockformat.lower() != 'gaussianfield':
            raise ValueError('Unrecognized mockformat {}!'.format(mockformat))
            return meta
        
        meta['FLUX_G'][:] = 10**((22.5 - (data['GR'][index] + data['RZ'][index] + data['MAG'][index]))/2.5) # g-band flux  
        meta['FLUX_R'][:] =  10**((22.5 - (data['MAG'][index] + data['RZ'][index]))/2.5) # r-band flux
        meta['FLUX_Z'][:] = 10**((22.5 - data['MAG'][index])/2.5) # z-band flux
        meta['FLUX_W1'][:] = 10**((22.5 - (data['RZ'][index] + data['MAG'][index] - data['RW1'][index]))/2.5)# wise flux 1
        meta['FLUX_W2'][:] = 10**((22.5 - (data['RZ'][index] + data['MAG'][index] - data['RW1'][index] - data['W1W2'][index]))/2.5)# wise flux 2
        
        return meta
Example #6
0
    def elg_test(self, data, index=None, mockformat='gaussianfield'):
        """Test script -- generate spectra for the ELG sample.

        Currently only the GaussianField mock sample is supported.  DATA needs
        to have Z, GR, RZ, VDISP, and SEED, which are assigned in
        mock.io.read_gaussianfield.  See also TemplateKDTree.elg().

        """
        objtype = 'ELG'
        if index is None:
            index = np.arange(len(data['Z']))

        input_meta = empty_metatable(nmodel=len(index), objtype=objtype)
        for inkey, datakey in zip(('SEED', 'MAG', 'REDSHIFT', 'VDISP'),
                                  ('SEED', 'MAG', 'Z', 'VDISP')):
            input_meta[inkey] = data[datakey][index]

        if mockformat.lower() == 'gaussianfield':
            alldata = np.vstack((data['Z'][index],
                                 data['GR'][index],
                                 data['RZ'][index])).T
            _, templateid = self.tree.query(objtype, alldata)
        else:
            raise ValueError('Unrecognized mockformat {}!'.format(mockformat))

        import matplotlib.pyplot as plt
        from scipy.interpolate import interp1d

        f1 = interp1d(np.squeeze(self.tree.elg_kcorr['REDSHIFT']), np.squeeze(self.tree.elg_kcorr['GR']), axis=0)
        gr = f1(data['Z'][index])
        plt.plot(np.squeeze(self.tree.elg_kcorr['REDSHIFT']), np.squeeze(self.tree.elg_kcorr['GR'])[:, 500])
        plt.scatter(data['Z'][index], gr[:, 500], marker='x', color='red', s=15)
        plt.show()

        def elg_colorbox(ax):
            """Draw the ELG selection box."""
            from matplotlib.patches import Polygon
            grlim = ax.get_ylim()
            coeff0, coeff1 = (1.15, -0.15), (-1.2, 1.6)
            rzmin, rzpivot = 0.3, (coeff1[1] - coeff0[1]) / (coeff0[0] - coeff1[0])
            verts = [(rzmin, grlim[0]),
                     (rzmin, np.polyval(coeff0, rzmin)),
                     (rzpivot, np.polyval(coeff1, rzpivot)),
                     ((grlim[0] - 0.1 - coeff1[1]) / coeff1[0], grlim[0] - 0.1)
                     ]
            ax.add_patch(Polygon(verts, fill=False, ls='--', color='k'))

        fig, ax = plt.subplots()
        ax.scatter(data['RZ'][index], data['GR'][index])
        ax.set_xlim(-0.5, 2) ; plt.ylim(-0.5, 2)
        elg_colorbox(ax)
        plt.show()
        import pdb ; pdb.set_trace()

        input_meta['TEMPLATEID'] = templateid
        flux, _, meta = self.elg_templates.make_templates(input_meta=input_meta,
                                                          nocolorcuts=True, novdisp=False,
                                                          verbose=self.verbose)

        return flux, meta
Example #7
0
    def qso(self, data, index=None, mockformat='gaussianfield'):
        """Generate magnitudes for the QSO or QSO/LYA samples.
        """
        from desisim.lya_spectra import get_spectra
        
        objtype = 'QSO'
        if index is None:
            index = np.arange(len(data['Z']))
        nobj = len(index)
        meta = empty_metatable(nmodel=nobj, objtype=objtype)

        if mockformat.lower() != 'gaussianfield':
            raise ValueError('Unrecognized mockformat {}!'.format(mockformat))
            return meta
        
        for inkey, datakey in zip(('SEED', 'MAG', 'REDSHIFT'),
                                      ('SEED', 'MAG', 'Z')):
            meta[inkey] = data[datakey][index]
        
        meta['FLUX_G'][:] = 10**((22.5 - data['MAG'][index])/2.5) # g-band flux
        meta['FLUX_R'][:] = 10**((22.5 - (data['MAG'][index] - data['GR'][index]))/2.5) # r-band flux
        meta['FLUX_Z'][:] = 10**((22.5 - (data['MAG'][index] - data['GR'][index] - data['RZ'][index]))/2.5) # z-band flux
        meta['FLUX_W1'][:] = 10**((22.5 - (data['MAG'][index] - data['GR'][index] - data['RW1'][index]))/2.5)# wise flux 1
        meta['FLUX_W2'][:] = 10**((22.5 - (data['MAG'][index] - data['GR'][index] - data['RW1'][index] - data['W1W2'][index]))/2.5)# wise flux 2
        
        lya = np.where( data['TEMPLATESUBTYPE'][index] == 'LYA' )[0]                 
        if len(lya) > 0:
            meta['SUBTYPE'][lya] = 'LYA'
                                       
        return meta
Example #8
0
    def mws_wd(self, data, index=None, mockformat='wd'):
        """Generate magnitudes for the MWS_WD sample.  Deal with DA vs DB white dwarfs
        separately.
        """
        objtype = 'WD'
        if index is None:
            index = np.arange(len(data['Z']))
        nobj = len(index)

        meta = empty_metatable(nmodel=nobj, objtype=objtype)
        for inkey, datakey in zip(('SEED', 'MAG', 'REDSHIFT', 'TEFF', 'LOGG', 'SUBTYPE'),
                                  ('SEED', 'MAG', 'Z', 'TEFF', 'LOGG', 'TEMPLATESUBTYPE')):
            meta[inkey] = data[datakey][index]

        if mockformat.lower() != 'wd':
            raise ValueError('Unrecognized mockformat {}!'.format(mockformat))
            return meta
            
        for subtype in ('DA', 'DB'):
            these = np.where(meta['SUBTYPE'] == subtype)[0]
            if len(these) > 0:
                meta['SUBTYPE'][these] = subtype    
                meta['TEMPLATEID'][:] = -1
    
        meta['FLUX_G'][:] = 10**((22.5 - data['MAG'][index])/2.5) # g-band flux
        return meta
Example #9
0
    def mws_wd(self, data, index=None, mockformat='wd'):
        """Generate spectra for the MWS_WD sample.  Deal with DA vs DB white dwarfs
        separately.

        """
        objtype = 'WD'
        if index is None:
            index = np.arange(len(data['Z']))
        nobj = len(index)

        input_meta = empty_metatable(nmodel=nobj, objtype=objtype)
        for inkey, datakey in zip(('SEED', 'MAG', 'REDSHIFT', 'TEFF', 'LOGG', 'SUBTYPE'),
                                  ('SEED', 'MAG', 'Z', 'TEFF', 'LOGG', 'TEMPLATESUBTYPE')):
            input_meta[inkey] = data[datakey][index]

        if mockformat.lower() == 'wd':
            meta = empty_metatable(nmodel=nobj, objtype=objtype)
            flux = np.zeros([nobj, len(self.wave)], dtype='f4')
            
            for subtype in ('DA', 'DB'):
                these = np.where(input_meta['SUBTYPE'] == subtype)[0]
                if len(these) > 0:
                    alldata = np.vstack((data['TEFF'][index][these],
                                         data['LOGG'][index][these])).T
                    _, templateid = self.tree.query(objtype, alldata, subtype=subtype)

                    input_meta['TEMPLATEID'][these] = templateid
                    
                    template_function = 'wd_{}_templates'.format(subtype.lower())
                    flux1, _, meta1 = getattr(self, template_function).make_templates(input_meta=input_meta[these],
                                                          verbose=self.verbose)
                    
                    meta[these] = meta1
                    flux[these, :] = flux1
            
        else:
            raise ValueError('Unrecognized mockformat {}!'.format(mockformat))

        return flux, meta
Example #10
0
    def sky(self, data, index=None, mockformat=None):
        """Generate spectra for SKY.

        """
        objtype = 'SKY'
        if index is None:
            index = np.arange(len(data['Z']))
        nobj = len(index)

        meta = empty_metatable(nmodel=nobj, objtype=objtype)
        for inkey, datakey in zip(('SEED', 'REDSHIFT'),
                                  ('SEED', 'Z')):
            meta[inkey] = data[datakey][index]

        return meta
Example #11
0
    def mws(self, data, index=None, mockformat='galaxia'):
        """Generate magnitudes for the MWS_NEARBY and MWS_MAIN samples.
        """
        objtype = 'STAR'
        if index is None:
            index = np.arange(len(data['Z']))

        meta = empty_metatable(nmodel=len(index), objtype=objtype)
        for inkey, datakey in zip(('SEED', 'MAG', 'REDSHIFT', 'TEFF', 'LOGG', 'FEH'),
                                  ('SEED', 'MAG', 'Z', 'TEFF', 'LOGG', 'FEH')):
            meta[inkey] = data[datakey][index]

        if not (mockformat.lower() in ['100pc','galaxia']):
            raise ValueError('Unrecognized mockformat {}!'.format(mockformat))
            return meta

        meta['FLUX_R'][:] = 10**((22.5 - data['MAG'][index])/2.5) # r-band flux
        return meta
Example #12
0
    def bgs(self, data, index=None, mockformat='durham_mxxl_hdf5'):
        """Generate magnitudes for BGS.
        Currently only the MXXL (durham_mxxl_hdf5) mock is supported.  DATA
        needs to have Z, SDSS_absmag_r01, SDSS_01gr, VDISP, and SEED, which are
        assigned in mock.io.read_durham_mxxl_hdf5.  
        """
        objtype = 'BGS'
        if index is None:
            index = np.arange(len(data['Z']))
        meta = empty_metatable(nmodel=len(index), objtype=objtype)
        for inkey, datakey in zip(('SEED', 'MAG', 'REDSHIFT', 'VDISP'),
                                  ('SEED', 'MAG', 'Z', 'VDISP')):
            meta[inkey] = data[datakey][index]

        if mockformat.lower() != 'durham_mxxl_hdf5':
            raise ValueError('Unrecognized mockformat {}!'.format(mockformat))
            return meta
        meta['FLUX_R'][:] = 10**((22.5 - data['MAG'][index])/2.5) # r-band flux
        meta['FLUX_G'][:] = 10**((22.5-(data['SDSS_01gr'][index] + data['MAG'][index]))/2.5) # g-band flux
        return meta
Example #13
0
def get_spectra(lyafile, nqso=None, wave=None, templateid=None, normfilter='sdss2010-g',
                seed=None, rand=None, qso=None, add_dlas=False, debug=False, nocolorcuts=True):
    """Generate a QSO spectrum which includes Lyman-alpha absorption.

    Args:
        lyafile (str): name of the Lyman-alpha spectrum file to read.
        nqso (int, optional): number of spectra to generate (starting from the
          first spectrum; if more flexibility is needed use TEMPLATEID).
        wave (numpy.ndarray, optional): desired output wavelength vector.
        templateid (int numpy.ndarray, optional): indices of the spectra
          (0-indexed) to read from LYAFILE (default is to read everything).  If
          provided together with NQSO, TEMPLATEID wins.
        normfilter (str, optional): normalization filter
        seed (int, optional): Seed for random number generator.
        rand (numpy.RandomState, optional): RandomState object used for the
          random number generation.  If provided together with SEED, this
          optional input superseeds the numpy.RandomState object instantiated by
          SEED.
        qso (desisim.templates.QSO, optional): object with which to generate
          individual spectra/templates.
        add_dlas (bool): Inject damped Lya systems into the Lya forest
          These are done according to the current best estimates for the incidence dN/dz
          (Prochaska et al. 2008, ApJ, 675, 1002)
          Set in calc_lz
          These are *not* inserted according to overdensity along the sightline
        nocolorcuts (bool, optional): Do not apply the fiducial rzW1W2 color-cuts
          cuts (default True).

    Returns (flux, wave, meta, dla_meta) where:

        * flux (numpy.ndarray): Array [nmodel, npix] of observed-frame spectra
          (erg/s/cm2/A).
        * wave (numpy.ndarray): Observed-frame [npix] wavelength array (Angstrom).
        * meta (astropy.Table): Table of meta-data [nmodel] for each output spectrum
          with columns defined in desisim.io.empty_metatable *plus* RA, DEC.
        * objmeta (astropy.Table): Table of additional object-specific meta-data
          [nmodel] for each output spectrum with columns defined in
          desisim.io.empty_metatable.
        * dla_meta (astropy.Table): Table of meta-data [ndla] for the DLAs injected
          into the spectra.  Only returned if add_dlas=True

    Note: `dla_meta` is only included if add_dlas=True.

    """
    from scipy.interpolate import interp1d

    import fitsio

    from speclite.filters import load_filters
    from desisim.templates import QSO
    from desisim.io import empty_metatable

    h = fitsio.FITS(lyafile)
    if templateid is None:
        if nqso is None:
            nqso = len(h)-1
        templateid = np.arange(nqso)
    else:
        templateid = np.asarray(templateid)
        nqso = len(np.atleast_1d(templateid))

    if rand is None:
        rand = np.random.RandomState(seed)
    templateseed = rand.randint(2**32, size=nqso)

    #heads = [head.read_header() for head in h[templateid + 1]]
    heads = []
    for indx in templateid:
        heads.append(h[indx + 1].read_header())

    zqso = np.array([head['ZQSO'] for head in heads])
    ra = np.array([head['RA'] for head in heads])
    dec = np.array([head['DEC'] for head in heads])
    mag_g = np.array([head['MAG_G'] for head in heads])

    # Hard-coded filtername!  Should match MAG_G!
    normfilt = load_filters(normfilter)

    if qso is None:
        qso = QSO(normfilter_south=normfilter, wave=wave)

    wave = qso.wave
    flux = np.zeros([nqso, len(wave)], dtype='f4')

    meta, objmeta = empty_metatable(objtype='QSO', nmodel=nqso)
    meta['TEMPLATEID'][:] = templateid
    meta['REDSHIFT'][:] = zqso
    meta['MAG'][:] = mag_g
    meta['MAGFILTER'][:] = normfilter
    meta['SEED'][:] = templateseed
    meta['RA'] = ra
    meta['DEC'] = dec

    # Lists for DLA meta data
    if add_dlas:
        dla_NHI, dla_z, dla_id = [], [], []

    # Loop on quasars
    for ii, indx in enumerate(templateid):
        flux1, _, meta1, objmeta1 = qso.make_templates(nmodel=1, redshift=np.atleast_1d(zqso[ii]),
                                                mag=np.atleast_1d(mag_g[ii]), seed=templateseed[ii],
                                                nocolorcuts=nocolorcuts, lyaforest=False)
        flux1 *= 1e-17
        for col in meta1.colnames:
            meta[col][ii] = meta1[col][0]
        for col in objmeta1.colnames:
            objmeta[col][ii] = objmeta1[col][0]

        # read lambda and forest transmission
        data = h[indx + 1].read()
        la = data['LAMBDA'][:]
        tr = data['FLUX'][:]

        if len(tr):
            # Interpolate the transmission at the spectral wavelengths, if
            # outside the forest, the transmission is 1.
            itr = interp1d(la, tr, bounds_error=False, fill_value=1.0)
            flux1 *= itr(wave)

        # Inject a DLA?
        if add_dlas:
            if np.min(wave/lambda_RF_LYA - 1) < zqso[ii]: # Any forest?
                dlas, dla_model = insert_dlas(wave, zqso[ii], seed=templateseed[ii])
                ndla = len(dlas)
                if ndla > 0:
                    flux1 *= dla_model
                    # Meta
                    dla_z += [idla['z'] for idla in dlas]
                    dla_NHI += [idla['N'] for idla in dlas]
                    dla_id += [indx]*ndla

        padflux, padwave = normfilt.pad_spectrum(flux1, wave, method='edge')
        normmaggies = np.array(normfilt.get_ab_maggies(padflux, padwave,
                                                       mask_invalid=True)[normfilter])

        factor = 10**(-0.4 * mag_g[ii]) / normmaggies
        flux1 *= factor
        for key in ('FLUX_G', 'FLUX_R', 'FLUX_Z', 'FLUX_W1', 'FLUX_W2'):
            meta[key][ii] *= factor
        flux[ii, :] = flux1[:]

    h.close()

    # Finish
    if add_dlas:
        ndla = len(dla_id)
        if ndla > 0:
            from astropy.table import Table
            dla_meta = Table()
            dla_meta['NHI'] = dla_NHI  # log NHI values
            dla_meta['z'] = dla_z
            dla_meta['ID'] = dla_id
        else:
            dla_meta = None
        return flux*1e17, wave, meta, objmeta, dla_meta
    else:
        return flux*1e17, wave, meta, objmeta
Example #14
0
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
Example #15
0
    def qso(self, data, index=None, mockformat='gaussianfield'):
        """Generate spectra for the QSO or QSO/LYA samples.

        Note: We need to make sure NORMFILTER matches!

        """
        from desisim.lya_spectra import get_spectra
        from desisim.lya_spectra import read_lya_skewers,apply_lya_transmission
        import fitsio

        log = get_logger()

        objtype = 'QSO'
        if index is None:
            index = np.arange(len(data['Z']))
        nobj = len(index)

        if mockformat.lower() == 'gaussianfield':
            input_meta = empty_metatable(nmodel=nobj, objtype=objtype)
            for inkey, datakey in zip(('SEED', 'MAG', 'REDSHIFT'),
                                      ('SEED', 'MAG', 'Z')):
                input_meta[inkey] = data[datakey][index]

            # Build the tracer and Lya forest QSO spectra separately.
            meta = empty_metatable(nmodel=nobj, objtype=objtype)
            flux = np.zeros([nobj, len(self.wave)], dtype='f4')

            lya = np.where( data['TEMPLATESUBTYPE'][index] == 'LYA' )[0]
            tracer = np.where( data['TEMPLATESUBTYPE'][index] == '' )[0]

            if len(tracer) > 0:
                flux1, _, meta1 = self.qso_templates.make_templates(input_meta=input_meta[tracer],
                                                                    lyaforest=False,
                                                                    nocolorcuts=True,
                                                                    verbose=self.verbose)
                meta[tracer] = meta1
                flux[tracer, :] = flux1

            if len(lya) > 0:


                ilya=index[lya].astype(int)
                nqso=ilya.size
                log.info("Generating spectra of %d lya QSOs"%nqso)

                if 'LYAHDU' in data :
                    # this is the old format with one HDU per spectrum
                    alllyafile = data['LYAFILES'][ilya]
                    alllyahdu = data['LYAHDU'][ilya]

                    for lyafile in sorted(set(alllyafile)):
                        these = np.where( lyafile == alllyafile )[0]

                        templateid = alllyahdu[these] - 1 # templateid is 0-indexed
                        flux1, _, meta1 = get_spectra(lyafile, templateid=templateid, normfilter=data['FILTERNAME'],
                                                      rand=self.rand, qso=self.lya_templates, nocolorcuts=True)
                        meta1['SUBTYPE'] = 'LYA'
                        meta[lya[these]] = meta1
                        flux[lya[these], :] = flux1
                else : # new format

                    # read skewers
                    skewer_wave=None
                    skewer_trans=None
                    skewer_meta=None

                    # all the files that contain at least one QSO skewer
                    alllyafile = data['LYAFILES'][ilya]
                    uniquelyafiles = sorted(set(alllyafile))

                    for lyafile in uniquelyafiles :
                        these = np.where( alllyafile == lyafile )[0]
                        objid_in_data=data['OBJID'][ilya][these]
                        objid_in_mock=(fitsio.read(lyafile, columns=['MOCKID'],upper=True,ext=1).astype(float)).astype(int)
                        o2i=dict()
                        for i,o in enumerate(objid_in_mock) :
                            o2i[o]=i
                        indices_in_mock_healpix=np.zeros(objid_in_data.size).astype(int)
                        for i,o in enumerate(objid_in_data) :
                            if not o in o2i :
                                log.error("No MOCKID={} in {}. It's a bug, should never happen".format(o,lyafile))
                                raise(KeyError("No MOCKID={} in {}. It's a bug, should never happen".format(o,lyafile)))
                            indices_in_mock_healpix[i]=o2i[o]

                        tmp_wave,tmp_trans,tmp_meta = read_lya_skewers(lyafile,indices=indices_in_mock_healpix)

                        if skewer_wave is None :
                            skewer_wave=tmp_wave
                            dw=skewer_wave[1]-skewer_wave[0] # this is just to check same wavelength
                            skewer_trans=np.zeros((nqso,skewer_wave.size)) # allocate skewer_array
                            skewer_meta=dict()
                            for k in tmp_meta.dtype.names :
                                skewer_meta[k]=np.zeros(nqso).astype(tmp_meta[k].dtype)
                        else :
                            # check wavelength is the same for all skewers
                            assert(np.max(np.abs(wave-tmp_wave))<0.001*dw)

                        skewer_trans[these] = tmp_trans
                        for k in skewer_meta.keys() :
                            skewer_meta[k][these]=tmp_meta[k]

                    # check we matched things correctly
                    assert(np.max(np.abs(skewer_meta["Z"]-data['Z'][ilya]))<0.000001)
                    assert(np.max(np.abs(skewer_meta["RA"]-data['RA'][ilya]))<0.000001)
                    assert(np.max(np.abs(skewer_meta["DEC"]-data['DEC'][ilya]))<0.000001)


                    # now we create a series of QSO spectra all at once
                    # this is faster than calling each one at a time
                    # we use the provided QSO template class

                    seed = self.rand.randint(2**32)
                    qso  = self.lya_templates
                    qso_flux, qso_wave, qso_meta = qso.make_templates(nmodel=nqso,
                                                                      redshift=data['Z'][ilya],
                                                                      mag=data['MAG'][ilya],
                                                                      seed=seed,
                                                                      lyaforest=False,
                                                                      nocolorcuts=True)

                    # apply transmission to QSOs
                    qso_flux = apply_lya_transmission(qso_wave,qso_flux,skewer_wave,skewer_trans)

                    # save this
                    qso_meta['SUBTYPE'] = 'LYA'
                    meta[lya] = qso_meta
                    flux[lya, :] = qso_flux

        else:
            raise ValueError('Unrecognized mockformat {}!'.format(mockformat))

        return flux, meta
Example #16
0
def get_targets(nspec,
                program,
                tileid=None,
                seed=None,
                specify_targets=dict(),
                specmin=0):
    """
    Generates a set of targets for the requested program

    Args:
        nspec: (int) number of targets to generate
        program: (str) program name DARK, BRIGHT, GRAY, MWS, BGS, LRG, ELG, ...

    Options:
      * tileid: (int) tileid, used for setting RA,dec
      * seed: (int) random number seed
      * specify_targets: (dict of dicts)  Define target properties like magnitude and redshift
                                 for each target class. Each objtype has its own key,value pair
                                 see simspec.templates.specify_galparams_dict() 
                                 or simsepc.templates.specify_starparams_dict()
      * specmin: (int) first spectrum number (0-indexed)

    Returns:
      * fibermap
      * targets as tuple of (flux, wave, meta)
    """
    if tileid is None:
        tile_ra, tile_dec = 0.0, 0.0
    else:
        tile_ra, tile_dec = io.get_tile_radec(tileid)

    program = program.upper()
    log.debug('Using random seed {}'.format(seed))
    np.random.seed(seed)

    #- Get distribution of target types
    true_objtype, target_objtype = sample_objtype(nspec, program)

    #- Get DESI wavelength coverage
    wavemin = desimodel.io.load_throughput('b').wavemin
    wavemax = desimodel.io.load_throughput('z').wavemax
    dw = 0.2
    wave = np.arange(round(wavemin, 1), wavemax, dw)
    nwave = len(wave)

    flux = np.zeros((nspec, len(wave)))
    meta = empty_metatable(nmodel=nspec, objtype='SKY')
    fibermap = empty_fibermap(nspec)

    for objtype in set(true_objtype):
        ii = np.where(true_objtype == objtype)[0]
        nobj = len(ii)

        fibermap['OBJTYPE'][ii] = target_objtype[ii]

        if objtype in specify_targets.keys():
            obj_kwargs = specify_targets[objtype]
        else:
            obj_kwargs = dict()

        # Simulate spectra
        if objtype == 'SKY':
            fibermap['DESI_TARGET'][ii] = desi_mask.SKY
            continue

        elif objtype == 'ELG':
            from desisim.templates import ELG
            elg = ELG(wave=wave)
            simflux, wave1, meta1 = elg.make_templates(nmodel=nobj,
                                                       seed=seed,
                                                       **obj_kwargs)
            fibermap['DESI_TARGET'][ii] = desi_mask.ELG

        elif objtype == 'LRG':
            from desisim.templates import LRG
            lrg = LRG(wave=wave)
            simflux, wave1, meta1 = lrg.make_templates(nmodel=nobj,
                                                       seed=seed,
                                                       **obj_kwargs)
            fibermap['DESI_TARGET'][ii] = desi_mask.LRG

        elif objtype == 'BGS':
            from desisim.templates import BGS
            bgs = BGS(wave=wave)
            simflux, wave1, meta1 = bgs.make_templates(nmodel=nobj,
                                                       seed=seed,
                                                       **obj_kwargs)
            fibermap['DESI_TARGET'][ii] = desi_mask.BGS_ANY
            fibermap['BGS_TARGET'][ii] = bgs_mask.BGS_BRIGHT

        elif objtype == 'QSO':
            from desisim.templates import QSO
            qso = QSO(wave=wave)
            simflux, wave1, meta1 = qso.make_templates(nmodel=nobj,
                                                       seed=seed,
                                                       lyaforest=False,
                                                       **obj_kwargs)
            fibermap['DESI_TARGET'][ii] = desi_mask.QSO

        # For a "bad" QSO simulate a normal star without color cuts, which isn't
        # right. We need to apply the QSO color-cuts to the normal stars to pull
        # out the correct population of contaminating stars.

        # Note by @moustakas: we can now do this using desisim/#150, but we are
        # going to need 'noisy' photometry (because the QSO color-cuts
        # explicitly avoid the stellar locus).
        elif objtype == 'QSO_BAD':
            from desisim.templates import STAR
            #from desitarget.cuts import isQSO
            #star = STAR(wave=wave, colorcuts_function=isQSO)
            star = STAR(wave=wave)
            simflux, wave1, meta1 = star.make_templates(nmodel=nobj,
                                                        seed=seed,
                                                        **obj_kwargs)
            fibermap['DESI_TARGET'][ii] = desi_mask.QSO

        elif objtype == 'STD':
            from desisim.templates import FSTD
            fstd = FSTD(wave=wave)
            simflux, wave1, meta1 = fstd.make_templates(nmodel=nobj,
                                                        seed=seed,
                                                        **obj_kwargs)
            fibermap['DESI_TARGET'][ii] = desi_mask.STD_FSTAR

        elif objtype == 'MWS_STAR':
            from desisim.templates import MWS_STAR
            mwsstar = MWS_STAR(wave=wave)
            # todo: mag ranges for different programs of STAR targets should be in desimodel
            if 'rmagrange' not in obj_kwargs.keys():
                obj_kwargs['rmagrange'] = (15.0, 20.0)
            simflux, wave1, meta1 = mwsstar.make_templates(nmodel=nobj,
                                                           seed=seed,
                                                           **obj_kwargs)
            fibermap['DESI_TARGET'][ii] = desi_mask.MWS_ANY
            #- MWS bit names changed after desitarget 0.6.0 so use number
            #- instead of name for now (bit 0 = mask 1 = MWS_MAIN currently)
            fibermap['MWS_TARGET'][ii] = 1

        else:
            raise ValueError('Unable to simulate OBJTYPE={}'.format(objtype))

        flux[ii] = simflux
        meta[ii] = meta1

        fibermap['FILTER'][ii, :6] = [
            'DECAM_G', 'DECAM_R', 'DECAM_Z', 'WISE_W1', 'WISE_W2'
        ]
        fibermap['MAG'][ii, 0] = 22.5 - 2.5 * np.log10(meta['FLUX_G'][ii])
        fibermap['MAG'][ii, 1] = 22.5 - 2.5 * np.log10(meta['FLUX_R'][ii])
        fibermap['MAG'][ii, 2] = 22.5 - 2.5 * np.log10(meta['FLUX_Z'][ii])
        fibermap['MAG'][ii, 3] = 22.5 - 2.5 * np.log10(meta['FLUX_W1'][ii])
        fibermap['MAG'][ii, 4] = 22.5 - 2.5 * np.log10(meta['FLUX_W2'][ii])

    #- Load fiber -> positioner mapping and tile information
    fiberpos = desimodel.io.load_fiberpos()

    #- Where are these targets?  Centered on positioners for now.
    x = fiberpos['X'][specmin:specmin + nspec]
    y = fiberpos['Y'][specmin:specmin + nspec]
    fp = FocalPlane(tile_ra, tile_dec)
    ra = np.zeros(nspec)
    dec = np.zeros(nspec)
    for i in range(nspec):
        ra[i], dec[i] = fp.xy2radec(x[i], y[i])

    #- Fill in the rest of the fibermap structure
    fibermap['FIBER'] = np.arange(nspec, dtype='i4')
    fibermap['POSITIONER'] = fiberpos['POSITIONER'][specmin:specmin + nspec]
    fibermap['SPECTROID'] = fiberpos['SPECTROGRAPH'][specmin:specmin + nspec]
    fibermap['TARGETID'] = np.random.randint(sys.maxsize, size=nspec)
    fibermap['TARGETCAT'] = np.zeros(nspec, dtype=(str, 20))
    fibermap['LAMBDAREF'] = np.ones(nspec, dtype=np.float32) * 5400
    fibermap['RA_TARGET'] = ra
    fibermap['DEC_TARGET'] = dec
    fibermap['X_TARGET'] = x
    fibermap['Y_TARGET'] = y
    fibermap['X_FVCOBS'] = fibermap['X_TARGET']
    fibermap['Y_FVCOBS'] = fibermap['Y_TARGET']
    fibermap['X_FVCERR'] = np.zeros(nspec, dtype=np.float32)
    fibermap['Y_FVCERR'] = np.zeros(nspec, dtype=np.float32)
    fibermap['RA_OBS'] = fibermap['RA_TARGET']
    fibermap['DEC_OBS'] = fibermap['DEC_TARGET']
    fibermap['BRICKNAME'] = brick.brickname(ra, dec)

    return fibermap, (flux, wave, meta)
Example #17
0
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
Example #18
0
def get_spectra(lyafile,
                nqso=None,
                wave=None,
                templateid=None,
                normfilter='sdss2010-g',
                seed=None,
                rand=None,
                qso=None,
                nocolorcuts=False):
    '''Generate a QSO spectrum which includes Lyman-alpha absorption.

    Args:
        lyafile (str): name of the Lyman-alpha spectrum file to read.
        nqso (int, optional): number of spectra to generate (starting from the
          first spectrum; if more flexibility is needed use TEMPLATEID).
        wave (numpy.ndarray, optional): desired output wavelength vector.
        templateid (int numpy.ndarray, optional): indices of the spectra
          (0-indexed) to read from LYAFILE (default is to read everything).  If
          provided together with NQSO, TEMPLATEID wins.
        normfilter (str, optional): normalization filter
        seed (int, optional): Seed for random number generator.
        rand (numpy.RandomState, optional): RandomState object used for the
          random number generation.  If provided together with SEED, this
          optional input superseeds the numpy.RandomState object instantiated by
          SEED.
        qso (desisim.templates.QSO, optional): object with which to generate
          individual spectra/templates.
        nocolorcuts (bool, optional): Do not apply the fiducial rzW1W2 color-cuts
          cuts (default False).

    Returns:
        flux (numpy.ndarray):
            Array [nmodel, npix] of observed-frame spectra (erg/s/cm2/A).

        wave (numpy.ndarray):
            Observed-frame [npix] wavelength array (Angstrom).

        meta (astropy.Table):
            Table of meta-data [nmodel] for each output spectrum
            with columns defined in desisim.io.empty_metatable *plus* RA, DEC.

    '''
    import numpy as np
    from scipy.interpolate import interp1d

    import fitsio

    from speclite.filters import load_filters
    from desisim.templates import QSO
    from desisim.io import empty_metatable

    h = fitsio.FITS(lyafile)
    if templateid is None:
        if nqso is None:
            nqso = len(h) - 1
        templateid = np.arange(nqso)
    else:
        templateid = np.asarray(templateid)
        nqso = len(templateid)

    if rand is None:
        rand = np.random.RandomState(seed)
    templateseed = rand.randint(2**32, size=nqso)

    #heads = [head.read_header() for head in h[templateid + 1]]
    heads = []
    for indx in templateid:
        heads.append(h[indx + 1].read_header())

    zqso = np.array([head['ZQSO'] for head in heads])
    ra = np.array([head['RA'] for head in heads])
    dec = np.array([head['DEC'] for head in heads])
    mag_g = np.array([head['MAG_G'] for head in heads])

    # Hard-coded filtername!
    normfilt = load_filters(normfilter)

    if qso is None:
        qso = QSO(normfilter=normfilter, wave=wave)

    wave = qso.wave
    flux = np.zeros([nqso, len(wave)], dtype='f4')

    meta = empty_metatable(objtype='QSO', nmodel=nqso)
    meta['TEMPLATEID'] = templateid
    meta['REDSHIFT'] = zqso
    meta['MAG'] = mag_g
    meta['SEED'] = templateseed
    meta['RA'] = ra
    meta['DEC'] = dec

    for ii, indx in enumerate(templateid):
        flux1, _, meta1 = qso.make_templates(nmodel=1,
                                             redshift=np.atleast_1d(zqso[ii]),
                                             mag=np.atleast_1d(mag_g[ii]),
                                             seed=templateseed[ii],
                                             nocolorcuts=nocolorcuts)
        flux1 *= 1e-17
        for col in meta1.colnames:
            meta[col][ii] = meta1[col][0]

        # read lambda and forest transmission
        data = h[indx + 1].read()
        la = data['LAMBDA'][:]
        tr = data['FLUX'][:]

        if len(tr):
            # Interpolate the transmission at the spectral wavelengths, if
            # outside the forest, the transmission is 1.
            itr = interp1d(la, tr, bounds_error=False, fill_value=1.0)
            flux1 *= itr(wave)

        padflux, padwave = normfilt.pad_spectrum(flux1, wave, method='edge')
        normmaggies = np.array(
            normfilt.get_ab_maggies(padflux, padwave,
                                    mask_invalid=True)[normfilter])
        factor = 10**(-0.4 * mag_g[ii]) / normmaggies
        flux1 *= factor
        for key in ('FLUX_G', 'FLUX_R', 'FLUX_Z', 'FLUX_W1', 'FLUX_W2'):
            meta[key][ii] *= factor

        flux[ii, :] = flux1[:]

    h.close()

    return flux, wave, meta
Example #19
0
def get_spectra(lyafile,
                nqso=None,
                wave=None,
                templateid=None,
                normfilter='sdss2010-g',
                seed=None,
                rand=None,
                qso=None,
                add_dlas=False,
                debug=False,
                nocolorcuts=True):
    """Generate a QSO spectrum which includes Lyman-alpha absorption.

    Args:
        lyafile (str): name of the Lyman-alpha spectrum file to read.
        nqso (int, optional): number of spectra to generate (starting from the
          first spectrum; if more flexibility is needed use TEMPLATEID).
        wave (numpy.ndarray, optional): desired output wavelength vector.
        templateid (int numpy.ndarray, optional): indices of the spectra
          (0-indexed) to read from LYAFILE (default is to read everything).  If
          provided together with NQSO, TEMPLATEID wins.
        normfilter (str, optional): normalization filter
        seed (int, optional): Seed for random number generator.
        rand (numpy.RandomState, optional): RandomState object used for the
          random number generation.  If provided together with SEED, this
          optional input superseeds the numpy.RandomState object instantiated by
          SEED.
        qso (desisim.templates.QSO, optional): object with which to generate
          individual spectra/templates.
        add_dlas (bool): Inject damped Lya systems into the Lya forest
          These are done according to the current best estimates for the incidence dN/dz
          (Prochaska et al. 2008, ApJ, 675, 1002)
          Set in calc_lz
          These are *not* inserted according to overdensity along the sightline
        nocolorcuts (bool, optional): Do not apply the fiducial rzW1W2 color-cuts
          cuts (default True).

    Returns (flux, wave, meta, dla_meta) where:

        * flux (numpy.ndarray): Array [nmodel, npix] of observed-frame spectra
          (erg/s/cm2/A).
        * wave (numpy.ndarray): Observed-frame [npix] wavelength array (Angstrom).
        * meta (astropy.Table): Table of meta-data [nmodel] for each output spectrum
          with columns defined in desisim.io.empty_metatable *plus* RA, DEC.
        * objmeta (astropy.Table): Table of additional object-specific meta-data
          [nmodel] for each output spectrum with columns defined in
          desisim.io.empty_metatable.
        * dla_meta (astropy.Table): Table of meta-data [ndla] for the DLAs injected
          into the spectra.  Only returned if add_dlas=True

    Note: `dla_meta` is only included if add_dlas=True.

    """
    from scipy.interpolate import interp1d

    import fitsio

    from speclite.filters import load_filters
    from desisim.templates import QSO
    from desisim.io import empty_metatable

    h = fitsio.FITS(lyafile)
    if templateid is None:
        if nqso is None:
            nqso = len(h) - 1
        templateid = np.arange(nqso)
    else:
        templateid = np.asarray(templateid)
        nqso = len(np.atleast_1d(templateid))

    if rand is None:
        rand = np.random.RandomState(seed)
    templateseed = rand.randint(2**32, size=nqso)

    #heads = [head.read_header() for head in h[templateid + 1]]
    heads = []
    for indx in templateid:
        heads.append(h[indx + 1].read_header())

    zqso = np.array([head['ZQSO'] for head in heads])
    ra = np.array([head['RA'] for head in heads])
    dec = np.array([head['DEC'] for head in heads])
    mag_g = np.array([head['MAG_G'] for head in heads])

    # Hard-coded filtername!  Should match MAG_G!
    normfilt = load_filters(normfilter)

    if qso is None:
        qso = QSO(normfilter_south=normfilter, wave=wave)

    wave = qso.wave
    flux = np.zeros([nqso, len(wave)], dtype='f4')

    meta, objmeta = empty_metatable(objtype='QSO', nmodel=nqso)
    meta['TEMPLATEID'][:] = templateid
    meta['REDSHIFT'][:] = zqso
    meta['MAG'][:] = mag_g
    meta['MAGFILTER'][:] = normfilter
    meta['SEED'][:] = templateseed
    meta['RA'] = ra
    meta['DEC'] = dec

    # Lists for DLA meta data
    if add_dlas:
        dla_NHI, dla_z, dla_id = [], [], []

    # Loop on quasars
    for ii, indx in enumerate(templateid):
        flux1, _, meta1, objmeta1 = qso.make_templates(
            nmodel=1,
            redshift=np.atleast_1d(zqso[ii]),
            mag=np.atleast_1d(mag_g[ii]),
            seed=templateseed[ii],
            nocolorcuts=nocolorcuts,
            lyaforest=False)
        flux1 *= 1e-17
        for col in meta1.colnames:
            meta[col][ii] = meta1[col][0]
        for col in objmeta1.colnames:
            objmeta[col][ii] = objmeta1[col][0]

        # read lambda and forest transmission
        data = h[indx + 1].read()
        la = data['LAMBDA'][:]
        tr = data['FLUX'][:]

        if len(tr):
            # Interpolate the transmission at the spectral wavelengths, if
            # outside the forest, the transmission is 1.
            itr = interp1d(la, tr, bounds_error=False, fill_value=1.0)
            flux1 *= itr(wave)

        # Inject a DLA?
        if add_dlas:
            if np.min(wave / lambda_RF_LYA - 1) < zqso[ii]:  # Any forest?
                dlas, dla_model = insert_dlas(wave,
                                              zqso[ii],
                                              seed=templateseed[ii])
                ndla = len(dlas)
                if ndla > 0:
                    flux1 *= dla_model
                    # Meta
                    dla_z += [idla['z'] for idla in dlas]
                    dla_NHI += [idla['N'] for idla in dlas]
                    dla_id += [indx] * ndla

        padflux, padwave = normfilt.pad_spectrum(flux1, wave, method='edge')
        normmaggies = np.array(
            normfilt.get_ab_maggies(padflux, padwave,
                                    mask_invalid=True)[normfilter])

        factor = 10**(-0.4 * mag_g[ii]) / normmaggies
        flux1 *= factor
        for key in ('FLUX_G', 'FLUX_R', 'FLUX_Z', 'FLUX_W1', 'FLUX_W2'):
            meta[key][ii] *= factor
        flux[ii, :] = flux1[:]

    h.close()

    # Finish
    if add_dlas:
        ndla = len(dla_id)
        if ndla > 0:
            from astropy.table import Table
            dla_meta = Table()
            dla_meta['NHI'] = dla_NHI  # log NHI values
            dla_meta['z'] = dla_z
            dla_meta['ID'] = dla_id
        else:
            dla_meta = None
        return flux * 1e17, wave, meta, objmeta, dla_meta
    else:
        return flux * 1e17, wave, meta, objmeta
Example #20
0
    def qso(self, data, index=None, mockformat='gaussianfield'):
        """Generate spectra for the QSO or QSO/LYA samples.

        Note: We need to make sure NORMFILTER matches!

        """
        from desisim.lya_spectra import get_spectra
        
        objtype = 'QSO'
        if index is None:
            index = np.arange(len(data['Z']))
        nobj = len(index)

        if mockformat.lower() == 'gaussianfield':
            input_meta = empty_metatable(nmodel=nobj, objtype=objtype)
            for inkey, datakey in zip(('SEED', 'MAG', 'REDSHIFT'),
                                      ('SEED', 'MAG', 'Z')):
                input_meta[inkey] = data[datakey][index]

            # Build the tracer and Lya forest QSO spectra separately.
            meta = empty_metatable(nmodel=nobj, objtype=objtype)
            flux = np.zeros([nobj, len(self.wave)], dtype='f4')

            lya = np.where( data['TEMPLATESUBTYPE'][index] == 'LYA' )[0]
            tracer = np.where( data['TEMPLATESUBTYPE'][index] == '' )[0]

            if len(tracer) > 0:
                flux1, _, meta1 = self.qso_templates.make_templates(input_meta=input_meta[tracer],
                                                                    lyaforest=False,
                                                                    nocolorcuts=True,
                                                                    verbose=self.verbose)
                meta[tracer] = meta1
                flux[tracer, :] = flux1

            if len(lya) > 0:
                alllyafile = data['LYAFILES'][index][lya]
                alllyahdu = data['LYAHDU'][index][lya]
                
                for lyafile in sorted(set(alllyafile)):
                    these = np.where( lyafile == alllyafile )[0]

                    templateid = alllyahdu[these] - 1 # templateid is 0-indexed
                    flux1, _, meta1 = get_spectra(lyafile, templateid=templateid, normfilter=data['FILTERNAME'],
                                                  rand=self.rand, qso=self.lya_templates, nocolorcuts=True)
                    meta1['SUBTYPE'] = 'LYA'
                    meta[lya[these]] = meta1
                    flux[lya[these], :] = flux1

        elif mockformat.lower() == 'lya':
            # Build spectra for Lyman-alpha QSOs. Deprecated!
            from desisim.lya_spectra import get_spectra
            from desitarget.mock.io import decode_rownum_filenum

            meta = empty_metatable(nmodel=nobj, objtype=objtype)
            flux = np.zeros([nobj, len(self.wave)], dtype='f4')
            
            rowindx, fileindx = decode_rownum_filenum(data['MOCKID'][index])
            for indx1 in set(fileindx):
                lyafile = data['FILES'][indx1]
                these = np.where(indx1 == fileindx)[0]
                templateid = rowindx[these].astype('int')
            
                flux1, _, meta1 = get_spectra(lyafile, templateid=templateid,
                                              normfilter=data['FILTERNAME'],
                                              rand=self.rand, qso=self.lya_templates)
                meta[these] = meta1
                flux[these, :] = flux1
        else:
            raise ValueError('Unrecognized mockformat {}!'.format(mockformat))

        return flux, meta
Example #21
0
def get_targets(nspec, program, tileid=None, seed=None, specify_targets=dict(), specmin=0):
    """
    Generates a set of targets for the requested program

    Args:
        nspec: (int) number of targets to generate
        program: (str) program name DARK, BRIGHT, GRAY, MWS, BGS, LRG, ELG, ...

    Options:
      * tileid: (int) tileid, used for setting RA,dec
      * seed: (int) random number seed
      * specify_targets: (dict of dicts)  Define target properties like magnitude and redshift
                                 for each target class. Each objtype has its own key,value pair
                                 see simspec.templates.specify_galparams_dict() 
                                 or simsepc.templates.specify_starparams_dict()
      * specmin: (int) first spectrum number (0-indexed)

    Returns:
      * fibermap
      * targets as tuple of (flux, wave, meta)
    """
    if tileid is None:
        tile_ra, tile_dec = 0.0, 0.0
    else:
        tile_ra, tile_dec = io.get_tile_radec(tileid)

    program = program.upper()
    log.debug('Using random seed {}'.format(seed))
    np.random.seed(seed)

    #- Get distribution of target types
    true_objtype, target_objtype = sample_objtype(nspec, program)

    #- Get DESI wavelength coverage
    try:
        params = desimodel.io.load_desiparams()
        wavemin = params['ccd']['b']['wavemin']
        wavemax = params['ccd']['z']['wavemax']
    except KeyError:
        wavemin = desimodel.io.load_throughput('b').wavemin
        wavemax = desimodel.io.load_throughput('z').wavemax
    dw = 0.2
    wave = np.arange(round(wavemin, 1), wavemax, dw)
    nwave = len(wave)

    flux = np.zeros( (nspec, len(wave)) )
    meta, _ = empty_metatable(nmodel=nspec, objtype='SKY')
    objmeta = dict()
    fibermap = empty_fibermap(nspec)

    targetid = np.random.randint(sys.maxsize, size=nspec).astype(np.int64)
    meta['TARGETID'] = targetid
    fibermap['TARGETID'] = targetid
    
    for objtype in set(true_objtype):
        ii = np.where(true_objtype == objtype)[0]
        nobj = len(ii)

        fibermap['OBJTYPE'][ii] = target_objtype[ii]

        if objtype in specify_targets.keys():
            obj_kwargs = specify_targets[objtype]
        else:
            obj_kwargs = dict()
                
        # Simulate spectra
        if objtype == 'SKY':
            fibermap['DESI_TARGET'][ii] = desi_mask.SKY
            continue

        elif objtype == 'ELG':
            from desisim.templates import ELG
            elg = ELG(wave=wave)
            simflux, wave1, meta1, objmeta1 = elg.make_templates(nmodel=nobj, seed=seed, **obj_kwargs)
            fibermap['DESI_TARGET'][ii] = desi_mask.ELG

        elif objtype == 'LRG':
            from desisim.templates import LRG
            lrg = LRG(wave=wave)
            simflux, wave1, meta1, objmeta1 = lrg.make_templates(nmodel=nobj, seed=seed, **obj_kwargs)
            fibermap['DESI_TARGET'][ii] = desi_mask.LRG

        elif objtype == 'BGS':
            from desisim.templates import BGS
            bgs = BGS(wave=wave)
            simflux, wave1, meta1, objmeta1 = bgs.make_templates(nmodel=nobj, seed=seed, **obj_kwargs)
            fibermap['DESI_TARGET'][ii] = desi_mask.BGS_ANY
            fibermap['BGS_TARGET'][ii] = bgs_mask.BGS_BRIGHT

        elif objtype == 'QSO':
            from desisim.templates import QSO
            qso = QSO(wave=wave)
            simflux, wave1, meta1, objmeta1 = qso.make_templates(nmodel=nobj, seed=seed, lyaforest=False, **obj_kwargs)
            fibermap['DESI_TARGET'][ii] = desi_mask.QSO

        # For a "bad" QSO simulate a normal star without color cuts, which isn't
        # right. We need to apply the QSO color-cuts to the normal stars to pull
        # out the correct population of contaminating stars.

        # Note by @moustakas: we can now do this using desisim/#150, but we are
        # going to need 'noisy' photometry (because the QSO color-cuts
        # explicitly avoid the stellar locus).
        elif objtype == 'QSO_BAD':
            from desisim.templates import STAR
            #from desitarget.cuts import isQSO
            #star = STAR(wave=wave, colorcuts_function=isQSO)
            star = STAR(wave=wave)
            simflux, wave1, meta1, objmeta1 = star.make_templates(nmodel=nobj, seed=seed, **obj_kwargs)
            fibermap['DESI_TARGET'][ii] = desi_mask.QSO

        elif objtype == 'STD':
            from desisim.templates import STD
            std = STD(wave=wave)
            simflux, wave1, meta1, objmeta1 = std.make_templates(nmodel=nobj, seed=seed, **obj_kwargs)
            #- Loop options for forwards/backwards compatibility
            for name in ['STD_FAINT', 'STD_FSTAR', 'STD']:
                if name in desi_mask.names():
                    fibermap['DESI_TARGET'][ii] |= desi_mask[name]
                    break

        elif objtype == 'MWS_STAR':
            from desisim.templates import MWS_STAR
            mwsstar = MWS_STAR(wave=wave)
            # TODO: mag ranges for different programs of STAR targets should be in desimodel
            if 'magrange' not in obj_kwargs.keys():
                obj_kwargs['magrange'] = (15.0,20.0)
            simflux, wave1, meta1, objmeta1 = mwsstar.make_templates(nmodel=nobj, seed=seed, **obj_kwargs)
            fibermap['DESI_TARGET'][ii] |= desi_mask.MWS_ANY
            #- MWS bit names changed after desitarget 0.6.0 so use number
            #- instead of name for now (bit 0 = mask 1 = MWS_MAIN currently)
            fibermap['MWS_TARGET'][ii] = 1

        else:
            raise ValueError('Unable to simulate OBJTYPE={}'.format(objtype))

        # Assign targetid
        meta1['TARGETID'] = targetid[ii]
        if hasattr(objmeta1, 'data'): # simqso.sqgrids.QsoSimPoints object
            objmeta1.data['TARGETID'] = targetid[ii]
        else:
            if len(objmeta1) > 0:
                objmeta1['TARGETID'] = targetid[ii]
                # We want the dict key tied to the "true" object type (e.g., STAR),
                # not, e.g., QSO_BAD.
                objmeta[meta1['OBJTYPE'][0]] = objmeta1

        flux[ii] = simflux
        meta[ii] = meta1

        for band in ['G', 'R', 'Z', 'W1', 'W2']:
            key = 'FLUX_'+band
            fibermap[key][ii] = meta[key][ii]
            #- TODO: FLUX_IVAR

    #- Load fiber -> positioner mapping and tile information
    fiberpos = desimodel.io.load_fiberpos()

    #- Where are these targets?  Centered on positioners for now.
    x = fiberpos['X'][specmin:specmin+nspec]
    y = fiberpos['Y'][specmin:specmin+nspec]
    fp = FocalPlane(tile_ra, tile_dec)
    ra = np.zeros(nspec)
    dec = np.zeros(nspec)
    for i in range(nspec):
        ra[i], dec[i] = fp.xy2radec(x[i], y[i])

    #- Fill in the rest of the fibermap structure
    fibermap['FIBER'] = np.arange(nspec, dtype='i4')
    fibermap['POSITIONER'] = fiberpos['POSITIONER'][specmin:specmin+nspec]
    fibermap['SPECTROID'] = fiberpos['SPECTROGRAPH'][specmin:specmin+nspec]
    fibermap['TARGETCAT'] = np.zeros(nspec, dtype=(str, 20))
    fibermap['LAMBDA_REF'] = np.ones(nspec, dtype=np.float32)*5400
    fibermap['TARGET_RA'] = ra
    fibermap['TARGET_DEC'] = dec
    fibermap['FIBERASSIGN_X'] = x
    fibermap['FIBERASSIGN_Y'] = y
    fibermap['FIBER_RA'] = fibermap['TARGET_RA']
    fibermap['FIBER_DEC'] = fibermap['TARGET_DEC']
    fibermap['BRICKNAME'] = brick.brickname(ra, dec)

    return fibermap, (flux, wave, meta, objmeta)
Example #22
0
def get_targets(nspec, program, tileid=None, seed=None, specify_targets=dict(), specmin=0):
    """
    Generates a set of targets for the requested program

    Args:
        nspec: (int) number of targets to generate
        program: (str) program name DARK, BRIGHT, GRAY, MWS, BGS, LRG, ELG, ...

    Options:
      * tileid: (int) tileid, used for setting RA,dec
      * seed: (int) random number seed
      * specify_targets: (dict of dicts)  Define target properties like magnitude and redshift
                                 for each target class. Each objtype has its own key,value pair
                                 see simspec.templates.specify_galparams_dict() 
                                 or simsepc.templates.specify_starparams_dict()
      * specmin: (int) first spectrum number (0-indexed)

    Returns:
      * fibermap
      * targets as tuple of (flux, wave, meta)
    """
    if tileid is None:
        tile_ra, tile_dec = 0.0, 0.0
    else:
        tile_ra, tile_dec = io.get_tile_radec(tileid)

    program = program.upper()
    log.debug('Using random seed {}'.format(seed))
    np.random.seed(seed)

    #- Get distribution of target types
    true_objtype, target_objtype = sample_objtype(nspec, program)

    #- Get DESI wavelength coverage
    try:
        params = desimodel.io.load_desiparams()
        wavemin = params['ccd']['b']['wavemin']
        wavemax = params['ccd']['z']['wavemax']
    except KeyError:
        wavemin = desimodel.io.load_throughput('b').wavemin
        wavemax = desimodel.io.load_throughput('z').wavemax
    dw = 0.2
    wave = np.arange(round(wavemin, 1), wavemax, dw)
    nwave = len(wave)

    flux = np.zeros( (nspec, len(wave)) )
    meta, _ = empty_metatable(nmodel=nspec, objtype='SKY')
    objmeta = dict()
    fibermap = empty_fibermap(nspec)

    targetid = np.random.randint(sys.maxsize, size=nspec).astype(np.int64)
    meta['TARGETID'] = targetid
    fibermap['TARGETID'] = targetid
    
    for objtype in set(true_objtype):
        ii = np.where(true_objtype == objtype)[0]
        nobj = len(ii)

        fibermap['OBJTYPE'][ii] = target_objtype[ii]

        if objtype in specify_targets.keys():
            obj_kwargs = specify_targets[objtype]
        else:
            obj_kwargs = dict()
                
        # Simulate spectra
        if objtype == 'SKY':
            fibermap['DESI_TARGET'][ii] = desi_mask.SKY
            continue

        elif objtype == 'ELG':
            from desisim.templates import ELG
            elg = ELG(wave=wave)
            simflux, wave1, meta1, objmeta1 = elg.make_templates(nmodel=nobj, seed=seed, **obj_kwargs)
            fibermap['DESI_TARGET'][ii] = desi_mask.ELG

        elif objtype == 'LRG':
            from desisim.templates import LRG
            lrg = LRG(wave=wave)
            simflux, wave1, meta1, objmeta1 = lrg.make_templates(nmodel=nobj, seed=seed, **obj_kwargs)
            fibermap['DESI_TARGET'][ii] = desi_mask.LRG

        elif objtype == 'BGS':
            from desisim.templates import BGS
            bgs = BGS(wave=wave)
            simflux, wave1, meta1, objmeta1 = bgs.make_templates(nmodel=nobj, seed=seed, **obj_kwargs)
            fibermap['DESI_TARGET'][ii] = desi_mask.BGS_ANY
            fibermap['BGS_TARGET'][ii] = bgs_mask.BGS_BRIGHT

        elif objtype == 'QSO':
            from desisim.templates import QSO
            qso = QSO(wave=wave)
            simflux, wave1, meta1, objmeta1 = qso.make_templates(nmodel=nobj, seed=seed, lyaforest=False, **obj_kwargs)
            fibermap['DESI_TARGET'][ii] = desi_mask.QSO

        # For a "bad" QSO simulate a normal star without color cuts, which isn't
        # right. We need to apply the QSO color-cuts to the normal stars to pull
        # out the correct population of contaminating stars.

        # Note by @moustakas: we can now do this using desisim/#150, but we are
        # going to need 'noisy' photometry (because the QSO color-cuts
        # explicitly avoid the stellar locus).
        elif objtype == 'QSO_BAD':
            from desisim.templates import STAR
            #from desitarget.cuts import isQSO
            #star = STAR(wave=wave, colorcuts_function=isQSO)
            star = STAR(wave=wave)
            simflux, wave1, meta1, objmeta1 = star.make_templates(nmodel=nobj, seed=seed, **obj_kwargs)
            fibermap['DESI_TARGET'][ii] = desi_mask.QSO

        elif objtype == 'STD':
            from desisim.templates import STD
            std = STD(wave=wave)
            simflux, wave1, meta1, objmeta1 = std.make_templates(nmodel=nobj, seed=seed, **obj_kwargs)
            #- Loop options for forwards/backwards compatibility
            for name in ['STD_FAINT', 'STD_FSTAR', 'STD']:
                if name in desi_mask.names():
                    fibermap['DESI_TARGET'][ii] |= desi_mask[name]
                    break

        elif objtype == 'MWS_STAR':
            from desisim.templates import MWS_STAR
            mwsstar = MWS_STAR(wave=wave)
            # TODO: mag ranges for different programs of STAR targets should be in desimodel
            if 'magrange' not in obj_kwargs.keys():
                obj_kwargs['magrange'] = (15.0,20.0)
            simflux, wave1, meta1, objmeta1 = mwsstar.make_templates(nmodel=nobj, seed=seed, **obj_kwargs)
            fibermap['DESI_TARGET'][ii] |= desi_mask.MWS_ANY
            #- MWS bit names changed after desitarget 0.6.0 so use number
            #- instead of name for now (bit 0 = mask 1 = MWS_MAIN currently)
            fibermap['MWS_TARGET'][ii] = 1

        else:
            raise ValueError('Unable to simulate OBJTYPE={}'.format(objtype))

        # Assign targetid
        meta1['TARGETID'] = targetid[ii]
        if hasattr(objmeta1, 'data'): # simqso.sqgrids.QsoSimPoints object
            objmeta1.data['TARGETID'] = targetid[ii]
        else:
            if len(objmeta1) > 0:
                objmeta1['TARGETID'] = targetid[ii]
                # We want the dict key tied to the "true" object type (e.g., STAR),
                # not, e.g., QSO_BAD.
                objmeta[meta1['OBJTYPE'][0]] = objmeta1

        flux[ii] = simflux
        meta[ii] = meta1

        for band in ['G', 'R', 'Z', 'W1', 'W2']:
            key = 'FLUX_'+band
            fibermap[key][ii] = meta[key][ii]
            #- TODO: FLUX_IVAR

    #- Load fiber -> positioner mapping and tile information
    fiberpos = desimodel.io.load_fiberpos()

    #- Where are these targets?  Centered on positioners for now.
    x = fiberpos['X'][specmin:specmin+nspec]
    y = fiberpos['Y'][specmin:specmin+nspec]
    fp = FocalPlane(tile_ra, tile_dec)
    ra = np.zeros(nspec)
    dec = np.zeros(nspec)
    for i in range(nspec):
        ra[i], dec[i] = fp.xy2radec(x[i], y[i])

    #- Fill in the rest of the fibermap structure
    fibermap['FIBER'] = np.arange(nspec, dtype='i4')
    fibermap['POSITIONER'] = fiberpos['POSITIONER'][specmin:specmin+nspec]
    fibermap['SPECTROID'] = fiberpos['SPECTROGRAPH'][specmin:specmin+nspec]
    fibermap['TARGETCAT'] = np.zeros(nspec, dtype=(str, 20))
    fibermap['LAMBDA_REF'] = np.ones(nspec, dtype=np.float32)*5400
    fibermap['TARGET_RA'] = ra
    fibermap['TARGET_DEC'] = dec
    fibermap['DESIGN_X'] = x
    fibermap['DESIGN_Y'] = y
    fibermap['FIBER_RA'] = fibermap['TARGET_RA']
    fibermap['FIBER_DEC'] = fibermap['TARGET_DEC']
    fibermap['BRICKNAME'] = brick.brickname(ra, dec)

    return fibermap, (flux, wave, meta, objmeta)