786934210

Build Type

github-actions

Committed by GitHub

Commit Message

<a href="https://github.com/desihub/desitarget/commit/<a class=hub.com/desihub/desitarget/commit/0ebc8987539d6f0642807d243d252b69be502382">0ebc89875<a href="https://github.com/desihub/desitarget/commit/0ebc8987539d6f0642807d243d252b69be502382">">Merge </a><a class="double-link" href="https://github.com/desihub/desitarget/commit/<a class="double-link" href="https://github.com/desihub/desitarget/commit/93aa111ef896d3b8f319284e40310274a79b15b2">93aa111ef</a>">93aa111ef</a><a href="https://github.com/desihub/desitarget/commit/0ebc8987539d6f0642807d243d252b69be502382"> into 3ec0b65a4">3ec0b65a4</a>

Pull Request Pull Request #714: Lyaretarget

Run Details

154 of 154 new or added lines in 1 file covered. (100.0%)

7216 of 12095 relevant lines covered (59.66%)

0.6 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

0.0

/py/desitarget/lyazcat.py

"""
desitarget.lyazcat
==============

Post-redrock ML processing for LyA Quasar object identification.
"""

import os
import numpy as np
from astropy.io import fits
import time

from desitarget.geomask import match
from desitarget.internal import sharedmem
from desispec.io import read_spectra

from quasarnp.io import load_model
from quasarnp.io import load_desi_coadd
from quasarnp.utils import process_preds
from squeze.model import Model
from squeze.common_functions import load_json
from squeze.candidates import Candidates
from squeze.desi_spectrum import DesiSpectrum
from squeze.spectra import Spectra

# ADM set up the DESI default logger.
from desiutil.log import get_logger
log = get_logger()

zcatdatamodel_names = np.array(['TARGETID', 'Z', 'ZWARN', 'DELTACHI2',
                                'Z_COMB', 'Z_COMB_PROB', 'Z_QN', 'Z_QN_CONF',
                                'IS_QSO_QN', 'Z_SQ', 'Z_SQ_CONF', 'Z_ABS',
                                'Z_ABS_CONF'])

zcatdatamodel_formats = np.array(['>i8', '>f8', '>i8', '>f8',
                                  '>f8', '>f8', '>f8', '>f8',
                                  '>i2', '>f8', '>f8', '>f8',
                                  '>f8'])
zcols_copy = np.array(['TARGETID', 'Z', 'ZWARN', 'DELTACHI2'])
n1_cols = np.array(['Z_QN', 'Z_QN_CONF', 'IS_QSO_QN', 'Z_SQ', 'Z_SQ_CONF',
                    'Z_ABS', 'Z_ABS_CONF'])


def tmark(istring):
    """A function to mark the time an operation starts or ends.
    
    Parameters
    ----------
    istring : :class:'str'
        The input string to print to the terminal.
        
    Notes
    -----
    - A string with the date and time in ISO 8061 standard followed
      by the 'istring'.
    """
    t0 = time.time()
    t_start = time.strftime('%Y-%m-%d | %H:%M:%S')
    log.info('\n{}: {}'.format(istring,t_start))

    
def make_new_zcat(zbestname):
    """Make the initial zcat array with redrock data.
    
    Parameters
    ----------
    zbestname : :class:`str`
        Full filename and path for the zbest file to process.
        
    Returns
    -------
    :class:`~numpy.array` or 'bool'
        A zcat in the official format (`zcatdatamodel`) compiled from
        the `tile', 'night', and 'petal_num', in `zcatdir`. If the zbest
        file for that petal doesn't exist, returns False.
    """
    tmark('    Making redrock zcat')
    try:
        zs = fits.open(zbestname)['ZBEST'].data
        # EBL write out the zcat as a file with the correct data model.
        zcat = np.zeros(len(zs), dtype={'names': zcatdatamodel_names, 'formats': zcatdatamodel_formats})
        for col in zcols_copy:
            zcat[col] = zs[col]
        for col in n1_cols:
            zcat[col][:] = -1

        return zcat
    except (FileNotFoundError, OSError):
        return False

    
def get_qn_model_fname(qnmodel_fname=None):
    """Convenience function to grab the $QN_MODEL_FNAME environment variable.

    Parameters
    ----------
    qnmodel_fname : :class:`str`, optional, defaults to $QN_MODEL_FNAME
        If `qnmodel_fname` is passed, it is returned from this function. If it's
        not passed, the $QN_MODEL_FNAME variable is returned.

    Returns
    -------
    :class:`str`
        not passed, the directory stored in the $QN_MODEL_FNAME environment
        variable is returned prepended to the default filename.
    """
    if qnmodel_fname is None:
        qnmodel_fname = os.environ.get('QN_MODEL_FILE')
        # EBL check that the $QN_MODEL_FNAME environment variable is set.
        if qnmodel_fname is None:
            msg = "Pass qnmodel_fname or set $QN_MODEL_FNAME environment variable!"
            log.critical(msg)
            raise ValueError(msg)

    return qnmodel_fname


def load_qn_model(model_filename):
    """Convenience function to load the QuasarNP model and line lists.
    
    Parameters
    ----------
    model_filename : :class:'str'
        The filename and path of the QuasarNP model. Either input by user or defaults
        to get_qn_model_fname().
    
    Returns
    -------
    :class:'~numpy.array'
        The QuasarNP model file loaded as an array.
    :class:'~numpy.array'
        An array of the emission line names to be used for quasarnp.process_preds().
    :class:'~numpy.array'
        An array of the BAL emission line names to be used by quasarnp.process_preds().
    """
    lines = ['LYA', 'CIV(1548)', 'CIII(1909)', 'MgII(2796)', 'Hbeta', 'Halpha']
    lines_bal = ['CIV(1548)']
    model = load_model(model_filename)
    
    return model, lines, lines_bal


def add_qn_data(zcat, coaddname, qnp_model, qnp_lines, qnp_lines_bal):
    """Apply the QuasarNP model to the input zcat and add data to columns.
    
    Parameters
    ----------
    zcat : :class:'~numpy.array'
        The structured array that was created by make_new_zcat()
    coaddname : :class:'str'
        The name of the coadd file corresponding to the zbest file used
        in make_new_zcat()
    qnp_model : :class:'h5.array'
        The array containing the pre-trained QuasarNP model.
    qnp_lines : :class:'list'
        A list containing the names of the emission lines that
        quasarnp.process_preds() should use.
    qnp_lines_bal : :class:'list'
        A list containing the names of the emission lines to check
        for BAL troughs.
        
    Returns
    -------
    :class:'~numpy.array'
        The zcat array with QuasarNP data included in the columns:
        
        * Z_QN        - The best QuasarNP redshift for the object
        * Z_QN_CONF   - The confidence of Z_QN
        * IS_QSO_QN   - A binary flag indicated object is a quasar
    """
    tmark('    Adding QuasarNP data')
    
    data, w = load_desi_coadd(coaddname)
    data = data[:, :, None]
    p = qnp_model.predict(data)
    c_line, z_line, zbest, *_ = process_preds(p, qnp_lines, qnp_lines_bal,
                                              verbose=False)

    cbest = np.array(c_line[c_line.argmax(axis=0), np.arange(len(zbest))])
    c_thresh = 0.5
    n_thresh = 1
    is_qso = np.sum(c_line > c_thresh, axis=0) >= n_thresh

    zcat['Z_QN'][w] = zbest
    zcat['Z_QN_CONF'][w] = cbest
    zcat['IS_QSO_QN'][w] = is_qso

    return zcat


def get_sq_model_fname(sqmodel_fname=None):
    """Convenience function to grab the $SQ_MODEL_FNAME environment variable.

    Parameters
    ----------
    sqmodel_fname : :class:`str`, optional, defaults to $SQ_MODEL_FNAME
        If `sqmodel_fname` is passed, it is returned from this function. If it's
        not passed, the $SQ_MODEL_FNAME environment variable is returned.

    Returns
    -------
    :class:`str`
        If `sqmodel_fname` is passed, it is returned from this function. If it's
        not passed, the directory stored in the $SQ_MODEL_FNAME environment
        variable is returned.
    """
    if sqmodel_fname is None:
        sqmodel_fname = os.environ.get('SQ_MODEL_FILE')
        # EBL check that the $SQ_MODEL_FNAME environment variable is set.
        if sqmodel_fname is None:
            msg = "Pass sqmodel_fname or set $SQ_MODEL_FNAME environment variable!"
            log.critical(msg)
            raise ValueError(msg)

    return sqmodel_fname


def load_sq_model(model_filename):
    """Convenience function for loading the SQUEzE model file.
    
    Parameters
    ----------
    model_filename : :class:'str'
        The filename and path of the SQUEzE model file. Either input by user
        or defaults to get_sq_model_fname().
    
    Returns
    -------
    :class:'~numpy.array'
        A numpy array of the SQUEzE model.
    
    Notes
    -----
    - The input model file needs to be in the json file format.
    """
    model = Model.from_json(load_json(model_filename))
    
    return model


def add_sq_data(zcat, coaddname, squeze_model):
    """Apply the SQUEzE model to the input zcat and add data to columns.
    
    Parameters
    ----------
    zcat : :class:'~numpy.array'
        The structured array that was created by make_new_zcat()
    coaddname : class:'str'
        The name of the coadd file corresponding to the zbest file used
        in make_new_zcat()
    squeze_model : :class:'numpy.array'
        The loaded SQUEzE model file
        
    Returns
    -------
    :class:'~numpy.array'
        The zcat array with SQUEzE data included in the columns:
        
        * Z_SQ        - The best redshift from SQUEzE for each object.
        * Z_SQ_CONF   - The confidence value of this redshift.
    """
    tmark('    Adding SQUEzE data')
    mdata = ['TARGETID']
    single_exposure = False
    sq_cols_keep = ['PROB', 'Z_TRY', 'TARGETID']

    tmark('      Reading spectra')
    desi_spectra = read_spectra(coaddname)
    # EBL Initialize squeze Spectra class
    squeze_spectra = Spectra()
    # EBL Get TARGETIDs
    targetid = np.unique(desi_spectra.fibermap['TARGETID'])
    # EBL Loop over TARGETIDs to build the Spectra objects
    for targid in targetid:
        # EBL Select objects
        pos = np.where(desi_spectra.fibermap['TARGETID'] == targid)
        # EBL Prepare column metadata
        metadata = {col.upper(): desi_spectra.fibermap[col][pos[0][0]] for col in mdata}
        # EBL Add the SPECID as the TARGETID
        metadata['SPECID'] = targid
        # EBL Extract the data
        flux = {}
        wave = {}
        ivar = {}
        mask = {}
        for band in desi_spectra.bands:
            flux[band] = desi_spectra.flux[band][pos]
            wave[band] = desi_spectra.wave[band]
            ivar[band] = desi_spectra.ivar[band][pos]
            mask[band] = desi_spectra.mask[band][pos]
            
        # EBL Format each spectrum for the model application
        spectrum = DesiSpectrum(flux, wave, ivar, mask, metadata, single_exposure)
        # EBL Append the spectrum to the Spectra object
        squeze_spectra.append(spectrum)
    
    # EBL Initialize candidate object. This takes a while with no feedback
    # so we want a time output for benchmarking purposes.
    tmark('      Initializing candidates')
    candidates = Candidates(mode='operation', model=squeze_model)
    # EBL Look for candidate objects. This also takes a while.
    tmark('      Looking for candidates')
    candidates.find_candidates(squeze_spectra.spectra_list(), save=False)
    # EBL Compute the probabilities of the line/model matches to the spectra
    tmark('      Computing probabilities')
    candidates.classify_candidates(save=False)
    # EBL Filter the results by removing the duplicate entries for each
    # TARGETID. Merge the remaining with the zcat data.
    tmark('      Merging SQUEzE data with zcat')
    data_frame = candidates.candidates()
    data_frame = data_frame[~data_frame['DUPLICATED']][sq_cols_keep]
    # EBL Strip the pandas data frame structure and put it into a numpy 
    # structured array first.
    sqdata_arr = np.zeros(len(data_frame), dtype=[('TARGETID', 'int64'),
                                                  ('Z_SQ', 'float64'),
                                                  ('Z_SQ_CONF', 'float64')])
    sqdata_arr['TARGETID'] = data_frame['TARGETID'].values
    sqdata_arr['Z_SQ'] = data_frame['Z_TRY'].values
    sqdata_arr['Z_SQ_CONF'] = data_frame['PROB'].values
    # EBL SQUEzE will reorder the objects, so match on TARGETID.
    zcat_args, sqdata_args = match(zcat['TARGETID'],sqdata_arr['TARGETID'])
    zcat['Z_SQ'][zcat_args] = sqdata_arr['Z_SQ'][sqdata_args]
    zcat['Z_SQ_CONF'][zcat_args] = sqdata_arr['Z_SQ_CONF'][sqdata_args]

    return zcat


def add_abs_data(zcat, coaddname):
    """Add the MgII absorption line finder data to the input zcat array.
    
    Parameters
    ----------
    zcat : :class:'~numpy.array'
        The structured array that was created by make_new_zcat()
    coaddname : class:'str'
        The name of the coadd file corresponding to the zbest file used
        in make_new_zcat()
        
    Returns
    -------
    :class:'~numpy.array'
        The zcat array with MgII Absorption data included in the columns:
        
        * Z_ABS        - The highest redshift of MgII absorption
        * Z_ABS_CONF   - The confidence value for this redshift.
    """
    tmark('    MgII Absorption data not yet added.')

    return zcat


def zcomb_selector(zcat, proc_flag=False):
    """Compare results from redrock, QuasarNP, SQUEzE, and MgII data.
    
    Parameters
    ----------
    zcat : :class:'~numpy.array'
        The structured array that was created by make_new_zcat()
    proc_flag : :class:'bool'
        Turn on extra comparison procedure.
        
    Returns
    -------
    :class:'~numpy.array'
        The zcat array with SQUEzE data included in the columns:
        
        * Z_COMB        - The best models-combined redshift for each object.
        * Z_COMB_PROB   - The combined probability value of that redshift.
    """
    zcat['Z_COMB'][:] = zcat['Z']
    zcat['Z_COMB_PROB'][:] = 0.95

    return zcat


def zcat_writer(zcat, outputdir, outputname, qn_flag=False, sq_flag=False, abs_flag=False):
    """Writes the zcat structured array out as a FITS file.
    
    Parameters
    ----------
    zcat : :class:'~numpy.array'
        The structured array that was created by make_new_zcat()
    outputdir : :class:'str'
        The directory where the zcat file will be written.
    outputname : :class:'str'
        The filename of the zqso output file.
    qn_flag : :class:'bool'
        Flag if QuasarNP data (or not) was added to the zcat file.
    sq_flag : :class:'bool'
        Flag if SQUEzE data (or not) was added to the zcat file.
    abs_flag : :class:'bool'
        Flag if MgII Absorption data (or not) was added to the zcat file.
        
    Returns
    -------
    :class:'str'
        The filename, with path, of the FITS file written out.
    """
    tmark('    Creating file...')
    full_outputname = os.path.join(outputdir, outputname)
    
    prim_hdu = fits.PrimaryHDU()
    prim_hdu.header['QN_ADDED'] = str(qn_flag)
    prim_hdu.header['SQ_ADDED'] = str(sq_flag)
    prim_hdu.header['AB_ADDED'] = str(abs_flag)
    
    data_hdu = fits.BinTableHDU.from_columns(zcat, name='ZCATALOG')
    data_out = fits.HDUList([prim_hdu, data_hdu])
    
    data_out.writeto(full_outputname)

    return full_outputname


def create_zcat(tile, night, petal_num, zcatdir, outputdir, qn_flag=False,
                qnp_model=None, qnp_lines=None, qnp_lines_bal=None,
                sq_flag=False, squeze_model=None, abs_flag=False):
    """This will create a single zqso file from a set of user inputs.
    
    Parameters
    ----------
    tile : :class:'str'
        The TILEID of the tile to process.
    night : :class:'str'
        The date associated with the observation of the 'tile' used.
        * Must be in YYYYMMDD format
    petal_num : :class:'int'
        If 'all_petals' isn't used, the single petal to create a zcat for.
    zcatdir : :class:'str'
        The location for the daily redrock output.
    outputdir : :class:'str'
        The filepath to write out the zcat file.
    qn_flag : :class:'bool', optional
        Flag to add QuasarNP data (or not) to the zcat file.
    qnp_model : :class:'h5 array', optional
        The QuasarNP model file to be used for line predictions.
    qnp_lines : :class:'list', optional
        The list of lines to use in the QuasarNP model to test against.
    qnp_lines : :class:'list', optional
        The list of BAL lines to use for QuasarNP to identify BALs.
    sq_flag : :class:'bool', optional
        Flag to add SQUEzE data (or not) to the zcat file.
    squeze_model : :class:'numpy.array', optional
        The numpy array for the SQUEzE model file.
    abs_flag : :class:'bool', optional
        Flag to add MgII Absorption data (or not) to the zcat file.
    
    Notes
    -----
    - Writes a FITS catalog that incorporates redrock, QuasarNP, SQUEzE, and MgII
      absorption redshifts and confidence values. This will write to the same
      directory of the zbest and coadd files unless a different output directory is
      passed.
    """
    # EBL Create the filepath for the input tile/night combination
    tiledir = os.path.join(zcatdir, tile)
    ymdir = os.path.join(tiledir, night)
        
    # EBL Create the filename tag that appends to zbest-*, coadd-*, and zqso-*
    # files.
    filename_tag = f'{petal_num}-{tile}-thru{night}.fits'
    zbestname = f'zbest-{filename_tag}'
    coaddname = f'coadd-{filename_tag}'
    outputname = f'zqso-{filename_tag}'
    
    zcat = make_new_zcat(os.path.join(ymdir, zbestname))
    if isinstance(zcat, bool):
        print('  !!!! Petal Number does not have a corresponding zbest file !!!!')
    else:
        if qn_flag :
            zcat = add_qn_data(zcat, os.path.join(ymdir, coaddname),
                               qnp_model, qnp_lines, qnp_lines_bal)
        if sq_flag:
            zcat = add_sq_data(zcat, os.path.join(ymdir, coaddname), squeze_model)
        if abs_flag:
            zcat = add_abs_data(zcat, os.path.join(ymdir, coaddname))
        
        fzcat = zcomb_selector(zcat)
    
        full_outputname = zcat_writer(fzcat, outputdir, outputname, qn_flag,
                                      sq_flag, abs_flag)
    
        tmark('    --{} written out correctly.'.format(full_outputname))
        print('='*79)

1	"""
2	desitarget.lyazcat
3	==============
4
5	Post-redrock ML processing for LyA Quasar object identification.
6	"""
7
8	import os	×
9	import numpy as np	×
10	from astropy.io import fits	×
11	import time	×
12
13	from desitarget.geomask import match	×
14	from desitarget.internal import sharedmem	×
15	from desispec.io import read_spectra	×
16
17	from quasarnp.io import load_model	×
18	from quasarnp.io import load_desi_coadd	×
19	from quasarnp.utils import process_preds	×
20	from squeze.model import Model	×
21	from squeze.common_functions import load_json	×
22	from squeze.candidates import Candidates	×
23	from squeze.desi_spectrum import DesiSpectrum	×
24	from squeze.spectra import Spectra	×
25
26	# ADM set up the DESI default logger.
27	from desiutil.log import get_logger	×
28	log = get_logger()	×
29
30	zcatdatamodel_names = np.array(['TARGETID', 'Z', 'ZWARN', 'DELTACHI2',	×
31	'Z_COMB', 'Z_COMB_PROB', 'Z_QN', 'Z_QN_CONF',
32	'IS_QSO_QN', 'Z_SQ', 'Z_SQ_CONF', 'Z_ABS',
33	'Z_ABS_CONF'])
34
35	zcatdatamodel_formats = np.array(['>i8', '>f8', '>i8', '>f8',	×
36	'>f8', '>f8', '>f8', '>f8',
37	'>i2', '>f8', '>f8', '>f8',
38	'>f8'])
39	zcols_copy = np.array(['TARGETID', 'Z', 'ZWARN', 'DELTACHI2'])	×
40	n1_cols = np.array(['Z_QN', 'Z_QN_CONF', 'IS_QSO_QN', 'Z_SQ', 'Z_SQ_CONF',	×
41	'Z_ABS', 'Z_ABS_CONF'])
42
43
44	def tmark(istring):	×
45	"""A function to mark the time an operation starts or ends.
46
47	Parameters
48	----------
49	istring : :class:'str'
50	The input string to print to the terminal.
51
52	Notes
53	-----
54	- A string with the date and time in ISO 8061 standard followed
55	by the 'istring'.
56	"""
57	t0 = time.time()	×
58	t_start = time.strftime('%Y-%m-%d \| %H:%M:%S')	×
59	log.info('\n{}: {}'.format(istring,t_start))	×
60
61
62	def make_new_zcat(zbestname):	×
63	"""Make the initial zcat array with redrock data.
64
65	Parameters
66	----------
67	zbestname : :class:`str`
68	Full filename and path for the zbest file to process.
69
70	Returns
71	-------
72	:class:`~numpy.array` or 'bool'
73	A zcat in the official format (`zcatdatamodel`) compiled from
74	the `tile', 'night', and 'petal_num', in `zcatdir`. If the zbest
75	file for that petal doesn't exist, returns False.
76	"""
77	tmark(' Making redrock zcat')	×
78	try:	×
79	zs = fits.open(zbestname)['ZBEST'].data	×
80	# EBL write out the zcat as a file with the correct data model.
81	zcat = np.zeros(len(zs), dtype={'names': zcatdatamodel_names, 'formats': zcatdatamodel_formats})	×
82	for col in zcols_copy:	×
83	zcat[col] = zs[col]	×
84	for col in n1_cols:	×
85	zcat[col][:] = -1	×
86
87	return zcat	×
88	except (FileNotFoundError, OSError):	×
89	return False	×
90
91
92	def get_qn_model_fname(qnmodel_fname=None):	×
93	"""Convenience function to grab the $QN_MODEL_FNAME environment variable.
94
95	Parameters
96	----------
97	qnmodel_fname : :class:`str`, optional, defaults to $QN_MODEL_FNAME
98	If `qnmodel_fname` is passed, it is returned from this function. If it's
99	not passed, the $QN_MODEL_FNAME variable is returned.
100
101	Returns
102	-------
103	:class:`str`
104	not passed, the directory stored in the $QN_MODEL_FNAME environment
105	variable is returned prepended to the default filename.
106	"""
107	if qnmodel_fname is None:	×
108	qnmodel_fname = os.environ.get('QN_MODEL_FILE')	×
109	# EBL check that the $QN_MODEL_FNAME environment variable is set.
110	if qnmodel_fname is None:	×
111	msg = "Pass qnmodel_fname or set $QN_MODEL_FNAME environment variable!"	×
112	log.critical(msg)	×
113	raise ValueError(msg)	×
114
115	return qnmodel_fname	×
116
117
118	def load_qn_model(model_filename):	×
119	"""Convenience function to load the QuasarNP model and line lists.
120
121	Parameters
122	----------
123	model_filename : :class:'str'
124	The filename and path of the QuasarNP model. Either input by user or defaults
125	to get_qn_model_fname().
126
127	Returns
128	-------
129	:class:'~numpy.array'
130	The QuasarNP model file loaded as an array.
131	:class:'~numpy.array'
132	An array of the emission line names to be used for quasarnp.process_preds().
133	:class:'~numpy.array'
134	An array of the BAL emission line names to be used by quasarnp.process_preds().
135	"""
136	lines = ['LYA', 'CIV(1548)', 'CIII(1909)', 'MgII(2796)', 'Hbeta', 'Halpha']	×
137	lines_bal = ['CIV(1548)']	×
138	model = load_model(model_filename)	×
139
140	return model, lines, lines_bal	×
141
142
143	def add_qn_data(zcat, coaddname, qnp_model, qnp_lines, qnp_lines_bal):	×
144	"""Apply the QuasarNP model to the input zcat and add data to columns.
145
146	Parameters
147	----------
148	zcat : :class:'~numpy.array'
149	The structured array that was created by make_new_zcat()
150	coaddname : :class:'str'
151	The name of the coadd file corresponding to the zbest file used
152	in make_new_zcat()
153	qnp_model : :class:'h5.array'
154	The array containing the pre-trained QuasarNP model.
155	qnp_lines : :class:'list'
156	A list containing the names of the emission lines that
157	quasarnp.process_preds() should use.
158	qnp_lines_bal : :class:'list'
159	A list containing the names of the emission lines to check
160	for BAL troughs.
161
162	Returns
163	-------
164	:class:'~numpy.array'
165	The zcat array with QuasarNP data included in the columns:
166
167	* Z_QN - The best QuasarNP redshift for the object
168	* Z_QN_CONF - The confidence of Z_QN
169	* IS_QSO_QN - A binary flag indicated object is a quasar
170	"""
171	tmark(' Adding QuasarNP data')	×
172
173	data, w = load_desi_coadd(coaddname)	×
174	data = data[:, :, None]	×
175	p = qnp_model.predict(data)	×
176	c_line, z_line, zbest, *_ = process_preds(p, qnp_lines, qnp_lines_bal,	×
177	verbose=False)
178
179	cbest = np.array(c_line[c_line.argmax(axis=0), np.arange(len(zbest))])	×
180	c_thresh = 0.5	×
181	n_thresh = 1	×
182	is_qso = np.sum(c_line > c_thresh, axis=0) >= n_thresh	×
183
184	zcat['Z_QN'][w] = zbest	×
185	zcat['Z_QN_CONF'][w] = cbest	×
186	zcat['IS_QSO_QN'][w] = is_qso	×
187
188	return zcat	×
189
190
191	def get_sq_model_fname(sqmodel_fname=None):	×
192	"""Convenience function to grab the $SQ_MODEL_FNAME environment variable.
193
194	Parameters
195	----------
196	sqmodel_fname : :class:`str`, optional, defaults to $SQ_MODEL_FNAME
197	If `sqmodel_fname` is passed, it is returned from this function. If it's
198	not passed, the $SQ_MODEL_FNAME environment variable is returned.
199
200	Returns
201	-------
202	:class:`str`
203	If `sqmodel_fname` is passed, it is returned from this function. If it's
204	not passed, the directory stored in the $SQ_MODEL_FNAME environment
205	variable is returned.
206	"""
207	if sqmodel_fname is None:	×
208	sqmodel_fname = os.environ.get('SQ_MODEL_FILE')	×
209	# EBL check that the $SQ_MODEL_FNAME environment variable is set.
210	if sqmodel_fname is None:	×
211	msg = "Pass sqmodel_fname or set $SQ_MODEL_FNAME environment variable!"	×
212	log.critical(msg)	×
213	raise ValueError(msg)	×
214
215	return sqmodel_fname	×
216
217
218	def load_sq_model(model_filename):	×
219	"""Convenience function for loading the SQUEzE model file.
220
221	Parameters
222	----------
223	model_filename : :class:'str'
224	The filename and path of the SQUEzE model file. Either input by user
225	or defaults to get_sq_model_fname().
226
227	Returns
228	-------
229	:class:'~numpy.array'
230	A numpy array of the SQUEzE model.
231
232	Notes
233	-----
234	- The input model file needs to be in the json file format.
235	"""
236	model = Model.from_json(load_json(model_filename))	×
237
238	return model	×
239
240
241	def add_sq_data(zcat, coaddname, squeze_model):	×
242	"""Apply the SQUEzE model to the input zcat and add data to columns.
243
244	Parameters
245	----------
246	zcat : :class:'~numpy.array'
247	The structured array that was created by make_new_zcat()
248	coaddname : class:'str'
249	The name of the coadd file corresponding to the zbest file used
250	in make_new_zcat()
251	squeze_model : :class:'numpy.array'
252	The loaded SQUEzE model file
253
254	Returns
255	-------
256	:class:'~numpy.array'
257	The zcat array with SQUEzE data included in the columns:
258
259	* Z_SQ - The best redshift from SQUEzE for each object.
260	* Z_SQ_CONF - The confidence value of this redshift.
261	"""
262	tmark(' Adding SQUEzE data')	×
263	mdata = ['TARGETID']	×
264	single_exposure = False	×
265	sq_cols_keep = ['PROB', 'Z_TRY', 'TARGETID']	×
266
267	tmark(' Reading spectra')	×
268	desi_spectra = read_spectra(coaddname)	×
269	# EBL Initialize squeze Spectra class
270	squeze_spectra = Spectra()	×
271	# EBL Get TARGETIDs
272	targetid = np.unique(desi_spectra.fibermap['TARGETID'])	×
273	# EBL Loop over TARGETIDs to build the Spectra objects
274	for targid in targetid:	×
275	# EBL Select objects
276	pos = np.where(desi_spectra.fibermap['TARGETID'] == targid)	×
277	# EBL Prepare column metadata
278	metadata = {col.upper(): desi_spectra.fibermap[col][pos[0][0]] for col in mdata}	×
279	# EBL Add the SPECID as the TARGETID
280	metadata['SPECID'] = targid	×
281	# EBL Extract the data
282	flux = {}	×
283	wave = {}	×
284	ivar = {}	×
285	mask = {}	×
286	for band in desi_spectra.bands:	×
287	flux[band] = desi_spectra.flux[band][pos]	×
288	wave[band] = desi_spectra.wave[band]	×
289	ivar[band] = desi_spectra.ivar[band][pos]	×
290	mask[band] = desi_spectra.mask[band][pos]	×
291
292	# EBL Format each spectrum for the model application
293	spectrum = DesiSpectrum(flux, wave, ivar, mask, metadata, single_exposure)	×
294	# EBL Append the spectrum to the Spectra object
295	squeze_spectra.append(spectrum)	×
296
297	# EBL Initialize candidate object. This takes a while with no feedback
298	# so we want a time output for benchmarking purposes.
299	tmark(' Initializing candidates')	×
300	candidates = Candidates(mode='operation', model=squeze_model)	×
301	# EBL Look for candidate objects. This also takes a while.
302	tmark(' Looking for candidates')	×
303	candidates.find_candidates(squeze_spectra.spectra_list(), save=False)	×
304	# EBL Compute the probabilities of the line/model matches to the spectra
305	tmark(' Computing probabilities')	×
306	candidates.classify_candidates(save=False)	×
307	# EBL Filter the results by removing the duplicate entries for each
308	# TARGETID. Merge the remaining with the zcat data.
309	tmark(' Merging SQUEzE data with zcat')	×
310	data_frame = candidates.candidates()	×
311	data_frame = data_frame[~data_frame['DUPLICATED']][sq_cols_keep]	×
312	# EBL Strip the pandas data frame structure and put it into a numpy
313	# structured array first.
314	sqdata_arr = np.zeros(len(data_frame), dtype=[('TARGETID', 'int64'),	×
315	('Z_SQ', 'float64'),
316	('Z_SQ_CONF', 'float64')])
317	sqdata_arr['TARGETID'] = data_frame['TARGETID'].values	×
318	sqdata_arr['Z_SQ'] = data_frame['Z_TRY'].values	×
319	sqdata_arr['Z_SQ_CONF'] = data_frame['PROB'].values	×
320	# EBL SQUEzE will reorder the objects, so match on TARGETID.
321	zcat_args, sqdata_args = match(zcat['TARGETID'],sqdata_arr['TARGETID'])	×
322	zcat['Z_SQ'][zcat_args] = sqdata_arr['Z_SQ'][sqdata_args]	×
323	zcat['Z_SQ_CONF'][zcat_args] = sqdata_arr['Z_SQ_CONF'][sqdata_args]	×
324
325	return zcat	×
326
327
328	def add_abs_data(zcat, coaddname):	×
329	"""Add the MgII absorption line finder data to the input zcat array.
330
331	Parameters
332	----------
333	zcat : :class:'~numpy.array'
334	The structured array that was created by make_new_zcat()
335	coaddname : class:'str'
336	The name of the coadd file corresponding to the zbest file used
337	in make_new_zcat()
338
339	Returns
340	-------
341	:class:'~numpy.array'
342	The zcat array with MgII Absorption data included in the columns:
343
344	* Z_ABS - The highest redshift of MgII absorption
345	* Z_ABS_CONF - The confidence value for this redshift.
346	"""
347	tmark(' MgII Absorption data not yet added.')	×
348
349	return zcat	×
350
351
352	def zcomb_selector(zcat, proc_flag=False):	×
353	"""Compare results from redrock, QuasarNP, SQUEzE, and MgII data.
354
355	Parameters
356	----------
357	zcat : :class:'~numpy.array'
358	The structured array that was created by make_new_zcat()
359	proc_flag : :class:'bool'
360	Turn on extra comparison procedure.
361
362	Returns
363	-------
364	:class:'~numpy.array'
365	The zcat array with SQUEzE data included in the columns:
366
367	* Z_COMB - The best models-combined redshift for each object.
368	* Z_COMB_PROB - The combined probability value of that redshift.
369	"""
370	zcat['Z_COMB'][:] = zcat['Z']	×
371	zcat['Z_COMB_PROB'][:] = 0.95	×
372
373	return zcat	×
374
375
376	def zcat_writer(zcat, outputdir, outputname, qn_flag=False, sq_flag=False, abs_flag=False):	×
377	"""Writes the zcat structured array out as a FITS file.
378
379	Parameters
380	----------
381	zcat : :class:'~numpy.array'
382	The structured array that was created by make_new_zcat()
383	outputdir : :class:'str'
384	The directory where the zcat file will be written.
385	outputname : :class:'str'
386	The filename of the zqso output file.
387	qn_flag : :class:'bool'
388	Flag if QuasarNP data (or not) was added to the zcat file.
389	sq_flag : :class:'bool'
390	Flag if SQUEzE data (or not) was added to the zcat file.
391	abs_flag : :class:'bool'
392	Flag if MgII Absorption data (or not) was added to the zcat file.
393
394	Returns
395	-------
396	:class:'str'
397	The filename, with path, of the FITS file written out.
398	"""
399	tmark(' Creating file...')	×
400	full_outputname = os.path.join(outputdir, outputname)	×
401
402	prim_hdu = fits.PrimaryHDU()	×
403	prim_hdu.header['QN_ADDED'] = str(qn_flag)	×
404	prim_hdu.header['SQ_ADDED'] = str(sq_flag)	×
405	prim_hdu.header['AB_ADDED'] = str(abs_flag)	×
406
407	data_hdu = fits.BinTableHDU.from_columns(zcat, name='ZCATALOG')	×
408	data_out = fits.HDUList([prim_hdu, data_hdu])	×
409
410	data_out.writeto(full_outputname)	×
411
412	return full_outputname	×
413
414
415	def create_zcat(tile, night, petal_num, zcatdir, outputdir, qn_flag=False,	×
416	qnp_model=None, qnp_lines=None, qnp_lines_bal=None,
417	sq_flag=False, squeze_model=None, abs_flag=False):
418	"""This will create a single zqso file from a set of user inputs.
419
420	Parameters
421	----------
422	tile : :class:'str'
423	The TILEID of the tile to process.
424	night : :class:'str'
425	The date associated with the observation of the 'tile' used.
426	* Must be in YYYYMMDD format
427	petal_num : :class:'int'
428	If 'all_petals' isn't used, the single petal to create a zcat for.
429	zcatdir : :class:'str'
430	The location for the daily redrock output.
431	outputdir : :class:'str'
432	The filepath to write out the zcat file.
433	qn_flag : :class:'bool', optional
434	Flag to add QuasarNP data (or not) to the zcat file.
435	qnp_model : :class:'h5 array', optional
436	The QuasarNP model file to be used for line predictions.
437	qnp_lines : :class:'list', optional
438	The list of lines to use in the QuasarNP model to test against.
439	qnp_lines : :class:'list', optional
440	The list of BAL lines to use for QuasarNP to identify BALs.
441	sq_flag : :class:'bool', optional
442	Flag to add SQUEzE data (or not) to the zcat file.
443	squeze_model : :class:'numpy.array', optional
444	The numpy array for the SQUEzE model file.
445	abs_flag : :class:'bool', optional
446	Flag to add MgII Absorption data (or not) to the zcat file.
447
448	Notes
449	-----
450	- Writes a FITS catalog that incorporates redrock, QuasarNP, SQUEzE, and MgII
451	absorption redshifts and confidence values. This will write to the same
452	directory of the zbest and coadd files unless a different output directory is
453	passed.
454	"""
455	# EBL Create the filepath for the input tile/night combination
456	tiledir = os.path.join(zcatdir, tile)	×
457	ymdir = os.path.join(tiledir, night)	×
458
459	# EBL Create the filename tag that appends to zbest-, coadd-, and zqso-*
460	# files.
461	filename_tag = f'{petal_num}-{tile}-thru{night}.fits'	×
462	zbestname = f'zbest-{filename_tag}'	×
463	coaddname = f'coadd-{filename_tag}'	×
464	outputname = f'zqso-{filename_tag}'	×
465
466	zcat = make_new_zcat(os.path.join(ymdir, zbestname))	×
467	if isinstance(zcat, bool):	×
468	print(' !!!! Petal Number does not have a corresponding zbest file !!!!')	×
469	else:
470	if qn_flag :	×
471	zcat = add_qn_data(zcat, os.path.join(ymdir, coaddname),	×
472	qnp_model, qnp_lines, qnp_lines_bal)
473	if sq_flag:	×
474	zcat = add_sq_data(zcat, os.path.join(ymdir, coaddname), squeze_model)	×
475	if abs_flag:	×
476	zcat = add_abs_data(zcat, os.path.join(ymdir, coaddname))	×
477
478	fzcat = zcomb_selector(zcat)	×
479
480	full_outputname = zcat_writer(fzcat, outputdir, outputname, qn_flag,	×
481	sq_flag, abs_flag)
482
483	tmark(' --{} written out correctly.'.format(full_outputname))	×
484	print('='*79)	×

desihub / desitarget / 786934210

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous