source: fact/tools/pyscripts/new_pyfact/pyfact.py@ 17804

#!/usr/bin/python -tt
#
# Werner Lustermann, Dominik Neise
# ETH Zurich, TU Dortmund
#
import os
import sys
import ctypes
from ctypes import *
import numpy as np
import pprint # for SlowData
from scipy import signal

import ROOT
hostname = ROOT.gSystem.HostName()

libz_path_dict = {
    # hostname  :   /path/to/libz.so
    'isdc'          : "/usr/lib64/libz.so",
    'neiseLenovo'   : "/usr/lib/libz.so",
    'factcontrol'   : "/usr/lib/libz.so",
    "max-K50AB"     : "/usr/lib/x86_64-linux-gnu/libz.so",
    "watz"          : "/usr/lib/x86_64-linux-gnu/libz.so",
    "grolsch"       : "/usr/lib/i386-linux-gnu/libz.so",
}
libz_loaded = False
for my_hostname in libz_path_dict:
    if my_hostname in hostname:
        ROOT.gSystem.Load(libz_path_dict[my_hostname])
        libz_loaded = True
if not libz_loaded:
    print """Warning - Warning - Warning - Warning - Warning - Warning - Warning
    I most probably need to load libz.so but I don't know where it is.
    
    Please edit pyfact.py around line 16-24 and insert your hostname and your
    path to your libz.so 
    Sorry for the inconvenience.
    """
    sys.exit(-1)

    
root_make_string = ROOT.gSystem.GetMakeSharedLib()
if not "-std=c++0x" in root_make_string:
    make_string = root_make_string.replace('$Opt', '$Opt -std=c++0x -D HAVE_ZLIB')
ROOT.gSystem.SetMakeSharedLib(make_string)
ROOT.gROOT.ProcessLine(".L extern_Mars_mcore/izstream.h+O")
ROOT.gROOT.ProcessLine(".L extern_Mars_mcore/fits.h+O")
ROOT.gROOT.ProcessLine(".L extern_Mars_mcore/zfits.h+O")
ROOT.gROOT.ProcessLine(".L extern_Mars_mcore/factfits.h+O")
if not "-std=c++0x" in root_make_string:
    make_string = root_make_string.replace('$Opt', "$Opt -std=c++0x -D HAVE_ZLIB -D'PACKAGE_NAME=\"PACKAGE_NAME\"' -D'PACKAGE_VERSION=\"PACKAGE_VERSION\"' -D'REVISION=\"REVISION\"' ")
ROOT.gSystem.SetMakeSharedLib(make_string)
ROOT.gROOT.ProcessLine(".L extern_Mars_mcore/DrsCalib.h+O")

ROOT.gInterpreter.GenerateDictionary("map<string,fits::Entry>","map;string;extern_Mars_mcore/fits.h")
ROOT.gInterpreter.GenerateDictionary("pair<string,fits::Entry>","map;string;extern_Mars_mcore/fits.h")
ROOT.gInterpreter.GenerateDictionary("map<string,fits::Table::Column>","map;string;extern_Mars_mcore/fits.h")
ROOT.gInterpreter.GenerateDictionary("pair<string,fits::Table::Column>","map;string;extern_Mars_mcore/fits.h")
ROOT.gInterpreter.GenerateDictionary("vector<DrsCalibrate::Step>","vector;extern_Mars_mcore/DrsCalib.h")

ROOT.gSystem.Load('extern_Mars_mcore/fits_h.so')
ROOT.gSystem.Load('extern_Mars_mcore/izstream_h.so')
ROOT.gSystem.Load('extern_Mars_mcore/zfits_h.so')
ROOT.gSystem.Load('extern_Mars_mcore/factfits_h.so')
ROOT.gSystem.Load('extern_Mars_mcore/DrsCalib_h.so')
from ROOT import *

class Fits( object ):
    """ General FITS file access
    
    Wrapper for factfits class from Mars/mcore/factfits.h 
    (factfits might not be suited to read any FITS file out there, but certainly
    all FACT FITS files can be read using this class)
    """
    __module__ = 'pyfact'
    def __init__(self, path):
        """
        """
        if not os.path.exists(path):
            raise IOError(path+' was not found')
        self.f = factfits(path)
        self._make_header()
        self._setup_columns()
    
    def _make_header(self):
        """
        """
        str_to_bool = { 'T':True, 'F':False}
        type_conversion = { 'I' : int, 'F' : float, 'T' : str, 'B' : str_to_bool.__getitem__}

        self.header = {}
        self.header_comments = {}
        for key,entry in self.f.GetKeys():
            try:
                self.header[key] = type_conversion[entry.type](entry.value)
            except KeyError:
                raise IOError("Error: entry type unknown.\n Is %s, but should be one of: [I,F,T,B]" % (entry.type) )
            self.header_comments[key] = entry.comment
    
    def _setup_columns(self):
        """
        """
        col_type_to_np_type_map = { 'L' : 'b1',  'A' : 'a1', 'B' : 'i1',
            'I' : 'i2', 'J' : 'i4', 'K' : 'i8', 'E' : 'f4', 'D' : 'f8'}
        self.cols = {}
        for key,col in self.f.GetColumns():
            self.cols[key] = np.zeros(col.num, col_type_to_np_type_map[col.type])
            if col.num != 0:
                self.f.SetPtrAddress(key, self.cols[key])

    def __iter__(self):
        return self

    def next(self, row=None):
        """ 
        """
        if row is None:
            if self.f.GetNextRow() == False:
                raise StopIteration
        else:
            row = int(row)
            if self.f.GetRow(row) == False:
                raise StopIteration
        return self

    
    
class RawData( Fits ):
    """ Special raw data FITS file access (with DRS4 calibration)
        
        During iteration the C++ method DrsCalibration::Apply is being called.
    """
    __module__='pyfact'
    def __init__(self, data_path, calib_path):
        """ -constructor-
        *data_path*  : fits or fits.gz file of the data including the path
        *calib_path* : fits or fits.gz file containing DRS calibration data
        """
        super(RawData, self).__init__(data_path)
        self.cols['CalibData'] = np.zeros( self.cols['Data'].shape, np.float32)
        self.cols['CalibData2D'] = self.cols['CalibData'].reshape( self.header['NPIX'], -1)
        if not self.cols['CalibData2D'].base is self.cols['CalibData']:
            print "Error seomthing went wrong!"
        self.drs_calibration = DrsCalibration()
        self.drs_calibration.ReadFitsImp( calib_path )
        
        self.drs_calibrate = DrsCalibrate()
        self.list_of_previous_start_cells = []

    def __iter__(self):
        """ iterator """
        return self

    def next(self, row=None):
        """ 
        """
        super(RawData, self).next(row)
        
        self.drs_calibration.Apply( self.cols['CalibData'], 
                                    self.cols['Data'], 
                                    self.cols['StartCellData'], 
                                    self.header['NROI'])
        
        for previous_start_cells in self.list_of_previous_start_cells:
            self.drs_calibrate.CorrectStep(
                                self.cols['CalibData'], 
                                self.header['NPIX'], 
                                self.header['NROI'], 
                                previous_start_cells, 
                                self.cols['StartCellData'], 
                                self.header['NROI']+10)
            self.drs_calibrate.CorrectStep(
                                self.cols['CalibData'], 
                                self.header['NPIX'], 
                                self.header['NROI'], 
                                previous_start_cells, 
                                self.cols['StartCellData'], 
                                3)
        self.list_of_previous_start_cells.append(self.cols['StartCellData'])
        if len(self.list_of_previous_start_cells) > 5:
            self.list_of_previous_start_cells.pop(0)

        for ch in range(self.header['NPIX']):
            self.drs_calibrate.RemoveSpikes3(self.cols['CalibData2D'][ch], self.header['NROI'])

        return self



if __name__ == '__main__':
    """ tests  """
    if len(sys.argv) < 3:
        print """ Usage:
        ----------------------------------------------------------------------
        To just build the shared object libs call:
        python pyfact.py
        
        To build (if necessary) and open an example file
        python -i pyfact.py /path/to/data_file.zfits /path/to/calib_file.drs.fits.gz
        
        Any of the 3 file 'types': fits.gz    zfits    fits
        should be supported.
        
        To clean all of automatically built files, do something like:
        rm *.so *.d AutoDict_* extern_Mars_mcore/*.so extern_Mars_mcore/*.d
        ----------------------------------------------------------------------
        """
        sys.exit(1)
    print "Example for calibrated raw-file"
    f = RawData(sys.argv[1], sys.argv[2])
    
    print "number of events:", f.header['NAXIS2']
    print "date of observation:", f.header['DATE']
    print "The files has these cols:", f.cols.keys()
    
    for counter,row in enumerate(f):
        print "Event Id:", row.cols['EventNum']
        print "shape of column 'StartCellData'", row.cols['StartCellData'].shape
        print "dtype of column 'Data'", row.cols['StartCellData'].dtype
        if counter > 3:
            break
    # get next row
    f.next()
    print "Event Id:", f.cols['EventNum']
    # get another row
    f.next(10)
    print "Event Id:", f.cols['EventNum']
    # Go back again
    f.next(3)
    print "Event Id:", f.cols['EventNum']
    
    import matplotlib.pyplot as plt
    plt.ion()
    for i in range(f.header['NPIX']):
        plt.cla()
        plt.title("Event %d"%(f.cols['EventNum']))
        plt.plot(f.cols['CalibData2D'][i], '.:', label='pixel %d'%i)
        plt.legend()
        answer = raw_input('anykey for next pixel; "q" to quit. :')
        if 'q' in answer.lower():
            break