nod3fits.py 13.8 KB
#!/usr/bin/env python

try:
   from Astropy.io import fits
   from astropy import __version__
except:
   import pyfits as fits
   from pyfits import __version__
fitsvers = __version__
fitsname = fits.__name__

import time, os
from numpy import array, float32, nanmin, nanmax, reshape, fliplr, flipud

import re, sys, os.path as osp
from guidata.qt.compat import getsavefilename, getopenfilename, getopenfilenames
from guidata.qt.QtGui import QMessageBox

from guiqwt.config import _
#from guiqwt.io import (IMAGE_SAVE_FILTERS, IMAGE_LOAD_FILTERS)

def nod3_header(header):
    import xml2nod3
    Nod3 = xml2nod3.NOD3_Fits()
    nod3keys, nod3header = Nod3.NOD3Header()
    for key in header.keys():
        if key not in nod3keys:
           del header[key]
    return header

def check_header(h):
    if type(h) == type(None):
       return h
    global header
    # patch for old nod2-fits structure
    for key in ('IDENT1', 'IDENT2', 'IDENT3', 'IDENT4'):
        if key in h: h[key] = "old nod2 header replacement"
    if "OBS LAT" in h: h.rename_key("OBS LAT", "SITELAT")
    if "OBSLAT" in h: h.rename_key("OBSLAT", "SITELAT")
    if "OBSEPOCH" in h: h.rename_key("OBSEPOCH", "DATE_OBS")
    if "DATE-OBS" in h: h.rename_key("DATE-OBS", "DATE_OBS")
    #if 'EQUINOX' in h and not 'EPOCH' in h: h.rename_key('EQUINOX', 'EPOCH')
    #if not 'PATLONG' in h: h.update('PATLONG', 0.0)
    if not 'PATLONG' in h: h['PATLONG'] = 0.0
    #if not 'PATLAT' in h: h.update('PATLAT', 0.0)
    if not 'PATLAT' in h: h['PATLAT'] = 0.0
    if abs(h['PATLONG']) > 10.0: h['PATLONG'] /= 3600.0
    if abs(h['PATLAT']) > 10.0: h['PATLAT'] /= 3600.0
    return h

def _add_all_supported_files(filters):
    extlist = re.findall(r'\*.[a-zA-Z0-9]*', filters)
    allfiles = '%s (%s)\n' % (_("All supported files"), ' '.join(extlist))
    return allfiles+filters

def exec_fits_open_dialog(parent, basedir='', app_name=None, mapnum=0, infile=None):
    """
    Executes an image*s* open dialog box (QFileDialog.getOpenFileNames)
        * parent: parent widget (None means no parent)
        * basedir: base directory ('' means current directory)
        * app_name (opt.): application name (used as a title for an eventual 
          error message box in case something goes wrong when saving image)
        * mapnum (opt): chooes the map number in the fitsfile
    
    Yields (filename, data) tuples if dialog is accepted, None otherwise
    """

    global header

    IMAGE_LOAD_FILTERS = '%s (*.fits *.FITS)' % (_(u"FitsImages"))
    IMAGE_LOAD_FILTERS = _add_all_supported_files(IMAGE_LOAD_FILTERS)

    saved_in, saved_out, saved_err = sys.stdin, sys.stdout, sys.stderr
    sys.stderr = None
    sys.stdout = None
    if not infile:
       filenames, _filter = getopenfilenames(parent, _("Open"), basedir,
                                          IMAGE_LOAD_FILTERS)
    else:
       filenames = [infile]
    sys.stdin, sys.stdout, sys.stderr = saved_in, saved_out, saved_err
    filenames = [str(fname) for fname in list(filenames)]
    SidData = {}
    ParData = {}
    for filename in filenames:
        hdulist = fits.open(filename)
        nhdu = len(hdulist)
        siddata = None
        pardata = None
        for n in range(nhdu):
            if hasattr(hdulist[n], 'header'):
               header = hdulist[n].header
               if 'CALUNIT' in header:
                  #header.update("BUNIT", header['CALUNIT'], "Map unit")
                  header["BUNIT"] = (header['CALUNIT'], "Map unit")
                  del header['CALUNIT']
               if 'SCANNUM' in header:
                  scannum = hdulist[n].header['SCANNUM']
               else:
                  scannum = 1
               if 'EXTNAME' in header:
               #if 'EXTNAME' in header:
                  extname = header['EXTNAME']
                  if extname == "SIDMAP":
                     mchan, header, siddata = fits_to_array(filename, hdulist=hdulist, mapnum=n)
                     SidData[scannum] = siddata
                  elif extname == "PARMAP":
                     mchan, header, pardata = fits_to_array(filename, hdulist=hdulist, mapnum=n)
                     ParData[scannum] = pardata
        for mapnum in range(nhdu):
            try:
                mchan, header, data = fits_to_array(filename, hdulist=hdulist, mapnum=mapnum)
                #header = check_header(header)
                #h = check_header(header)
                if mchan > 1:
                   if 'EXTNAME' in header:
                      header['extname'] = h['extname'][:-2] + str("C1")
                #if 'EQUINOX' in header: header.rename_key('EQUINOX', 'EPOCH')
                if 'EXTNAME' in header:
                #if 'EXTNAME' in header:
                   extname = header['extname']
                else:
                   extname = None
            except Exception, msg:
                import traceback
                traceback.print_exc()
                QMessageBox.critical(parent,
                     _('Error') if app_name is None else app_name,
                     (_(u"%s could not be opened:") % osp.basename(filename))+\
                     #"\n"+str(msg))
                     "\n"+str('not a NOD3 fits file'))
                return
            if 'SCANNUM' in header: scannum = header['SCANNUM']
            if extname != "SIDMAP" and extname != "PARMAP":
               try:
                  siddata = SidData[scannum]
               except:
                  try: siddata = SidData[1]
                  except: siddata = None
               try:
                  pardata = ParData[scannum]
               except:
                  try: siddata = SidData[1]
                  except: pardata = None
               header = check_history(header)
               yield filename, header, data, siddata, pardata
            for chan in range(1, mchan):
                nchan, h, data = fits_to_array(filename, hdulist=hdulist, mapnum=mapnum, chan=chan)
                if 'EXTNAME' in header:
                   header['extname'] = h['extname'][:-2] + str("C%d" % (chan+1))
                header = check_history(header)
                yield filename, header, data, siddata, pardata

def check_history(header):
    if "CTYPE1" in header:
       header["CTYPE1"] = header["CTYPE1"].replace("DES", "CAR")
       header["CTYPE2"] = header["CTYPE2"].replace("DES", "CAR")
    if 'HISTORY' in header:
       def get_history():
           return header['HISTORY']
       header.get_history = get_history
    elif not hasattr(header, 'get_history'):
       header.add_history("MPIfR NOD3 package")
       def get_history():
           return header['HISTORY']
       header.get_history = get_history
    return header
    
def exec_fits_save_dialog(parent, rows, outdir='', app_name=None, dim=2, nod3=False):
    """
    Executes an image save dialog box (QFileDialog.getSaveFileName)
        * data: image pixel array data
        * parent: parent widget (None means no parent)
        * basedir: base directory ('' means current directory)
        * app_name (opt.): application name (used as a title for an eventual 
          error message box in case something goes wrong when saving image)
    
    Returns filename if dialog is accepted, None otherwise
    """

    IMAGE_SAVE_FILTERS = '%s (*.fits *.FITS)' % (_(u"FitsImages"))
    IMAGE_SAVE_FILTERS = _add_all_supported_files(IMAGE_SAVE_FILTERS)

    saved_in, saved_out, saved_err = sys.stdin, sys.stdout, sys.stderr
    sys.stdout = None
    filename, _filter = getsavefilename(parent, _("Save as"), outdir,
                                        IMAGE_SAVE_FILTERS)
    sys.stdin, sys.stdout, sys.stderr = saved_in, saved_out, saved_err

    if hasattr(parent, 'SidOut'):
       if parent.SidOut: sidout = True
       else: sidout = False
    else: sidout = False
    if hasattr(parent, 'ParOut'):
       if parent.ParOut: parout = True
       else: parout = False
    else: parout = False

    if filename:
       filename = str(filename)
       try:
          irow = 0
          if dim == 2:
             for row in rows:
                 obj = parent.objects[row]
                 #obj.header.update('NOD3VERS', parent.VERSION, 'NOD3 software version')
                 obj.header['NOD3VERS'] = (parent.VERSION, 'NOD3 software version')
                 if nod3: obj.header = nod3_header(obj.header)
                 array_to_fits(obj.data, obj.header, filename, irow)
                 irow += 1
          else:
             obj = parent.objects[-1]
             #obj.header.update('NOD3VERS', parent.VERSION, 'NOD3 software version')
             obj.header['NOD3VERS'] = (parent.VERSION, 'NOD3 software version')
             if nod3: obj.header = nod3_header(obj.header)
             array_to_fits(obj.data, obj.header, filename, irow)
          if 'SCANNUM' in obj.header:
             scans = [obj.header['SCANNUM']]
          else:
             scans = [0]
          if hasattr(obj, 'sidmap') and sidout:
             ParSid_to_fits([obj.sidmap], obj.header, filename, scans, name='sidmap')
          if hasattr(obj, 'parmap') and parout:
             ParSid_to_fits([obj.parmap], obj.header, filename, scans, name='parmap')
          parent.SidOut = False
          parent.ParOut = False
          return filename
       except Exception, msg:
          import traceback
          traceback.print_exc()
          QMessageBox.critical(parent,
               _('Error') if app_name is None else app_name,
               (_(u"%s could not be written:") % osp.basename(filename))+\
               "\n"+str(msg))
          parent.SidOut = False
          parent.ParOut = False
          return

def header_clear(header):
    for item in header.items():
        #header.__delitem__(item[0])
        del header[item[0]]
    return header

def ParSid_to_fits(Data, header, outfile, scans, name):
    if outfile.lower().find(".fits") < 0:
       outfile += ".fits"
    for m in range(len(scans)):
        rows, cols = Data[m].shape
        scannum = scans[m]
        data = array(Data[m], dtype=float32)
        try: header.clear()
        except: header = header_clear(header)
        header.update('XTENSION', 'IMAGE   ', 'Image extension')
        header['XTENSION'] = ('IMAGE   ', 'Image extension')
        #header.update("BITPIX", -32, 'array data type')
        header["BITPIX"] = (-32, 'array data type')
        #header.update('NAXIS', 2, 'number of array dimensions')
        header['NAXIS'] = (2, 'number of array dimensions')
        #header.update('NAXIS1', cols)
        header['NAXIS1'] = cols
        #header.update('NAXIS2', rows)
        header['NAXIS2'] = rows
        #header.update('SCANNUM', scannum)
        header['SCANNUM'] = scannum
        #header.update('EXTNAME', name.upper(), 'Name of image extension')
        header['EXTNAME'] = (name.upper(), 'Name of image extension')
        fits.append(outfile, data, header)

def array_to_fits(Data, header, outfile, irow):
    if outfile.lower().find(".fits") < 0:
       outfile += ".fits"
    data = array(Data, dtype=float32)
    header["BITPIX"] = -32
    if "DATAMAX" in header:
    #if "DATAMAX" in header:
       header["DATAMAX"] = nanmax(data)
    else:
       #header.update("DATAMAX", nanmax(data), "MAX VALUE OF ARRAY")
       header["DATAMAX"] = (nanmax(data), "MAX VALUE OF ARRAY")
    if "PATLONG" in header:
       #if abs(header["PATLONG"]) > 30.0: header.update("PATLONG", header["PATLONG"]/3600.0)
       if abs(header["PATLONG"]) > 30.0: header["PATLONG"] = header["PATLONG"]/3600.0
    if "PATLAT" in header:
       #if abs(header["PATLAT"]) > 30.0: header.update("PATLAT", header["PATLAT"]/3600.0)
       if abs(header["PATLAT"]) > 30.0: header["PATLAT"] = header["PATLAT"]/3600.0
    if "DATAMIN" in header:
    #if "DATAMIN" in header:
       header["DATAMIN"] = nanmin(data)
    else:
       #header.update("DATAMIN", nanmin(data), "MIN VALUE OF ARRAY")
       header["DATAMIN"] = (nanmin(data), "MIN VALUE OF ARRAY")
    #header.update("DATE", time.strftime("%m/%d/%y", time.gmtime()), "DATE FITS DATA WRITTEN  (MM/DD/YY)")
    header["DATE"] = (time.strftime("%m/%d/%y", time.gmtime()), "DATE FITS DATA WRITTEN  (MM/DD/YY)")
    try:
       #header.update("USER", os.getenv('USER'), "NOD3 user account")
       header["USER"] = (os.getenv('USER'), "NOD3 user account")
    except: pass 
    #if hasattr(os, 'getlogin'):
    #   header.update("USER", os.getlogin(), "NOD3 user account")
    #header.update("CREATOR", "MPIfR NOD3", "NOD3 single dish data reduction sofware")
    header["CREATOR"] = ("MPIfR NOD3", "NOD3 single dish data reduction sofware")
    N = 1
    while str("NAXIS%d" % (N)) in header:
          N += 1
    header['NAXIS'] = N-1
    if 'BLANK' in header:
       #header.__delitem__('BLANK')
       del header['BLANK']
    rows, cols = data.shape
    shape = [rows, cols]
    for n in range(header['NAXIS']-2):
        shape = [1] + shape
    data = data.reshape(shape)
    if irow == 0: 
       if os.path.exists(outfile): os.remove(outfile)
       fits.writeto(outfile, data, header)
    else:
       fits.append(outfile, data, header)

def fits_to_array(fitsfile, hdulist=False, mapnum=0, chan=0):
    global header
    #try:
    if 1:
       if not hdulist:
          hdulist = fits.open(fitsfile, do_not_scale_image_data=True)
          #hdulist = fits.open(fitsfile, do_not_scale_image_data=False)
       header = hdulist[mapnum].header
       try:
          data = 1.0*hdulist[mapnum].data
       except:
          return None, None, None
       nchan = 1
       if len(data.shape) == 3:
          nchan = data.shape[0]
          data = data[chan]
       elif len(data.shape) == 4:
          data = data[0][0]
       elif len(data.shape) == 5:
          data = data[0][0][0]
       elif len(data.shape) == 6:
          data = data[0][0][0][0]
       else: None, None, None
       return nchan, header, data
    #except:
    #   return None, None, None
    

def main():
    infile = sys.argv[1]
    nchan, header, data = fits_to_array(infile)
    array_to_fits(data, header, 'new.fits', 0)
    array_to_fits(data, header, 'new.fits', 1)
    
if __name__ == '__main__':
   main()