<commit_before>from desisim.templates import QSO
from desisim.io import empty_metatable
import scipy as sp
from scipy import interpolate
from speclite import filters
import numpy as np
import astropy.table 

def get_spectra(infile, first=0, nqso=None, seed=None):
 
    '''
    read input lyman-alpha absorption files,
    return normalized spectra including forest absorption.

    Args:
        infile: name of the lyman-alpha spectra file

    Options:
        first: first spectrum to read
        nqso: number of spectra to read
        seed: random seed for template generation

    returns (flux, wave, meta)
    spectra with forest absorption normalized according to the magnitude
    
    Note:
        meta is metadata table from QSO continuum template generation
        plus RA and DEC columns
    '''

    import fitsio
    h = fitsio.FITS(infile)
    if nqso is None:
        nqso = len(h)-1

    if first<0:
        raise ValueError("first must be >= 0")

    heads = [head.read_header() for head in h[first+1:first+1+nqso]]

    zqso = [head["ZQSO"] for head in heads]
    ra = [head["RA"] for head in heads]
    dec = [head["DEC"] for head in heads]
    mag_g = [head["MAG_G"] for head in heads]

    assert(len(zqso) == nqso)

    rand = sp.random.RandomState(seed)
    seed = rand.randint(2**32, size=nqso)

    input_meta = empty_metatable(objtype='QSO', nmodel=1)

    filter_name = 'sdss2010-g'

    normfilt = filters.load_filters(filter_name)
    qso = QSO(normfilter=filter_name)

    flux = np.zeros([nqso, len(qso.wave)], dtype='f4')
    meta = None
    for i,head in enumerate(h[first+1:first+1+nqso]):
        f, wave, meta_qso = qso.make_templates(nmodel=1,
                                               redshift=np.array([zqso[i]]), mag=np.array([mag_g[i]]), seed=seed[i])

        meta_qso['TEMPLATEID'] = first + i + 1
        if meta is None:
            meta = meta_qso.copy()
        else:
            meta = astropy.table.vstack([meta, meta_qso])

        # read lambda and forest transmission
        la = head["LAMBDA"][:]
        tr = head["FLUX"][:]
        if len(tr):
            # will interpolate the transmission at the spectral wavelengths, 
            # if outside the forest, the transmission is 1
            itr = interpolate.interp1d(la,tr,bounds_error=False,fill_value=1)
            f *= itr(wave)

        padflux, padwave = normfilt.pad_spectrum(f, wave, method='edge')
        normmaggies = sp.array(normfilt.get_ab_maggies(padflux, padwave, 
                               mask_invalid=True)[filter_name])
        f *= 10**(-0.4 * mag_g[i]) / normmaggies
        flux[i,:] = f[:]

    h.close()

    # Add RA,DEC to output meta
    meta['RA'] = ra
    meta['DEC'] = dec
    
    return flux,wave,meta

<commit_msg>refactor lya_spectra to support optional inputs<commit_after>"""
desisim.lya_spectra
===================

Function to simulate a QSO spectrum including Lyman-alpha absorption.
"""

from __future__ import division, print_function

def get_spectra(lyafile, templateid=None, wave=None, normfilter='sdss2010-g',
                rand=None, qso=None):
    '''Generate a QSO spectrum which includes Lyman-alpha absorption.

    Args:
        lyafile (str): name of the Lyman-alpha spectrum file to read.
        templateid (int numpy.ndarray, optional): indices of the spectra
          (0-indexed) to read from LYAFILE (default is to read everything).
        wave (numpy.ndarray, optional): desired output wavelength vector.
        normfilter (str, optional): normalization filter
        rand (numpy.RandomState, optional): RandomState object used for the
          random number generation.
        qso (desisim.templates.QSO, optional): object with which to generate
          individual spectra/templates.

    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:
        nqso = len(h)-1
    else:
        nqso = len(templateid)

    if rand is None:
        rand = np.random.RandomState()
    seed = 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'] = seed
    meta['RA'] = ra
    meta['DEC'] = dec
    
    for ii, indx in enumerate(templateid):
        flux1, _, meta1 = qso.make_templates(nmodel=1, redshift=np.array([zqso[ii]]),
                                             mag=np.array([mag_g[ii]]), seed=seed[ii])

        # 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])
        flux1 *= 10**(-0.4 * mag_g[ii]) / normmaggies
        flux[ii, :] = flux1[:]

    h.close()

    return flux, wave, meta


{'fraction_non_alphanumeric': 0.11364350833071361, 'fraction_numerical': 0.008802263439170073, 'mean_word_length': 2.533037201554692, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '87870', 'n_tokens_mistral': 2090, 'n_tokens_neox': 1991, 'n_words': 560}