wlframes.py

"""
Last modified: 2017-03-07

Reduction of the UNe lamps:
- check the exposure time???
- check the flux ratio in chosen areas???
- subtract the masterdark
- remove bad pixel using the bad pixel mask
- if more than one lamps, combine them and adjust ron and gain
- save the masterlamp in the calibration database
- straigthen the masterlamp
- extract the central 10 pixels of each order for quality check??
- UNe wavelength calibration (TO DO: FP)
- TO DO: UNe line diagnostics (checking the ratio between specific emission lines)
"""

from astropy import units as u
from astropy.io import fits
import warnings
from astropy.utils.exceptions import AstropyWarning


import ccdproc
from drslib.config import CONFIG
from drslib import db, varie
from drslib import metadata
import numpy as np
import math, os, subprocess, shutil


class GBWls():
    def __init__(self, wls, dbconn):
        self.wls = wls
        self.dbconn = dbconn
        self.quality = []
        self.messages = []
        self.wllist = []
        self.wlcorr = []
        self.mlamp = ''


    def qualitycheck(self,unefp):
        """
        Quality check for UNe and FP lamps still to be defined.
        --> It will be done after the default wavelength calibration,
        here everything is defined as ok.
        """
        for frame in self.wls:

            wl = ccdproc.CCDData.read(frame, unit=u.adu)

            self.wllist.append(wl)
            self.quality.append('OK')
        return

    def masterlamp(self,unefp):

        if self.wllist:
            exptime = float(self.wllist[0].header[CONFIG['KEYS']['EXPTIME']])
            darkname = 'dark' + str(int(exptime))
            use_dark = True

            try:
                masterdark = db.extract_dbfile(self.dbconn,darkname)
            except:
                masterdark = False

            if not masterdark:
                self.messages.append('No masterdark found for this night, it will be taken from the calibration database.')

                try:
                    db.copy_dbfile(self.dbconn,darkname)
                    masterdark = db.extract_dbfile(self.dbconn,darkname)
                except:
                    for key in CONFIG['DARKLIST']:
                        try:
                            darkname = 'dark' + str(int(key))
                            db.copy_dbfile(self.dbconn,darkname)
                            masterdark = db.extract_dbfile(self.dbconn,darkname)
                            self.messages.append('No masterdark in the calibration database with the same exposure time as the lamp. The %s sec masterdark will be used instead' % (str(int(key))))
                            break
                        except:
                            self.messages.append('There are no masterdark in the calibration database. The masterdark will not be used.')
                            use_dark = False

            if use_dark:
                mdark = ccdproc.CCDData.read(masterdark, unit=u.adu)
            badpix = ccdproc.CCDData.read(CONFIG['BADPIX_MASK'], unit=u.adu)
            bad_mask=badpix.data
            inverse_mask=np.logical_not(bad_mask)


            for wl in self.wllist:

                # subdtract the masterdak
                if use_dark:
                    subwl = ccdproc.subtract_dark(wl,mdark,exposure_time=CONFIG['KEYS']['EXPTIME'],exposure_unit=u.second)
                    self.messages.append('Masterdark subtracted.')
                else:
                    subwl = wl

                # mask the lamp using the badpix mask
                wl_corrected = varie.badpix(subwl.data,bad_mask,inverse_mask)
                self.messages.append('Bad pixel correction done.')

                wl_corr = ccdproc.CCDData(wl_corrected,unit=u.adu)

                self.wlcorr.append(wl_corr)

            if len(self.wlcorr)>1:

                combinewl = ccdproc.Combiner(self.wlcorr)


                # mask the pixels using ron/gain
                #combinewl.sigma_clipping(func=np.ma.mean, dev_func=varie.stdcombine)
                combinewl.sigma_clipping(func=np.ma.median, dev_func=varie.stdcombine)

                #combineflat.sigma_clipping(func=np.ma.mean)
                mwl = combinewl.average_combine()
            else:
                mwl = self.wlcorr[0]

            mwl.data = np.asarray(mwl.data, dtype='float32')

            hea = self.wllist[0].header
            if unefp == 'une':
                wlname = hea[CONFIG['KEYS']['IMANAME']].replace('.fts','_UNE.fits')

            else:
                wlname = hea[CONFIG['KEYS']['IMANAME']].replace('.fts','_FP.fits')
                #wlname = hea[CONFIG['KEYS']['IMANAME']].replace('.fts','_LFC.fits')

            #wlname = wlname.replace(' ', '').replace('-', '_')
            self.mlamp = os.path.join(CONFIG['CALIB_DIR'],wlname)

            mwl.header = hea
            mwl.header[CONFIG['KEYS']['FILENAME']] = wlname
            mwl.header[CONFIG['KEYS']['NCOMBINE']] = len(self.wllist)
            mwl.header[CONFIG['RON_EFF'][0]] = (math.sqrt(len(self.wllist))*CONFIG['RON'],CONFIG['RON_EFF'][1])
            mwl.header[CONFIG['GAIN_EFF'][0]] = (len(self.wllist)*CONFIG['GAIN'],CONFIG['GAIN_EFF'][1])

            mwl.header[CONFIG['DRS_VERSION'][0]] = (CONFIG['VERSION'], CONFIG['DRS_VERSION'][1])

            # Add metadata to header
            mwl.header = metadata.add_metadata(mwl.header)

            hdu = fits.PrimaryHDU(data=mwl.data,header=mwl.header)
            masterwl = fits.HDUList([hdu])
            masterwl.writeto(self.mlamp,clobber=True)
            mwl = None

            #db.insert_dbfile(self.dbconn,unefp,self.mlamp)
            #self.messages.append('%s masterlamp %s was created and inserted in the calibration database.' % (unefp,str(os.path.basename(self.mlamp))))


        else:
            self.messages.append('There are no useful %s lamps. The reference lamp will be taken from the calibration database.' % (unefp))

            if db.check_dbfile(self.dbconn,unefp):
                try:
                    db.copy_dbfile(self.dbconn,unefp)
                    self.mlamp = db.extract_dbfile(self.dbconn,unefp)
                except:
                    self.messages.append('There are no %s lamps in the calibration database.' %(unefp))
                    return False

            # Since masterlamp has been copied from calibration database
            # also str and extr will be copied
            self.copy_straighten(unefp)
            return False

        return True


    def copy_straighten(self,unefp):
        """
        The masterlamp has been copied from calibration database
        This function also copy str and extr file from calibration database
        """

        # Get masterlamp name from calibdb
        # lamp is a fullpath name
        lamp = self.mlamp

        # Generate str and extr fullpath name from calibdb
        straight = lamp.replace('.fits','_str.fits')
        extracted = lamp.replace('.fits','_ext.fits')

        # Insert str and extr in calibdb for this night
        strname = '_'.join((unefp,'str'))
        db.insert_dbfile(self.dbconn,strname,straight)

        calname = '_'.join((unefp,'calib'))
        db.insert_dbfile(self.dbconn,calname,extracted)

        # Copy straight and extracted to RED_CALIB directory
        red_straight = os.path.join(CONFIG['RED_CALIB'],os.path.basename(straight))
        shutil.copyfile(straight,red_straight)
        red_extr = os.path.join(CONFIG['RED_CALIB'],os.path.basename(extracted))
        shutil.copyfile(extracted,red_extr)

        return


    def extract(self,unefp):
        # straighten the lamp's orders
        # lamp = db.extract_dbfile(self.dbconn,unefp)
        lamp = self.mlamp

        #print lamp
        straight = lamp.replace('.fits','_str.fits')
        extracted = lamp.replace('.fits','_ext.fits')

        args = [CONFIG['STRAIGHT'],lamp,straight]
        args.extend(CONFIG['STRAIGHT_OPT'])
        # search for shift defined in the straighten options in config.py
        dy = True
        for opt in CONFIG['STRAIGHT_OPT']:
            try:
                dy = opt.rindex('DY=')
                ypos = int(opt[dy-2:])
                shift = CONFIG['SHIFT_Y'] + ypos
                dy = False
            except:
                pass

        if dy:
            try:
                shift = db.extract_dbfile(self.dbconn,'shiftY')
            except:
                try:
                    cal_flat = db.extract_dbfile(self.dbconn,'flat')
                    mflat = ccdproc.CCDData.read(cal_flat, unit=u.adu)
                    shift = varie.shiftY(mflat.data)
                    db.insert_dbfile(self.dbconn,'shiftY',shift)
                except:
                    db.copy_dbfile(self.dbconn,'shiftY')
                    shift = db.extract_dbfile(self.dbconn,'shiftY')
            if not shift:
                shift = CONFIG['SHIFT_Y']

            shiftY = [''.join(('DY=',str(shift - CONFIG['SHIFT_Y'])))]
            #print shiftY
            args.extend(shiftY)

        subprocess.call(args)

        #self.messages.append('%s: orders straightened.' % (str(os.path.basename(lamp))))

        #strname = '_'.join((unefp,'str'))
        #db.insert_dbfile(self.dbconn,strname,straight)
        #self.messages.append('Straightened %s masterlamp %s was created and inserted in the calibration database.' % (unefp,str(os.path.basename(straight))))

        wl = ccdproc.CCDData.read(straight, unit=u.adu)

        # Update FILENAME in header then
        # add metadata to header and save straight file
        wl.header[CONFIG['KEYS']['FILENAME']] = os.path.basename(straight)
        wl.header = metadata.add_metadata(wl.header)

        hdu = fits.PrimaryHDU(data=wl.data, header=wl.header)
        str_fits = fits.HDUList([hdu])
        str_fits.writeto(straight, overwrite=True)


        wldata = wl.data
        wl.header[CONFIG['WLFIT'][0]] = (CONFIG['WLFIT_FUNC'],CONFIG['WLFIT'][1])

        ordini = np.zeros((CONFIG['N_ORD'],CONFIG['XCCD']))

        # read the lines to use in the wavelength calibration
        select_lines, all_lines = varie.UNe_linelist() 

        # extract the central 10 pixel of each order
        # wavelength calibration, to check the lamp and
        # to be used as default if later ad hoc calibration fails
        for x in xrange(CONFIG['N_ORD']):
            row = x*CONFIG['W_ORD'] + 19
 
            ordini[x] = np.sum(wldata[row-CONFIG['WEXT']:row+CONFIG['WEXT']],axis=0)
            if unefp == 'une':
                calib_failed, coeffs, comments = varie.UNe_calibrate(ordini[x],x+32,select_lines[x+32],all_lines[x+32],True,True)
                for comment in comments:
                    self.messages.append(comment)

                keyfail = ''.join((CONFIG['CAL_FAILED'][0],str(x+32)))

                if calib_failed:

                    if self.quality[0] == 'OK':
                        self.quality = ['FAILED' for q in self.quality]
                        #db.remove_dbfile(CONFIG['DATE'],'UNE')
                        #return

                    wl.header[keyfail] = (False,CONFIG['CAL_FAILED'][1])
                    self.messages.append(' *** WARNING ***')
                    self.messages.append('The default wavelength calibration for the order %s will be taken from the database and as such it will not be optimal for the night.' % str(x+32))
                    try:
                        wcalib = db.extract_dbfile(self.dbconn,'une_calib')
                    except:
                        #wcalib = False

                    #if not wcalib:
                        #self.messages.append('No UNe calibration found for this night, it will be taken from the calibration database.')
                        try:
                            db.copy_dbfile(self.dbconn,'une_calib')
                            wcalib = db.extract_dbfile(self.dbconn,'une_calib')
                        except:
                            wcalib = False

                    try:
                        wlc = ccdproc.CCDData.read(wcalib, unit=u.adu)
                        wl.header = wlc.header
                    except:
                        self.messages.append(' *** WARNING ***')
                        self.messages.append('There is no default wavelength calibration in the database, please observe a new lamp, otherwise the spectra will not be reduced.')

                else:
                    wl.header[keyfail] = (True,CONFIG['CAL_FAILED'][1])
                    for key in coeffs:
                        keyword = ''.join((CONFIG['WLCOEFFS'][key][0],str(x+32)))
                        wl.header[keyword] = (coeffs[key],CONFIG['WLCOEFFS'][key][1])

            else:
                # calibrate FP using UNe
                try:
                    wcalib = db.extract_dbfile(self.dbconn,'une_calib')
                except:
                    wcalib = False

                if not wcalib:
                    self.messages.append('No UNe calibration found for this night, it will be taken from the calibration database.')
                    try:
                        db.copy_dbfile(self.dbconn,'une_calib')
                        wcalib = db.extract_dbfile(self.dbconn,'une_calib')
                    except:
                        return

                #print wcalib
                wlc = ccdproc.CCDData.read(wcalib, unit=u.adu)

                for key in CONFIG['WLCOEFFS']:
                    keyword = ''.join((CONFIG['WLCOEFFS'][key][0],str(x+32)))
                    try:
                        wl.header[keyword] = (wlc.header[keyword],CONFIG['WLCOEFFS'][key][1])
                    except:
                        pass


        # save the extracted lamp as FITS file
        wlextract = ccdproc.CCDData(ordini, unit=u.adu)
        wlextract.data = np.asarray(wlextract.data, dtype='float32')
        wlextract.header = wl.header
        wlextract.header[CONFIG['KEY_WEXT'][0]] = (CONFIG['WEXT']*2,CONFIG['KEY_WEXT'][1])

        # Update FILENAME in header then
        # add metadata to header and save extracted file
        wlextract.header[CONFIG['KEYS']['FILENAME']] = os.path.basename(extracted)
        wlextract.header = metadata.add_metadata(wlextract.header)

        hdu = fits.PrimaryHDU(data=wlextract.data, header=wlextract.header)
        wlextr = fits.HDUList([hdu])
        wlextr.writeto(extracted, overwrite=True)

        # Copy straight and extracted to RED_CALIB directory
        red_straight = os.path.join(CONFIG['RED_CALIB'],os.path.basename(straight))
        shutil.copyfile(straight,red_straight)
        red_extr = os.path.join(CONFIG['RED_CALIB'],os.path.basename(extracted))
        shutil.copyfile(extracted,red_extr)

        if self.quality[0] == 'OK':
            db.insert_dbfile(self.dbconn,unefp,self.mlamp)
            self.messages.append('%s masterlamp %s was created and inserted in the calibration database.' % (unefp,str(os.path.basename(self.mlamp))))

            strname = '_'.join((unefp,'str'))
            db.insert_dbfile(self.dbconn,strname,straight)
            self.messages.append('Straightened %s masterlamp %s was created and inserted in the calibration database.' % (unefp,str(os.path.basename(straight))))

            calname = '_'.join((unefp,'calib'))
            db.insert_dbfile(self.dbconn,calname,extracted)
            self.messages.append('%s wavelength calibration %s was created and inserted in the calibration database.' % (unefp,str(os.path.basename(extracted))))

        else:
            os.remove(self.mlamp)
            os.remove(straight)
            os.remove(extracted)
            try: db.remove_dbfile1(self.mlamp)
            except: pass
            self.messages.append('%s frame %s failed the quality check and it will not be used.' % (unefp,str(os.path.basename(extracted))))


        wlextract = None
        ordini = None

        return


    def une_process(self):
        warnings.simplefilter('ignore', category=AstropyWarning)
        self.qualitycheck('une')
        if self.masterlamp('une'):
            self.extract('une')
        #self.wls[:] = []
        return

    def fp_process(self):
        warnings.simplefilter('ignore', category=AstropyWarning)
        self.qualitycheck('fp')
        if self.masterlamp('fp'):
            self.extract('fp')
        #self.wls[:] = []
        return