<commit_before>
from __future__ import absolute_import, division


from desispec.io import read_frame
from desispec.io import read_fiberflat
from desispec.io import read_sky
from desispec.io import write_qa_frame
from desispec.io.fluxcalibration import read_stdstar_models
from desispec.io.fluxcalibration import write_flux_calibration
from desispec.io.qa import load_qa_frame
from desispec.fiberflat import apply_fiberflat
from desispec.sky import subtract_sky
from desispec.fluxcalibration import compute_flux_calibration, isStdStar
from desiutil.log import get_logger
from desispec.qa import qa_plots

import argparse
import os
import os.path
import numpy as np
import sys


def parse(options=None):
    parser = argparse.ArgumentParser(description="Compute the flux calibration for a DESI frame using precomputed spectro-photometrically calibrated stellar models.")

    parser.add_argument('--infile', type = str, default = None, required=True,
                        help = 'path of DESI exposure frame fits file')
    parser.add_argument('--fiberflat', type = str, default = None, required=True,
                        help = 'path of DESI fiberflat fits file')
    parser.add_argument('--sky', type = str, default = None, required=True,
                        help = 'path of DESI sky fits file')
    parser.add_argument('--models', type = str, default = None, required=True,
                        help = 'path of spetro-photometric stellar spectra fits file')
    parser.add_argument('--chi2cut', type = float, default = 0., required=False,
                        help = 'apply a reduced chi2 cut for the selection of stars')
    parser.add_argument('--chi2cut-nsig', type = float, default = 0., required=False,
                        help = 'discard n-sigma outliers from the reduced chi2 of the standard star fit')
    parser.add_argument('--min-color', type = float, default = 0.25, required=False,
                        help = 'only consider stars with g-r greater than this')
    parser.add_argument('--delta-color-cut', type = float, default = 0.3, required=False,
                        help = 'discard model stars with different broad-band color from imaging')
    parser.add_argument('--outfile', type = str, default = None, required=True,
                        help = 'path of DESI flux calbration fits file')
    parser.add_argument('--qafile', type=str, default=None, required=False,
                        help='path of QA file.')
    parser.add_argument('--qafig', type = str, default = None, required=False,
                        help = 'path of QA figure file')
    parser.add_argument('--highest-throughput', type = int, default = 0, required=False,
                        help = 'use this number of stars ranked by highest throughput to normalize transmission (for DESI commissioning)')
    
    args = None
    if options is None:
        args = parser.parse_args()
    else:
        args = parser.parse_args(options)
    return args


def main(args) :

    log=get_logger()

    cmd = ['desi_compute_fluxcalibration',]
    for key, value in args.__dict__.items():
        if value is not None:
            cmd += ['--'+key, str(value)]
    cmd = ' '.join(cmd)
    log.info(cmd)

    log.info("read frame")
    # read frame
    frame = read_frame(args.infile)

    log.info("apply fiberflat")
    # read fiberflat
    fiberflat = read_fiberflat(args.fiberflat)

    # apply fiberflat
    apply_fiberflat(frame, fiberflat)

    log.info("subtract sky")
    # read sky
    skymodel=read_sky(args.sky)

    # subtract sky
    subtract_sky(frame, skymodel)

    log.info("compute flux calibration")

    # read models
    model_flux,model_wave,model_fibers,model_metadata=read_stdstar_models(args.models)

    ok=np.ones(len(model_metadata),dtype=bool)
    
    if args.chi2cut > 0 :
        log.info("Apply cut CHI2DOF<{}".format(args.chi2cut))
        ok &= (model_metadata["CHI2DOF"]<args.chi2cut)
    if args.delta_color_cut > 0 :
        log.info("Apply cut |delta color|<{}".format(args.delta_color_cut))
        ok &= (np.abs(model_metadata["MODEL_G-R"]-model_metadata["DATA_G-R"])<args.delta_color_cut)
    log.info("Apply cut DATA_G-R>{}".format(args.min_color))
    ok &= (model_metadata["DATA_G-R"]>args.min_color)

    if args.chi2cut_nsig > 0 :
        # automatically reject stars that ar chi2 outliers
        mchi2=np.median(model_metadata["CHI2DOF"])
        rmschi2=np.std(model_metadata["CHI2DOF"])
        maxchi2=mchi2+args.chi2cut_nsig*rmschi2
        log.info("Apply cut CHI2DOF<{} based on chi2cut_nsig={}".format(maxchi2,args.chi2cut_nsig))
        ok &= (model_metadata["CHI2DOF"]<=maxchi2)
    
    ok=np.where(ok)[0]
    if ok.size == 0 :
        log.error("cuts discarded all stars")
        sys.exit(12)
    nstars=model_flux.shape[0]
    nbad=nstars-ok.size
    if nbad>0 :
        log.warning("discarding %d star(s) out of %d because of cuts"%(nbad,nstars))
        model_flux=model_flux[ok]
        model_fibers=model_fibers[ok]
        model_metadata=model_metadata[:][ok]
    
    # check that the model_fibers are actually standard stars
    fibermap = frame.fibermap

    ## check whether star fibers from args.models are consistent with fibers from fibermap
    ## if not print the OBJTYPE from fibermap for the fibers numbers in args.models and exit
    fibermap_std_indices = np.where(isStdStar(fibermap))[0]
    if np.any(~np.in1d(model_fibers%500, fibermap_std_indices)):
        if 'DESI_TARGET' in fibermap:
            colname = 'DESI_TARGET'
        else:
            colname = 'SV1_DESI_TARGET'  #- TODO: could become SV2_DESI_TARGET

        for i in model_fibers%500:
            log.error("inconsistency with spectrum {}, OBJTYPE='{}', {}={} in fibermap".format(
                (i, fibermap["OBJTYPE"][i], colname, fibermap[colname][i])))
        sys.exit(12)

    # Make sure the fibers of interest aren't entirely masked.
    if np.sum(np.sum(frame.ivar[model_fibers%500, :] == 0, axis=1) == frame.nwave) == len(model_fibers):
        log.warning('All standard-star spectra are masked!')
        return
        
    fluxcalib = compute_flux_calibration(frame, model_wave, model_flux, model_fibers%500, highest_throughput_nstars = args.highest_throughput)

    # QA
    if (args.qafile is not None):
        log.info("performing fluxcalib QA")
        # Load
        qaframe = load_qa_frame(args.qafile, frame_meta=frame.meta, flavor=frame.meta['FLAVOR'])
        # Run
        #import pdb; pdb.set_trace()<commit_msg>fix error message<commit_after>
from __future__ import absolute_import, division


from desispec.io import read_frame
from desispec.io import read_fiberflat
from desispec.io import read_sky
from desispec.io import write_qa_frame
from desispec.io.fluxcalibration import read_stdstar_models
from desispec.io.fluxcalibration import write_flux_calibration
from desispec.io.qa import load_qa_frame
from desispec.fiberflat import apply_fiberflat
from desispec.sky import subtract_sky
from desispec.fluxcalibration import compute_flux_calibration, isStdStar
from desiutil.log import get_logger
from desispec.qa import qa_plots
from desitarget.targets import main_cmx_or_sv

import argparse
import os
import os.path
import numpy as np
import sys


def parse(options=None):
    parser = argparse.ArgumentParser(description="Compute the flux calibration for a DESI frame using precomputed spectro-photometrically calibrated stellar models.")

    parser.add_argument('--infile', type = str, default = None, required=True,
                        help = 'path of DESI exposure frame fits file')
    parser.add_argument('--fiberflat', type = str, default = None, required=True,
                        help = 'path of DESI fiberflat fits file')
    parser.add_argument('--sky', type = str, default = None, required=True,
                        help = 'path of DESI sky fits file')
    parser.add_argument('--models', type = str, default = None, required=True,
                        help = 'path of spetro-photometric stellar spectra fits file')
    parser.add_argument('--chi2cut', type = float, default = 0., required=False,
                        help = 'apply a reduced chi2 cut for the selection of stars')
    parser.add_argument('--chi2cut-nsig', type = float, default = 0., required=False,
                        help = 'discard n-sigma outliers from the reduced chi2 of the standard star fit')
    parser.add_argument('--min-color', type = float, default = 0.25, required=False,
                        help = 'only consider stars with g-r greater than this')
    parser.add_argument('--delta-color-cut', type = float, default = 0.3, required=False,
                        help = 'discard model stars with different broad-band color from imaging')
    parser.add_argument('--outfile', type = str, default = None, required=True,
                        help = 'path of DESI flux calbration fits file')
    parser.add_argument('--qafile', type=str, default=None, required=False,
                        help='path of QA file.')
    parser.add_argument('--qafig', type = str, default = None, required=False,
                        help = 'path of QA figure file')
    parser.add_argument('--highest-throughput', type = int, default = 0, required=False,
                        help = 'use this number of stars ranked by highest throughput to normalize transmission (for DESI commissioning)')
    
    args = None
    if options is None:
        args = parser.parse_args()
    else:
        args = parser.parse_args(options)
    return args


def main(args) :

    log=get_logger()

    cmd = ['desi_compute_fluxcalibration',]
    for key, value in args.__dict__.items():
        if value is not None:
            cmd += ['--'+key, str(value)]
    cmd = ' '.join(cmd)
    log.info(cmd)

    log.info("read frame")
    # read frame
    frame = read_frame(args.infile)

    log.info("apply fiberflat")
    # read fiberflat
    fiberflat = read_fiberflat(args.fiberflat)

    # apply fiberflat
    apply_fiberflat(frame, fiberflat)

    log.info("subtract sky")
    # read sky
    skymodel=read_sky(args.sky)

    # subtract sky
    subtract_sky(frame, skymodel)

    log.info("compute flux calibration")

    # read models
    model_flux,model_wave,model_fibers,model_metadata=read_stdstar_models(args.models)

    ok=np.ones(len(model_metadata),dtype=bool)
    
    if args.chi2cut > 0 :
        log.info("Apply cut CHI2DOF<{}".format(args.chi2cut))
        ok &= (model_metadata["CHI2DOF"]<args.chi2cut)
    if args.delta_color_cut > 0 :
        log.info("Apply cut |delta color|<{}".format(args.delta_color_cut))
        ok &= (np.abs(model_metadata["MODEL_G-R"]-model_metadata["DATA_G-R"])<args.delta_color_cut)
    log.info("Apply cut DATA_G-R>{}".format(args.min_color))
    ok &= (model_metadata["DATA_G-R"]>args.min_color)

    if args.chi2cut_nsig > 0 :
        # automatically reject stars that ar chi2 outliers
        mchi2=np.median(model_metadata["CHI2DOF"])
        rmschi2=np.std(model_metadata["CHI2DOF"])
        maxchi2=mchi2+args.chi2cut_nsig*rmschi2
        log.info("Apply cut CHI2DOF<{} based on chi2cut_nsig={}".format(maxchi2,args.chi2cut_nsig))
        ok &= (model_metadata["CHI2DOF"]<=maxchi2)
    
    ok=np.where(ok)[0]
    if ok.size == 0 :
        log.error("cuts discarded all stars")
        sys.exit(12)
    nstars=model_flux.shape[0]
    nbad=nstars-ok.size
    if nbad>0 :
        log.warning("discarding %d star(s) out of %d because of cuts"%(nbad,nstars))
        model_flux=model_flux[ok]
        model_fibers=model_fibers[ok]
        model_metadata=model_metadata[:][ok]
    
    # check that the model_fibers are actually standard stars
    fibermap = frame.fibermap

    ## check whether star fibers from args.models are consistent with fibers from fibermap
    ## if not print the OBJTYPE from fibermap for the fibers numbers in args.models and exit
    fibermap_std_indices = np.where(isStdStar(fibermap))[0]
    if np.any(~np.in1d(model_fibers%500, fibermap_std_indices)):
        target_colnames, target_masks, survey = main_cmx_or_sv(fibermap)
        colname =  target_colnames[0]
        for i in model_fibers%500:
            log.error("inconsistency with spectrum {}, OBJTYPE={}, {}={} in fibermap".format(
                i, fibermap["OBJTYPE"][i], colname, fibermap[colname][i]))
        sys.exit(12)

    # Make sure the fibers of interest aren't entirely masked.
    if np.sum(np.sum(frame.ivar[model_fibers%500, :] == 0, axis=1) == frame.nwave) == len(model_fibers):
        log.warning('All standard-star spectra are masked!')
        return
        
    fluxcalib = compute_flux_calibration(frame, model_wave, model_flux, model_fibers%500, highest_throughput_nstars = args.highest_throughput)

    # QA
    if (args.qafile is not None):
        log.info("performing fluxcalib QA")
        # Load
        qaframe = load_qa_frame(args.qafile, frame_meta=frame.meta, flavor=frame.meta['FLAVOR'])
        # Run
        #import pdb; pdb.set_trace()
{'fraction_non_alphanumeric': 0.10838650165270197, 'fraction_numerical': 0.009454992697363364, 'mean_word_length': 3.0353598014888337, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 15, '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': '5310153', 'n_tokens_mistral': 4232, 'n_tokens_neox': 4024, 'n_words': 1020}