<commit_before>'''
redrock wrapper tools for DESI
'''
from __future__ import absolute_import, division, print_function

import os, sys
import warnings
if sys.version_info[0] > 2:
    basestring = str

import numpy as np
from astropy.io import fits
from astropy.table import Table, vstack
import desispec.io
from desispec.resolution import Resolution

from ..dataobj import (Target, MultiprocessingSharedSpectrum, 
    SimpleSpectrum, MPISharedTargets)

from .. import io
from .. import zfind


def write_zbest(outfile, zbest, fibermap):
    '''
    Write zbest and fibermap Tables to outfile

    Adds blank BRICKNAME and SUBTYPE columns if needed
    Adds zbest.meta['EXTNAME'] = 'ZBEST'
    Adds fibermap.meta['EXTNAME'] = 'FIBERMAP'
    '''
    ntargets = len(zbest)
    if 'BRICKNAME' not in zbest.colnames:
        zbest['BRICKNAME'] = np.zeros(ntargets, dtype='S8')

    if 'SUBTYPE' not in zbest.colnames:
        zbest['SUBTYPE'] = np.zeros(ntargets, dtype='S8')

    zbest.meta['EXTNAME'] = 'ZBEST'
    fibermap.meta['EXTNAME'] = 'FIBERMAP'

    hx = fits.HDUList()
    hx.append(fits.PrimaryHDU())
    hx.append(fits.convenience.table_to_hdu(zbest))
    hx.append(fits.convenience.table_to_hdu(fibermap))
    hx.writeto(outfile, overwrite=True)

def read_spectra(spectrafiles, targetids=None, spectrum_class=SimpleSpectrum):
    '''
    Read targets from a list of spectra files.

    Args:
        spectrafiles : list of input spectra files, or string glob to match

    Options:
        targetids : list of target ids. If set, only those target spectra will be read.
        spectrum_class :  The spectrum_class argument is needed to use the same read_spectra
        routine for the two parallelism approaches for redrock. The multiprocessing version
        uses a shared memory at the spectrum class initialization, see
        the redrock.dataobj.MultiprocessingSharedSpectrum class, whereas the MPI version
        implements the shared memory after all spectra have been read by the root process,
        and so the MPI version used another more simple spectrum class (see redrock.dataobj.SimpleSpectrum).
    
    Returns tuple of (targets, fibermap) where
        targets is a list of Target objects and
        fibermap is the Table from the input spectra files

    Note that due to multiple exposures of the same targets,
    len(targets) != len(fibermap)
    '''
    if isinstance(spectrafiles, basestring):
        import glob
        spectrafiles = glob.glob(spectrafiles)

    assert len(spectrafiles) > 0

    input_spectra = list()
    input_targetids = set()
    input_fibermaps = list()

    #- Ignore warnings about zdc2 bricks lacking bricknames in header
    for infile in spectrafiles:
        sp = desispec.io.read_spectra(infile)
        if hasattr(sp, 'fmap'):
            sp.fibermap = sp.fmap   #- for future compatibility
        input_spectra.append(sp)
        input_targetids.update(sp.fibermap['TARGETID'])
        input_fibermaps.append(sp.fibermap)

    if targetids is None:
        targetids = input_targetids

    targets = list()
    for targetid in targetids:
        spectra = list()
        for sp in input_spectra:
            ii = np.where(sp.fibermap['TARGETID'] == targetid)[0]
            if np.count_nonzero(ii) == 0:
                continue
            for x in sp.bands:          #- typically 'b', 'r', 'z'
                wave = sp.wave[x]                
                flux = sp.flux[x][ii]
                ivar = sp.ivar[x][ii]*(sp.mask[x][ii]==0)
                Rdata = sp.resolution_data[x][ii]

                for i in range(flux.shape[0]):
                    ifiber = ii[i]      #- index into sp.fibermap
                    if np.all(flux[i] == 0):
                        continue

                    if np.all(ivar[i] == 0):
                        continue

                    R = Resolution(Rdata[i])

                    meta = dict()
                    meta['NIGHT'] = sp.fibermap['NIGHT'][ifiber]
                    meta['EXPID'] = sp.fibermap['EXPID'][ifiber]

                    #- for backwards compatibility, allow missing TILEID
                    if 'TILEID' in sp.fibermap.dtype.names:
                        meta['TILEID'] = sp.fibermap['TILEID'][ifiber]
                    else:
                        meta['TILEID'] = -1

                    if 'HPXPIXEL' in sp.fibermap.dtype.names:
                        meta['HPXPIXEL'] = sp.fibermap['HPXPIXEL'][ifiber]
                        meta['HPXNSIDE'] = sp.meta['HPXNSIDE']
                    else:
                        import desimodel.footprint
                        meta['HPXNSIDE'] = 64
                        ra = sp.fibermap['RA_TARGET'][ifiber]
                        dec = sp.fibermap['DEC_TARGET'][ifiber]
                        meta['HPXPIXEL'] = desimodel.footprint.radec2pix(64, ra, dec)

                    spectra.append(spectrum_class(wave, flux[i], ivar[i], R, meta=meta))

        #- end of for targetid in targetids loop

        if len(spectra) > 0:
            numexp = len(set([sp.meta['EXPID'] for sp in spectra]))
            if spectra[0].meta['TILEID'] >= 0:
                numtile = len(set([sp.meta['TILEID'] for sp in spectra]))
            else:
                numtile = -1

            meta = dict(NUMEXP=numexp, NUMTILE=numtile)
            targets.append(Target(targetid, spectra, meta=meta))
        else:
            print('ERROR: Target {} on {} has no good spectra'.format(targetid, os.path.basename(brickfiles[0])))

    fibermap = vstack(input_fibermaps)

    return targets, fibermap


def rrdesi(options=None, comm=None):
    import optparse
    from astropy.io import fits
    import time

    # Note, in the pure multiprocessing case (comm == None), "rank"
    # will always be set to zero, which is fine since we are outside
    # any areas using multiprocessing and this is just used to control
    # program flow.
    rank = 0
    nproc = 1
    if comm is not None:
        rank = comm.rank
        nproc = comm.size

    start_time = time.time()
    pid = os.getpid()

    parser = optparse.OptionParser(usage = "%prog [options] spectra1 spectra2...")
    parser.add_option("-t", "--templates", type="string",  help="template file or directory")
    parser.add_option("-o", "--output", type="string",  help="output file")
    parser.add_option("--zbest", type="string",  help="output zbest fits file")
    parser.add_option("-n", "--ntargets", type=int,  help="number of targets to process")
    parser.add_option("--mintarget", type=int,  help="first target to include", default=0)
    parser.add_option("--ncpu", type=int,  help="number of cpu cores for multiprocessing", default=None)
    parser.add_option("--debug", help="debug with ipython", action="store_true")
    ### parser.add_option("--coadd", help="use coadd instead of individual spectra", action="store_true")
    parser.add_option("--allspec", help="use individual spectra instead of coadd", action="store_true")

    if options is None:
        opts, infiles = parser.parse_args()
    else:
        opts, infiles = parser.parse_args(options)

    if (opts.output is None) and (opts.zbest is None):
        print('ERROR: --output or --zbest required')
        sys.exit(1)

    if len(infiles) == 0:
        print('ERROR: must provide input spectra/brick files')
        sys.exit(1)

    templates = None
    targets   = None
    meta      = None
    if rank == 0:
        t0 = time.time()
        print('INFO: reading targets')
        sys.stdout.flush()
        
        spectrum_class = SimpleSpectrum
        if opts.ncpu is None or opts.ncpu > 1:
            spectrum_class = MultiprocessingSharedSpectrum

        targets, fibermap = read_spectra(infiles,spectrum_class=spectrum_class)
            
        if not opts.allspec:
            for t in targets:
                t._all_spectra = t.spectra
                t.spectra = t.coadd

        if opts.ntargets is not None:
            targets = targets[opts.mintarget:opts.mintarget+opts.ntargets]
            targetids = [t.id for t in targets]
            ii = np.in1d(fibermap['TARGETID'], targetids)
            fibermap = fibermap[ii]<commit_msg>fix bug in error printout<commit_after>'''
redrock wrapper tools for DESI
'''
from __future__ import absolute_import, division, print_function

import os, sys
import warnings
if sys.version_info[0] > 2:
    basestring = str

import numpy as np
from astropy.io import fits
from astropy.table import Table, vstack
import desispec.io
from desispec.resolution import Resolution

from ..dataobj import (Target, MultiprocessingSharedSpectrum, 
    SimpleSpectrum, MPISharedTargets)

from .. import io
from .. import zfind


def write_zbest(outfile, zbest, fibermap):
    '''
    Write zbest and fibermap Tables to outfile

    Adds blank BRICKNAME and SUBTYPE columns if needed
    Adds zbest.meta['EXTNAME'] = 'ZBEST'
    Adds fibermap.meta['EXTNAME'] = 'FIBERMAP'
    '''
    ntargets = len(zbest)
    if 'BRICKNAME' not in zbest.colnames:
        zbest['BRICKNAME'] = np.zeros(ntargets, dtype='S8')

    if 'SUBTYPE' not in zbest.colnames:
        zbest['SUBTYPE'] = np.zeros(ntargets, dtype='S8')

    zbest.meta['EXTNAME'] = 'ZBEST'
    fibermap.meta['EXTNAME'] = 'FIBERMAP'

    hx = fits.HDUList()
    hx.append(fits.PrimaryHDU())
    hx.append(fits.convenience.table_to_hdu(zbest))
    hx.append(fits.convenience.table_to_hdu(fibermap))
    hx.writeto(outfile, overwrite=True)

def read_spectra(spectrafiles, targetids=None, spectrum_class=SimpleSpectrum):
    '''
    Read targets from a list of spectra files.

    Args:
        spectrafiles : list of input spectra files, or string glob to match

    Options:
        targetids : list of target ids. If set, only those target spectra will be read.
        spectrum_class :  The spectrum_class argument is needed to use the same read_spectra
        routine for the two parallelism approaches for redrock. The multiprocessing version
        uses a shared memory at the spectrum class initialization, see
        the redrock.dataobj.MultiprocessingSharedSpectrum class, whereas the MPI version
        implements the shared memory after all spectra have been read by the root process,
        and so the MPI version used another more simple spectrum class (see redrock.dataobj.SimpleSpectrum).
    
    Returns tuple of (targets, fibermap) where
        targets is a list of Target objects and
        fibermap is the Table from the input spectra files

    Note that due to multiple exposures of the same targets,
    len(targets) != len(fibermap)
    '''
    if isinstance(spectrafiles, basestring):
        import glob
        spectrafiles = glob.glob(spectrafiles)

    assert len(spectrafiles) > 0

    input_spectra = list()
    input_targetids = set()
    input_fibermaps = list()

    #- Ignore warnings about zdc2 bricks lacking bricknames in header
    for infile in spectrafiles:
        sp = desispec.io.read_spectra(infile)
        if hasattr(sp, 'fmap'):
            sp.fibermap = sp.fmap   #- for future compatibility
        input_spectra.append(sp)
        input_targetids.update(sp.fibermap['TARGETID'])
        input_fibermaps.append(sp.fibermap)

    if targetids is None:
        targetids = input_targetids

    targets = list()
    for targetid in targetids:
        spectra = list()
        for sp in input_spectra:
            ii = np.where(sp.fibermap['TARGETID'] == targetid)[0]
            if np.count_nonzero(ii) == 0:
                continue
            for x in sp.bands:          #- typically 'b', 'r', 'z'
                wave = sp.wave[x]                
                flux = sp.flux[x][ii]
                ivar = sp.ivar[x][ii]*(sp.mask[x][ii]==0)
                Rdata = sp.resolution_data[x][ii]

                for i in range(flux.shape[0]):
                    ifiber = ii[i]      #- index into sp.fibermap
                    if np.all(flux[i] == 0):
                        continue

                    if np.all(ivar[i] == 0):
                        continue

                    R = Resolution(Rdata[i])

                    meta = dict()
                    meta['NIGHT'] = sp.fibermap['NIGHT'][ifiber]
                    meta['EXPID'] = sp.fibermap['EXPID'][ifiber]

                    #- for backwards compatibility, allow missing TILEID
                    if 'TILEID' in sp.fibermap.dtype.names:
                        meta['TILEID'] = sp.fibermap['TILEID'][ifiber]
                    else:
                        meta['TILEID'] = -1

                    if 'HPXPIXEL' in sp.fibermap.dtype.names:
                        meta['HPXPIXEL'] = sp.fibermap['HPXPIXEL'][ifiber]
                        meta['HPXNSIDE'] = sp.meta['HPXNSIDE']
                    else:
                        import desimodel.footprint
                        meta['HPXNSIDE'] = 64
                        ra = sp.fibermap['RA_TARGET'][ifiber]
                        dec = sp.fibermap['DEC_TARGET'][ifiber]
                        meta['HPXPIXEL'] = desimodel.footprint.radec2pix(64, ra, dec)

                    spectra.append(spectrum_class(wave, flux[i], ivar[i], R, meta=meta))

        #- end of for targetid in targetids loop

        if len(spectra) > 0:
            numexp = len(set([sp.meta['EXPID'] for sp in spectra]))
            if spectra[0].meta['TILEID'] >= 0:
                numtile = len(set([sp.meta['TILEID'] for sp in spectra]))
            else:
                numtile = -1

            meta = dict(NUMEXP=numexp, NUMTILE=numtile)
            targets.append(Target(targetid, spectra, meta=meta))
        else:
            print('ERROR: Target {} on {} has no good spectra'.format(targetid, os.path.basename(spectrafiles[0])))

    fibermap = vstack(input_fibermaps)

    return targets, fibermap


def rrdesi(options=None, comm=None):
    import optparse
    from astropy.io import fits
    import time

    # Note, in the pure multiprocessing case (comm == None), "rank"
    # will always be set to zero, which is fine since we are outside
    # any areas using multiprocessing and this is just used to control
    # program flow.
    rank = 0
    nproc = 1
    if comm is not None:
        rank = comm.rank
        nproc = comm.size

    start_time = time.time()
    pid = os.getpid()

    parser = optparse.OptionParser(usage = "%prog [options] spectra1 spectra2...")
    parser.add_option("-t", "--templates", type="string",  help="template file or directory")
    parser.add_option("-o", "--output", type="string",  help="output file")
    parser.add_option("--zbest", type="string",  help="output zbest fits file")
    parser.add_option("-n", "--ntargets", type=int,  help="number of targets to process")
    parser.add_option("--mintarget", type=int,  help="first target to include", default=0)
    parser.add_option("--ncpu", type=int,  help="number of cpu cores for multiprocessing", default=None)
    parser.add_option("--debug", help="debug with ipython", action="store_true")
    ### parser.add_option("--coadd", help="use coadd instead of individual spectra", action="store_true")
    parser.add_option("--allspec", help="use individual spectra instead of coadd", action="store_true")

    if options is None:
        opts, infiles = parser.parse_args()
    else:
        opts, infiles = parser.parse_args(options)

    if (opts.output is None) and (opts.zbest is None):
        print('ERROR: --output or --zbest required')
        sys.exit(1)

    if len(infiles) == 0:
        print('ERROR: must provide input spectra files')
        sys.exit(1)

    templates = None
    targets   = None
    meta      = None
    if rank == 0:
        t0 = time.time()
        print('INFO: reading targets')
        sys.stdout.flush()
        
        spectrum_class = SimpleSpectrum
        if opts.ncpu is None or opts.ncpu > 1:
            spectrum_class = MultiprocessingSharedSpectrum

        targets, fibermap = read_spectra(infiles,spectrum_class=spectrum_class)
            
        if not opts.allspec:
            for t in targets:
                t._all_spectra = t.spectra
                t.spectra = t.coadd

        if opts.ntargets is not None:
            targets = targets[opts.mintarget:opts.mintarget+opts.ntargets]
            targetids = [t.id for t in targets]
            ii = np.in1d(fibermap['TARGETID'], targetids)
            fibermap = fibermap[ii]
{'fraction_non_alphanumeric': 0.10559423476242824, 'fraction_numerical': 0.00427507023329669, 'mean_word_length': 2.245144669044788, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3921746', 'n_tokens_mistral': 5044, 'n_tokens_neox': 4863, 'n_words': 1343}