-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathquick_spectra.py
224 lines (194 loc) · 8.37 KB
/
quick_spectra.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
import numpy as np
import os
import matplotlib.pyplot as plt
from scio import scio
import argparse
import pytz
import datetime as dt
def _parse_slice(s):
a = [int(e) if e.strip() else None for e in s.split(":")]
return slice(*a)
def get_slice(s,acctime):
tstart, tstop, tstep = s.start, s.stop, s.step
if tstart is not None:
tstart = int(np.floor(tstart*60/acctime))
if tstop is not None:
tstop = int(np.floor(tstop*60/acctime))
if tstep is not None:
tstep = int(np.floor(tstep*60/acctime))
return slice(tstart, tstop, tstep)
def get_acctime(fpath):
dat = np.fromfile(fpath,dtype='uint32')
diff = np.diff(dat)
acctime = np.mean(diff[(diff>0)&(diff<100)]) #sometimes timestamps are 0, which causes diff to be huge. could also use np. median
return acctime
def get_vmin_vmax(data_arr):
'''
Automatically gets vmin and vmax for colorbar
'''
# print("shape of passed array", data_arr.shape, data_arr.dtype)
xx=np.ravel(data_arr).copy()
med = np.percentile(xx,50)
# print(med, "median")
u=np.percentile(xx,99)
b=np.percentile(xx,1)
xx_clean=xx[(xx<=u)&(xx>=b)] # remove some outliers for better plotting
stddev = np.std(xx_clean)
vmin= max(med - 2*stddev,10**7)
vmax = med + 2*stddev
if(vmin>vmax):
vmin=10**(np.log10(vmin)-1)
print(vmin,vmax)
# print("vmin, vmax are", vmin, vmax)
return vmin,vmax
if __name__ == "__main__":
"Example usage: python quick_spectra.py ~/data_auto_cross/16171/1617100000"
parser = argparse.ArgumentParser()
parser.add_argument("data_dir", type=str, help="Auto/cross-spectra location. Ex: ~/data_auto_cross/16171/161700000")
parser.add_argument("-o", "--output_dir", type=str, default="./", help="Output directory for plots")
parser.add_argument("-l", "--logplot", action="store_true", help="Plot in logscale")
parser.add_argument("-c", "--common", action="store_true", help="Common colorbar for both pols")
parser.add_argument("-s", "--show", action="store_true", help="Show final plot")
parser.add_argument("-tz", "--timezone", type=str, default='US/Eastern', help="Valid timezone of the telescope recognized by pytz. E.g. US/Eastern. Default is US/Eastern.")
parser.add_argument("-sl", "--tslice", type=_parse_slice, help="Slice on time axis to restrict plot to. Format: -sl=tmin:tmax for timin, tmax in minutes")
parser.add_argument("-vmi", "--vmin", dest='vmin', default = None, type=float, help="minimum for colorbar. if nothing is specified, vmin is automatically set")
parser.add_argument("-vma", "--vmax", dest='vmax', default = None, type=float, help="maximum for colorbar. if nothing is specified, vmax is automatically set")
parser.add_argument("-fma", "--fmax", dest='fmax', default = None, type=float, help="maximum for frequency to plot")
parser.add_argument("-fmi", "--fmin", dest='fmin', default = None, type=float, help="minimum for frequency to plot")
args = parser.parse_args()
#data_dir = pathlib.Path(args.data_dir)
#output_dir = pathlib.Path(args.output_dir)
pol00 = scio.read(os.path.join(args.data_dir, "pol00.scio.bz2"))
pol11 = scio.read(os.path.join(args.data_dir, "pol11.scio.bz2"))
pol01r = scio.read(os.path.join(args.data_dir, "pol01r.scio.bz2"))
pol01i = scio.read(os.path.join(args.data_dir, "pol01i.scio.bz2"))
acctime = get_acctime(os.path.join(args.data_dir, "time_gps_start.raw"))
fmin, fmax = 0, 125
if args.fmin:
fmin = args.fmin
if args.fmax:
fmax = args.fmax
cstart = int(np.floor(fmin/(250/4096)))
cend = int(np.floor(fmax/(250/4096)))
tstart = 0 #for myext below
if args.tslice is not None:
tstart=int(args.tslice.start)
if args.tslice:
#convert tslice in minutes to samps
tslice = get_slice(args.tslice, acctime)
pol00 = pol00[tslice, cstart:cend]
pol11 = pol11[tslice, cstart:cend]
pol01r = pol01r[tslice, cstart:cend]
pol01i = pol01i[tslice, cstart:cend]
else:
# Remove starting row since it's sometimes bad :(
pol00 = pol00[1:,cstart:cend]
pol11 = pol11[1:,cstart:cend]
pol01r = pol01r[1:,cstart:cend]
pol01i = pol01i[1:,cstart:cend]
# Add real and image for pol01
pol01 = pol01r + 1J*pol01i
freq = np.linspace(fmin,fmax, np.shape(pol00)[1])
pol00_med = np.median(pol00, axis=0)
pol11_med = np.median(pol11, axis=0)
pol00_mean = np.mean(pol00, axis=0)
pol11_mean = np.mean(pol11, axis=0)
pol00_max = np.max(pol00, axis=0)
pol11_max = np.max(pol11, axis=0)
pol00_min = np.min(pol00, axis=0)
pol11_min = np.min(pol11, axis=0)
if args.vmin is None and args.vmax is None:
vmin,vmax=get_vmin_vmax(pol00)
vmin2,vmax2=get_vmin_vmax(pol11)
else:
vmin=10**args.vmin
vmax=10**args.vmax
vmin2=vmin
vmax2=vmax
pmax = np.max(pol00)
axrange = [fmin, fmax, 0, pmax]
if(args.common):
vmin_common = min(vmin,vmin2)
vmax_common = max(vmax,vmax2)
vmin=vmin_common
vmax=vmax_common
vmin2=vmin_common
vmax2=vmax_common
if args.logplot == True:
pol00 = np.log10(pol00)
pol11 = np.log10(pol11)
pol00_med = np.log10(pol00_med)
pol11_med = np.log10(pol11_med)
pol00_mean = np.log10(pol00_mean)
pol11_mean = np.log10(pol11_mean)
pol00_max = np.log10(pol00_max)
pol11_max = np.log10(pol11_max)
pol00_min = np.log10(pol00_min)
pol11_min = np.log10(pol11_min)
vmin = np.log10(vmin)
vmax = np.log10(vmax)
vmin2 = np.log10(vmin2)
vmax2 = np.log10(vmax2)
pmax = np.log10(pmax)
axrange = [fmin, fmax, 6.5, pmax]
print("Estimated accumulation time from timestamp file: ", acctime)
tot_minutes = int(np.ceil(acctime * pol00.shape[0]/60))
myext = np.array([fmin,fmax, tstart+tot_minutes, tstart])
plt.figure(figsize=(18,10), dpi=200)
plt.subplot(2,3,1)
plt.imshow(pol00, vmin=vmin, vmax=vmax, aspect='auto', extent=myext)
plt.title('pol00')
cb00 = plt.colorbar()
plt.xlabel('Frequency (MHz)')
cb00.ax.plot([0, 1], [7.0]*2, 'w')
cb00.ax.set_ylabel('Uncalibrated log(Power)')
plt.subplot(2,3,4)
plt.imshow(pol11, vmin=vmin2, vmax=vmax2, aspect='auto', extent=myext)
plt.title('pol11')
cb=plt.colorbar()
cb.ax.set_ylabel('Uncalibrated log(Power)')
plt.xlabel('Frequency (MHz)')
plt.subplot(2,3,2)
plt.title('Basic stats for frequency bins')
plt.plot(freq, pol00_max, 'r-', label='Max')
plt.plot(freq, pol00_min, 'b-', label='Min')
plt.plot(freq, pol00_mean, 'k-', label='Mean')
plt.plot(freq, pol00_med, color='#666666', linestyle='-', label='Median')
plt.xlabel('Frequency (MHz)')
plt.ylabel('pol00')
plt.ylim(vmin,vmax)
plt.xlim(fmin, fmax)
plt.subplot(2,3,5)
plt.plot(freq, pol11_max, 'r-', label='Max')
plt.plot(freq, pol11_min, 'b-', label='Min')
plt.plot(freq, pol11_mean, 'k-', label='Mean')
plt.plot(freq, pol11_med, color='#666666', linestyle='-', label='Median')
plt.xlabel('Frequency (MHz)')
plt.ylabel('pol11')
plt.ylim(vmin2,vmax2)
plt.xlim(fmin, fmax)
plt.legend(loc='lower right', fontsize='small')
plt.subplot(2,3,3)
plt.imshow(np.log10(np.abs(pol01)), vmin=3, vmax=8, aspect='auto', extent=myext)
plt.title('pol01 magnitude')
cb=plt.colorbar()
cb.ax.set_ylabel('Uncalibrated log(Power)')
plt.xlabel('Frequency (MHz)')
plt.subplot(2,3,6)
plt.imshow(np.angle(pol01), vmin=-np.pi, vmax=np.pi, aspect='auto', extent=myext, cmap='RdBu')
plt.title('pol01 phase')
cb=plt.colorbar()
cb.ax.set_ylabel('Phase (rad)')
plt.xlabel('Frequency (MHz)')
args.data_dir=os.path.abspath(args.data_dir)
timestamp = args.data_dir.split('/')[-1]
mytz = pytz.timezone(args.timezone)
utctime = dt.datetime.fromtimestamp(int(timestamp),tz=pytz.utc) # this might be a safer way compared to passing tz directly to fromtimestamp
localtimestr = utctime.astimezone(tz=mytz).strftime("%b-%d %H:%M:%S")
plt.suptitle(f"Minutes since {localtimestr} localtime. File ctime {timestamp}")
outfile = os.path.normpath(args.output_dir + '/' + timestamp + '_quick' + '.png')
plt.savefig(outfile)
print('Wrote ' + outfile)
if args.show == True:
plt.show()
plt.close()