Example #1
0
def main(args) :

    log=get_logger()

    log.info("starting")

    # read exposure to load data and get range of spectra
    frame = read_frame(args.infile)
    specmin, specmax = np.min(frame.fibers), np.max(frame.fibers)

    if args.cosmics_nsig>0 : # Reject cosmics
        reject_cosmic_rays_1d(frame,args.cosmics_nsig)

    # read fiberflat
    fiberflat = read_fiberflat(args.fiberflat)

    # apply fiberflat to sky fibers
    apply_fiberflat(frame, fiberflat)

    # compute sky model
    skymodel = compute_sky(frame,add_variance=(not args.no_extra_variance),\
                           angular_variation_deg=args.angular_variation_deg,\
                           chromatic_variation_deg=args.chromatic_variation_deg,\
                           adjust_wavelength=args.adjust_wavelength,\
                           adjust_lsf=args.adjust_lsf)

    # QA
    if (args.qafile is not None) or (args.qafig is not None):
        log.info("performing skysub QA")
        # Load
        qaframe = load_qa_frame(args.qafile, frame_meta=frame.meta, flavor=frame.meta['FLAVOR'])
        # Run
        qaframe.run_qa('SKYSUB', (frame, skymodel))
        # Write
        if args.qafile is not None:
            write_qa_frame(args.qafile, qaframe)
            log.info("successfully wrote {:s}".format(args.qafile))
        # Figure(s)
        if args.qafig is not None:
            qa_plots.frame_skyres(args.qafig, frame, skymodel, qaframe)

    # record inputs
    frame.meta['IN_FRAME'] = shorten_filename(args.infile)
    frame.meta['FIBERFLT'] = shorten_filename(args.fiberflat)

    # write result
    write_sky(args.outfile, skymodel, frame.meta)
    log.info("successfully wrote %s"%args.outfile)
Example #2
0
def main(args):

    log = get_logger()

    log.info("degxx={} degxy={} degyx={} degyy={}".format(
        args.degxx, args.degxy, args.degyx, args.degyy))

    # read preprocessed image
    image = read_image(args.image)
    log.info("read image {}".format(args.image))
    if image.mask is not None:
        image.ivar *= (image.mask == 0)

    tset = fit_trace_shifts(image=image, args=args)
    tset.meta['IN_PSF'] = shorten_filename(args.psf)
    tset.meta['IN_IMAGE'] = shorten_filename(args.image)

    if args.outpsf is not None:
        write_traces_in_psf(args.psf, args.outpsf, tset)
        log.info("wrote modified PSF in %s" % args.outpsf)
Example #3
0
def main_mpi(args, comm=None, timing=None):
    freeze_iers()
    nproc = 1
    rank = 0
    if comm is not None:
        nproc = comm.size
        rank = comm.rank

    mark_start = time.time()

    log = get_logger()

    psf_file = args.psf
    input_file = args.input

    # these parameters are interpreted as the *global* spec range,
    # to be divided among processes.
    specmin = args.specmin
    nspec = args.nspec

    #- Load input files and broadcast

    # FIXME: after we have fixed the serialization
    # of the PSF, read and broadcast here, to reduce
    # disk contention.

    img = None
    if rank == 0:
        img = io.read_image(input_file)
    if comm is not None:
        img = comm.bcast(img, root=0)

    psf = load_psf(psf_file)

    mark_read_input = time.time()

    # get spectral range
    if nspec is None:
        nspec = psf.nspec

    if args.fibermap is not None:
        fibermap = io.read_fibermap(args.fibermap)
    else:
        try:
            fibermap = io.read_fibermap(args.input)
        except (AttributeError, IOError, KeyError):
            fibermap = None

    if fibermap is not None:
        fibermap = fibermap[specmin:specmin + nspec]
        if nspec > len(fibermap):
            log.warning(
                "nspec {} > len(fibermap) {}; reducing nspec to {}".format(
                    nspec, len(fibermap), len(fibermap)))
            nspec = len(fibermap)
        fibers = fibermap['FIBER']
    else:
        fibers = np.arange(specmin, specmin + nspec)

    specmax = specmin + nspec

    #- Get wavelength grid from options
    if args.wavelength is not None:
        raw_wstart, raw_wstop, raw_dw = [
            float(tmp) for tmp in args.wavelength.split(',')
        ]
    else:
        raw_wstart = np.ceil(psf.wmin_all)
        raw_wstop = np.floor(psf.wmax_all)
        raw_dw = 0.7

    raw_wave = np.arange(raw_wstart, raw_wstop + raw_dw / 2.0, raw_dw)
    nwave = len(raw_wave)
    bundlesize = args.bundlesize

    if args.barycentric_correction:
        if ('RA' in img.meta) or ('TARGTRA' in img.meta):
            barycentric_correction_factor = \
                    barycentric_correction_multiplicative_factor(img.meta)
        #- Early commissioning has RA/TARGTRA in fibermap but not HDU 0
        elif fibermap is not None and \
                (('RA' in fibermap.meta) or ('TARGTRA' in fibermap.meta)):
            barycentric_correction_factor = \
                    barycentric_correction_multiplicative_factor(fibermap.meta)
        else:
            msg = 'Barycentric corr requires (TARGT)RA in HDU 0 or fibermap'
            log.critical(msg)
            raise KeyError(msg)
    else:
        barycentric_correction_factor = 1.

    # Explictly define the correct wavelength values to avoid confusion of reference frame
    # If correction applied, otherwise divide by 1 and use the same raw values
    wstart = raw_wstart / barycentric_correction_factor
    wstop = raw_wstop / barycentric_correction_factor
    dw = raw_dw / barycentric_correction_factor
    wave = raw_wave / barycentric_correction_factor

    #- Confirm that this PSF covers these wavelengths for these spectra
    psf_wavemin = np.max(psf.wavelength(list(range(specmin, specmax)), y=-0.5))
    psf_wavemax = np.min(
        psf.wavelength(list(range(specmin, specmax)), y=psf.npix_y - 0.5))
    if psf_wavemin - 5 > wstart:
        raise ValueError(
            'Start wavelength {:.2f} < min wavelength {:.2f} for these fibers'.
            format(wstart, psf_wavemin))
    if psf_wavemax + 5 < wstop:
        raise ValueError(
            'Stop wavelength {:.2f} > max wavelength {:.2f} for these fibers'.
            format(wstop, psf_wavemax))

    if rank == 0:
        #- Print parameters
        log.info("extract:  input = {}".format(input_file))
        log.info("extract:  psf = {}".format(psf_file))
        log.info("extract:  specmin = {}".format(specmin))
        log.info("extract:  nspec = {}".format(nspec))
        log.info("extract:  wavelength = {},{},{}".format(wstart, wstop, dw))
        log.info("extract:  nwavestep = {}".format(args.nwavestep))
        log.info("extract:  regularize = {}".format(args.regularize))

    if barycentric_correction_factor != 1.:
        img.meta['HELIOCOR'] = barycentric_correction_factor

    #- Augment input image header for output
    img.meta['NSPEC'] = (nspec, 'Number of spectra')
    img.meta['WAVEMIN'] = (raw_wstart, 'First wavelength [Angstroms]')
    img.meta['WAVEMAX'] = (raw_wstop, 'Last wavelength [Angstroms]')
    img.meta['WAVESTEP'] = (raw_dw, 'Wavelength step size [Angstroms]')
    img.meta['SPECTER'] = (specter.__version__,
                           'https://github.com/desihub/specter')
    img.meta['IN_PSF'] = (io.shorten_filename(psf_file), 'Input spectral PSF')
    img.meta['IN_IMG'] = io.shorten_filename(input_file)
    depend.add_dependencies(img.meta)

    #- Check if input PSF was itself a traceshifted version of another PSF
    orig_psf = None
    if rank == 0:
        try:
            psfhdr = fits.getheader(psf_file, 'PSF')
            orig_psf = psfhdr['IN_PSF']
        except KeyError:
            #- could happen due to PSF format not having "PSF" extension,
            #- or due to PSF header not having 'IN_PSF' keyword.  Either is OK
            pass

    if comm is not None:
        orig_psf = comm.bcast(orig_psf, root=0)

    if orig_psf is not None:
        img.meta['ORIG_PSF'] = orig_psf

    #- If not using MPI, use a single call to each of these and then end this function call
    #  Otherwise, continue on to splitting things up for the different ranks
    if comm is None:
        _extract_and_save(img, psf, specmin, nspec, specmin, wave, raw_wave,
                          fibers, fibermap, args.output, args.model,
                          bundlesize, args, log)

        #- This is it if we aren't running MPI, so return
        return
    #else:
    #    # Continue to the MPI section, which could go under this else statment
    #    # But to save on indentation we'll just pass on to the rest of the function
    #    # since the alternative has already returned
    #    pass

    # Now we divide our spectra into bundles
    checkbundles = set()
    checkbundles.update(
        np.floor_divide(np.arange(specmin, specmax),
                        bundlesize * np.ones(nspec)).astype(int))
    bundles = sorted(checkbundles)
    nbundle = len(bundles)

    bspecmin = {}
    bnspec = {}

    for b in bundles:
        if specmin > b * bundlesize:
            bspecmin[b] = specmin
        else:
            bspecmin[b] = b * bundlesize
        if (b + 1) * bundlesize > specmax:
            bnspec[b] = specmax - bspecmin[b]
        else:
            bnspec[b] = bundlesize

    # Now we assign bundles to processes
    mynbundle = int(nbundle // nproc)
    myfirstbundle = 0
    leftover = nbundle % nproc
    if rank < leftover:
        mynbundle += 1
        myfirstbundle = rank * mynbundle
    else:
        myfirstbundle = ((mynbundle + 1) * leftover) + (mynbundle *
                                                        (rank - leftover))

    # get the root output file
    outpat = re.compile(r'(.*)\.fits')
    outmat = outpat.match(args.output)
    if outmat is None:
        raise RuntimeError(
            "extraction output file should have .fits extension")
    outroot = outmat.group(1)

    outdir = os.path.normpath(os.path.dirname(outroot))

    if rank == 0:
        if not os.path.isdir(outdir):
            os.makedirs(outdir)

    if comm is not None:
        comm.barrier()

    mark_preparation = time.time()
    time_total_extraction = 0.0
    time_total_write_output = 0.0
    failcount = 0

    for b in range(myfirstbundle, myfirstbundle + mynbundle):
        mark_iteration_start = time.time()
        outbundle = "{}_{:02d}.fits".format(outroot, b)
        outmodel = "{}_model_{:02d}.fits".format(outroot, b)

        log.info('extract:  Rank {} extracting {} spectra {}:{} at {}'.format(
            rank,
            os.path.basename(input_file),
            bspecmin[b],
            bspecmin[b] + bnspec[b],
            time.asctime(),
        ))
        sys.stdout.flush()

        #- The actual extraction
        try:
            mark_extraction = _extract_and_save(img, psf, bspecmin[b],
                                                bnspec[b], specmin, wave,
                                                raw_wave, fibers, fibermap,
                                                outbundle, outmodel,
                                                bundlesize, args, log)

            mark_write_output = time.time()

            time_total_extraction += mark_extraction - mark_iteration_start
            time_total_write_output += mark_write_output - mark_extraction
        except:
            # Log the error and increment the number of failures
            log.error(
                "extract:  FAILED bundle {}, spectrum range {}:{}".format(
                    b, bspecmin[b], bspecmin[b] + bnspec[b]))
            exc_type, exc_value, exc_traceback = sys.exc_info()
            lines = traceback.format_exception(exc_type, exc_value,
                                               exc_traceback)
            log.error(''.join(lines))
            failcount += 1
            sys.stdout.flush()

    if comm is not None:
        failcount = comm.allreduce(failcount)

    if failcount > 0:
        # all processes throw
        raise RuntimeError("some extraction bundles failed")

    time_merge = None
    if rank == 0:
        mark_merge_start = time.time()
        mergeopts = ['--output', args.output, '--force', '--delete']
        mergeopts.extend(
            ["{}_{:02d}.fits".format(outroot, b) for b in bundles])
        mergeargs = mergebundles.parse(mergeopts)
        mergebundles.main(mergeargs)

        if args.model is not None:
            model = None
            for b in bundles:
                outmodel = "{}_model_{:02d}.fits".format(outroot, b)
                if model is None:
                    model = fits.getdata(outmodel)
                else:
                    #- TODO: test and warn if models overlap for pixels with
                    #- non-zero values
                    model += fits.getdata(outmodel)

                os.remove(outmodel)

            fits.writeto(args.model, model)
        mark_merge_end = time.time()
        time_merge = mark_merge_end - mark_merge_start

    # Resolve difference timer data

    if type(timing) is dict:
        timing["read_input"] = mark_read_input - mark_start
        timing["preparation"] = mark_preparation - mark_read_input
        timing["total_extraction"] = time_total_extraction
        timing["total_write_output"] = time_total_write_output
        timing["merge"] = time_merge
Example #4
0
def preproc(rawimage,
            header,
            primary_header,
            bias=True,
            dark=True,
            pixflat=True,
            mask=True,
            bkgsub=False,
            nocosmic=False,
            cosmics_nsig=6,
            cosmics_cfudge=3.,
            cosmics_c2fudge=0.5,
            ccd_calibration_filename=None,
            nocrosstalk=False,
            nogain=False,
            overscan_per_row=False,
            use_overscan_row=False,
            use_savgol=None,
            nodarktrail=False,
            remove_scattered_light=False,
            psf_filename=None,
            bias_img=None,
            model_variance=False):
    '''
    preprocess image using metadata in header

    image = ((rawimage-bias-overscan)*gain)/pixflat

    Args:
        rawimage : 2D numpy array directly from raw data file
        header : dict-like metadata, e.g. from FITS header, with keywords
            CAMERA, BIASSECx, DATASECx, CCDSECx
            where x = A, B, C, D for each of the 4 amplifiers
            (also supports old naming convention 1, 2, 3, 4).
        primary_header: dict-like metadata fit keywords EXPTIME, DOSVER
            DATE-OBS is also required if bias, pixflat, or mask=True

    Optional bias, pixflat, and mask can each be:
        False: don't apply that step
        True: use default calibration data for that night
        ndarray: use that array
        filename (str or unicode): read HDU 0 and use that

    Optional overscan features:
        overscan_per_row : bool,  Subtract the overscan_col values
            row by row from the data.
        use_overscan_row : bool,  Subtract off the overscan_row
            from the data (default: False).  Requires ORSEC in
            the Header
        use_savgol : bool,  Specify whether to use Savitsky-Golay filter for
            the overscan.   (default: False).  Requires use_overscan_row=True
            to have any effect.

    Optional variance model if model_variance=True
    Optional background subtraction with median filtering if bkgsub=True

    Optional disabling of cosmic ray rejection if nocosmic=True
    Optional disabling of dark trail correction if nodarktrail=True

    Optional bias image (testing only) may be provided by bias_img=

    Optional tuning of cosmic ray rejection parameters:
        cosmics_nsig: number of sigma above background required
        cosmics_cfudge: number of sigma inconsistent with PSF required
        cosmics_c2fudge:  fudge factor applied to PSF

    Optional fit and subtraction of scattered light

    Returns Image object with member variables:
        pix : 2D preprocessed image in units of electrons per pixel
        ivar : 2D inverse variance of image
        mask : 2D mask of image (0=good)
        readnoise : 2D per-pixel readnoise of image
        meta : metadata dictionary
        TODO: define what keywords are included

    preprocessing includes the following steps:
        - bias image subtraction
        - overscan subtraction (from BIASSEC* keyword defined regions)
        - readnoise estimation (from BIASSEC* keyword defined regions)
        - gain correction (from GAIN* keywords)
        - pixel flat correction
        - cosmic ray masking
        - propagation of input known bad pixel mask
        - inverse variance estimation

    Notes:

    The bias image is subtracted before any other calculation to remove any
    non-uniformities in the overscan regions prior to calculating overscan
    levels and readnoise.

    The readnoise is an image not just one number per amp, because the pixflat
    image also affects the interpreted readnoise.

    The inverse variance is estimated from the readnoise and the image itself,
    and thus is biased.
    '''
    log = get_logger()

    header = header.copy()
    depend.setdep(header, 'DESI_SPECTRO_CALIB',
                  os.getenv('DESI_SPECTRO_CALIB'))

    for key in ['DESI_SPECTRO_REDUX', 'SPECPROD']:
        if key in os.environ:
            depend.setdep(header, key, os.environ[key])

    cfinder = None

    if ccd_calibration_filename is not False:
        cfinder = CalibFinder([header, primary_header],
                              yaml_file=ccd_calibration_filename)

    #- TODO: Check for required keywords first

    #- Subtract bias image
    camera = header['CAMERA'].lower()

    #- convert rawimage to float64 : this is the output format of read_image
    rawimage = rawimage.astype(np.float64)

    # Savgol
    if cfinder and cfinder.haskey("USE_ORSEC"):
        use_overscan_row = cfinder.value("USE_ORSEC")
    if cfinder and cfinder.haskey("SAVGOL"):
        use_savgol = cfinder.value("SAVGOL")

    # Set bias image, as desired
    if bias_img is None:
        bias = get_calibration_image(cfinder, "BIAS", bias, header)
    else:
        bias = bias_img

    #- Check if this file uses amp names 1,2,3,4 (old) or A,B,C,D (new)
    amp_ids = get_amp_ids(header)
    #- Double check that we have the necessary keywords
    missing_keywords = list()
    for prefix in ['CCDSEC', 'BIASSEC']:
        for amp in amp_ids:
            key = prefix + amp
            if not key in header:
                log.error('No {} keyword in header'.format(key))
                missing_keywords.append(key)

    if len(missing_keywords) > 0:
        raise KeyError("Missing keywords {}".format(
            ' '.join(missing_keywords)))

    #- Output arrays
    ny = 0
    nx = 0
    for amp in amp_ids:
        yy, xx = parse_sec_keyword(header['CCDSEC%s' % amp])
        ny = max(ny, yy.stop)
        nx = max(nx, xx.stop)
    image = np.zeros((ny, nx))

    readnoise = np.zeros_like(image)

    #- Load dark
    if cfinder and cfinder.haskey("DARK") and (dark is not False):

        #- Exposure time
        if cfinder and cfinder.haskey("EXPTIMEKEY"):
            exptime_key = cfinder.value("EXPTIMEKEY")
            log.info("Using exposure time keyword %s for dark normalization" %
                     exptime_key)
        else:
            exptime_key = "EXPTIME"
        exptime = primary_header[exptime_key]
        log.info(
            "Use exptime = {} sec to compute the dark current".format(exptime))

        dark_filename = cfinder.findfile("DARK")
        depend.setdep(header, 'CCD_CALIB_DARK',
                      shorten_filename(dark_filename))
        log.info(f'Using DARK model from {dark_filename}')
        # dark is multipled by exptime, or we use the non-linear dark model in the routine
        dark = read_dark(filename=dark_filename, exptime=exptime)

        if dark.shape == image.shape:
            log.info("dark is trimmed")
            trimmed_dark_in_electrons = dark
            dark_is_trimmed = True
        elif dark.shape == rawimage.shape:
            log.info("dark is not trimmed")
            trimmed_dark_in_electrons = np.zeros_like(image)
            dark_is_trimmed = False
        else:
            message = "incompatible dark shape={} when raw shape={} and preproc shape={}".format(
                dark.shape, rawimage.shape, image.shape)
            log.error(message)
            raise ValueError(message)

    else:
        dark = False

    if bias is not False:  #- it's an array
        if bias.shape == rawimage.shape:
            log.info("subtracting bias")
            rawimage = rawimage - bias
        else:
            raise ValueError('shape mismatch bias {} != rawimage {}'.format(
                bias.shape, rawimage.shape))

    #- Load mask
    mask = get_calibration_image(cfinder, "MASK", mask, header)

    if mask is False:
        mask = np.zeros(image.shape, dtype=np.int32)
    else:
        if mask.shape != image.shape:
            raise ValueError('shape mismatch mask {} != image {}'.format(
                mask.shape, image.shape))

    for amp in amp_ids:
        # Grab the sections
        ov_col = parse_sec_keyword(header['BIASSEC' + amp])
        if 'ORSEC' + amp in header.keys():
            ov_row = parse_sec_keyword(header['ORSEC' + amp])
        elif use_overscan_row:
            log.error('No ORSEC{} keyword; not using overscan_row'.format(amp))
            use_overscan_row = False

        if nogain:
            gain = 1.
        else:
            #- Initial teststand data may be missing GAIN* keywords; don't crash
            if 'GAIN' + amp in header:
                gain = header['GAIN' + amp]  #- gain = electrons / ADU
            else:
                if cfinder and cfinder.haskey('GAIN' + amp):
                    gain = float(cfinder.value('GAIN' + amp))
                    log.info('Using GAIN{}={} from calibration data'.format(
                        amp, gain))
                else:
                    gain = 1.0
                    log.warning(
                        'Missing keyword GAIN{} in header and nothing in calib data; using {}'
                        .format(amp, gain))

        #- Record what gain value was actually used
        header['GAIN' + amp] = gain

        #- Add saturation level
        if 'SATURLEV' + amp in header:
            saturlev_adu = header['SATURLEV' + amp]  # in ADU
        else:
            if cfinder and cfinder.haskey('SATURLEV' + amp):
                saturlev_adu = float(cfinder.value('SATURLEV' + amp))
                log.info('Using SATURLEV{}={} from calibration data'.format(
                    amp, saturlev_adu))
            else:
                saturlev_adu = 2**16 - 1  # 65535 is the max value in the images
                log.warning(
                    'Missing keyword SATURLEV{} in header and nothing in calib data; using {} ADU'
                    .format(amp, saturlev_adu))
        header['SATULEV' +
               amp] = (saturlev_adu,
                       "saturation or non lin. level, in ADU, inc. bias")

        # Generate the overscan images
        raw_overscan_col = rawimage[ov_col].copy()

        if use_overscan_row:
            raw_overscan_row = rawimage[ov_row].copy()
            overscan_row = np.zeros_like(raw_overscan_row)

            # Remove overscan_col from overscan_row
            raw_overscan_squared = rawimage[ov_row[0], ov_col[1]].copy()
            for row in range(raw_overscan_row.shape[0]):
                o, r = _overscan(raw_overscan_squared[row])
                overscan_row[row] = raw_overscan_row[row] - o

        # Now remove the overscan_col
        nrows = raw_overscan_col.shape[0]
        log.info("nrows in overscan=%d" % nrows)
        overscan_col = np.zeros(nrows)
        rdnoise = np.zeros(nrows)
        if (cfinder and cfinder.haskey('OVERSCAN' + amp)
                and cfinder.value("OVERSCAN" + amp).upper()
                == "PER_ROW") or overscan_per_row:
            log.info(
                "Subtracting overscan per row for amplifier %s of camera %s" %
                (amp, camera))
            for j in range(nrows):
                if np.isnan(np.sum(overscan_col[j])):
                    log.warning(
                        "NaN values in row %d of overscan of amplifier %s of camera %s"
                        % (j, amp, camera))
                    continue
                o, r = _overscan(raw_overscan_col[j])
                #log.info("%d %f %f"%(j,o,r))
                overscan_col[j] = o
                rdnoise[j] = r
        else:
            log.info(
                "Subtracting average overscan for amplifier %s of camera %s" %
                (amp, camera))
            o, r = _overscan(raw_overscan_col)
            overscan_col += o
            rdnoise += r
            if bias is not False:
                jj = parse_sec_keyword(header['DATASEC' + amp])
                o, biasnoise = _overscan(bias[jj])
                new_rdnoise = np.sqrt(rdnoise**2 + biasnoise**2)
                log.info(
                    "Master bias noise for AMP %s = %4.3f ADU, rdnoise %4.3f -> %4.3f ADU"
                    % (amp, biasnoise, np.mean(rdnoise), np.mean(new_rdnoise)))
                rdnoise = new_rdnoise
        rdnoise *= gain
        median_rdnoise = np.median(rdnoise)
        median_overscan = np.median(overscan_col)
        log.info("Median rdnoise and overscan= %f %f" %
                 (median_rdnoise, median_overscan))

        kk = parse_sec_keyword(header['CCDSEC' + amp])
        for j in range(nrows):
            readnoise[kk][j] = rdnoise[j]

        header['OVERSCN' + amp] = (median_overscan, 'ADUs (gain not applied)')
        if gain != 1:
            rdnoise_message = 'electrons (gain is applied)'
            gain_message = 'e/ADU (gain applied to image)'
        else:
            rdnoise_message = 'ADUs (gain not applied)'
            gain_message = 'gain not applied to image'
        header['OBSRDN' + amp] = (median_rdnoise, rdnoise_message)
        header['GAIN' + amp] = (gain, gain_message)

        #- Warn/error if measured readnoise is very different from expected if exists
        if 'RDNOISE' + amp in header:
            expected_readnoise = header['RDNOISE' + amp]
            if median_rdnoise < 0.5 * expected_readnoise:
                log.error(
                    'Amp {} measured readnoise {:.2f} < 0.5 * expected readnoise {:.2f}'
                    .format(amp, median_rdnoise, expected_readnoise))
            elif median_rdnoise < 0.9 * expected_readnoise:
                log.warning(
                    'Amp {} measured readnoise {:.2f} < 0.9 * expected readnoise {:.2f}'
                    .format(amp, median_rdnoise, expected_readnoise))
            elif median_rdnoise > 2.0 * expected_readnoise:
                log.error(
                    'Amp {} measured readnoise {:.2f} > 2 * expected readnoise {:.2f}'
                    .format(amp, median_rdnoise, expected_readnoise))
            elif median_rdnoise > 1.2 * expected_readnoise:
                log.warning(
                    'Amp {} measured readnoise {:.2f} > 1.2 * expected readnoise {:.2f}'
                    .format(amp, median_rdnoise, expected_readnoise))
        #else:
        #    log.warning('Expected readnoise keyword {} missing'.format('RDNOISE'+amp))

        log.info("Measured readnoise for AMP %s = %f" % (amp, median_rdnoise))

        #- subtract overscan from data region and apply gain
        jj = parse_sec_keyword(header['DATASEC' + amp])

        data = rawimage[jj].copy()
        # Subtract columns
        for k in range(nrows):
            data[k] -= overscan_col[k]

        saturlev_elec = gain * (saturlev_adu - np.mean(overscan_col))
        header['SATUELE' +
               amp] = (saturlev_elec,
                       "saturation or non lin. level, in electrons")

        # And now the rows
        if use_overscan_row:
            # Savgol?
            if use_savgol:
                log.info("Using savgol")
                collapse_oscan_row = np.zeros(overscan_row.shape[1])
                for col in range(overscan_row.shape[1]):
                    o, _ = _overscan(overscan_row[:, col])
                    collapse_oscan_row[col] = o
                oscan_row = _savgol_clipped(collapse_oscan_row, niter=0)
                oimg_row = np.outer(np.ones(data.shape[0]), oscan_row)
                data -= oimg_row
            else:
                o, r = _overscan(overscan_row)
                data -= o

        #- apply saturlev (defined in ADU), prior to multiplication by gain
        saturated = (rawimage[jj] >= saturlev_adu)
        mask[kk][saturated] |= ccdmask.SATURATED

        #- ADC to electrons
        image[kk] = data * gain

        if dark is not False:
            if not dark_is_trimmed:
                trimmed_dark_in_electrons[kk] = dark[jj] * gain

    if not nocrosstalk:
        #- apply cross-talk

        # the ccd looks like :
        # C D
        # A B
        # for cross talk, we need a symmetric 4x4 flip_matrix
        # of coordinates ABCD giving flip of both axis
        # when computing crosstalk of
        #    A   B   C   D
        #
        # A  AA  AB  AC  AD
        # B  BA  BB  BC  BD
        # C  CA  CB  CC  CD
        # D  DA  DB  DC  BB
        # orientation_matrix_defines change of orientation
        #
        fip_axis_0 = np.array([[1, 1, -1, -1], [1, 1, -1, -1], [-1, -1, 1, 1],
                               [-1, -1, 1, 1]])
        fip_axis_1 = np.array([[1, -1, 1, -1], [-1, 1, -1, 1], [1, -1, 1, -1],
                               [-1, 1, -1, 1]])

        for a1 in range(len(amp_ids)):
            amp1 = amp_ids[a1]
            ii1 = parse_sec_keyword(header['CCDSEC' + amp1])
            a1flux = image[ii1]
            #a1mask=mask[ii1]

            for a2 in range(len(amp_ids)):
                if a1 == a2:
                    continue
                amp2 = amp_ids[a2]
                if cfinder is None: continue
                if not cfinder.haskey("CROSSTALK%s%s" % (amp1, amp2)): continue
                crosstalk = cfinder.value("CROSSTALK%s%s" % (amp1, amp2))
                if crosstalk == 0.: continue
                log.info("Correct for crosstalk=%f from AMP %s into %s" %
                         (crosstalk, amp1, amp2))
                a12flux = crosstalk * a1flux.copy()
                #a12mask=a1mask.copy()
                if fip_axis_0[a1, a2] == -1:
                    a12flux = a12flux[::-1]
                    #a12mask=a12mask[::-1]
                if fip_axis_1[a1, a2] == -1:
                    a12flux = a12flux[:, ::-1]
                    #a12mask=a12mask[:,::-1]
                ii2 = parse_sec_keyword(header['CCDSEC' + amp2])
                image[ii2] -= a12flux
                # mask[ii2]  |= a12mask (not sure we really need to propagate the mask)

    #- Poisson noise variance (prior to dark subtraction and prior to pixel flat field)
    #- This is biasing, but that's what we have for now
    poisson_var = image.clip(0)

    #- subtract dark after multiplication by gain
    if dark is not False:
        log.info("subtracting dark")
        image -= trimmed_dark_in_electrons
        # measure its noise
        new_readnoise = np.zeros(readnoise.shape)
        for amp in amp_ids:
            kk = parse_sec_keyword(header['CCDSEC' + amp])
            o, darknoise = _overscan(trimmed_dark_in_electrons[kk])
            new_readnoise[kk] = np.sqrt(readnoise[kk]**2 + darknoise**2)
            log.info(
                "Master dark noise for AMP %s = %4.3f elec, rdnoise %4.3f -> %4.3f elec"
                % (amp, darknoise, np.mean(
                    readnoise[kk]), np.mean(new_readnoise[kk])))
        readnoise = new_readnoise

    #- Correct for dark trails if any
    if not nodarktrail and cfinder is not None:
        for amp in amp_ids:
            if cfinder.haskey("DARKTRAILAMP%s" % amp):
                amplitude = cfinder.value("DARKTRAILAMP%s" % amp)
                width = cfinder.value("DARKTRAILWIDTH%s" % amp)
                ii = _parse_sec_keyword(header["CCDSEC" + amp])
                log.info(
                    "Removing dark trails for amplifier %s with width=%3.1f and amplitude=%5.4f"
                    % (amp, width, amplitude))
                correct_dark_trail(image,
                                   ii,
                                   left=((amp == "B") | (amp == "D")),
                                   width=width,
                                   amplitude=amplitude)

    #- Divide by pixflat image
    pixflat = get_calibration_image(cfinder, "PIXFLAT", pixflat, header)
    if pixflat is not False:
        if pixflat.shape != image.shape:
            raise ValueError('shape mismatch pixflat {} != image {}'.format(
                pixflat.shape, image.shape))

        almost_zero = 0.001

        if np.all(pixflat > almost_zero):
            image /= pixflat
            readnoise /= pixflat
            poisson_var /= pixflat**2
        else:
            good = (pixflat > almost_zero)
            image[good] /= pixflat[good]
            readnoise[good] /= pixflat[good]
            poisson_var[good] /= pixflat[good]**2
            mask[~good] |= ccdmask.PIXFLATZERO

        lowpixflat = (0 < pixflat) & (pixflat < 0.1)
        if np.any(lowpixflat):
            mask[lowpixflat] |= ccdmask.PIXFLATLOW

    #- Inverse variance, estimated directly from the data (BEWARE: biased!)
    var = poisson_var + readnoise**2
    ivar = np.zeros(var.shape)
    ivar[var > 0] = 1.0 / var[var > 0]

    #- Ridiculously high readnoise is bad
    mask[readnoise > 100] |= ccdmask.BADREADNOISE

    if bkgsub:
        bkg = _background(image, header)
        image -= bkg

    img = Image(image,
                ivar=ivar,
                mask=mask,
                meta=header,
                readnoise=readnoise,
                camera=camera)

    #- update img.mask to mask cosmic rays
    if not nocosmic:
        cosmics.reject_cosmic_rays(img,
                                   nsig=cosmics_nsig,
                                   cfudge=cosmics_cfudge,
                                   c2fudge=cosmics_c2fudge)
        mask = img.mask

    xyset = None

    if model_variance:

        psf = None
        if psf_filename is None:
            psf_filename = cfinder.findfile("PSF")

        depend.setdep(header, 'CCD_CALIB_PSF', shorten_filename(psf_filename))
        xyset = read_xytraceset(psf_filename)

        fiberflat = None
        with_spectral_smoothing = True
        with_sky_model = True

        if with_sky_model:
            log.debug("Will use a sky model to model the spectra")
            fiberflat_filename = cfinder.findfile("FIBERFLAT")
            depend.setdep(header, 'CCD_CALIB_FIBERFLAT',
                          shorten_filename(fiberflat_filename))
            if fiberflat_filename is not None:
                fiberflat = read_fiberflat(fiberflat_filename)

        log.info("compute an image model after dark correction and pixel flat")
        nsig = 5.
        mimage = compute_image_model(
            img,
            xyset,
            fiberflat=fiberflat,
            with_spectral_smoothing=with_spectral_smoothing,
            with_sky_model=with_sky_model,
            spectral_smoothing_nsig=nsig,
            psf=psf)

        # here we bring back original image for large outliers
        # this allows to have a correct ivar for cosmic rays and bright sources
        eps = 0.1
        out = (((ivar > 0) * (image - mimage)**2 /
                (1. / (ivar + (ivar == 0)) + (0.1 * mimage)**2)) > nsig**2)
        # out &= (image>mimage) # could request this to be conservative on the variance ... but this could cause other issues
        mimage[out] = image[out]

        log.info("use image model to compute variance")
        if bkgsub:
            mimage += bkg
        if pixflat is not False:
            # undo pixflat
            mimage *= pixflat
        if dark is not False:
            mimage += dark
        poisson_var = mimage.clip(0)
        if pixflat is not False:
            if np.all(pixflat > almost_zero):
                poisson_var /= pixflat**2
            else:
                poisson_var[good] /= pixflat[good]**2
        var = poisson_var + readnoise**2
        ivar[var > 0] = 1.0 / var[var > 0]

        # regenerate img object
        img = Image(image,
                    ivar=ivar,
                    mask=mask,
                    meta=header,
                    readnoise=readnoise,
                    camera=camera)

    if remove_scattered_light:
        if xyset is None:
            if psf_filename is None:
                psf_filename = cfinder.findfile("PSF")
                depend.setdep(header, 'SCATTERED_LIGHT_PSF',
                              shorten_filename(psf_filename))
            xyset = read_xytraceset(psf_filename)
        img.pix -= model_scattered_light(img, xyset)

    #- Extend header with primary header keywords too
    addkeys(img.meta, primary_header)

    return img
Example #5
0
def get_calibration_image(cfinder, keyword, entry, header=None):
    """Reads a calibration file

    Args:
        cfinder : None or CalibFinder object
        keyword :  BIAS, MASK, or PIXFLAT
        entry : boolean or filename or image
                if entry==False return False
                if entry==True use calibration filename from calib. config and read it
                if entry==str use this for the filename
                if entry==image return input

    Options:
        header : if not None, update header['CAL...'] = calib provenance

    returns:
       2D numpy array with calibration image
    """
    log = get_logger()

    #- set the header to something so that we don't have to keep checking it
    if header is None:
        header = dict()

    calkey = 'CCD_CALIB_{}'.format(keyword.upper())
    if entry is False:
        depend.setdep(header, calkey, 'None')
        return False  # we don't want do anything

    filename = None
    if entry is True:
        # we have to find the filename
        if cfinder is None:
            log.error("no calibration data was found")
            raise ValueError("no calibration data was found")
        if cfinder.haskey(keyword):
            filename = cfinder.findfile(keyword)
            depend.setdep(header, calkey, shorten_filename(filename))
        else:
            depend.setdep(header, calkey, 'None')
            return False  # we say in the calibration data we don't need this
    elif isinstance(entry, str):
        filename = entry
        depend.setdep(header, calkey, shorten_filename(filename))
    else:
        depend.setdep(header, calkey, 'Unknown image')
        return entry  # it's expected to be an image array

    log.info("Using %s %s" % (keyword, filename))
    if keyword == "BIAS":
        return read_bias(filename=filename)
    elif keyword == "MASK":
        return read_mask(filename=filename)
    elif keyword == "PIXFLAT":
        return read_pixflat(filename=filename)
    elif keyword == "DARK":
        raise ValueError("Dark are now treated separately.")
    else:
        log.error("Don't known how to read %s in %s" % (keyword, path))
        raise ValueError("Don't known how to read %s in %s" % (keyword, path))
    return False
Example #6
0
def main(args=None, comm=None):
    if args is None:
        args = parse()
    # elif isinstance(args, (list, tuple)):
    #     args = parse(args)

    log = get_logger()

    start_mpi_connect = time.time()
    if comm is not None:
        #- Use the provided comm to determine rank and size
        rank = comm.rank
        size = comm.size
    else:
        #- Check MPI flags and determine the comm, rank, and size given the arguments
        comm, rank, size = assign_mpi(do_mpi=args.mpi,
                                      do_batch=args.batch,
                                      log=log)
    stop_mpi_connect = time.time()

    #- Start timer; only print log messages from rank 0 (others are silent)
    timer = desiutil.timer.Timer(silent=(rank > 0))

    #- Fill in timing information for steps before we had the timer created
    if args.starttime is not None:
        timer.start('startup', starttime=args.starttime)
        timer.stop('startup', stoptime=start_imports)

    timer.start('imports', starttime=start_imports)
    timer.stop('imports', stoptime=stop_imports)

    timer.start('mpi_connect', starttime=start_mpi_connect)
    timer.stop('mpi_connect', stoptime=stop_mpi_connect)

    #- Freeze IERS after parsing args so that it doesn't bother if only --help
    timer.start('freeze_iers')
    desiutil.iers.freeze_iers()
    timer.stop('freeze_iers')

    #- Preflight checks
    timer.start('preflight')
    if rank > 0:
        #- Let rank 0 fetch these, and then broadcast
        args, hdr, camhdr = None, None, None
    else:
        args, hdr, camhdr = update_args_with_headers(args)

    ## Make sure badamps is formatted properly
    if comm is not None and rank == 0 and args.badamps is not None:
        args.badamps = validate_badamps(args.badamps)

    if comm is not None:
        args = comm.bcast(args, root=0)
        hdr = comm.bcast(hdr, root=0)
        camhdr = comm.bcast(camhdr, root=0)

    known_obstype = [
        'SCIENCE', 'ARC', 'FLAT', 'ZERO', 'DARK', 'TESTARC', 'TESTFLAT',
        'PIXFLAT', 'SKY', 'TWILIGHT', 'OTHER'
    ]
    if args.obstype not in known_obstype:
        raise RuntimeError('obstype {} not in {}'.format(
            args.obstype, known_obstype))

    timer.stop('preflight')

    #-------------------------------------------------------------------------
    #- Create and submit a batch job if requested

    if args.batch:
        #exp_str = '{:08d}'.format(args.expid)
        jobdesc = args.obstype.lower()
        if args.obstype == 'SCIENCE':
            # if not doing pre-stdstar fitting or stdstar fitting and if there is
            # no flag stopping flux calibration, set job to poststdstar
            if args.noprestdstarfit and args.nostdstarfit and (
                    not args.nofluxcalib):
                jobdesc = 'poststdstar'
            # elif told not to do std or post stdstar but the flag for prestdstar isn't set,
            # then perform prestdstar
            elif (not args.noprestdstarfit
                  ) and args.nostdstarfit and args.nofluxcalib:
                jobdesc = 'prestdstar'
            #elif (not args.noprestdstarfit) and (not args.nostdstarfit) and (not args.nofluxcalib):
            #    jobdesc = 'science'
        scriptfile = create_desi_proc_batch_script(night=args.night, exp=args.expid, cameras=args.cameras,\
                                                jobdesc=jobdesc, queue=args.queue, runtime=args.runtime,\
                                                batch_opts=args.batch_opts, timingfile=args.timingfile,
                                                system_name=args.system_name)
        err = 0
        if not args.nosubmit:
            err = subprocess.call(['sbatch', scriptfile])
        sys.exit(err)

    #-------------------------------------------------------------------------
    #- Proceeding with running

    #- What are we going to do?
    if rank == 0:
        log.info('----------')
        log.info('Input {}'.format(args.input))
        log.info('Night {} expid {}'.format(args.night, args.expid))
        log.info('Obstype {}'.format(args.obstype))
        log.info('Cameras {}'.format(args.cameras))
        log.info('Output root {}'.format(desispec.io.specprod_root()))
        log.info('----------')

    #- Create output directories if needed
    if rank == 0:
        preprocdir = os.path.dirname(
            findfile('preproc', args.night, args.expid, 'b0'))
        expdir = os.path.dirname(
            findfile('frame', args.night, args.expid, 'b0'))
        os.makedirs(preprocdir, exist_ok=True)
        os.makedirs(expdir, exist_ok=True)

    #- Wait for rank 0 to make directories before proceeding
    if comm is not None:
        comm.barrier()

    #-------------------------------------------------------------------------
    #- Preproc
    #- All obstypes get preprocessed

    timer.start('fibermap')

    #- Assemble fibermap for science exposures
    fibermap = None
    fibermap_ok = None
    if rank == 0 and args.obstype == 'SCIENCE':
        fibermap = findfile('fibermap', args.night, args.expid)
        if not os.path.exists(fibermap):
            tmp = findfile('preproc', args.night, args.expid, 'b0')
            preprocdir = os.path.dirname(tmp)
            fibermap = os.path.join(preprocdir, os.path.basename(fibermap))

            log.info('Creating fibermap {}'.format(fibermap))
            cmd = 'assemble_fibermap -n {} -e {} -o {}'.format(
                args.night, args.expid, fibermap)
            if args.badamps is not None:
                cmd += ' --badamps={}'.format(args.badamps)
            runcmd(cmd, inputs=[], outputs=[fibermap])

        fibermap_ok = os.path.exists(fibermap)

        #- Some commissioning files didn't have coords* files that caused assemble_fibermap to fail
        #- these are well known failures with no other solution, so for those, just force creation
        #- of a fibermap with null coordinate information
        if not fibermap_ok and int(args.night) < 20200310:
            log.info(
                "Since night is before 20200310, trying to force fibermap creation without coords file"
            )
            cmd += ' --force'
            runcmd(cmd, inputs=[], outputs=[fibermap])
            fibermap_ok = os.path.exists(fibermap)

    #- If assemble_fibermap failed and obstype is SCIENCE, exit now
    if comm is not None:
        fibermap_ok = comm.bcast(fibermap_ok, root=0)

    if args.obstype == 'SCIENCE' and not fibermap_ok:
        sys.stdout.flush()
        if rank == 0:
            log.critical(
                'assemble_fibermap failed for science exposure; exiting now')

        sys.exit(13)

    #- Wait for rank 0 to make fibermap if needed
    if comm is not None:
        fibermap = comm.bcast(fibermap, root=0)

    timer.stop('fibermap')

    if not (args.obstype in ['SCIENCE'] and args.noprestdstarfit):
        timer.start('preproc')
        for i in range(rank, len(args.cameras), size):
            camera = args.cameras[i]
            outfile = findfile('preproc', args.night, args.expid, camera)
            outdir = os.path.dirname(outfile)
            cmd = "desi_preproc -i {} -o {} --outdir {} --cameras {}".format(
                args.input, outfile, outdir, camera)
            if args.scattered_light:
                cmd += " --scattered-light"
            if fibermap is not None:
                cmd += " --fibermap {}".format(fibermap)
            if not args.obstype in ['ARC']:  # never model variance for arcs
                if not args.no_model_pixel_variance:
                    cmd += " --model-variance"
            runcmd(cmd, inputs=[args.input], outputs=[outfile])

        timer.stop('preproc')
        if comm is not None:
            comm.barrier()

    #-------------------------------------------------------------------------
    #- Get input PSFs
    timer.start('findpsf')
    input_psf = dict()
    if rank == 0:
        for camera in args.cameras:
            if args.psf is not None:
                input_psf[camera] = args.psf
            elif args.calibnight is not None:
                # look for a psfnight psf for this calib night
                psfnightfile = findfile('psfnight', args.calibnight,
                                        args.expid, camera)
                if not os.path.isfile(psfnightfile):
                    log.error("no {}".format(psfnightfile))
                    raise IOError("no {}".format(psfnightfile))
                input_psf[camera] = psfnightfile
            else:
                # look for a psfnight psf
                psfnightfile = findfile('psfnight', args.night, args.expid,
                                        camera)
                if os.path.isfile(psfnightfile):
                    input_psf[camera] = psfnightfile
                elif args.most_recent_calib:
                    nightfile = find_most_recent(args.night,
                                                 file_type='psfnight')
                    if nightfile is None:
                        input_psf[camera] = findcalibfile(
                            [hdr, camhdr[camera]], 'PSF')
                    else:
                        input_psf[camera] = nightfile
                else:
                    input_psf[camera] = findcalibfile([hdr, camhdr[camera]],
                                                      'PSF')
            log.info("Will use input PSF : {}".format(input_psf[camera]))

    if comm is not None:
        input_psf = comm.bcast(input_psf, root=0)

    timer.stop('findpsf')

    #-------------------------------------------------------------------------
    #- Traceshift

    if ( args.obstype in ['FLAT', 'TESTFLAT', 'SKY', 'TWILIGHT']     )   or \
    ( args.obstype in ['SCIENCE'] and (not args.noprestdstarfit) ):

        timer.start('traceshift')

        if rank == 0 and args.traceshift:
            log.info('Starting traceshift at {}'.format(time.asctime()))

        for i in range(rank, len(args.cameras), size):
            camera = args.cameras[i]
            preprocfile = findfile('preproc', args.night, args.expid, camera)
            inpsf = input_psf[camera]
            outpsf = findfile('psf', args.night, args.expid, camera)
            if not os.path.isfile(outpsf):
                if args.traceshift:
                    cmd = "desi_compute_trace_shifts"
                    cmd += " -i {}".format(preprocfile)
                    cmd += " --psf {}".format(inpsf)
                    cmd += " --outpsf {}".format(outpsf)
                    cmd += " --degxx 2 --degxy 0"
                    if args.obstype in ['FLAT', 'TESTFLAT', 'TWILIGHT']:
                        cmd += " --continuum"
                    else:
                        cmd += " --degyx 2 --degyy 0"
                    if args.obstype in ['SCIENCE', 'SKY']:
                        cmd += ' --sky'
                else:
                    cmd = "ln -s {} {}".format(inpsf, outpsf)
                runcmd(cmd, inputs=[preprocfile, inpsf], outputs=[outpsf])
            else:
                log.info("PSF {} exists".format(outpsf))

        timer.stop('traceshift')
        if comm is not None:
            comm.barrier()

    #-------------------------------------------------------------------------
    #- PSF
    #- MPI parallelize this step

    if args.obstype in ['ARC', 'TESTARC']:

        timer.start('arc_traceshift')

        if rank == 0:
            log.info('Starting traceshift before specex PSF fit at {}'.format(
                time.asctime()))

        for i in range(rank, len(args.cameras), size):
            camera = args.cameras[i]
            preprocfile = findfile('preproc', args.night, args.expid, camera)
            inpsf = input_psf[camera]
            outpsf = findfile('psf', args.night, args.expid, camera)
            outpsf = replace_prefix(outpsf, "psf", "shifted-input-psf")
            if not os.path.isfile(outpsf):
                cmd = "desi_compute_trace_shifts"
                cmd += " -i {}".format(preprocfile)
                cmd += " --psf {}".format(inpsf)
                cmd += " --outpsf {}".format(outpsf)
                cmd += " --degxx 0 --degxy 0 --degyx 0 --degyy 0"
                cmd += ' --arc-lamps'
                runcmd(cmd, inputs=[preprocfile, inpsf], outputs=[outpsf])
            else:
                log.info("PSF {} exists".format(outpsf))

        timer.stop('arc_traceshift')
        if comm is not None:
            comm.barrier()

        timer.start('psf')

        if rank == 0:
            log.info('Starting specex PSF fitting at {}'.format(
                time.asctime()))

        if rank > 0:
            cmds = inputs = outputs = None
        else:
            cmds = dict()
            inputs = dict()
            outputs = dict()
            for camera in args.cameras:
                preprocfile = findfile('preproc', args.night, args.expid,
                                       camera)
                tmpname = findfile('psf', args.night, args.expid, camera)
                inpsf = replace_prefix(tmpname, "psf", "shifted-input-psf")
                outpsf = replace_prefix(tmpname, "psf", "fit-psf")

                log.info("now run specex psf fit")

                cmd = 'desi_compute_psf'
                cmd += ' --input-image {}'.format(preprocfile)
                cmd += ' --input-psf {}'.format(inpsf)
                cmd += ' --output-psf {}'.format(outpsf)

                # look for fiber blacklist
                cfinder = CalibFinder([hdr, camhdr[camera]])
                blacklistkey = "FIBERBLACKLIST"
                if not cfinder.haskey(blacklistkey) and cfinder.haskey(
                        "BROKENFIBERS"):
                    log.warning(
                        "BROKENFIBERS yaml keyword deprecated, please use FIBERBLACKLIST"
                    )
                    blacklistkey = "BROKENFIBERS"

                if cfinder.haskey(blacklistkey):
                    blacklist = cfinder.value(blacklistkey)
                    cmd += ' --broken-fibers {}'.format(blacklist)
                    if rank == 0:
                        log.warning('broken fibers: {}'.format(blacklist))

                if not os.path.exists(outpsf):
                    cmds[camera] = cmd
                    inputs[camera] = [preprocfile, inpsf]
                    outputs[camera] = [
                        outpsf,
                    ]

        if comm is not None:
            cmds = comm.bcast(cmds, root=0)
            inputs = comm.bcast(inputs, root=0)
            outputs = comm.bcast(outputs, root=0)
            #- split communicator by 20 (number of bundles)
            group_size = 20
            if (rank == 0) and (size % group_size != 0):
                log.warning(
                    'MPI size={} should be evenly divisible by {}'.format(
                        size, group_size))

            group = rank // group_size
            num_groups = (size + group_size - 1) // group_size
            comm_group = comm.Split(color=group)

            if rank == 0:
                log.info(
                    f'Fitting PSFs with {num_groups} sub-communicators of size {group_size}'
                )

            for i in range(group, len(args.cameras), num_groups):
                camera = args.cameras[i]
                if camera in cmds:
                    cmdargs = cmds[camera].split()[1:]
                    cmdargs = desispec.scripts.specex.parse(cmdargs)
                    if comm_group.rank == 0:
                        print('RUNNING: {}'.format(cmds[camera]))
                        t0 = time.time()
                        timestamp = time.asctime()
                        log.info(
                            f'MPI group {group} ranks {rank}-{rank+group_size-1} fitting PSF for {camera} at {timestamp}'
                        )
                    try:
                        desispec.scripts.specex.main(cmdargs, comm=comm_group)
                    except Exception as e:
                        if comm_group.rank == 0:
                            log.error(
                                f'FAILED: MPI group {group} ranks {rank}-{rank+group_size-1} camera {camera}'
                            )
                            log.error('FAILED: {}'.format(cmds[camera]))
                            log.error(e)

                    if comm_group.rank == 0:
                        specex_time = time.time() - t0
                        log.info(
                            f'specex fit for {camera} took {specex_time:.1f} seconds'
                        )

            comm.barrier()

        else:
            log.warning(
                'fitting PSFs without MPI parallelism; this will be SLOW')
            for camera in args.cameras:
                if camera in cmds:
                    runcmd(cmds[camera],
                           inputs=inputs[camera],
                           outputs=outputs[camera])

        if comm is not None:
            comm.barrier()

        # loop on all cameras and interpolate bad fibers
        for camera in args.cameras[rank::size]:
            t0 = time.time()
            log.info(f'Rank {rank} interpolating {camera} PSF over bad fibers')
            # look for fiber blacklist
            cfinder = CalibFinder([hdr, camhdr[camera]])
            blacklistkey = "FIBERBLACKLIST"
            if not cfinder.haskey(blacklistkey) and cfinder.haskey(
                    "BROKENFIBERS"):
                log.warning(
                    "BROKENFIBERS yaml keyword deprecated, please use FIBERBLACKLIST"
                )
                blacklistkey = "BROKENFIBERS"

            if cfinder.haskey(blacklistkey):
                fiberblacklist = cfinder.value(blacklistkey)
                tmpname = findfile('psf', args.night, args.expid, camera)
                inpsf = replace_prefix(tmpname, "psf", "fit-psf")
                outpsf = replace_prefix(tmpname, "psf",
                                        "fit-psf-fixed-blacklisted")
                if os.path.isfile(inpsf) and not os.path.isfile(outpsf):
                    cmd = 'desi_interpolate_fiber_psf'
                    cmd += ' --infile {}'.format(inpsf)
                    cmd += ' --outfile {}'.format(outpsf)
                    cmd += ' --fibers {}'.format(fiberblacklist)
                    log.info(
                        'For camera {} interpolating PSF for broken fibers: {}'
                        .format(camera, fiberblacklist))
                    runcmd(cmd, inputs=[inpsf], outputs=[outpsf])
                    if os.path.isfile(outpsf):
                        os.rename(
                            inpsf,
                            inpsf.replace("fit-psf",
                                          "fit-psf-before-blacklisted-fix"))
                        subprocess.call('cp {} {}'.format(outpsf, inpsf),
                                        shell=True)

            dt = time.time() - t0
            log.info(
                f'Rank {rank} {camera} PSF interpolation took {dt:.1f} sec')

        timer.stop('psf')

    #-------------------------------------------------------------------------
    #- Merge PSF of night if applicable

    #if args.obstype in ['ARC']:
    if False:
        if rank == 0:
            for camera in args.cameras:
                psfnightfile = findfile('psfnight', args.night, args.expid,
                                        camera)
                if not os.path.isfile(
                        psfnightfile
                ):  # we still don't have a psf night, see if we can compute it ...
                    psfs = glob.glob(
                        findfile('psf', args.night, args.expid,
                                 camera).replace("psf", "fit-psf").replace(
                                     str(args.expid), "*"))
                    log.info(
                        "Number of PSF for night={} camera={} = {}".format(
                            args.night, camera, len(psfs)))
                    if len(psfs) > 4:  # lets do it!
                        log.info("Computing psfnight ...")
                        dirname = os.path.dirname(psfnightfile)
                        if not os.path.isdir(dirname):
                            os.makedirs(dirname)
                        desispec.scripts.specex.mean_psf(psfs, psfnightfile)
                if os.path.isfile(psfnightfile):  # now use this one
                    input_psf[camera] = psfnightfile

    #-------------------------------------------------------------------------
    #- Extract
    #- This is MPI parallel so handle a bit differently

    # maybe add ARC and TESTARC too
    if ( args.obstype in ['FLAT', 'TESTFLAT', 'SKY', 'TWILIGHT']     )   or \
    ( args.obstype in ['SCIENCE'] and (not args.noprestdstarfit) ):

        timer.start('extract')
        if rank == 0:
            log.info('Starting extractions at {}'.format(time.asctime()))

        if rank > 0:
            cmds = inputs = outputs = None
        else:
            cmds = dict()
            inputs = dict()
            outputs = dict()
            for camera in args.cameras:
                cmd = 'desi_extract_spectra'

                #- Based on data from SM1-SM8, looking at central and edge fibers
                #- with in mind overlapping arc lamps lines
                if camera.startswith('b'):
                    cmd += ' -w 3600.0,5800.0,0.8'
                elif camera.startswith('r'):
                    cmd += ' -w 5760.0,7620.0,0.8'
                elif camera.startswith('z'):
                    cmd += ' -w 7520.0,9824.0,0.8'

                preprocfile = findfile('preproc', args.night, args.expid,
                                       camera)
                psffile = findfile('psf', args.night, args.expid, camera)
                framefile = findfile('frame', args.night, args.expid, camera)
                cmd += ' -i {}'.format(preprocfile)
                cmd += ' -p {}'.format(psffile)
                cmd += ' -o {}'.format(framefile)
                cmd += ' --psferr 0.1'

                if args.obstype == 'SCIENCE' or args.obstype == 'SKY':
                    if rank == 0:
                        log.info('Include barycentric correction')
                    cmd += ' --barycentric-correction'

                if not os.path.exists(framefile):
                    cmds[camera] = cmd
                    inputs[camera] = [preprocfile, psffile]
                    outputs[camera] = [
                        framefile,
                    ]

        #- TODO: refactor/combine this with PSF comm splitting logic
        if comm is not None:
            cmds = comm.bcast(cmds, root=0)
            inputs = comm.bcast(inputs, root=0)
            outputs = comm.bcast(outputs, root=0)

            #- split communicator by 20 (number of bundles)
            extract_size = 20
            if (rank == 0) and (size % extract_size != 0):
                log.warning(
                    'MPI size={} should be evenly divisible by {}'.format(
                        size, extract_size))

            extract_group = rank // extract_size
            num_extract_groups = (size + extract_size - 1) // extract_size
            comm_extract = comm.Split(color=extract_group)

            for i in range(extract_group, len(args.cameras),
                           num_extract_groups):
                camera = args.cameras[i]
                if camera in cmds:
                    cmdargs = cmds[camera].split()[1:]
                    extract_args = desispec.scripts.extract.parse(cmdargs)
                    if comm_extract.rank == 0:
                        print('RUNNING: {}'.format(cmds[camera]))

                    desispec.scripts.extract.main_mpi(extract_args,
                                                      comm=comm_extract)

            comm.barrier()

        else:
            log.warning(
                'running extractions without MPI parallelism; this will be SLOW'
            )
            for camera in args.cameras:
                if camera in cmds:
                    runcmd(cmds[camera],
                           inputs=inputs[camera],
                           outputs=outputs[camera])

        timer.stop('extract')
        if comm is not None:
            comm.barrier()

    #-------------------------------------------------------------------------
    #- Fiberflat

    if args.obstype in ['FLAT', 'TESTFLAT']:
        timer.start('fiberflat')
        if rank == 0:
            log.info('Starting fiberflats at {}'.format(time.asctime()))

        for i in range(rank, len(args.cameras), size):
            camera = args.cameras[i]
            framefile = findfile('frame', args.night, args.expid, camera)
            fiberflatfile = findfile('fiberflat', args.night, args.expid,
                                     camera)
            cmd = "desi_compute_fiberflat"
            cmd += " -i {}".format(framefile)
            cmd += " -o {}".format(fiberflatfile)
            runcmd(cmd, inputs=[
                framefile,
            ], outputs=[
                fiberflatfile,
            ])

        timer.stop('fiberflat')
        if comm is not None:
            comm.barrier()

    #-------------------------------------------------------------------------
    #- Average and auto-calib fiberflats of night if applicable

    #if args.obstype in ['FLAT']:
    if False:
        if rank == 0:
            fiberflatnightfile = findfile('fiberflatnight', args.night,
                                          args.expid, args.cameras[0])
            fiberflatdirname = os.path.dirname(fiberflatnightfile)
            if not os.path.isfile(fiberflatnightfile) and len(
                    args.cameras
            ) >= 6:  # we still don't have them, see if we can compute them, but need at least 2 spectros ...
                flats = glob.glob(
                    findfile('fiberflat', args.night, args.expid,
                             "b0").replace(str(args.expid),
                                           "*").replace("b0", "*"))
                log.info("Number of fiberflat for night {} = {}".format(
                    args.night, len(flats)))
                if len(flats) >= 3 * 4 * len(
                        args.cameras
                ):  # lets do it! (3 exposures x 4 lamps x N cameras)
                    log.info(
                        "Computing fiberflatnight per lamp and camera ...")
                    tmpdir = os.path.join(fiberflatdirname, "tmp")
                    if not os.path.isdir(tmpdir):
                        os.makedirs(tmpdir)

                    log.info(
                        "First average measurements per camera and per lamp")
                    average_flats = dict()
                    for camera in args.cameras:
                        # list of flats for this camera
                        flats_for_this_camera = []
                        for flat in flats:
                            if flat.find(camera) >= 0:
                                flats_for_this_camera.append(flat)
                        #log.info("For camera {} , flats = {}".format(camera,flats_for_this_camera))
                        #sys.exit(12)

                        # average per lamp (and camera)
                        average_flats[camera] = list()
                        for lampbox in range(4):
                            ofile = os.path.join(
                                tmpdir,
                                "fiberflatnight-camera-{}-lamp-{}.fits".format(
                                    camera, lampbox))
                            if not os.path.isfile(ofile):
                                log.info(
                                    "Average flat for camera {} and lamp box #{}"
                                    .format(camera, lampbox))
                                pg = "CALIB DESI-CALIB-0{} LEDs only".format(
                                    lampbox)

                                cmd = "desi_average_fiberflat --program '{}' --outfile {} -i ".format(
                                    pg, ofile)
                                for flat in flats_for_this_camera:
                                    cmd += " {} ".format(flat)
                                runcmd(cmd,
                                       inputs=flats_for_this_camera,
                                       outputs=[
                                           ofile,
                                       ])
                                if os.path.isfile(ofile):
                                    average_flats[camera].append(ofile)
                            else:
                                log.info("Will use existing {}".format(ofile))
                                average_flats[camera].append(ofile)

                    log.info(
                        "Auto-calibration across lamps and spectro  per camera arm (b,r,z)"
                    )
                    for camera_arm in ["b", "r", "z"]:
                        cameras_for_this_arm = []
                        flats_for_this_arm = []
                        for camera in args.cameras:
                            if camera[0].lower() == camera_arm:
                                cameras_for_this_arm.append(camera)
                                if camera in average_flats:
                                    for flat in average_flats[camera]:
                                        flats_for_this_arm.append(flat)
                        cmd = "desi_autocalib_fiberflat --night {} --arm {} -i ".format(
                            args.night, camera_arm)
                        for flat in flats_for_this_arm:
                            cmd += " {} ".format(flat)
                        runcmd(cmd, inputs=flats_for_this_arm, outputs=[])
                    log.info("Done with fiber flats per night")

        if comm is not None:
            comm.barrier()

    #-------------------------------------------------------------------------
    #- Get input fiberflat
    if args.obstype in ['SCIENCE', 'SKY'] and (not args.nofiberflat):
        timer.start('find_fiberflat')
        input_fiberflat = dict()
        if rank == 0:
            for camera in args.cameras:
                if args.fiberflat is not None:
                    input_fiberflat[camera] = args.fiberflat
                elif args.calibnight is not None:
                    # look for a fiberflatnight for this calib night
                    fiberflatnightfile = findfile('fiberflatnight',
                                                  args.calibnight, args.expid,
                                                  camera)
                    if not os.path.isfile(fiberflatnightfile):
                        log.error("no {}".format(fiberflatnightfile))
                        raise IOError("no {}".format(fiberflatnightfile))
                    input_fiberflat[camera] = fiberflatnightfile
                else:
                    # look for a fiberflatnight fiberflat
                    fiberflatnightfile = findfile('fiberflatnight', args.night,
                                                  args.expid, camera)
                    if os.path.isfile(fiberflatnightfile):
                        input_fiberflat[camera] = fiberflatnightfile
                    elif args.most_recent_calib:
                        nightfile = find_most_recent(
                            args.night, file_type='fiberflatnight')
                        if nightfile is None:
                            input_fiberflat[camera] = findcalibfile(
                                [hdr, camhdr[camera]], 'FIBERFLAT')
                        else:
                            input_fiberflat[camera] = nightfile
                    else:
                        input_fiberflat[camera] = findcalibfile(
                            [hdr, camhdr[camera]], 'FIBERFLAT')
                log.info("Will use input FIBERFLAT: {}".format(
                    input_fiberflat[camera]))

        if comm is not None:
            input_fiberflat = comm.bcast(input_fiberflat, root=0)

        timer.stop('find_fiberflat')

    #-------------------------------------------------------------------------
    #- Apply fiberflat and write fframe file

    if args.obstype in ['SCIENCE', 'SKY'] and args.fframe and \
    ( not args.nofiberflat ) and (not args.noprestdstarfit):
        timer.start('apply_fiberflat')
        if rank == 0:
            log.info('Applying fiberflat at {}'.format(time.asctime()))

        for i in range(rank, len(args.cameras), size):
            camera = args.cameras[i]
            fframefile = findfile('fframe', args.night, args.expid, camera)
            if not os.path.exists(fframefile):
                framefile = findfile('frame', args.night, args.expid, camera)
                fr = desispec.io.read_frame(framefile)
                flatfilename = input_fiberflat[camera]
                if flatfilename is not None:
                    ff = desispec.io.read_fiberflat(flatfilename)
                    fr.meta['FIBERFLT'] = desispec.io.shorten_filename(
                        flatfilename)
                    apply_fiberflat(fr, ff)

                    fframefile = findfile('fframe', args.night, args.expid,
                                          camera)
                    desispec.io.write_frame(fframefile, fr)
                else:
                    log.warning(
                        "Missing fiberflat for camera {}".format(camera))

        timer.stop('apply_fiberflat')
        if comm is not None:
            comm.barrier()

    #-------------------------------------------------------------------------
    #- Select random sky fibers (inplace update of frame file)
    #- TODO: move this to a function somewhere
    #- TODO: this assigns different sky fibers to each frame of same spectrograph

    if (args.obstype in [
            'SKY', 'SCIENCE'
    ]) and (not args.noskysub) and (not args.noprestdstarfit):
        timer.start('picksky')
        if rank == 0:
            log.info('Picking sky fibers at {}'.format(time.asctime()))

        for i in range(rank, len(args.cameras), size):
            camera = args.cameras[i]
            framefile = findfile('frame', args.night, args.expid, camera)
            orig_frame = desispec.io.read_frame(framefile)

            #- Make a copy so that we can apply fiberflat
            fr = deepcopy(orig_frame)

            if np.any(fr.fibermap['OBJTYPE'] == 'SKY'):
                log.info('{} sky fibers already set; skipping'.format(
                    os.path.basename(framefile)))
                continue

            #- Apply fiberflat then select random fibers below a flux cut
            flatfilename = input_fiberflat[camera]
            if flatfilename is None:
                log.error("No fiberflat for {}".format(camera))
                continue
            ff = desispec.io.read_fiberflat(flatfilename)
            apply_fiberflat(fr, ff)
            sumflux = np.sum(fr.flux, axis=1)
            fluxcut = np.percentile(sumflux, 30)
            iisky = np.where(sumflux < fluxcut)[0]
            iisky = np.random.choice(iisky, size=100, replace=False)

            #- Update fibermap or original frame and write out
            orig_frame.fibermap['OBJTYPE'][iisky] = 'SKY'
            orig_frame.fibermap['DESI_TARGET'][iisky] |= desi_mask.SKY

            desispec.io.write_frame(framefile, orig_frame)

        timer.stop('picksky')
        if comm is not None:
            comm.barrier()

    #-------------------------------------------------------------------------
    #- Sky subtraction
    if args.obstype in [
            'SCIENCE', 'SKY'
    ] and (not args.noskysub) and (not args.noprestdstarfit):
        timer.start('skysub')
        if rank == 0:
            log.info('Starting sky subtraction at {}'.format(time.asctime()))

        for i in range(rank, len(args.cameras), size):
            camera = args.cameras[i]
            framefile = findfile('frame', args.night, args.expid, camera)
            hdr = fitsio.read_header(framefile, 'FLUX')
            fiberflatfile = input_fiberflat[camera]
            if fiberflatfile is None:
                log.error("No fiberflat for {}".format(camera))
                continue
            skyfile = findfile('sky', args.night, args.expid, camera)

            cmd = "desi_compute_sky"
            cmd += " -i {}".format(framefile)
            cmd += " --fiberflat {}".format(fiberflatfile)
            cmd += " --o {}".format(skyfile)
            if args.no_extra_variance:
                cmd += " --no-extra-variance"
            if not args.no_sky_wavelength_adjustment:
                cmd += " --adjust-wavelength"
            if not args.no_sky_lsf_adjustment: cmd += " --adjust-lsf"

            runcmd(cmd, inputs=[framefile, fiberflatfile], outputs=[
                skyfile,
            ])

            #- sframe = flatfielded sky-subtracted but not flux calibrated frame
            #- Note: this re-reads and re-does steps previously done for picking
            #- sky fibers; desi_proc is about human efficiency,
            #- not I/O or CPU efficiency...
            sframefile = desispec.io.findfile('sframe', args.night, args.expid,
                                              camera)
            if not os.path.exists(sframefile):
                frame = desispec.io.read_frame(framefile)
                fiberflat = desispec.io.read_fiberflat(fiberflatfile)
                sky = desispec.io.read_sky(skyfile)
                apply_fiberflat(frame, fiberflat)
                subtract_sky(frame, sky, apply_throughput_correction=True)
                frame.meta['IN_SKY'] = shorten_filename(skyfile)
                frame.meta['FIBERFLT'] = shorten_filename(fiberflatfile)
                desispec.io.write_frame(sframefile, frame)

        timer.stop('skysub')
        if comm is not None:
            comm.barrier()

    #-------------------------------------------------------------------------
    #- Standard Star Fitting

    if args.obstype in ['SCIENCE',] and \
            (not args.noskysub ) and \
            (not args.nostdstarfit) :

        timer.start('stdstarfit')
        if rank == 0:
            log.info('Starting flux calibration at {}'.format(time.asctime()))

        #- Group inputs by spectrograph
        framefiles = dict()
        skyfiles = dict()
        fiberflatfiles = dict()
        night, expid = args.night, args.expid  #- shorter
        for camera in args.cameras:
            sp = int(camera[1])
            if sp not in framefiles:
                framefiles[sp] = list()
                skyfiles[sp] = list()
                fiberflatfiles[sp] = list()

            framefiles[sp].append(findfile('frame', night, expid, camera))
            skyfiles[sp].append(findfile('sky', night, expid, camera))
            fiberflatfiles[sp].append(input_fiberflat[camera])

        #- Hardcoded stdstar model version
        starmodels = os.path.join(os.getenv('DESI_BASIS_TEMPLATES'),
                                  'stdstar_templates_v2.2.fits')

        #- Fit stdstars per spectrograph (not per-camera)
        spectro_nums = sorted(framefiles.keys())
        ## for sp in spectro_nums[rank::size]:
        for i in range(rank, len(spectro_nums), size):
            sp = spectro_nums[i]

            stdfile = findfile('stdstars', night, expid, spectrograph=sp)
            cmd = "desi_fit_stdstars"
            cmd += " --frames {}".format(' '.join(framefiles[sp]))
            cmd += " --skymodels {}".format(' '.join(skyfiles[sp]))
            cmd += " --fiberflats {}".format(' '.join(fiberflatfiles[sp]))
            cmd += " --starmodels {}".format(starmodels)
            cmd += " --outfile {}".format(stdfile)
            cmd += " --delta-color 0.1"
            if args.maxstdstars is not None:
                cmd += " --maxstdstars {}".format(args.maxstdstars)

            inputs = framefiles[sp] + skyfiles[sp] + fiberflatfiles[sp]
            runcmd(cmd, inputs=inputs, outputs=[stdfile])

        timer.stop('stdstarfit')
        if comm is not None:
            comm.barrier()

    # -------------------------------------------------------------------------
    # - Flux calibration

    if args.obstype in ['SCIENCE'] and \
                (not args.noskysub) and \
                (not args.nofluxcalib):
        timer.start('fluxcalib')

        night, expid = args.night, args.expid  #- shorter
        #- Compute flux calibration vectors per camera
        for camera in args.cameras[rank::size]:
            framefile = findfile('frame', night, expid, camera)
            skyfile = findfile('sky', night, expid, camera)
            spectrograph = int(camera[1])
            stdfile = findfile('stdstars',
                               night,
                               expid,
                               spectrograph=spectrograph)
            calibfile = findfile('fluxcalib', night, expid, camera)

            fiberflatfile = input_fiberflat[camera]

            cmd = "desi_compute_fluxcalibration"
            cmd += " --infile {}".format(framefile)
            cmd += " --sky {}".format(skyfile)
            cmd += " --fiberflat {}".format(fiberflatfile)
            cmd += " --models {}".format(stdfile)
            cmd += " --outfile {}".format(calibfile)
            cmd += " --delta-color-cut 0.1"

            inputs = [framefile, skyfile, fiberflatfile, stdfile]
            runcmd(cmd, inputs=inputs, outputs=[
                calibfile,
            ])

        timer.stop('fluxcalib')
        if comm is not None:
            comm.barrier()

    #-------------------------------------------------------------------------
    #- Applying flux calibration

    if args.obstype in [
            'SCIENCE',
    ] and (not args.noskysub) and (not args.nofluxcalib):

        night, expid = args.night, args.expid  #- shorter

        timer.start('applycalib')
        if rank == 0:
            log.info('Starting cframe file creation at {}'.format(
                time.asctime()))

        for camera in args.cameras[rank::size]:
            framefile = findfile('frame', night, expid, camera)
            skyfile = findfile('sky', night, expid, camera)
            spectrograph = int(camera[1])
            stdfile = findfile('stdstars',
                               night,
                               expid,
                               spectrograph=spectrograph)
            calibfile = findfile('fluxcalib', night, expid, camera)
            cframefile = findfile('cframe', night, expid, camera)

            fiberflatfile = input_fiberflat[camera]

            cmd = "desi_process_exposure"
            cmd += " --infile {}".format(framefile)
            cmd += " --fiberflat {}".format(fiberflatfile)
            cmd += " --sky {}".format(skyfile)
            cmd += " --calib {}".format(calibfile)
            cmd += " --outfile {}".format(cframefile)
            cmd += " --cosmics-nsig 6"
            if args.no_xtalk:
                cmd += " --no-xtalk"

            inputs = [framefile, fiberflatfile, skyfile, calibfile]
            runcmd(cmd, inputs=inputs, outputs=[
                cframefile,
            ])

        if comm is not None:
            comm.barrier()

        timer.stop('applycalib')

    #-------------------------------------------------------------------------
    #- Wrap up

    # if rank == 0:
    #     report = timer.report()
    #     log.info('Rank 0 timing report:\n' + report)

    if comm is not None:
        timers = comm.gather(timer, root=0)
    else:
        timers = [
            timer,
        ]

    if rank == 0:
        stats = desiutil.timer.compute_stats(timers)
        log.info('Timing summary statistics:\n' + json.dumps(stats, indent=2))

        if args.timingfile:
            if os.path.exists(args.timingfile):
                with open(args.timingfile) as fx:
                    previous_stats = json.load(fx)

                #- augment previous_stats with new entries, but don't overwrite old
                for name in stats:
                    if name not in previous_stats:
                        previous_stats[name] = stats[name]

                stats = previous_stats

            tmpfile = args.timingfile + '.tmp'
            with open(tmpfile, 'w') as fx:
                json.dump(stats, fx, indent=2)
            os.rename(tmpfile, args.timingfile)

    if rank == 0:
        log.info('All done at {}'.format(time.asctime()))
Example #7
0
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)

    # Set fibermask flagged spectra to have 0 flux and variance
    frame = get_fiberbitmasked_frame(frame,
                                     bitmask='flux',
                                     ivar_framemask=True)

    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))
        good = (model_metadata["CHI2DOF"] < args.chi2cut)
        bad = ~good
        ok &= good
        if np.any(bad):
            log.info(" discard {} stars with CHI2DOF= {}".format(
                np.sum(bad), list(model_metadata["CHI2DOF"][bad])))

    legacy_filters = ('G-R', 'R-Z')
    gaia_filters = ('GAIA-BP-RP', 'GAIA-G-RP')
    model_column_list = model_metadata.columns.names
    if args.color is None:
        if 'MODEL_G-R' in model_column_list:
            color = 'G-R'
        elif 'MODEL_GAIA-BP-RP' in model_column_list:
            log.info('Using Gaia filters')
            color = 'GAIA-BP-RP'
        else:
            log.error(
                "Can't find either G-R or BP-RP color in the model file.")
            sys.exit(15)
    else:
        if args.color not in legacy_filters and args.color not in gaia_filters:
            log.error(
                'Color name {} is not allowed, must be one of {} {}'.format(
                    args.color, legacy_filters, gaia_filters))
            sys.exit(14)
        color = args.color
        if color not in model_column_list:
            # This should't happen
            log.error(
                'The color {} was not computed in the models'.format(color))
            sys.exit(16)

    if args.delta_color_cut > 0:
        log.info("apply cut |delta color|<{}".format(args.delta_color_cut))
        good = (np.abs(model_metadata["MODEL_" + color] -
                       model_metadata["DATA_" + color]) < args.delta_color_cut)
        bad = ok & (~good)
        ok &= good
        if np.any(bad):
            vals = model_metadata["MODEL_" +
                                  color][bad] - model_metadata["DATA_" +
                                                               color][bad]
            log.info(" discard {} stars with dcolor= {}".format(
                np.sum(bad), list(vals)))

    if args.min_color is not None:
        log.info("apply cut DATA_{}>{}".format(color, args.min_color))
        good = (model_metadata["DATA_{}".format(color)] > args.min_color)
        bad = ok & (~good)
        ok &= good
        if np.any(bad):
            vals = model_metadata["DATA_{}".format(color)][bad]
            log.info(" discard {} stars with {}= {}".format(
                np.sum(bad), color, list(vals)))

    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))
        good = (model_metadata["CHI2DOF"] <= maxchi2)
        bad = ok & (~good)
        ok &= good
        if np.any(bad):
            log.info(" discard {} stars with CHI2DOF={}".format(
                np.sum(bad), list(model_metadata["CHI2DOF"][bad])))

    ok = np.where(ok)[0]
    if ok.size == 0:
        log.error("selection 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,
        exposure_seeing_fwhm=args.seeing_fwhm)

    # QA
    if (args.qafile is not None):

        from desispec.io import write_qa_frame
        from desispec.io.qa import load_qa_frame
        from desispec.qa import qa_plots

        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()
        qaframe.run_qa('FLUXCALIB', (frame, fluxcalib))
        # Write
        if args.qafile is not None:
            write_qa_frame(args.qafile, qaframe)
            log.info("successfully wrote {:s}".format(args.qafile))
        # Figure(s)
        if args.qafig is not None:
            qa_plots.frame_fluxcalib(args.qafig, qaframe, frame, fluxcalib)

    # record inputs
    frame.meta['IN_FRAME'] = shorten_filename(args.infile)
    frame.meta['IN_SKY'] = shorten_filename(args.sky)
    frame.meta['FIBERFLT'] = shorten_filename(args.fiberflat)
    frame.meta['STDMODEL'] = shorten_filename(args.models)

    # write result
    write_flux_calibration(args.outfile, fluxcalib, header=frame.meta)

    log.info("successfully wrote %s" % args.outfile)
Example #8
0
def main(args):
    log = get_logger()

    if (args.fiberflat is None) and (args.sky is None) and (args.calib is
                                                            None):
        log.critical('no --fiberflat, --sky, or --calib; nothing to do ?!?')
        sys.exit(12)

    if (not args.no_tsnr) and (args.calib is None):
        log.critical(
            'need --fiberflat --sky and --calib to compute template SNR')
        sys.exit(12)

    frame = read_frame(args.infile)

    if not args.no_tsnr:
        # tsnr alpha calc. requires uncalibrated + no substraction rame.
        uncalibrated_frame = copy.deepcopy(frame)

    #- Raw scores already added in extraction, but just in case they weren't
    #- it is harmless to rerun to make sure we have them.
    compute_and_append_frame_scores(frame, suffix="RAW")

    if args.cosmics_nsig > 0 and args.sky == None:  # Reject cosmics (otherwise do it after sky subtraction)
        log.info("cosmics ray 1D rejection")
        reject_cosmic_rays_1d(frame, args.cosmics_nsig)

    if args.fiberflat != None:
        log.info("apply fiberflat")
        # read fiberflat
        fiberflat = read_fiberflat(args.fiberflat)

        # apply fiberflat to all fibers
        apply_fiberflat(frame, fiberflat)
        compute_and_append_frame_scores(frame, suffix="FFLAT")
    else:
        fiberflat = None

    if args.no_xtalk:
        zero_ivar = (not args.no_zero_ivar)
    else:
        zero_ivar = False

    if args.sky != None:

        # read sky
        skymodel = read_sky(args.sky)

        if args.cosmics_nsig > 0:

            # use a copy the frame (not elegant but robust)
            copied_frame = copy.deepcopy(frame)

            # first subtract sky without throughput correction
            subtract_sky(copied_frame,
                         skymodel,
                         apply_throughput_correction=False,
                         zero_ivar=zero_ivar)

            # then find cosmics
            log.info("cosmics ray 1D rejection after sky subtraction")
            reject_cosmic_rays_1d(copied_frame, args.cosmics_nsig)

            # copy mask
            frame.mask = copied_frame.mask

            # and (re-)subtract sky, but just the correction term
            subtract_sky(frame,
                         skymodel,
                         apply_throughput_correction=(
                             not args.no_sky_throughput_correction),
                         zero_ivar=zero_ivar)

        else:
            # subtract sky
            subtract_sky(frame,
                         skymodel,
                         apply_throughput_correction=(
                             not args.no_sky_throughput_correction),
                         zero_ivar=zero_ivar)

        compute_and_append_frame_scores(frame, suffix="SKYSUB")

    if not args.no_xtalk:
        log.info("fiber crosstalk correction")
        correct_fiber_crosstalk(frame, fiberflat)

        if not args.no_zero_ivar:
            frame.ivar *= (frame.mask == 0)

    if args.calib != None:
        log.info("calibrate")
        # read calibration
        fluxcalib = read_flux_calibration(args.calib)
        # apply calibration
        apply_flux_calibration(frame, fluxcalib)

        # Ensure that ivars are set to 0 for all values if any designated
        # fibermask bit is set. Also flips a bits for each frame.mask value using specmask.BADFIBER
        frame = get_fiberbitmasked_frame(
            frame, bitmask="flux", ivar_framemask=(not args.no_zero_ivar))
        compute_and_append_frame_scores(frame, suffix="CALIB")

    if not args.no_tsnr:
        log.info("calculating tsnr")
        results, alpha = calc_tsnr2(uncalibrated_frame,
                                    fiberflat=fiberflat,
                                    skymodel=skymodel,
                                    fluxcalib=fluxcalib,
                                    alpha_only=args.alpha_only)

        frame.meta['TSNRALPH'] = alpha

        comments = {k: "from calc_frame_tsnr" for k in results.keys()}
        append_frame_scores(frame, results, comments, overwrite=True)

    # record inputs
    frame.meta['IN_FRAME'] = shorten_filename(args.infile)
    frame.meta['FIBERFLT'] = shorten_filename(args.fiberflat)
    frame.meta['IN_SKY'] = shorten_filename(args.sky)
    frame.meta['IN_CALIB'] = shorten_filename(args.calib)

    # save output
    write_frame(args.outfile, frame, units='10**-17 erg/(s cm2 Angstrom)')
    log.info("successfully wrote %s" % args.outfile)