| 1 | #!/usr/bin/python -itt
|
|---|
| 2 | import numpy as np
|
|---|
| 3 | import pprint
|
|---|
| 4 | import ctypes
|
|---|
| 5 |
|
|---|
| 6 | from ROOT import gSystem
|
|---|
| 7 | gSystem.Load('pyfits_h.so')
|
|---|
| 8 | from ROOT import *
|
|---|
| 9 |
|
|---|
| 10 | class FactFits( fits ):
|
|---|
| 11 | """ -Fact Fits File-
|
|---|
| 12 | A Python wrapper for the fits-class implemented in pyfits.h
|
|---|
| 13 | provides easy access to the fits file meta data.
|
|---|
| 14 | * dictionary of file metadata - self.meta
|
|---|
| 15 | * dict of table metadata - self.columns
|
|---|
| 16 | * variable table column access, thus possibly increased speed while looping
|
|---|
| 17 | """
|
|---|
| 18 | def __init__(self, path):
|
|---|
| 19 | """ creates meta and columns dictionaries
|
|---|
| 20 | """
|
|---|
| 21 | self.path = path
|
|---|
| 22 | try:
|
|---|
| 23 | fits.__init__(self,path)
|
|---|
| 24 | except IOError:
|
|---|
| 25 | print 'problem accessing data file: ', data_file_name
|
|---|
| 26 | raise # stop ! no data
|
|---|
| 27 |
|
|---|
| 28 | self.meta = self._make_meta_dict()
|
|---|
| 29 | self.columns = self._make_columns_dict()
|
|---|
| 30 |
|
|---|
| 31 | self.treat_meta_dict()
|
|---|
| 32 |
|
|---|
| 33 |
|
|---|
| 34 | # list of columns, which are already registered
|
|---|
| 35 | # see method register()
|
|---|
| 36 | self._registered_cols = []
|
|---|
| 37 | # dict of column data, this is used, in order to be able to remove
|
|---|
| 38 | # the ctypes of
|
|---|
| 39 | self._table_cols = {}
|
|---|
| 40 |
|
|---|
| 41 | # I need to count the rows, since the normal loop mechanism seems not to work.
|
|---|
| 42 | self._current_row = 0
|
|---|
| 43 |
|
|---|
| 44 | self.stacked_cols = {}
|
|---|
| 45 |
|
|---|
| 46 | def _make_meta_dict(self):
|
|---|
| 47 | # intermediate variables for file metadata dict generation
|
|---|
| 48 | keys=self.GetPy_KeyKeys()
|
|---|
| 49 | values=self.GetPy_KeyValues()
|
|---|
| 50 | comments=self.GetPy_KeyComments()
|
|---|
| 51 | types=self.GetPy_KeyTypes()
|
|---|
| 52 |
|
|---|
| 53 | if len(keys) != len(values):
|
|---|
| 54 | raise TypeError('len(keys)',len(keys),' != len(values)', len(values))
|
|---|
| 55 | if len(keys) != len(types):
|
|---|
| 56 | raise TypeError('len(keys)',len(keys),' != len(types)', len(types))
|
|---|
| 57 | if len(keys) != len(comments):
|
|---|
| 58 | raise TypeError('len(keys)',len(keys),' != len(comments)', len(comments))
|
|---|
| 59 |
|
|---|
| 60 | meta_dict = {}
|
|---|
| 61 | for i in range(len(keys)):
|
|---|
| 62 | type = types[i]
|
|---|
| 63 | if type == 'I':
|
|---|
| 64 | value = int(values[i])
|
|---|
| 65 | elif type == 'F':
|
|---|
| 66 | value = float(values[i])
|
|---|
| 67 | elif type == 'B':
|
|---|
| 68 | if values[i] == 'T':
|
|---|
| 69 | value = True
|
|---|
| 70 | elif values[i] == 'F':
|
|---|
| 71 | value = False
|
|---|
| 72 | else:
|
|---|
| 73 | raise TypeError("meta-type is 'B', but meta-value is neither 'T' nor 'F'. meta-value:",values[i])
|
|---|
| 74 | elif type == 'T':
|
|---|
| 75 | value = values[i]
|
|---|
| 76 | else:
|
|---|
| 77 | raise TypeError("unknown meta-type: known meta types are: I,F,B and T. meta-type:",type)
|
|---|
| 78 | meta_dict[keys[i]]=(value, comments[i])
|
|---|
| 79 | return meta_dict
|
|---|
| 80 |
|
|---|
| 81 |
|
|---|
| 82 | def _make_columns_dict(self):
|
|---|
| 83 | # intermediate variables for file table-metadata dict generation
|
|---|
| 84 | keys=self.GetPy_ColumnKeys()
|
|---|
| 85 | #offsets=self.GetPy_ColumnOffsets() #not needed on python level...
|
|---|
| 86 | nums=self.GetPy_ColumnNums()
|
|---|
| 87 | sizes=self.GetPy_ColumnSizes()
|
|---|
| 88 | types=self.GetPy_ColumnTypes()
|
|---|
| 89 | units=self.GetPy_ColumnUnits()
|
|---|
| 90 |
|
|---|
| 91 | # zip the values
|
|---|
| 92 | values = zip(nums,sizes,types,units)
|
|---|
| 93 | # create the columns dictionary
|
|---|
| 94 | columns = dict(zip(keys ,values))
|
|---|
| 95 | return columns
|
|---|
| 96 |
|
|---|
| 97 | def stack(self, on=True):
|
|---|
| 98 | self.next()
|
|---|
| 99 | for col in self._registered_cols:
|
|---|
| 100 | if isinstance( self.dict[col], type(np.array('')) ):
|
|---|
| 101 | self.stacked_cols[col] = self.dict[col]
|
|---|
| 102 | else:
|
|---|
| 103 | # elif isinstance(self.dict[col], ctypes._SimpleCData):
|
|---|
| 104 | self.stacked_cols[col] = np.array(self.dict[col])
|
|---|
| 105 | # else:
|
|---|
| 106 | # raise TypeError("I don't know how to stack "+col+". It is of type: "+str(type(self.dict[col])))
|
|---|
| 107 |
|
|---|
| 108 | def register(self, input_str):
|
|---|
| 109 | columns = self.columns
|
|---|
| 110 | if input_str.lower() == 'all':
|
|---|
| 111 | for col in columns:
|
|---|
| 112 | self._register(col)
|
|---|
| 113 | else:
|
|---|
| 114 | #check if colname is in columns:
|
|---|
| 115 | if input_str not in columns:
|
|---|
| 116 | error_msg = 'colname:'+ input_str +' is not a column in the binary table.\n'
|
|---|
| 117 | error_msg+= 'possible colnames are\n'
|
|---|
| 118 | for key in columns:
|
|---|
| 119 | error_msg += key+'\n'
|
|---|
| 120 | raise KeyError(error_msg)
|
|---|
| 121 | else:
|
|---|
| 122 | self._register(input_str)
|
|---|
| 123 |
|
|---|
| 124 | # 'private' method, do not use
|
|---|
| 125 | def _register( self, colname):
|
|---|
| 126 | columns = self.columns
|
|---|
| 127 | local = None
|
|---|
| 128 |
|
|---|
| 129 | number_of_elements = int(columns[colname][0])
|
|---|
| 130 | size_of_elements_in_bytes = int(columns[colname][1])
|
|---|
| 131 | ctypecode_of_elements = columns[colname][2]
|
|---|
| 132 | physical_unit_of_elements = columns[colname][3]
|
|---|
| 133 |
|
|---|
| 134 | # snippet from the C++ source code, or header file to be precise:
|
|---|
| 135 | #case 'L': gLog << "bool(8)"; break;
|
|---|
| 136 | #case 'B': gLog << "byte(8)"; break;
|
|---|
| 137 | #case 'I': gLog << "short(16)"; break;
|
|---|
| 138 | #case 'J': gLog << "int(32)"; break;
|
|---|
| 139 | #case 'K': gLog << "int(64)"; break;
|
|---|
| 140 | #case 'E': gLog << "float(32)"; break;
|
|---|
| 141 | #case 'D': gLog << "double(64)"; break;
|
|---|
| 142 |
|
|---|
| 143 |
|
|---|
| 144 |
|
|---|
| 145 | # the fields inside the columns can either contain single numbers,
|
|---|
| 146 | # or whole arrays of numbers as well.
|
|---|
| 147 | # we treat single elements differently...
|
|---|
| 148 | if number_of_elements == 1:
|
|---|
| 149 | # allocate some memory for a single number according to its type
|
|---|
| 150 | if ctypecode_of_elements == 'J': # J is for a 4byte int, i.e. an unsigned long
|
|---|
| 151 | local = ctypes.c_ulong()
|
|---|
| 152 | un_c_type = long
|
|---|
| 153 | elif ctypecode_of_elements == 'I': # I is for a 2byte int, i.e. an unsinged int
|
|---|
| 154 | local = ctypes.c_ushort()
|
|---|
| 155 | un_c_type = int
|
|---|
| 156 | elif ctypecode_of_elements == 'B': # B is for a byte
|
|---|
| 157 | local = ctypes.c_ubyte()
|
|---|
| 158 | un_c_type = int
|
|---|
| 159 | elif ctypecode_of_elements == 'D':
|
|---|
| 160 | local = ctypes.c_double()
|
|---|
| 161 | un_c_type = float
|
|---|
| 162 | elif ctypecode_of_elements == 'E':
|
|---|
| 163 | local = ctypes.c_float()
|
|---|
| 164 | un_c_type = float
|
|---|
| 165 | elif ctypecode_of_elements == 'A':
|
|---|
| 166 | local = ctypes.c_uchar()
|
|---|
| 167 | un_c_type = chr
|
|---|
| 168 | elif ctypecode_of_elements == 'K':
|
|---|
| 169 | local = ctypes.c_ulonglong()
|
|---|
| 170 | un_c_type = long
|
|---|
| 171 | else:
|
|---|
| 172 | raise TypeError('unknown ctypecode_of_elements:',ctypecode_of_elements)
|
|---|
| 173 | else:
|
|---|
| 174 | if ctypecode_of_elements == 'B': # B is for a byte
|
|---|
| 175 | nptype = np.int8
|
|---|
| 176 | elif ctypecode_of_elements == 'A': # A is for a char .. but I don't know how to handle it
|
|---|
| 177 | nptype = np.int8
|
|---|
| 178 | elif ctypecode_of_elements == 'I': # I is for a 2byte int
|
|---|
| 179 | nptype = np.int16
|
|---|
| 180 | elif ctypecode_of_elements == 'J': # J is for a 4byte int
|
|---|
| 181 | nptype = np.int32
|
|---|
| 182 | elif ctypecode_of_elements == 'K': # B is for a byte
|
|---|
| 183 | nptype = np.int64
|
|---|
| 184 | elif ctypecode_of_elements == 'E': # B is for a byte
|
|---|
| 185 | nptype = np.float32
|
|---|
| 186 | elif ctypecode_of_elements == 'D': # B is for a byte
|
|---|
| 187 | nptype = np.float64
|
|---|
| 188 | else:
|
|---|
| 189 | raise TypeError('unknown ctypecode_of_elements:',ctypecode_of_elements)
|
|---|
| 190 | local = np.zeros( number_of_elements, nptype)
|
|---|
| 191 |
|
|---|
| 192 | # Set the Pointer Address
|
|---|
| 193 | self.SetPtrAddress(colname, local)
|
|---|
| 194 | self._table_cols[colname] = local
|
|---|
| 195 | if number_of_elements > 1:
|
|---|
| 196 | self.__dict__[colname] = local
|
|---|
| 197 | self.dict[colname] = local
|
|---|
| 198 | else:
|
|---|
| 199 | # remove any traces of ctypes:
|
|---|
| 200 | self.__dict__[colname] = local.value
|
|---|
| 201 | self.dict[colname] = local.value
|
|---|
| 202 | self._registered_cols.append(colname)
|
|---|
| 203 |
|
|---|
| 204 |
|
|---|
| 205 | def treat_meta_dict(self):
|
|---|
| 206 | """make 'interesting' meta information available like normal members.
|
|---|
| 207 | non interesting are:
|
|---|
| 208 | TFORM, TUNIT, and TTYPE
|
|---|
| 209 | since these are available via the columns dict.
|
|---|
| 210 | """
|
|---|
| 211 |
|
|---|
| 212 | self.number_of_rows = self.meta['NAXIS2'][0]
|
|---|
| 213 | self.number_of_columns = self.meta['TFIELDS'][0]
|
|---|
| 214 |
|
|---|
| 215 | # there are some information in the meta dict, which are alsways there:
|
|---|
| 216 | # there are regarded as not interesting:
|
|---|
| 217 | uninteresting_meta = {}
|
|---|
| 218 | uninteresting_meta['arraylike'] = {}
|
|---|
| 219 | uninteresting = ['NAXIS', 'NAXIS1', 'NAXIS2',
|
|---|
| 220 | 'TFIELDS',
|
|---|
| 221 | 'XTENSION','EXTNAME','EXTREL',
|
|---|
| 222 | 'BITPIX', 'PCOUNT', 'GCOUNT',
|
|---|
| 223 | 'ORIGIN',
|
|---|
| 224 | 'PACKAGE', 'COMPILED', 'CREATOR',
|
|---|
| 225 | 'TELESCOP','TIMESYS','TIMEUNIT','VERSION']
|
|---|
| 226 | for key in uninteresting:
|
|---|
| 227 | if key in self.meta:
|
|---|
| 228 | uninteresting_meta[key]=self.meta[key]
|
|---|
| 229 | del self.meta[key]
|
|---|
| 230 |
|
|---|
| 231 | # the table meta data contains
|
|---|
| 232 |
|
|---|
| 233 | # loop over keys:
|
|---|
| 234 | # * try to find array-like keys
|
|---|
| 235 | arraylike = {}
|
|---|
| 236 | singlelike = []
|
|---|
| 237 | for key in self.meta:
|
|---|
| 238 | stripped = key.rstrip('1234567890')
|
|---|
| 239 | if stripped == key:
|
|---|
| 240 | singlelike.append(key)
|
|---|
| 241 | else:
|
|---|
| 242 | if stripped not in arraylike:
|
|---|
| 243 | arraylike[stripped] = 0
|
|---|
| 244 | else:
|
|---|
| 245 | arraylike[stripped] += 1
|
|---|
| 246 | newmeta = {}
|
|---|
| 247 | for key in singlelike:
|
|---|
| 248 | newmeta[key.lower()] = meta[key]
|
|---|
| 249 | for key in arraylike:
|
|---|
| 250 | uninteresting_meta['arraylike'][key.lower()] = []
|
|---|
| 251 | for i in range(arraylike[key]+1):
|
|---|
| 252 | if key+str(i) in meta:
|
|---|
| 253 | uninteresting_meta['arraylike'][key.lower()].append(meta[key+str(i)])
|
|---|
| 254 | self.ui_meta = uninteresting_meta
|
|---|
| 255 | # make newmeta self
|
|---|
| 256 | for key in newmeta:
|
|---|
| 257 | self.__dict__[key]=newmeta[key]
|
|---|
| 258 |
|
|---|
| 259 | dict = self.__dict__.copy()
|
|---|
| 260 | del dict['meta']
|
|---|
| 261 | del dict['ui_meta']
|
|---|
| 262 | self.dict = dict
|
|---|
| 263 |
|
|---|
| 264 | def __iter__(self):
|
|---|
| 265 | """ iterator """
|
|---|
| 266 | return self
|
|---|
| 267 |
|
|---|
| 268 | def next(self):
|
|---|
| 269 | """ used by __iter__ """
|
|---|
| 270 | # Here one might check, if looping makes any sense, and if not
|
|---|
| 271 | # one could stop looping or so...
|
|---|
| 272 | # like this:
|
|---|
| 273 | #
|
|---|
| 274 | # if len(self._registered_cols) == 0:
|
|---|
| 275 | # print 'warning: looping without any registered columns'
|
|---|
| 276 | if self._current_row < self.number_of_rows:
|
|---|
| 277 | if self.GetNextRow() == False:
|
|---|
| 278 | raise StopIteration
|
|---|
| 279 | for col in self._registered_cols:
|
|---|
| 280 | if isinstance(self._table_cols[col], ctypes._SimpleCData):
|
|---|
| 281 | self.__dict__[col] = self._table_cols[col].value
|
|---|
| 282 | self.dict[col] = self._table_cols[col].value
|
|---|
| 283 |
|
|---|
| 284 | for col in self.stacked_cols:
|
|---|
| 285 | if isinstance(self.dict[col], type(np.array(''))):
|
|---|
| 286 | self.stacked_cols[col] = np.vstack( (self.stacked_cols[col],self.dict[col]) )
|
|---|
| 287 | else:
|
|---|
| 288 | # elif isinstance(self.dict[col], ctypes._SimpleCData):
|
|---|
| 289 | self.stacked_cols[col] = np.vstack( (self.stacked_cols[col],np.array(self.dict[col])) )
|
|---|
| 290 | #else:
|
|---|
| 291 | #raise TypeError("I don't know how to stack "+col+". It is of type: "+str(type(self.dict[col])))
|
|---|
| 292 |
|
|---|
| 293 |
|
|---|
| 294 |
|
|---|
| 295 |
|
|---|
| 296 | self._current_row += 1
|
|---|
| 297 | else:
|
|---|
| 298 | raise StopIteration
|
|---|
| 299 | return self
|
|---|
| 300 |
|
|---|
| 301 | def show(self):
|
|---|
| 302 | pprint.pprint(self.dict)
|
|---|
| 303 |
|
|---|
| 304 | if __name__ == '__main__':
|
|---|
| 305 | import sys
|
|---|
| 306 | if len(sys.argv) == 1:
|
|---|
| 307 | print 'usage:', sys.argv[0], 'fits-file-name'
|
|---|
| 308 |
|
|---|
| 309 | file = FactFits(sys.argv[1])
|
|---|
| 310 | print '-'*70
|
|---|
| 311 | print "opened :", sys.argv[1], " as 'file'"
|
|---|
| 312 | print
|
|---|
| 313 | print '-'*70
|
|---|
| 314 | print 'type file.show() to look at its contents'
|
|---|
| 315 | print "type file.register( columnname ) or file.register('all') in order to register columns"
|
|---|
| 316 | print
|
|---|
| 317 | print " due column-registration you declare, that you would like to retrieve the contents of one of the columns"
|
|---|
| 318 | print " after column-registration, the 'file' has new member variables, they are named like the columns"
|
|---|
| 319 | print " PLEASE NOTE: immediatly after registration, the members exist, but they are empty."
|
|---|
| 320 | print " the values are assigned only, when you call file.next() or when you loop over the 'file'"
|
|---|
| 321 | print
|
|---|
| 322 | print "in order to loop over it, just go like this:"
|
|---|
| 323 | print "for row in file:"
|
|---|
| 324 | print " print row.columnname_one, row.columnname_two"
|
|---|
| 325 | print
|
|---|
| 326 | print ""
|
|---|
| 327 | print '-'*70
|
|---|
| 328 | |
|---|
