-
Notifications
You must be signed in to change notification settings - Fork 0
/
Data.py
227 lines (218 loc) · 9.1 KB
/
Data.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
from SaveLoadable import MetaSaveLoader
import os, fnmatch, re
from fitstools import *
import tarfile
from Exceptions import *
from numpy import mean
#class RawData(SaveLoadable):
class RawData(object):
"""A class that collect basic data infomation, and provide methods for unpacking and inspecting the zipped/tarred data files and saving the collected informations."""
__metaclass__ = MetaSaveLoader
TarFile=''
ObsID=''
def __init__(self, Location):
"""Way to loading the data using the path to the zipped/tarred file."""
if os.access(Location, os.R_OK):
#print 'File %s readable;' % (Location)
self.TarFile = Location
else:
raise FileError(Location)
self.rootpath = os.path.split(Location)[0]
self.inspect()
def unzip(self):
"""Function that unzip/untar the zipped/tarred data file."""
tar = tarfile.open(self.TarFile)
tar.extractall()
tar.close()
def inspect(self):
"""Function that inspect the data file and collect relative informations like target, telescope, instrument, mode, PI, etc."""
tar = tarfile.open(self.TarFile, 'r')
ObsID=[]
for tarinfo in tar:
FirstBlock = os.path.split(tarinfo.name)[0].split('/')[0]
try:
FirstBlock = int(FirstBlock)
if any([FirstBlock == ID for ID in ObsID]):pass
else: ObsID.append(FirstBlock)
except:
raise PatternError(FirstBlock, 'ObsID')
tar.close()
print 'Find Possible Observation IDs: ', ObsID
self.ObsID = ObsID
class chandraData(object):
__metaclass__ = MetaSaveLoader
"""A class that point to an unpacked chandra Data directory, and automaticly look up the important files such as the event list file."""
def __init__(self, Path):
if os.access(Path, os.R_OK):
#print 'File %s readable;' % (Location)
self.path = Path
if Path.split('/')[-1] == '':
LastPart = Path.split('/')[-2]
else:
LastPart = Path.split('/')[-1]
try:
LastPart = int(LastPart)
self.ObsID = LastPart
if os.path.exists(Path+'/primary/'):pass
else: FileError('/primary/')
except:
raise PatternError('Base name of the path %s' % (LastPart), 'ObsID')
else:
raise FileError(Path)
def findfile(file, dir, pattern, name):
if fnmatch.fnmatch(file, pattern):
if os.path.splitext(file)[1] == '.gz':
os.system('gunzip '+Path+dir+file)
self.__dict__[name] = Path+dir+file[:-3]
else:
self.__dict__[name] = Path+dir+file
return True
else: return False
for file in os.listdir(Path+'/primary/'):
if findfile(file, '/primary/', '*_evt2.fits*', 'evtfile'):print 'event list found.'
elif findfile(file, '/primary/', '*_asol*', 'asolfile'):print 'asolfile found'
else:pass
for file in os.listdir(Path+'/secondary/'):
if findfile(file, '/secondary/', '*_pbk*', 'pbkfile'):print 'pbk file found.'
else:pass
class region(object):
def __init__(self, regionstr):
p = re.compile('(?P<shape>circle|annulus|ellipse|box)\(\s*(?P<X>\d+\.*\d+)\s*,\s*(?P<Y>\d+\.*\d+)\s*,\s*(?P<R>\d+\.*\d+)(\s*,\s*(?P<Rout>\d+\.*\d+)(\s*,\s*(?P<Angle>\d+\.*\d+)){0,1}){0,1}\)', re.U)
for m in p.finditer(regionstr):
self.shape = m.group('shape')
self.X = float(m.group('X'))
self.Y = float(m.group('Y'))
self.R = float(m.group('R'))
try:
self.Rout = float(m.group('Rout'))
self.Angle = float(m.group('Angle'))
except:pass
def area(self):
PI = 3.14159265
if self.shape == 'circle':
return PI*self.R**2
if self.shape == 'annulus':
return PI*(self.Rout**2 - self.R**2)
if self.shape == 'box':
return self.R*self.Rout
if self.shape == 'ellipse':
return PI*self.R*self.Rout
def __str__(self):
if self.shape == 'ellipse' or self.shape == 'box':
return '%s(%f,%f,%f,%f,%f)' % (self.shape, self.X, self.Y, self.R, self.Rout, self.Angle)
elif self.shape == 'annulus':
return '%s(%f,%f,%f,%f)' % (self.shape, self.X, self.Y, self.R, self.Rout)
else:
return '%s(%f,%f,%f)' % (self.shape, self.X, self.Y, self.R)
def isinside(self, x, y):
if self.shape == 'circle':
if (x-self.X)**2 + (y-self.Y)**2 < self.R**2:
return True
else:
return False
elif self.shape == 'annulus':
rsquared = (x-self.X)**2 + (y-self.Y)**2
if rsquared < self.Rout**2 and rsquared > self.R**2:
return True
else:
return False
class psregion(object):
'''A class that provides APIs for accessing the region files of point sources.'''
def __init__(self, file=None, text=None):
#if not file[0] == '/':
#self.file = os.getcwd()+'/'+file
#else:
#self.file = file
if text == None:
text = open(file,'r').read()
#text = file
self.header = ''
self.regions = []
regionstr = ''
for lines in text.split('\n'):
if lines.find('(') == -1 or lines.find(')') == -1 or not lines.find('#') == -1:
self.header +=lines
else:
self.regions.append(region(lines))
#regionstr +=lines
def area(self):
return sum([x.area() for x in self.regions])
def __str__(self):
return self.header + '\n' + '\n'.join([str(x) for x in self.regions])
def save(self, file):
if file == None:file = self.file
f = open(file, 'w')
f.write(str(self))
def filter(self, table, X_index, Y_index):
newtable = []
for row in table:
x = row[X_index]
y = row[Y_index]
if any([reg.isinside(x, y) for reg in self.regions]):
newtable.append(row)
return newtable
def _tableandindices_(evtfile):
file = fitsfile(evtfile)
table = file.gettable()
cols = file.colinfo()
return (table, cols.names)
def improve_position(evtfile, region):
'''Improve the region file set around a point source by the mean position of all the events in the region.'''
(table, names) = _tableandindices_(evtfile)
(X_index, Y_index) = names.index('x'), names.index('y')
print 'Before improving position, the region reads:\n', str(region)
old_Radius = region.R
newtable = table
for region_enlarge_factor in [1.5, 1.]:
region.R = region_enlarge_factor * old_Radius
newtable = region.filter(newtable, X_index, Y_index)
region.X = mean([row[X_index] for row in newtable])
region.Y = mean([row[Y_index] for row in newtable])
print 'After improving position, the region reads:\n', str(region)
print 'But you still have to save this region using the save() method of region object.'
def filterevents(evtfile, Emin=0.3, Emax=10.0, region=None, outfile=None, cols=None, id='energy'):
'''Filter for event table or creat filtered file.'''
if not outfile == None: #outfile='src%s-%s.fits' % (Emin, Emax)
if region == None: regionfmt = ''
else:
regionfmt = '[sky=region(%s)]' % (region)
if cols == None: columnfmt = ''
else:
columnfmt = '[col %s]' % ','.join(cols)
if os.access(outfile,os.R_OK) and outfile.find('*') == -1: os.system('rm %s' % outfile)
cmdline = 'dmcopy "%(evtfile)s[%(id)s=%(Emin).0f:%(Emax).0f]%(region)s%(column)s" %(outfile)s clobber=yes ' % {
'evtfile':evtfile,
'id':id,
'Emin':Emin*1000,
'Emax':Emax*1000,
'region':regionfmt,
'column':columnfmt,
'outfile':outfile}
print cmdline
os.system(cmdline)
return fitsfile(outfile).gettable()
else:
(table, names) = _tableandindices_(evtfile)
try:
(X_index, Y_index, E_index) = [names.index(item) for item in ['x','y','energy']]
except:
(X_index, Y_index, E_index) = [names.index(item) for item in ['X','Y','ENERGY']]
if not region == None:
text = open(region, 'r').read()
psreg = psregion(text=text)
table = psreg.filter(table, X_index, Y_index)
newtable = []
for row in table:
if row[E_index] > Emin*1000 and row[E_index] < Emax*1000:
newtable.append(row)
del(table)
table = newtable
if cols == None: return table
else:
indics = [names.index(item) for item in cols]
newtable= []
for row in table:
newtable.append([row[i] for i in indics])
del(table)
table = newtable
return table