/
test_turbo_seti.py
261 lines (213 loc) · 8.56 KB
/
test_turbo_seti.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
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
"""
Testing for turbo_seti various functions.
TODO: more description of actual coverage versus ideal coverage.
"""
import os
import time
import pylab as plt
import numpy as np
import pytest
from blimpy import Waterfall
from turbo_seti import FindDoppler, seti_event
from turbo_seti import find_event, plot_event
import pyximport
pyximport.install(setup_args={"include_dirs":np.get_include()}, reload_support=True)
from turbo_seti.find_doppler import data_handler, helper_functions, taylor_tree
HERE = os.path.split(os.path.abspath(__file__))[0]
VOYAH5 = 'Voyager1.single_coarse.fine_res.h5'
VOYAH5FLIPPED = 'Voyager1.single_coarse.fine_res.flipped.h5'
def find_doppler(filename_fil):
""" Run turboseti doppler search on filename with default params """
t0 = time.time()
print("\n===== find_doppler =====")
print("Searching %s" % filename_fil)
filename_dat = filename_fil.replace('.h5', '.dat')
filename_log = filename_fil.replace('.h5', 'log')
if os.path.exists(filename_dat):
os.remove(filename_dat)
if os.path.exists(filename_log):
os.remove(filename_log)
snr = 5
coarse_chans = ''
obs_info = None
n_coarse_chan = 1
max_drift = 1.0
find_seti_event = FindDoppler(filename_fil, max_drift=max_drift, snr=snr, out_dir=HERE,
coarse_chans=coarse_chans, obs_info=obs_info, n_coarse_chan=n_coarse_chan)
find_seti_event.search()
t_taken = time.time() - t0
print("Time taken: %2.2fs" % t_taken)
def plot_hit(fil_filename, dat_filename, hit_id, bw=None, offset=0):
"""Plot a candidate from a .dat file
Args:
fil_filename(str): Path to filterbank file to plot
dat_filename(str): Path to turbosSETI generated .dat output file of events
hit_id(int): ID of hit in the dat file to plot (TopHitNum)
offset(float, optional): Offset drift line on plot. Default 0.
bw: (Default value = None)
Returns:
"""
# Load hit details
dat = find_event.make_table(dat_filename)
hit = dat.iloc[hit_id]
f0 = hit['Freq']
if bw is None:
bw_mhz = np.abs(hit['FreqStart'] - hit['FreqEnd'])
else:
bw_mhz = bw * 1e-6
fil = Waterfall(fil_filename, f_start=f0 - bw_mhz / 2, f_stop=f0 + bw_mhz / 2)
t_duration = (fil.n_ints_in_file - 1) * fil.header['tsamp']
fil.plot_waterfall()
plot_event.overlay_drift(f0, f0, f0, hit['DriftRate'], t_duration, offset)
def plot_hits(filename_fil, filename_dat):
""" Plot the hits in a .dat file. """
print("\n===== plot_hits =====")
table = find_event.read_dat(filename_dat)
print(table)
plt.figure(figsize=(10, 8))
N_hit = len(table)
if N_hit > 10:
print("Warning: More than 10 hits found. Only plotting first 10")
N_hit = 10
for ii in range(N_hit):
plt.subplot(N_hit, 1, ii+1)
plot_hit(filename_fil, filename_dat, ii)
plt.tight_layout()
plt.savefig(filename_dat.replace('.dat', '.png'))
def validate_voyager_hits(filename_dat):
""" This checks voyager hits against known values.
Known values:
# --------------------------
# ID Drift_Rate SNR Unc_Freq Corr_Freq Index freq_start freq_end ...
# --------------------------
001 -0.392226 30.612128 8419.319368 8419.319368 739933 8419.319779 8419.318963 ...
002 -0.373093 245.707984 8419.297028 8419.297028 747929 8419.297439 8419.296623 ...
003 -0.392226 31.220652 8419.274374 8419.274374 756037 8419.274785 8419.273969 ...
(from flipped)
003 -0.392226 30.612118 8419.319366 8419.319366 308642 8419.318955 8419.319771 ...
002 -0.373093 245.707905 8419.297025 8419.297025 300646 8419.296614 8419.297430 ...
001 -0.392226 31.220642 8419.274372 8419.274372 292538 8419.273961 8419.274777 ...
"""
print("\n===== validate_voyager_hits =====")
h = find_event.read_dat(filename_dat)
print(h)
valid_data = [
{
'Freq': 8419.319368,
'FreqStart': 8419.319779,
'FreqEnd': 8419.318963,
'SNR': 30.612128,
'DriftRate': -0.392226,
},
{
'Freq': 8419.297028,
'FreqStart': 8419.297439,
'FreqEnd': 8419.296623,
'SNR': 245.707984,
'DriftRate': -0.373093,
},
{
'Freq': 8419.274374,
'FreqStart': 8419.274785,
'FreqEnd': 8419.273969,
'SNR': 31.220652,
'DriftRate': -0.392226,
}
]
atols = {'Freq': 0.000005, 'FreqStart': 0.00001, 'FreqEnd': 0.00001, 'SNR': 0.001, 'DriftRate': 0.02}
for vd in valid_data:
hmax = h[np.isclose(h['Freq'], vd['Freq'], rtol=0.000001)].iloc[0]
for key in vd.keys():
print(key, hmax[key], vd[key])
if key in ('FreqStart', 'FreqEnd'):
upper = np.isclose(hmax[key], vd['FreqStart'], atol=atols[key], rtol=0)
lower = np.isclose(hmax[key], vd['FreqEnd'], atol=atols[key], rtol=0)
assert upper or lower
else:
assert np.isclose(hmax[key], vd[key], atol=atols[key], rtol=0)
return hmax
def test_find_doppler_voyager():
""" Run turboseti on Voyager data """
print("\n===== test_find_doppler_voyager =====")
filename_fil = os.path.join(HERE, VOYAH5)
filename_dat = filename_fil.replace('.h5', '.dat')
find_doppler(filename_fil)
validate_voyager_hits(filename_dat)
plot_hits(filename_fil, filename_dat)
def test_find_doppler_voyager_flipped():
""" Run turboseti on Voyager data (flipped in frequency) """
print("\n===== test_find_doppler_voyager_flipped =====")
filename_fil = os.path.join(HERE, VOYAH5FLIPPED)
filename_dat = filename_fil.replace('.h5', '.dat')
find_doppler(filename_fil)
validate_voyager_hits(filename_dat)
plot_hits(filename_fil, filename_dat)
def test_find_doppler_voyager_filterbank():
""" Run turboseti on Voyager data (filterbank version) """
print("\n===== test_find_doppler_voyager_filterbank =====")
filename_fil = os.path.join(HERE, VOYAH5)
find_doppler(filename_fil)
def test_turboSETI_entry_point():
""" Test the command line utility turboSETI """
print("\n===== test_turboSETI_entry_point =====")
filename_fil = os.path.join(HERE, VOYAH5FLIPPED)
args = [filename_fil, ]
seti_event.main(args)
def test_make_waterfall_plots():
""" Some basic plotting tests
TODO: Improve these tests (and the functions for that matter!
"""
print("\n===== test_plotting =====")
filename_fil = os.path.join(HERE, VOYAH5)
# Test make_waterfall_plots -- needs 6x files
filenames_list = [filename_fil] * 6
target = 'Voyager'
drate = -0.392226
fmid = 8419.274785
f_start = 8419.274374 - 600e-6
f_stop = 8419.274374 + 600e-6
source_name_list = ['test_make_waterfall_plots'] * 6
filter_level = "1"
plot_event.make_waterfall_plots(filenames_list,
target,
f_start,
f_stop,
drate,
fmid,
filter_level,
source_name_list)
def test_data_handler():
""" Basic data handler test """
print("\n===== test_data_handler =====")
with pytest.raises(OSError): # not AttributeError
data_handler.DATAHandle(filename='made_up_not_existing_file.h5')
def test_dask():
""" Test dask capability on Voyager data """
filename_fil = os.path.join(HERE, VOYAH5)
FD = FindDoppler(datafile=filename_fil, max_drift=2, out_dir=HERE)
print("===== test_dask ===== n_partitions=None")
FD.search()
print("===== test_dask ===== n_partitions=2")
FD.search(n_partitions=2)
print("===== test_dask ===== n_partitions=2, progress_bar='n'")
FD.search(n_partitions=2, progress_bar='n')
print("===== test_dask ===== End")
def test_bitrev():
before = 32769
nbits = 7
out_c = taylor_tree.bitrev(before, nbits)
out_p = helper_functions.bitrev(before, nbits)
assert out_c == out_p
before = 32770
out_c = taylor_tree.bitrev(before, nbits)
out_p = helper_functions.bitrev(before, nbits)
assert out_c == out_p
if __name__ == "__main__":
test_make_waterfall_plots()
test_turboSETI_entry_point()
test_find_doppler_voyager()
test_find_doppler_voyager_flipped()
test_find_doppler_voyager_filterbank()
test_data_handler()
test_dask()
test_bitrev()