Ejemplo n.º 1
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)
Ejemplo n.º 2
0
 def test_input_wave(self):
     '''Confirm that we can specify the wavelength array.'''
     print('In function test_input_wave, seed = {}'.format(self.seed))
     lrg = LRG(minwave=self.wavemin, maxwave=self.wavemax, cdelt=self.dwave)
     flux, wave, meta = lrg.make_templates(self.nspec, seed=self.seed)
     self._check_output_size(flux, wave, meta)
Ejemplo n.º 3
0
def combspec(ELGmag, BGSmag, ELGseed, BGSseed, ELGrShifts=None, BGSrShifts=None, nrShifts=None, returnmeta=False, sepflux=False):

    BGSmags = []
    ELGmags = []
    i = 0

    if ELGrShifts is not None and BGSrShifts is not None:

        while i < len(BGSrShifts):

            BGSmags.append(BGSmag)
            ELGmags.append(ELGmag)
            i += 1

        makeBGS = LRG()
        makeELG = ELG()

        fluxBGS, waveBGS, metaBGS, objmetaBGS = makeBGS.make_templates(seed=BGSseed, nmodel=len(BGSrShifts), redshift=BGSrShifts, mag=BGSmags, nocolorcuts=True)
        fluxELG, waveELG, metaELG, objmetaELG = makeELG.make_templates(seed=ELGseed, nmodel=len(ELGrShifts), redshift=ELGrShifts, mag=ELGmags, nocolorcuts=True)

        wave = waveBGS

        flux = fluxBGS + fluxELG

        if returnmeta == True and sepflux == True:
            return wave, flux, metaBGS, metaELG, fluxBGS, fluxELG

        elif returnmeta == True:
            return wave, flux, metaBGS, metaELG

        elif sepflux == True:
            return wave, flux, fluxBGS, fluxELG

        else:
            return wave, flux

    if nrShifts is not None:

        while i < nrShifts:

            BGSmags.append(BGSmag)
            ELGmags.append(ELGmag)
            i += 1

        makeBGS = BGS()
        makeELG = ELG()

        fluxBGS, waveBGS, metaBGS, objmetaBGS = makeBGS.make_templates(seed=BGSseed, nmodel=nrShifts, mag=BGSmags, nocolorcuts=True)
        fluxELG, waveELG, metaELG, objmetaELG = makeELG.make_templates(seed=ELGseed, nmodel=nrShifts, mag=ELGmags, nocolorcuts=True)

        wave = waveBGS
        flux = fluxBGS + fluxELG

        if returnmeta == True and sepflux == True:
            return wave, flux, metaBGS, metaELG, fluxBGS, fluxELG

        elif returnmeta == True:
            return wave, flux, metaBGS, metaELG

        elif sepflux == True:
            return wave, flux, fluxBGS, fluxELG

        else:
            return wave, flux
Ejemplo n.º 4
0
class MockSpectra(object):
    """Generate spectra for each type of mock.  Currently just choose the closest
    template; we can get fancier later.

    ToDo (@moustakas): apply Galactic extinction.

    """
    def __init__(self, wavemin=None, wavemax=None, dw=0.2, nproc=1,
                 rand=None, verbose=False):

        from desimodel.io import load_throughput
        
        self.tree = TemplateKDTree(nproc=nproc)

        # Build a default (buffered) wavelength vector.
        if wavemin is None:
            wavemin = load_throughput('b').wavemin - 10.0
        if wavemax is None:
            wavemax = load_throughput('z').wavemax + 10.0
            
        self.wavemin = wavemin
        self.wavemax = wavemax
        self.dw = dw
        self.wave = np.arange(round(wavemin, 1), wavemax, dw)

        self.rand = rand
        self.verbose = verbose

        #self.__normfilter = 'decam2014-r' # default normalization filter

        # Initialize the templates once:
        from desisim.templates import BGS, ELG, LRG, QSO, STAR, WD
        self.bgs_templates = BGS(wave=self.wave, normfilter='sdss2010-r') # Need to generalize this!
        self.elg_templates = ELG(wave=self.wave, normfilter='decam2014-r')
        self.lrg_templates = LRG(wave=self.wave, normfilter='decam2014-z')
        self.qso_templates = QSO(wave=self.wave, normfilter='decam2014-g')
        self.lya_templates = QSO(wave=self.wave, normfilter='decam2014-g')
        self.star_templates = STAR(wave=self.wave, normfilter='decam2014-r')
        self.wd_da_templates = WD(wave=self.wave, normfilter='decam2014-g', subtype='DA')
        self.wd_db_templates = WD(wave=self.wave, normfilter='decam2014-g', subtype='DB')
        
    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
        TemplateKDTree.bgs().

        """
        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

    def elg(self, data, index=None, mockformat='gaussianfield'):
        """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))

        if False:
            import matplotlib.pyplot as plt
            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

    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

    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

    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

    def mws_nearby(self, data, index=None, mockformat='100pc'):
        """Generate spectra for the MWS_NEARBY sample.

        """
        flux, meta = self.mws(data, index=index, mockformat=mockformat)
        return flux, meta

    def mws_main(self, data, index=None, mockformat='galaxia'):
        """Generate spectra for the MWS_MAIN sample.

        """
        flux, meta = self.mws(data, index=index, mockformat=mockformat)
        return flux, meta

    def faintstar(self, data, index=None, mockformat='galaxia'):
        """Generate spectra for the FAINTSTAR (faint stellar) sample.

        """
        flux, meta = self.mws(data, index=index, mockformat=mockformat)
        return flux, meta

    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

    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

    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]
        flux = np.zeros((nobj, len(self.wave)), dtype='i1')

        return flux, meta
Ejemplo n.º 5
0
def get_targets(nspec, tileid=None):
    """
    Returns:
        fibermap
        truth table

    TODO (@moustakas): Deal with the random seed correctly.

    TODO: document this better
    """
    if tileid is None:
        tile_ra, tile_dec = 0.0, 0.0
    else:
        tile_ra, tile_dec = io.get_tile_radec(tileid)

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

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

    truth = dict()
    truth['FLUX'] = np.zeros((nspec, len(wave)))
    truth['REDSHIFT'] = np.zeros(nspec, dtype='f4')
    truth['TEMPLATEID'] = np.zeros(nspec, dtype='i4')
    truth['OIIFLUX'] = np.zeros(nspec, dtype='f4')
    truth['D4000'] = np.zeros(nspec, dtype='f4')
    truth['VDISP'] = np.zeros(nspec, dtype='f4')
    truth['OBJTYPE'] = np.zeros(nspec, dtype='S10')
    #- Note: unlike other elements, first index of WAVE isn't spectrum index
    truth['WAVE'] = wave

    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]
        truth['OBJTYPE'][ii] = true_objtype[ii]

        # Simulate spectra
        if objtype == 'SKY':
            continue

        elif objtype == 'ELG':
            from desisim.templates import ELG
            elg = ELG(wave=wave)
            simflux, wave1, meta = elg.make_templates(nmodel=nobj)

        elif objtype == 'LRG':
            from desisim.templates import LRG
            lrg = LRG(wave=wave)
            simflux, wave1, meta = lrg.make_templates(nmodel=nobj)

        elif objtype == 'QSO':
            from desisim.templates import QSO
            qso = QSO(wave=wave)
            simflux, wave1, meta = qso.make_templates(nmodel=nobj)

        # 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.
        elif objtype == 'QSO_BAD':
            from desisim.templates import STAR
            star = STAR(wave=wave)
            simflux, wave1, meta = star.make_templates(nmodel=nobj)

        elif objtype == 'STD':
            from desisim.templates import STAR
            star = STAR(wave=wave, FSTD=True)
            simflux, wave1, meta = star.make_templates(nmodel=nobj)

        truth['FLUX'][ii] = 1e17 * simflux
        truth['UNITS'] = '1e-17 erg/s/cm2/A'
        truth['TEMPLATEID'][ii] = meta['TEMPLATEID']
        truth['REDSHIFT'][ii] = meta['REDSHIFT']

        # Pack in the photometry.  TODO: Include WISE.
        magg = meta['GMAG']
        magr = meta['RMAG']
        magz = meta['ZMAG']
        fibermap['MAG'][ii, 0:3] = np.vstack([magg, magr, magz]).T
        fibermap['FILTER'][ii, 0:3] = ['DECAM_G', 'DECAM_R', 'DECAM_Z']

        if objtype == 'ELG':
            truth['OIIFLUX'][ii] = meta['OIIFLUX']
            truth['D4000'][ii] = meta['D4000']
            truth['VDISP'][ii] = meta['VDISP']

        if objtype == 'LRG':
            truth['D4000'][ii] = meta['D4000']
            truth['VDISP'][ii] = meta['VDISP']

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

    #- Where are these targets?  Centered on positioners for now.
    x = fiberpos['X'][0:nspec]
    y = fiberpos['Y'][0: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'][0:nspec]
    fibermap['SPECTROID'] = fiberpos['SPECTROGRAPH'][0:nspec]
    fibermap['TARGETID'] = np.random.randint(sys.maxint, size=nspec)
    fibermap['TARGETCAT'] = np.zeros(nspec, dtype='|S20')
    fibermap['LAMBDAREF'] = np.ones(nspec, dtype=np.float32) * 5400
    fibermap['TARGET_MASK0'] = np.zeros(nspec, dtype='i8')
    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, truth
Ejemplo n.º 6
0
def get_targets(nspec, tileid=None):
    """
    Returns:
        fibermap
        truth table

    TODO (@moustakas): Deal with the random seed correctly. 
    
    TODO: document this better
    """
    if tileid is None:
        tile_ra, tile_dec = 0.0, 0.0
    else:
        tile_ra, tile_dec = io.get_tile_radec(tileid)

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

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

    truth = dict()
    truth["FLUX"] = np.zeros((nspec, len(wave)))
    truth["REDSHIFT"] = np.zeros(nspec, dtype="f4")
    truth["TEMPLATEID"] = np.zeros(nspec, dtype="i4")
    truth["OIIFLUX"] = np.zeros(nspec, dtype="f4")
    truth["D4000"] = np.zeros(nspec, dtype="f4")
    truth["OBJTYPE"] = np.zeros(nspec, dtype="S10")
    # - Note: unlike other elements, first index of WAVE isn't spectrum index
    truth["WAVE"] = wave

    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]
        truth["OBJTYPE"][ii] = true_objtype[ii]

        # Simulate spectra
        if objtype == "SKY":
            continue

        elif objtype == "ELG":
            from desisim.templates import ELG

            elg = ELG(nmodel=nobj, wave=wave)
            simflux, wave1, meta = elg.make_templates()

        elif objtype == "LRG":
            from desisim.templates import LRG

            lrg = LRG(nmodel=nobj, wave=wave)
            simflux, wave1, meta = lrg.make_templates()

        elif objtype == "QSO":
            from desisim.templates import QSO

            qso = QSO(nmodel=nobj, wave=wave)
            simflux, wave1, meta = qso.make_templates()

        # 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.
        elif objtype == "QSO_BAD":
            from desisim.templates import STAR

            star = STAR(nmodel=nobj, wave=wave)
            simflux, wave1, meta = star.make_templates()

        elif objtype == "STD":
            from desisim.templates import STAR

            star = STAR(nmodel=nobj, wave=wave, FSTD=True)
            simflux, wave1, meta = star.make_templates()

        truth["FLUX"][ii] = simflux
        truth["TEMPLATEID"][ii] = meta["TEMPLATEID"]
        truth["REDSHIFT"][ii] = meta["REDSHIFT"]

        # Pack in the photometry.  TODO: Include WISE.
        magg = meta["GMAG"]
        magr = meta["RMAG"]
        magz = meta["ZMAG"]
        fibermap["MAG"][ii, 0:3] = np.vstack([magg, magr, magz]).T
        fibermap["FILTER"][ii, 0:3] = ["DECAM_G", "DECAM_R", "DECAM_Z"]

        if objtype == "ELG":
            truth["OIIFLUX"][ii] = meta["OIIFLUX"]
            truth["D4000"][ii] = meta["D4000"]

        if objtype == "LRG":
            truth["D4000"][ii] = meta["D4000"]

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

    # - Where are these targets?  Centered on positioners for now.
    x = fiberpos["X"][0:nspec]
    y = fiberpos["Y"][0: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"][0:nspec]
    fibermap["SPECTROID"] = fiberpos["SPECTROGRAPH"][0:nspec]
    fibermap["TARGETID"] = np.random.randint(sys.maxint, size=nspec)
    fibermap["TARGETCAT"] = np.zeros(nspec, dtype="|S20")
    fibermap["LAMBDAREF"] = np.ones(nspec, dtype=np.float32) * 5400
    fibermap["TARGET_MASK0"] = np.zeros(nspec, dtype="i8")
    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, truth
Ejemplo n.º 7
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)
Ejemplo n.º 8
0
class MockSpectra(object):
    """Generate spectra for each type of mock.  Currently just choose the closest
    template; we can get fancier later.

    ToDo (@moustakas): apply Galactic extinction.

    """
    def __init__(self, wavemin=None, wavemax=None, dw=0.2, nproc=1,
                 rand=None, verbose=False):

        from lvmmodel.io import load_throughput

        self.tree = TemplateKDTree(nproc=nproc)

        # Build a default (buffered) wavelength vector.
        if wavemin is None:
            wavemin = load_throughput('b').wavemin - 10.0
        if wavemax is None:
            wavemax = load_throughput('z').wavemax + 10.0

        self.wavemin = wavemin
        self.wavemax = wavemax
        self.dw = dw
        self.wave = np.arange(round(wavemin, 1), wavemax, dw)

        self.rand = rand
        self.verbose = verbose

        #self.__normfilter = 'decam2014-r' # default normalization filter

        # Initialize the templates once:
        from desisim.templates import BGS, ELG, LRG, QSO, STAR, WD
        self.bgs_templates = BGS(wave=self.wave, normfilter='sdss2010-r') # Need to generalize this!
        self.elg_templates = ELG(wave=self.wave, normfilter='decam2014-r')
        self.lrg_templates = LRG(wave=self.wave, normfilter='decam2014-z')
        self.qso_templates = QSO(wave=self.wave, normfilter='decam2014-g')
        self.lya_templates = QSO(wave=self.wave, normfilter='decam2014-g')
        self.star_templates = STAR(wave=self.wave, normfilter='decam2014-r')
        self.wd_da_templates = WD(wave=self.wave, normfilter='decam2014-g', subtype='DA')
        self.wd_db_templates = WD(wave=self.wave, normfilter='decam2014-g', subtype='DB')

    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
        TemplateKDTree.bgs().

        """
        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

    def elg(self, data, index=None, mockformat='gaussianfield'):
        """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))

        if False:
            import matplotlib.pyplot as plt
            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

    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

    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

    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

    def mws_nearby(self, data, index=None, mockformat='100pc'):
        """Generate spectra for the MWS_NEARBY sample.

        """
        flux, meta = self.mws(data, index=index, mockformat=mockformat)
        return flux, meta

    def mws_main(self, data, index=None, mockformat='galaxia'):
        """Generate spectra for the MWS_MAIN sample.

        """
        flux, meta = self.mws(data, index=index, mockformat=mockformat)
        return flux, meta

    def faintstar(self, data, index=None, mockformat='galaxia'):
        """Generate spectra for the FAINTSTAR (faint stellar) sample.

        """
        flux, meta = self.mws(data, index=index, mockformat=mockformat)
        return flux, meta

    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

    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

    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]
        flux = np.zeros((nobj, len(self.wave)), dtype='i1')

        return flux, meta
Ejemplo n.º 9
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)