<commit_before>"""
desispec.zmtl
=============

Post-redrock ML processing for LyA Quasar object identification for
creating zmtl files to be used by the merged target list (MTL) updates.
"""

import os
import numpy as np
import time
import fitsio

from operator import itemgetter
from itertools import groupby
from astropy.modeling import fitting
from astropy.modeling import models
from astropy.table import Table
from scipy.signal import medfilt
from astropy.convolution import Gaussian1DKernel
from astropy.convolution import convolve

from quasarnp.io import load_model
from quasarnp.io import load_desi_coadd
from quasarnp.utils import process_preds

from desitarget.geomask import match, match_to
from desitarget.internal import sharedmem
from desitarget.io import write_with_units
from desiutil.depend import add_dependencies
from desitarget.targets import main_cmx_or_sv, switch_main_cmx_or_sv
from desitarget.targetmask import zwarn_mask

from desispec.io import read_spectra, findfile
from desispec.exposure_qa import get_qa_params
from desispec.maskbits import fibermask

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

# ADM data models for the various afterburners.
zmtldatamodel = np.array([], [('RA', '>f8'), ('DEC', '>f8'), ('TARGETID','>i8'),
                              ('DESI_TARGET', '>i8'), ('BGS_TARGET', '>i8'),
                              ('MWS_TARGET', '>i8'), ('SCND_TARGET', '>i8'),
                              ('Z', '>f8'), ('ZWARN', '>i8'),
                              ('SPECTYPE', '<U6'), ('DELTACHI2', '>f8'),
                              ('NUMOBS', '>i4'), ('ZTILEID', '>i4')])
qndm = [('Z_QN', '>f8'), ('Z_QN_CONF', '>f8'), ('IS_QSO_QN', '>i2')]
sqdm = [('Z_SQ', '>f8'), ('Z_SQ_CONF', '>f8')]
absdm = [('Z_ABS', '>f8'), ('Z_ABS_CONF', '>f8')]
combdm = [('Z_COMB', '>f8'), ('Z_COMB_PROB', '>f8')]


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_zmtl(zbestname, qn_flag=False, sq_flag=False, abs_flag=False,
                  zcomb_flag=False):
    """Make the initial zmtl array with redrock data.

    Parameters
    ----------
    zbestname : :class:`str`
        Full filename and path for the zbest file to process.
    qn_flag : :class:`bool'` optional
        Flag to add QuasarNP data (or not) to the zmtl file.
    sq_flag : :class:`bool`, optional
        Flag to add SQUEzE data (or not) to the zmtl file.
    abs_flag : :class:`bool`, optional
        Flag to add MgII Absorption data (or not) to the zmtl file.
    zcomb_flag : :class:`bool`, optional
        Flag if a combined redshift (or not) was added to the zmtl file.

    Returns
    -------
    :class:`~numpy.array` or `bool`
        A zmtl in the official format (`zmtldatamodel`) compiled from
        the `tile', 'night', and 'petal_num', in `zmtldir`. If the zbest
        file for that petal doesn't exist, returns ``False``.
    """
    tmark('    Making redrock zmtl')

    # ADM read in the zbest and fibermap extensions for the RR
    # ADM redshift catalog, if they exist.
    try:
        zs = fitsio.read(zbestname, "ZBEST")
        fms = fitsio.read(zbestname, "FIBERMAP")
        log.info(f'Read {zbestname}')
    except (FileNotFoundError, OSError):
        log.error(f'Missing {zbestname}')
        return False

    # ADM recover the information for unique targets based on the
    # ADM first entry for each TARGETID.
    _, ii = np.unique(fms['TARGETID'], return_index=True)
    fms = fms[ii]

    # ADM check for some glitches.
    if len(zs) != len(set(zs["TARGETID"])):
        msg = "a target is duplicated in file {}!!!".format(zbestname)
        log.critical(msg)
        raise ValueError(msg)
    # ADM check for some glitches.
    if len(zs) != len(fms):
        msg = "TARGETID mismatch for extensions in file {}!!!".format(zbestname)
        log.critical(msg)
        raise ValueError(msg)

    # ADM Strictly match the targets in the z catalog and fibermap.
    zid = match_to(fms["TARGETID"], zs["TARGETID"])

    # ADM set up the output zmtl file, which differs depending on which
    # ADM afterburners were specified.
    dtswitched = switch_main_cmx_or_sv(zmtldatamodel, fms)
    dt = dtswitched.dtype.descr
    for flag, dm in zip([qn_flag, sq_flag, abs_flag, zcomb_flag],
                        [qndm, sqdm, absdm, combdm]):
        if flag:
            dt += dm
    zmtl = np.full(len(zs), -1, dtype=dt)

    # ADM add the columns from the original zbest file.
    zmtl["RA"] = fms[zid]["TARGET_RA"]
    zmtl["DEC"] = fms[zid]["TARGET_DEC"]
    zmtl["ZTILEID"] = fms[zid]["TILEID"]
    zmtl["NUMOBS"] = zs["NUMTILE"]

    # ADM also add the appropriate bit-columns.
    Mxcols, _, _, = main_cmx_or_sv(fms, scnd=True)
    for col in Mxcols:
        if col in fms.dtype.names:
            zmtl[col] = fms[zid][col]
        # SB fail on missing required columns ...
        elif col in zmtldatamodel.dtype.names:
            msg = f'Input fibermap missing {col}, which is required by zmtl datamodel'
            log.critical(msg)
            raise ValueError(msg)
        # SB ... but only log error about unexpectedly missing optional columns
        else:
            log.error(f'Input fibermap missing optional {col}; leaving it blank')

    # ADM write out the unwritten columns.
    allcols = set(dtswitched.dtype.names)
    usedcols = set(['RA', 'DEC', 'NUMOBS', 'ZTILEID'] + Mxcols)
    for col in allcols - usedcols:
        zmtl[col] = zs[col]<commit_msg>zmtl update for new FIBERMAP<commit_after>"""
desispec.zmtl
=============

Post-redrock ML processing for LyA Quasar object identification for
creating zmtl files to be used by the merged target list (MTL) updates.
"""

import os
import numpy as np
import time
import fitsio

from operator import itemgetter
from itertools import groupby
from astropy.modeling import fitting
from astropy.modeling import models
from astropy.table import Table
from scipy.signal import medfilt
from astropy.convolution import Gaussian1DKernel
from astropy.convolution import convolve

from quasarnp.io import load_model
from quasarnp.io import load_desi_coadd
from quasarnp.utils import process_preds

from desitarget.geomask import match, match_to
from desitarget.internal import sharedmem
from desitarget.io import write_with_units
from desiutil.depend import add_dependencies
from desitarget.targets import main_cmx_or_sv, switch_main_cmx_or_sv
from desitarget.targetmask import zwarn_mask

from desispec.io import read_spectra, findfile
from desispec.exposure_qa import get_qa_params
from desispec.maskbits import fibermask

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

# ADM data models for the various afterburners.
zmtldatamodel = np.array([], [('RA', '>f8'), ('DEC', '>f8'), ('TARGETID','>i8'),
                              ('DESI_TARGET', '>i8'), ('BGS_TARGET', '>i8'),
                              ('MWS_TARGET', '>i8'), ('SCND_TARGET', '>i8'),
                              ('Z', '>f8'), ('ZWARN', '>i8'),
                              ('SPECTYPE', '<U6'), ('DELTACHI2', '>f8'),
                              ('NUMOBS', '>i4'), ('ZTILEID', '>i4')])
qndm = [('Z_QN', '>f8'), ('Z_QN_CONF', '>f8'), ('IS_QSO_QN', '>i2')]
sqdm = [('Z_SQ', '>f8'), ('Z_SQ_CONF', '>f8')]
absdm = [('Z_ABS', '>f8'), ('Z_ABS_CONF', '>f8')]
combdm = [('Z_COMB', '>f8'), ('Z_COMB_PROB', '>f8')]


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_zmtl(zbestname, qn_flag=False, sq_flag=False, abs_flag=False,
                  zcomb_flag=False):
    """Make the initial zmtl array with redrock data.

    Parameters
    ----------
    zbestname : :class:`str`
        Full filename and path for the zbest file to process.
    qn_flag : :class:`bool'` optional
        Flag to add QuasarNP data (or not) to the zmtl file.
    sq_flag : :class:`bool`, optional
        Flag to add SQUEzE data (or not) to the zmtl file.
    abs_flag : :class:`bool`, optional
        Flag to add MgII Absorption data (or not) to the zmtl file.
    zcomb_flag : :class:`bool`, optional
        Flag if a combined redshift (or not) was added to the zmtl file.

    Returns
    -------
    :class:`~numpy.array` or `bool`
        A zmtl in the official format (`zmtldatamodel`) compiled from
        the `tile', 'night', and 'petal_num', in `zmtldir`. If the zbest
        file for that petal doesn't exist, returns ``False``.
    """
    tmark('    Making redrock zmtl')

    # ADM read in the zbest and fibermap extensions for the RR
    # ADM redshift catalog, if they exist.
    try:
        zs = fitsio.read(zbestname, "ZBEST")
        fms = fitsio.read(zbestname, "FIBERMAP")
        log.info(f'Read {zbestname}')
    except (FileNotFoundError, OSError):
        log.error(f'Missing {zbestname}')
        return False

    # SB check assumption that FIBERMAP is row-matched to ZBEST
    if len(zs) != len(fms) or np.any(zs['TARGETID'] != fms['TARGETID']):
        msg = "ZBEST and FIBERMAP TARGETIDs aren't row-matched"
        log.critical(msg)
        raise ValueError(msg)

    # ADM check for some glitches.
    if len(zs) != len(set(zs["TARGETID"])):
        msg = "a target is duplicated in file {}!!!".format(zbestname)
        log.critical(msg)
        raise ValueError(msg)
    # ADM check for some glitches.
    if len(zs) != len(fms):
        msg = "TARGETID mismatch for extensions in file {}!!!".format(zbestname)
        log.critical(msg)
        raise ValueError(msg)

    # ADM Strictly match the targets in the z catalog and fibermap.
    zid = match_to(fms["TARGETID"], zs["TARGETID"])

    # ADM set up the output zmtl file, which differs depending on which
    # ADM afterburners were specified.
    dtswitched = switch_main_cmx_or_sv(zmtldatamodel, fms)
    dt = dtswitched.dtype.descr
    for flag, dm in zip([qn_flag, sq_flag, abs_flag, zcomb_flag],
                        [qndm, sqdm, absdm, combdm]):
        if flag:
            dt += dm
    zmtl = np.full(len(zs), -1, dtype=dt)

    # ADM add the columns from the original zbest file.
    zmtl["RA"] = fms[zid]["TARGET_RA"]
    zmtl["DEC"] = fms[zid]["TARGET_DEC"]
    zmtl["ZTILEID"] = fms[zid]["TILEID"]
    zmtl["NUMOBS"] = fms["COADD_NUMTILE"]

    # ADM also add the appropriate bit-columns.
    Mxcols, _, _, = main_cmx_or_sv(fms, scnd=True)
    for col in Mxcols:
        if col in fms.dtype.names:
            zmtl[col] = fms[zid][col]
        # SB fail on missing required columns ...
        elif col in zmtldatamodel.dtype.names:
            msg = f'Input fibermap missing {col}, which is required by zmtl datamodel'
            log.critical(msg)
            raise ValueError(msg)
        # SB ... but only log error about unexpectedly missing optional columns
        else:
            log.error(f'Input fibermap missing optional {col}; leaving it blank')

    # ADM write out the unwritten columns.
    allcols = set(dtswitched.dtype.names)
    usedcols = set(['RA', 'DEC', 'NUMOBS', 'ZTILEID'] + Mxcols)
    for col in allcols - usedcols:
        zmtl[col] = zs[col]
{'fraction_non_alphanumeric': 0.12709284627092846, 'fraction_numerical': 0.0050735667174023336, 'mean_word_length': 2.8499348958333335, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 45, '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': '4696926', 'n_tokens_mistral': 4069, 'n_tokens_neox': 4001, 'n_words': 1211}