<commit_before># = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
import numpy as np
import sys
import pyfits
import astropy.io.fits as afits
import collections
import tdose_utilities as tu
import tdose_extract_spectra as tes
import tdose_build_mock_cube as tbmc
import pdb
import scipy.ndimage.filters as snf
import matplotlib as mpl
mpl.use('Agg') # prevent pyplot from opening window; enables closing ssh session with detached screen running TDOSE
import matplotlib.pyplot as plt
# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
def extract_spectra(model_cube_file,source_association_dictionary=None,nameext='tdose_spectrum',outputdir='./',clobber=False,
                    variance_cube_file=None,variance_cube_ext='ERROR',source_model_cube_file=None,source_cube_ext='DATA',
                    model_cube_ext='DATA',layer_scale_ext='WAVESCL',data_cube_file=None,verbose=True):
    """
    Assemble the spectra determined by the wavelength layer scaling of the normalized models
    when generating the source model cube

    --- INPUT ---
    model_cube_file                     Model cube to base extraction on (using header info and layer scales)
    source_association_dictionary       Source association dictionary defining what sources should be combined into
                                        objects (individual spectra).
    nameext                             The name extension to use for saved spectra
    outputdir                           Directory to save spectra to
    clobber                             Overwrite spectra if they already exists
    variance_cube_file                  File containing variance cube of data to be used to estimate nois on 1D spectrum
    variance_cube_ext                   Extension of variance cube to use
    source_model_cube_file              The source model cube defining the individual sources
    source_cube_ext                     Extension of source model cube file that contins source models
    model_cube_ext                      Extension of model cube file that contains model
    layer_scale_ext                     Extension of model cube file that contains the layer scales
    data_cube_file                      File containing original data cube used for extraction of aperture spectra
    verbose

    --- EXAMPLE OF USE ---


    """
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if verbose: print ' - Loading data needed for spectral assembly'
    model_cube        = pyfits.open(model_cube_file)[model_cube_ext].data
    model_cube_hdr    = pyfits.open(model_cube_file)[model_cube_ext].header
    layer_scale_arr   = pyfits.open(model_cube_file)[layer_scale_ext].data
    if variance_cube_file is not None:
        stddev_cube       = np.sqrt(pyfits.open(variance_cube_file)[variance_cube_ext].data) # turn varinace into standard deviation
        source_model_cube = pyfits.open(source_model_cube_file)[source_cube_ext].data
    else:
        stddev_cube       = None
        source_model_cube = None
    Nsources = layer_scale_arr.shape[0]

    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if data_cube_file is not None:
        if verbose: print ' - Loading data cube '
        data_cube  = pyfits.open(data_cube_file)[model_cube_ext].data
    else:
        data_cube  = None

    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if source_association_dictionary is None:
        if verbose: print ' - Building default source association dictionary ' \
                          '(determining what sources are combined into objects), i.e., one source per object '
        sourcIDs_dic = collections.OrderedDict()
        for oo in xrange(Nsources):
            sourcIDs_dic[str(oo)] = [oo]
    else:
        sourcIDs_dic = source_association_dictionary
    Nobj = len(sourcIDs_dic.keys())
    if verbose: print ' - Found '+str(Nobj)+' objects to generate spectra for in source_association_dictionary '

    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if verbose: print ' - Assembling wavelength vector for spectra '
    wavelengths     =  np.arange(model_cube_hdr['NAXIS3'])*model_cube_hdr['CD3_3']+model_cube_hdr['CRVAL3']
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    specfiles = []
    for oo, key in enumerate(sourcIDs_dic.keys()):
        obj_cube_hdr = model_cube_hdr.copy()
        try:
            specid       = str("%.10d" % int(key))
        except:
            specid       = str(key)
        specname     = outputdir+nameext+'_'+specid+'.fits'
        specfiles.append(specname)
        sourceIDs    = sourcIDs_dic[key]

        obj_cube_hdr.append(('OBJID   ',specid         ,'ID of object'),end=True)
        obj_cube_hdr.append(('SRCIDS  ',str(sourceIDs) ,'IDs of sources combined in object'),end=True)

        if verbose:
            infostr = ' - Extracting spectrum '+str("%6.f" % (oo+1))+' / '+str("%6.f" % Nobj)
            sys.stdout.write("%s\r" % infostr)
            sys.stdout.flush()

        sourceoutput = tes.extract_spectrum(sourceIDs,layer_scale_arr,wavelengths,noise_cube=stddev_cube,
                                            source_model_cube=source_model_cube, data_cube=data_cube,
                                            specname=specname,obj_cube_hdr=obj_cube_hdr,clobber=clobber,verbose=True)

    if verbose: print '\n - Done extracting spectra. Returning list of fits files generated'
    return specfiles
# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
def extract_spectrum(sourceIDs,layer_scale_arr,wavelengths,noise_cube=None,source_model_cube=None,
                     specname='tdose_extract_spectra_extractedspec.fits',obj_cube_hdr=None,data_cube=None,
                     clobber=False,verbose=True):
    """
    Extracting a spectrum based on the layer scale image from the model cube provided a list of sources to combine.
    Noise is estimated from the noise cube (of the data)

    If all layer_scales are 1 a data_cube for the extractions is expected

    --- INPUT ---
    sourceIDs         The source IDs to combine into spectrum
    layer_scale_arr   Layer scale array (or image) produced when generating the model cube
                      fractional flux belonging to the source in each pixel
    wavelengths       Wavelength vector to use for extracted 1D spectrum.
    noise_cube        Cube with uncertainties (sqrt(variance)) of data cube to be used for estimating 1D uncertainties
                      To estimate S/N and 1D noise, providing a source model cube is required
    source_model_cube Source model cube containing the model cube for each individual source seperately
                      Needed in order to estimate noise from noise-cube
    specname          Name of file to save spectrum to
    obj_cube_hdr      Provide a template header to save the object cube (from combining the individual source cube)
                      as an extension to the extracted spectrum
    data_cube         In case all layers scales are 1, it is assumed that the source_model_cube contains a mask for the
                      spectral extraction, which will then be performed on this data_cube.
    clobber           To overwrite existing files set clobber=True
    verbose           Toggle verbosity

    --- EXAMPLE OF USE ---


    """
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if verbose: print ' - Checking shape of wavelengths and layer_scale_arr'
    if wavelengths.shape[0] != layer_scale_arr.shape[1]:
        sys.exit(' ---> Shape of wavelength vector ('+str(wavelengths.shape)+
                 ') and wavelength dimension of layer scale array ('+
                 layer_scale_arr.shape[1].shape+') do not match.')
    else:
        if verbose: print '   dimensions match; proceeding...'
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if verbose: print ' - Checking all sources have spectra in layer_scale_arr'
    maxsource = np.max(sourceIDs)
    if maxsource >= layer_scale_arr.shape[0]:
        sys.exit(' ---> Sources in list '+str(str(sourceIDs))+
                 ' not available among '+str(layer_scale_arr.shape[0])+' sources in layer_scale_arr.')
    else:
        if verbose: print '   All sources exist in layer_scale_arr; proceeding...'
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if verbose: print ' - Assembling object spectrum from source scaling'
    source_ent = np.asarray(sourceIDs).astype(int)

    if (layer_scale_arr == 1).all():
        if verbose: print ' - All layer scales are 1; assuming source model cube contain mask for spectral extraction'
        object_cube  = np.sum(np.abs(source_model_cube[source_ent,:,:]),axis=0)
        if data_cube is None:
            sys.exit(' ---> Did not find a data cube to extrac spectra from as expected')

        object_mask     = (object_cube == 0) # masking all zeros in object mask
        invalid_mask    = np.ma.masked_invalid(data_cube).mask
        comb_mask       = (invalid_mask | object_mask)
        spec_1D_masked  = np.sum(np.sum(  np.ma.array(data_cube,mask=comb_mask)  ,axis=1),axis=1)
        spec_1D         = spec_1D_masked.filled(fill_value=0.0)

        if noise_cube is not None:
            if verbose: print '   Calculating noise as d_spec_k = sqrt( SUMij d_pix_ij**2 ), i.e., as the sqrt of variances summed'
            invalid_mask_noise = np.ma.masked_invalid(noise_cube).mask
            comb_mask          = (comb_mask | invalid_mask_noise)
            variance_1D_masked = np.ma.array(noise_cube,mask=comb_mask)**2
            noise_1D_masked    = np.sqrt( np.sum( np.sum( variance_1D_masked, axis=1), axis=1) )
            noise_1D           = noise_1D_masked.filled(fill_value=0.0)

            if verbose: print '   Generating S/N vector'
            SN_1D         = spec_1D / noise_1D
        else:
            if verbose: print ' - No "noise_cube" provided. Setting all errors and S/N values to NaN'
            SN_1D    = np.zeros(spec_1D.shape)*np.NaN
            noise_1D = np.zeros(spec_1D.shape)*np.NaN
    else:
        if verbose: print ' - Some layer scales are different from 1; hence assembling spectra using layer scales'
        if len(source_ent) < 1:
            spec_1D   = layer_scale_arr[source_ent,:]
        else:
            spec_1D   = np.sum( layer_scale_arr[source_ent,:],axis=0)

        # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        if noise_cube is not None:
            if verbose: print ' - Estimate S/N at each wavelength for 1D spectrum (see Eq. 16 of Kamann+2013)'
            if verbose: print '   Estimating fraction of flux in each pixel wrt. total flux in each layer'<commit_msg>Change noise_1D fill valur from 0 to NaN to indicate masking as opposed to value know without uncertainty<commit_after># = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
import numpy as np
import sys
import pyfits
import astropy.io.fits as afits
import collections
import tdose_utilities as tu
import tdose_extract_spectra as tes
import tdose_build_mock_cube as tbmc
import pdb
import scipy.ndimage.filters as snf
import matplotlib as mpl
mpl.use('Agg') # prevent pyplot from opening window; enables closing ssh session with detached screen running TDOSE
import matplotlib.pyplot as plt
# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
def extract_spectra(model_cube_file,source_association_dictionary=None,nameext='tdose_spectrum',outputdir='./',clobber=False,
                    variance_cube_file=None,variance_cube_ext='ERROR',source_model_cube_file=None,source_cube_ext='DATA',
                    model_cube_ext='DATA',layer_scale_ext='WAVESCL',data_cube_file=None,verbose=True):
    """
    Assemble the spectra determined by the wavelength layer scaling of the normalized models
    when generating the source model cube

    --- INPUT ---
    model_cube_file                     Model cube to base extraction on (using header info and layer scales)
    source_association_dictionary       Source association dictionary defining what sources should be combined into
                                        objects (individual spectra).
    nameext                             The name extension to use for saved spectra
    outputdir                           Directory to save spectra to
    clobber                             Overwrite spectra if they already exists
    variance_cube_file                  File containing variance cube of data to be used to estimate nois on 1D spectrum
    variance_cube_ext                   Extension of variance cube to use
    source_model_cube_file              The source model cube defining the individual sources
    source_cube_ext                     Extension of source model cube file that contins source models
    model_cube_ext                      Extension of model cube file that contains model
    layer_scale_ext                     Extension of model cube file that contains the layer scales
    data_cube_file                      File containing original data cube used for extraction of aperture spectra
    verbose

    --- EXAMPLE OF USE ---


    """
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if verbose: print ' - Loading data needed for spectral assembly'
    model_cube        = pyfits.open(model_cube_file)[model_cube_ext].data
    model_cube_hdr    = pyfits.open(model_cube_file)[model_cube_ext].header
    layer_scale_arr   = pyfits.open(model_cube_file)[layer_scale_ext].data
    if variance_cube_file is not None:
        stddev_cube       = np.sqrt(pyfits.open(variance_cube_file)[variance_cube_ext].data) # turn varinace into standard deviation
        source_model_cube = pyfits.open(source_model_cube_file)[source_cube_ext].data
    else:
        stddev_cube       = None
        source_model_cube = None
    Nsources = layer_scale_arr.shape[0]

    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if data_cube_file is not None:
        if verbose: print ' - Loading data cube '
        data_cube  = pyfits.open(data_cube_file)[model_cube_ext].data
    else:
        data_cube  = None

    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if source_association_dictionary is None:
        if verbose: print ' - Building default source association dictionary ' \
                          '(determining what sources are combined into objects), i.e., one source per object '
        sourcIDs_dic = collections.OrderedDict()
        for oo in xrange(Nsources):
            sourcIDs_dic[str(oo)] = [oo]
    else:
        sourcIDs_dic = source_association_dictionary
    Nobj = len(sourcIDs_dic.keys())
    if verbose: print ' - Found '+str(Nobj)+' objects to generate spectra for in source_association_dictionary '

    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if verbose: print ' - Assembling wavelength vector for spectra '
    wavelengths     =  np.arange(model_cube_hdr['NAXIS3'])*model_cube_hdr['CD3_3']+model_cube_hdr['CRVAL3']
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    specfiles = []
    for oo, key in enumerate(sourcIDs_dic.keys()):
        obj_cube_hdr = model_cube_hdr.copy()
        try:
            specid       = str("%.10d" % int(key))
        except:
            specid       = str(key)
        specname     = outputdir+nameext+'_'+specid+'.fits'
        specfiles.append(specname)
        sourceIDs    = sourcIDs_dic[key]

        obj_cube_hdr.append(('OBJID   ',specid         ,'ID of object'),end=True)
        obj_cube_hdr.append(('SRCIDS  ',str(sourceIDs) ,'IDs of sources combined in object'),end=True)

        if verbose:
            infostr = ' - Extracting spectrum '+str("%6.f" % (oo+1))+' / '+str("%6.f" % Nobj)
            sys.stdout.write("%s\r" % infostr)
            sys.stdout.flush()

        sourceoutput = tes.extract_spectrum(sourceIDs,layer_scale_arr,wavelengths,noise_cube=stddev_cube,
                                            source_model_cube=source_model_cube, data_cube=data_cube,
                                            specname=specname,obj_cube_hdr=obj_cube_hdr,clobber=clobber,verbose=True)

    if verbose: print '\n - Done extracting spectra. Returning list of fits files generated'
    return specfiles
# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
def extract_spectrum(sourceIDs,layer_scale_arr,wavelengths,noise_cube=None,source_model_cube=None,
                     specname='tdose_extract_spectra_extractedspec.fits',obj_cube_hdr=None,data_cube=None,
                     clobber=False,verbose=True):
    """
    Extracting a spectrum based on the layer scale image from the model cube provided a list of sources to combine.
    Noise is estimated from the noise cube (of the data)

    If all layer_scales are 1 a data_cube for the extractions is expected

    --- INPUT ---
    sourceIDs         The source IDs to combine into spectrum
    layer_scale_arr   Layer scale array (or image) produced when generating the model cube
                      fractional flux belonging to the source in each pixel
    wavelengths       Wavelength vector to use for extracted 1D spectrum.
    noise_cube        Cube with uncertainties (sqrt(variance)) of data cube to be used for estimating 1D uncertainties
                      To estimate S/N and 1D noise, providing a source model cube is required
    source_model_cube Source model cube containing the model cube for each individual source seperately
                      Needed in order to estimate noise from noise-cube
    specname          Name of file to save spectrum to
    obj_cube_hdr      Provide a template header to save the object cube (from combining the individual source cube)
                      as an extension to the extracted spectrum
    data_cube         In case all layers scales are 1, it is assumed that the source_model_cube contains a mask for the
                      spectral extraction, which will then be performed on this data_cube.
    clobber           To overwrite existing files set clobber=True
    verbose           Toggle verbosity

    --- EXAMPLE OF USE ---


    """
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if verbose: print ' - Checking shape of wavelengths and layer_scale_arr'
    if wavelengths.shape[0] != layer_scale_arr.shape[1]:
        sys.exit(' ---> Shape of wavelength vector ('+str(wavelengths.shape)+
                 ') and wavelength dimension of layer scale array ('+
                 layer_scale_arr.shape[1].shape+') do not match.')
    else:
        if verbose: print '   dimensions match; proceeding...'
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if verbose: print ' - Checking all sources have spectra in layer_scale_arr'
    maxsource = np.max(sourceIDs)
    if maxsource >= layer_scale_arr.shape[0]:
        sys.exit(' ---> Sources in list '+str(str(sourceIDs))+
                 ' not available among '+str(layer_scale_arr.shape[0])+' sources in layer_scale_arr.')
    else:
        if verbose: print '   All sources exist in layer_scale_arr; proceeding...'
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    if verbose: print ' - Assembling object spectrum from source scaling'
    source_ent = np.asarray(sourceIDs).astype(int)

    if (layer_scale_arr == 1).all():
        if verbose: print ' - All layer scales are 1; assuming source model cube contain mask for spectral extraction'
        object_cube  = np.sum(np.abs(source_model_cube[source_ent,:,:]),axis=0)
        if data_cube is None:
            sys.exit(' ---> Did not find a data cube to extrac spectra from as expected')

        object_mask     = (object_cube == 0) # masking all zeros in object mask
        invalid_mask    = np.ma.masked_invalid(data_cube).mask
        comb_mask       = (invalid_mask | object_mask)
        spec_1D_masked  = np.sum(np.sum(  np.ma.array(data_cube,mask=comb_mask)  ,axis=1),axis=1)
        spec_1D         = spec_1D_masked.filled(fill_value=0.0)

        if noise_cube is not None:
            if verbose: print '   Calculating noise as d_spec_k = sqrt( SUMij d_pix_ij**2 ), i.e., as the sqrt of variances summed'
            invalid_mask_noise = np.ma.masked_invalid(noise_cube).mask
            comb_mask          = (comb_mask | invalid_mask_noise)
            variance_1D_masked = np.ma.array(noise_cube,mask=comb_mask)**2
            noise_1D_masked    = np.sqrt( np.sum( np.sum( variance_1D_masked, axis=1), axis=1) )
            noise_1D           = noise_1D_masked.filled(fill_value=np.nan)

            if verbose: print '   Generating S/N vector'
            SN_1D         = spec_1D / noise_1D
        else:
            if verbose: print ' - No "noise_cube" provided. Setting all errors and S/N values to NaN'
            SN_1D    = np.zeros(spec_1D.shape)*np.NaN
            noise_1D = np.zeros(spec_1D.shape)*np.NaN
    else:
        if verbose: print ' - Some layer scales are different from 1; hence assembling spectra using layer scales'
        if len(source_ent) < 1:
            spec_1D   = layer_scale_arr[source_ent,:]
        else:
            spec_1D   = np.sum( layer_scale_arr[source_ent,:],axis=0)

        # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        if noise_cube is not None:
            if verbose: print ' - Estimate S/N at each wavelength for 1D spectrum (see Eq. 16 of Kamann+2013)'
            if verbose: print '   Estimating fraction of flux in each pixel wrt. total flux in each layer'
{'fraction_non_alphanumeric': 0.11768885730781188, 'fraction_numerical': 0.005506950304214594, 'mean_word_length': 1.8736600306278715, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 192, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3274724', 'n_tokens_mistral': 7251, 'n_tokens_neox': 6879, 'n_words': 1742}