def draw_n(self, tchans, fchans, same, diff):
random_set = np.random.uniform((same)) *fchans/2
unique_set = np.random.uniform((diff)) *fchans/2
frame = stg.Frame(fchans=fchans*u.pixel,
tchans=tchans*u.pixel,
df=2.7939677238464355*u.Hz,
dt=18.25361108*u.s,
fch1=6095.214842353016*u.MHz)
for i in range (same):
frame = self.draw(frame=frame, seed=random_set[i], drift=0.5, noise_toggle=False)
for t in range (diff):
frame = self.draw(frame=frame, seed=unique_set[i], drift=0.5, noise_toggle=False)
frame2 = stg.Frame(fchans=fchans*u.pixel,
tchans=tchans*u.pixel,
df=2.7939677238464355*u.Hz,
dt=18.25361108*u.s,
fch1=6095.214842353016*u.MHz)
unique_set = np.random.uniform((diff))*fchans/2
for k in range (same):
frame2 = self.draw(frame=frame2, seed=random_set[k], drift=0.5, noise_toggle=False)
for h in range (diff):
frame2 = self.draw(frame=frame2, seed=unique_set[h], drift=0.5, noise_toggle=False)
return frame.get_data(), frame2.get_data()
def draw_2(self, fchans, tchans):
random_num = int(random()*fchans/2)
random_num_2 = int(random()*fchans/2)
frame = stg.Frame(fchans=fchans*u.pixel,
tchans=tchans*u.pixel,
df=2.7939677238464355*u.Hz,
dt=18.25361108*u.s,
fch1=6095.214842353016*u.MHz)
frame = self.draw(frame = frame, seed= random_num, drift=0.5, noise=True)
frame = self.draw(frame = frame, seed = random_num_2, drift=0.1, noise=False)
unique = int(random()*fchans/2)
driftrate_unique = int(random()*2)-1
frame = self.draw(frame = frame, seed=unique, drift=driftrate_unique, noise=False)
frame2 = stg.Frame(fchans=fchans*u.pixel,
tchans=tchans*u.pixel,
df=2.7939677238464355*u.Hz,
dt=18.25361108*u.s,
fch1=6095.214842353016*u.MHz)
frame2 = self.draw(frame = frame, seed= random_num, drift=0.5, noise=True)
frame2 = self.draw(frame = frame, seed = random_num_2, drift=0.1, noise=False)
unique = int(random()*fchans/2)
driftrate_unique = int(random()*2)-1
frame2 = self.draw(frame = frame, seed=unique, drift=driftrate_unique, noise=False)
return frame.get_data(),frame2.get_data()
def test_h5_io(frame_setup_no_data):
frame = copy.deepcopy(frame_setup_no_data)
fil_fn = 'temp.h5'
frame.save_hdf5(fil_fn)
temp_frame = stg.Frame(waterfall=fil_fn)
assert_allclose(temp_frame.get_data(), frame.get_data())
wf = bl.Waterfall(fil_fn)
temp_frame = stg.Frame(waterfall=wf)
assert_allclose(temp_frame.get_data(), frame.get_data())
os.remove(fil_fn)
def generate_frame(snr):
frame = stg.Frame(fchans=512,
tchans=16,
df=2.7939677238464355*u.Hz,
dt=18.25361108*u.s,
fch1=6095.214842353016*u.MHz)
frame.add_noise_from_obs(share_index=True)
noise_mean, noise_std = frame.get_noise_stats()
start_index = np.random.randint(0, frame.fchans)
end_index = np.random.randint(0, frame.fchans)
drift_rate = frame.get_drift_rate(start_index, end_index)
width = np.random.uniform(10, 30)
signal = frame.add_constant_signal(f_start=frame.get_frequency(start_index),
drift_rate=drift_rate*u.Hz/u.s,
level=frame.get_intensity(snr=snr),
width=width*u.Hz,
f_profile_type='gaussian')
frame_info = {
'noise_mean': noise_mean,
'noise_std': noise_std,
'snr': snr,
'start_index': start_index,
'end_index': end_index,
'width': width
}
return frame, frame_info
def test_setigen():
""" Generate a setigen frame and convert it into a DataArray"""
metadata = {'fch1': 6095.214842353016*u.MHz,
'dt': 18.25361108*u.s,
'df': 2.7939677238464355*u.Hz}
frame = stg.Frame(fchans=2**12*u.pixel,
tchans=128*u.pixel,
df=metadata['df'],
dt=metadata['dt'],
fch1=metadata['fch1'])
test_tones = [
{'f_start': frame.get_frequency(index=500), 'drift_rate': 0.01*u.Hz/u.s, 'snr': 100, 'width': 20*u.Hz},
{'f_start': frame.get_frequency(index=800), 'drift_rate': -0.10*u.Hz/u.s, 'snr': 100, 'width': 20*u.Hz},
{'f_start': frame.get_frequency(index=2048), 'drift_rate': 0.00*u.Hz/u.s, 'snr': 20, 'width': 6*u.Hz},
{'f_start': frame.get_frequency(index=3000), 'drift_rate': 0.07*u.Hz/u.s, 'snr': 50, 'width': 3*u.Hz}
]
noise = frame.add_noise(x_mean=100, x_std=5, noise_type='gaussian')
for tone in test_tones:
signal = frame.add_signal(stg.constant_path(f_start=tone['f_start'],
drift_rate=tone['drift_rate']),
stg.constant_t_profile(level=frame.get_intensity(snr=tone['snr'])),
stg.gaussian_f_profile(width=tone['width']),
stg.constant_bp_profile(level=1))
d = from_setigen(frame)
def frame_setup_no_data():
frame = stg.Frame(fchans=1024 * u.pixel,
tchans=32 * u.pixel,
df=2.7939677238464355 * u.Hz,
dt=18.253611008 * u.s,
fch1=6095.214842353016 * u.MHz)
return frame
def test_normalize_mask(plot=False):
""" Test normalization works when data are masked """
n_int, n_ifs, n_chan = 16, 1, 8192
np.random.seed(1234)
frame = stg.Frame(fchans=1024 * u.pixel,
tchans=32 * u.pixel,
df=2.7939677238464355 * u.Hz,
dt=18.253611008 * u.s,
fch1=6095.214842353016 * u.MHz)
noise = frame.add_noise(x_mean=10, noise_type='chi2')
signal = frame.add_signal(
stg.constant_path(f_start=frame.get_frequency(index=200),
drift_rate=1 * u.Hz / u.s),
stg.constant_t_profile(level=frame.get_intensity(snr=1000)),
stg.gaussian_f_profile(width=20 * u.Hz),
stg.constant_bp_profile(level=1))
d = from_setigen(frame)
d.data = d.data.astype('float32')
if plot:
frame.plot()
# Over 20% of data will be masked
mask = d.data.mean(axis=0).squeeze() > 15
d_norm = normalize(d, mask=mask, return_space='cpu')
#print(d_norm.mean(), d_norm.std())
plt.figure()
if plot:
plt.plot(d_norm.mean(axis=0).squeeze())
mask_cpu = cp.asnumpy(mask)
assert np.isclose(d_norm[..., ~mask_cpu].mean(), 0, atol=1e-6)
assert np.isclose(d_norm[..., ~mask_cpu].std(), 1.0, atol=1e-6)
print("Mask test passed!")
def frame_setup_no_data_ascending():
frame = stg.Frame(fchans=1024 * u.pixel,
tchans=32 * u.pixel,
df=2.7939677238464355 * u.Hz,
dt=18.253611008 * u.s,
fch1=6095211984.124035,
ascending=True)
return frame
def test_signal_from_floats():
frame = stg.Frame(tchans=3, fchans=3, dt=1, df=1, fch1=3)
frame.add_signal(path=3,
t_profile=2,
f_profile=stg.box_f_profile(width=1),
bp_profile=3)
data = np.array([[0., 0., 6.], [0., 0., 6.], [0., 0., 6.]])
assert_allclose(data, frame.get_data(use_db=False))
def test_signal_from_arrays():
frame = stg.Frame(tchans=3, fchans=3, dt=1, df=1, fch1=3)
frame.add_signal(path=[2, 1, 3],
t_profile=[1, 0.5, 1],
f_profile=stg.box_f_profile(width=1),
bp_profile=[1, 0.5, 1])
data = np.array([[0., 0.5, 0.], [0.5, 0., 0.], [0., 0., 1.]])
assert_allclose(data, frame.get_data(use_db=False))
def generate(self, total_num_samples, data, intensity = 0.7, test=False, labels= True):
start_time = time.time()
num_channels = 32
# Prepare the target set
fchans = num_channels
tchans = 32
df = 2.7939677238464355*u.Hz
dt = 18.25361108*u.s
fch1 = 6095.214842353016*u.MHz
tune_base = intensity
generated_signals = np.zeros((total_num_samples,data.shape[1],1, num_channels), dtype=float)
for i in range(0,total_num_samples):
base = data[i,:,0,:]
start = int(random()*(fchans-15))
artifical_radio = np.zeros((32,32))
period = random()
drifrate = (random()-0.5)*10**int(random()*-2) *u.Hz/u.s
choose = random()
amplitude = random()
tune = tune_base * (random()+0.3)
for blur in range(0,2):
frame = stg.Frame(fchans=fchans,
tchans=tchans,
df=df,
dt=dt,
fch1=fch1)
signal = frame.add_signal(stg.constant_path(f_start=frame.fs[start+10],
drift_rate=drifrate),
stg.constant_t_profile(level=tune*blur),
stg.gaussian_f_profile(width=(20-5*blur)*u.Hz),
stg.constant_bp_profile(level=tune*blur))
artifical_radio = np.add(frame.get_data(), artifical_radio)
artifical_radio = np.add(base, artifical_radio)
artifical_radio = np.reshape(artifical_radio, (tchans,1,fchans))
generated_signals[i,:,:,:]= artifical_radio
if i %int(total_num_samples/2) ==0 and test:
fig = plt.figure(figsize=(10, 6))
plt.imshow(generated_signals[i,:,0,:], aspect='auto')
plt.colorbar()
# Label the dataset
supervised_true = np.concatenate((np.ones((generated_signals.shape[0],1),dtype='int64'),np.zeros((generated_signals.shape[0],1),dtype='int64')), axis=1)
supervised_false = np.concatenate((np.zeros((total_num_samples,1),dtype='int64'),np.ones((total_num_samples,1),dtype='int64')), axis=1)
label = np.concatenate((supervised_true, supervised_false))
supervised_dataset = np.concatenate((generated_signals, data[0:total_num_samples,:,:,:]))
print(label.shape)
print(supervised_dataset.shape)
print("Synethtic Generation Execution Time: "+ str(time.time()-start_time))
X_train_supervised, X_test_supervised, y_train_supervised, y_test_supervised = train_test_split(supervised_dataset, label, test_size=0.2, random_state=2)
return X_train_supervised, X_test_supervised, y_train_supervised, y_test_supervised
def generate_frame(sig_db, rfi_num=0):
"""
Create frame with chi-squared synthetic noise, with 1 drifting signal
and either 0 or 1 non-drifting 'RFI' signals.
"""
frame = stg.Frame(fchans=1024,
tchans=32,
df=1.3969838619232178*u.Hz,
dt=1.4316557653333333*u.s,
fch1=6095.214842353016*u.MHz)
frame.add_noise_from_obs()
noise_mean, noise_std = frame.get_noise_stats()
start_index = np.random.randint(0, frame.fchans)
end_index = np.random.randint(0, frame.fchans)
drift_rate = frame.get_drift_rate(start_index, end_index)
width = np.random.uniform(1, 30)
frame.add_constant_signal(f_start=frame.get_frequency(start_index),
drift_rate=drift_rate*u.Hz/u.s,
level=frame.get_intensity(snr=db_to_snr(sig_db)),
width=width*u.Hz,
f_profile_type='gaussian')
if rfi_num == 1:
rfi_snr = db_to_snr(25)
rfi_start_index = rfi_end_index = np.random.randint(0, frame.fchans)
rfi_width = np.random.uniform(1, 30)
frame.add_constant_signal(f_start=frame.get_frequency(rfi_start_index),
drift_rate=0*u.Hz/u.s,
level=frame.get_intensity(snr=rfi_snr),
width=rfi_width*u.Hz,
f_profile_type='gaussian')
else:
rfi_start_index = rfi_end_index = -1
frame_info = {
'noise_mean': noise_mean,
'noise_std': noise_std,
'sig_db': sig_db,
'start_index': start_index,
'end_index': end_index,
'width': width,
'rfi_num': rfi_num,
'rfi_start_index': rfi_start_index,
'rfi_end_index': rfi_end_index,
}
return frame, frame_info
def test_hitsearch_multi():
""" Test hit search routine with multiple signals """
metadata = {'fch1': 6095.214842353016*u.MHz,
'dt': 18.25361108*u.s,
'df': 2.7939677238464355*u.Hz}
frame = stg.Frame(fchans=2**12*u.pixel,
tchans=32*u.pixel,
df=metadata['df'],
dt=metadata['dt'],
fch1=metadata['fch1'])
test_tones = [
{'f_start': frame.get_frequency(index=500), 'drift_rate': 0.50*u.Hz/u.s, 'snr': 100, 'width': 20*u.Hz},
{'f_start': frame.get_frequency(index=800), 'drift_rate': -0.40*u.Hz/u.s, 'snr': 100, 'width': 20*u.Hz},
{'f_start': frame.get_frequency(index=2048), 'drift_rate': 0.00*u.Hz/u.s, 'snr': 20, 'width': 6*u.Hz},
{'f_start': frame.get_frequency(index=3000), 'drift_rate': 0.07*u.Hz/u.s, 'snr': 50, 'width': 3*u.Hz}
]
frame.add_noise(x_mean=0, x_std=5, noise_type='gaussian')
for tone in test_tones:
frame.add_signal(stg.constant_path(f_start=tone['f_start'],
drift_rate=tone['drift_rate']),
stg.constant_t_profile(level=frame.get_intensity(snr=tone['snr'])),
stg.gaussian_f_profile(width=tone['width']),
stg.constant_bp_profile(level=1))
frame.save_fil(filename=synthetic_fil)
darray = from_fil(synthetic_fil)
fig = plt.figure(figsize=(10, 6)) # <============ fig is UNUSED
dedopp, md, hits = run_pipeline(darray.data, metadata, max_dd=1.0, min_dd=None, threshold=100,
n_boxcar=5, merge_boxcar_trials=True)
print(hits.sort_values('snr', ascending=False))
plt.figure(figsize=(10, 4))
plt.subplot(1,2,1)
imshow_waterfall(darray.data, md, 'channel', 'timestep')
plt.subplot(1,2,2)
imshow_dedopp(dedopp, md, 'channel', 'driftrate')
overlay_hits(hits, 'channel', 'driftrate')
plt.savefig('docs/figs/test_hitsearch_multi.png')
plt.show()
def generate_fil_file(outpath, flag_fascending, flag_sign_drift_rate):
r'''
Using setigen, generate a filterbank file.
Parameters:
outpath - full path of where to store the resultant filterbank file.
flag_fascending - use an ascending (+1) or descending (-1) sequence of frequencies
flag_sign_drift_rate - use a positive (+1) or negative (-1) drift rate
'''
if DEBUGGING:
print('generate_fil_file: flag_fascending={}, flag_sign_drift_rate={}'.
format(flag_fascending, flag_sign_drift_rate))
# Set up setigne parameters
stg_parms = SetigenParms()
if flag_sign_drift_rate < 0:
stg_parms.drift_rate_1 = -stg_parms.drift_rate_1
stg_parms.drift_rate_2 = -stg_parms.drift_rate_2
stg_parms.drift_rate_3 = -stg_parms.drift_rate_3
stg_parms.drift_rate_4 = -stg_parms.drift_rate_4
stg_parms.drift_rate_5 = -stg_parms.drift_rate_5
# Instantiate a setigen Frame object
frame = stg.Frame(fchans=stg_parms.fchans,
tchans=stg_parms.tchans,
df=stg_parms.df,
dt=stg_parms.dt,
fch1=stg_parms.fch1,
ascending=(flag_fascending > 0))
# Add noise to stg object.
frame.add_noise(x_mean=0, x_std=stg_parms.noise_std, noise_type='gaussian')
# Signal 1 will be detected.
signal_1_intensity = frame.get_intensity(snr=stg_parms.snr_1)
frame.add_constant_signal(f_start=frame.get_frequency(
stg_parms.signal_start_1),
drift_rate=stg_parms.drift_rate_1,
level=signal_1_intensity,
width=stg_parms.width_1,
f_profile_type='gaussian')
# Signal 2 will be detected.
signal_2_intensity = frame.get_intensity(snr=stg_parms.snr_2)
frame.add_constant_signal(f_start=frame.get_frequency(
stg_parms.signal_start_2),
drift_rate=stg_parms.drift_rate_2,
level=signal_2_intensity,
width=stg_parms.width_2,
f_profile_type='gaussian')
# Signal 3 is a symmetric signal with three Gaussians
# that will fall below the SNR requirements.
signal_3_intensity = frame.get_intensity(snr=stg_parms.snr_3)
frame.add_signal(
stg.constant_path(f_start=frame.get_frequency(
stg_parms.signal_start_3),
drift_rate=stg_parms.drift_rate_3),
stg.constant_t_profile(level=1),
stg.multiple_gaussian_f_profile(width=stg_parms.width_3),
stg.constant_bp_profile(level=signal_3_intensity))
# Signal 4 is a symmetric signal with three Gaussians
# that will be drifting too quickly.
signal_4_intensity = frame.get_intensity(snr=stg_parms.snr_4)
frame.add_signal(
stg.constant_path(f_start=frame.get_frequency(
stg_parms.signal_start_4),
drift_rate=stg_parms.drift_rate_4),
stg.constant_t_profile(level=1),
stg.multiple_gaussian_f_profile(width=stg_parms.width_4),
stg.constant_bp_profile(level=signal_4_intensity))
# Signal 5 is similar to signal 4 but drifting in the opposite direction.
signal_5_intensity = frame.get_intensity(snr=stg_parms.snr_5)
frame.add_signal(
stg.constant_path(f_start=frame.get_frequency(
stg_parms.signal_start_5),
drift_rate=stg_parms.drift_rate_5),
stg.constant_t_profile(level=1),
stg.multiple_gaussian_f_profile(width=stg_parms.width_5),
stg.constant_bp_profile(level=signal_5_intensity))
# Save the frame as a filterbank file.
frame.save_fil(filename=outpath)
print("generate_fil_file: generated {}".format(outpath))
del frame
def test_dedoppler():
""" Basic tests of the dedoppler functionality """
# zero drift test, no normalization
test_data = np.ones(shape=(32, 1, 1024), dtype='float32')
test_data[:, :, 511] = 10
metadata_in = {
'frequency_start': 1000 * u.MHz,
'time_step': 1.0 * u.s,
'frequency_step': 1.0 * u.Hz
}
dedopp, metadata = dedoppler(test_data,
metadata_in,
boxcar_size=1,
max_dd=1.0)
print("type(dedopp):", type(dedopp))
print("dedopp.data:", dedopp.data)
print("np.max(dedopp.data):", np.max(dedopp.data),
", np.sum(test_data[:, :, 511]):", np.sum(test_data[:, :, 511]))
assert np.max(dedopp.data) == np.sum(test_data[:, :, 511])
for dr_test in (0.0, 0.1, 0.5, -0.25, -0.5):
# single drifting tone
frame = stg.Frame(fchans=2**10 * u.pixel,
tchans=32 * u.pixel,
df=metadata_in['frequency_step'],
dt=metadata_in['time_step'],
fch1=metadata_in['frequency_start'])
tone = {
'f_start': frame.get_frequency(index=500),
'drift_rate': dr_test * u.Hz / u.s,
'snr': 500,
'width': metadata_in['frequency_step']
}
frame.add_noise(x_mean=1, noise_type='chi2')
frame.add_signal(
stg.constant_path(f_start=tone['f_start'],
drift_rate=tone['drift_rate']),
stg.constant_t_profile(level=frame.get_intensity(snr=tone['snr'])),
stg.gaussian_f_profile(width=tone['width']),
stg.constant_bp_profile(level=1))
frame.save_fil(filename=synthetic_fil)
darray = from_fil(synthetic_fil)
dedopp, metadata = dedoppler(darray,
boxcar_size=1,
max_dd=1.0,
return_space='cpu')
# Manual dedoppler search -- just find max channel (only works if S/N is good)
manual_dd_tot = 0
for ii in range(darray.data.shape[0]):
manual_dd_tot += np.max(darray.data[ii])
imshow_dedopp(dedopp, show_colorbar=False)
maxpixel = np.argmax(dedopp.data)
mdrift, mchan = (maxpixel // 1024, maxpixel % 1024)
optimal_drift = metadata['drift_rates'][mdrift].value
maxpixel_val = np.max(dedopp.data)
frac_recovered = (maxpixel_val / manual_dd_tot)
print(
f"Inserted drift rate: {tone['drift_rate']} \tSUM: {manual_dd_tot:2.2f}"
)
print(
f"Recovered drift rate: {optimal_drift} Hz / s \tSUM: {maxpixel_val:2.2f}\n"
)
# Channel should detected at +/- 1 chan
print(mdrift, mchan)
#assert np.abs(mchan - 500) <= 1
# Drift rate should be detected +/- 1 drift resolution
assert np.abs(optimal_drift - dr_test) <= 1.01 * np.abs(
metadata['drift_rate_step'].value)
# Recovered signal sum should be close to manual method
assert 1.001 >= frac_recovered >= 0.825
# Finish off figure plotting
plt.colorbar()
plt.savefig(os.path.join(test_fig_dir, 'test_dedoppler.png'))
plt.show()
def test_hitsearch():
""" Test the hit search routines """
n_timesteps = 32
n_chan = 4096
signal_bw = 16
# Create test data
metadata = {
'frequency_start': 1000 * u.MHz,
'time_step': 1.0 * u.s,
'frequency_step': 1.0 * u.Hz
}
frame = stg.Frame(fchans=n_chan * u.pixel,
tchans=n_timesteps * u.pixel,
df=metadata['frequency_step'],
dt=metadata['time_step'],
fch1=metadata['frequency_start'])
frame.add_noise(x_mean=0, x_std=1, noise_type='gaussian')
frame.save_fil(filename=synthetic_fil)
darray = from_fil(synthetic_fil)
# Add a signal with bandwidth into the data, SNR of 1000
for ii in range(signal_bw):
darray.data[:, :, n_chan // 2 + ii] = 1000 / signal_bw
print("--- Run dedoppler() then hitsearch() ---")
dedopp, md = dedoppler(darray, boxcar_size=16, max_dd=1.0)
hits0 = hitsearch(dedopp, threshold=1000)
print(hits0)
# Output should be
#driftrate f_start snr driftrate_idx channel_idx boxcar_size
#0 0.0 1000.002056 32000.0 32 2056 16
# Note that SNR here is unnormalized, so actually peak value -- as we didn't renormalize
h0 = hits0.iloc[0]
assert h0['snr'] == 32000.0
assert h0['channel_idx'] == 2056
assert h0['driftrate_idx'] == 32 or h0[
'driftrate_idx'] == 33 ## Not sure why new algorithm puts centroid to the side?
assert len(hits0) == 1
print("--- run_pipeline with w/o merge --- ")
hits = run_pipeline(darray,
metadata,
max_dd=1.0,
min_dd=None,
threshold=100,
n_boxcar=7,
merge_boxcar_trials=False)
for rid, hit in hits.iterrows():
assert (np.abs(hit['channel_idx'] - 2048) < np.max(
(signal_bw, hit['boxcar_size'])))
print(hits)
print("--- run merge_hits --- ")
print(hits.dtypes)
merged_hits = merge_hits(hits)
assert len(merged_hits == 1)
print(merged_hits)
print("--- run_pipeline with merge --- ")
hits2 = run_pipeline(darray.data,
metadata,
max_dd=1.0,
min_dd=None,
threshold=100,
n_boxcar=7,
merge_boxcar_trials=True)
hits2
print(hits2)
assert hits2.iloc[0]['boxcar_size'] == signal_bw
assert len(hits2) == len(merged_hits) == 1
plt.figure(figsize=(10, 4))
plt.subplot(1, 2, 1)
imshow_waterfall(darray, xaxis='channel', yaxis='timestep')
plt.subplot(1, 2, 2)
imshow_dedopp(dedopp, xaxis='channel', yaxis='driftrate')
plt.savefig(os.path.join(test_fig_dir, 'test_hitsearch.png'))
plt.show()
def gen_fil(arg_path):
r''' Generate a Filterbank file '''
# Define time and frequency arrays, essentially labels for the 2D data array
fchans = 1048576
tchans = 16
df = 1.0 * u.Hz
dt = 1.0 * u.s
fch1 = 6095.214842353016 * u.MHz
noise_std = 0.05 # Gaussian standard deviation
sig_snr_1 = 100.0
sig_width_1 = 1.1 * u.Hz
drate_1 = 1.6 * u.Hz / u.s
f_start_1 = 0
sig_snr_2 = 200.0
sig_width_2 = 1.2 * u.Hz
drate_2 = 1.3 * u.Hz / u.s
f_start_2 = fchans * 0.1
sig_snr_3 = 300.0
sig_width_3 = 1.3 * u.Hz
drate_3 = 2.6 * u.Hz / u.s
f_start_3 = fchans * 0.2
sig_snr_4 = 400.0
sig_width_4 = 1.4 * u.Hz
drate_4 = 3.2 * u.Hz / u.s
f_start_4 = fchans * 0.3
# Generate the frame.
frame = stg.Frame(fchans=fchans, tchans=tchans, df=df, dt=dt, fch1=fch1)
# Add noise.
frame.add_noise(x_mean=0, x_std=noise_std, noise_type='gaussian')
# Add signal 1.
signal_intensity = frame.get_intensity(snr=sig_snr_1)
frame.add_constant_signal(f_start=frame.get_frequency(f_start_1),
drift_rate=drate_1,
level=signal_intensity,
width=sig_width_1,
f_profile_type='gaussian')
# Add signal 2.
signal_intensity = frame.get_intensity(snr=sig_snr_2)
frame.add_constant_signal(f_start=frame.get_frequency(f_start_2),
drift_rate=drate_2,
level=signal_intensity,
width=sig_width_2,
f_profile_type='gaussian')
# Add signal 3.
signal_intensity = frame.get_intensity(snr=sig_snr_3)
frame.add_constant_signal(f_start=frame.get_frequency(f_start_3),
drift_rate=drate_3,
level=signal_intensity,
width=sig_width_3,
f_profile_type='gaussian')
# Add signal 4.
signal_intensity = frame.get_intensity(snr=sig_snr_4)
frame.add_constant_signal(f_start=frame.get_frequency(f_start_4),
drift_rate=drate_4,
level=signal_intensity,
width=sig_width_4,
f_profile_type='gaussian')
# Save Filterbank file.
frame.save_fil(arg_path)
def main():
experiment_path = '/datax/scratch/bbrzycki/data/nb-localization/training/'
output_dir = experiment_path + 'turboseti/'
turbo_rmse_dict = {}
wrong_num_signals = 0
timing_array = []
turbo_start = time.time()
for db in range(0, 30, 5):
for j in range(4000):
print(db, j)
fn = '{:02}db_{:06d}.npy'.format(db, j)
npy_fn = '/datax/scratch/bbrzycki/data/nb-localization/1sig/test/{}'.format(
fn)
fil_fn = output_dir + '{:02}db_{:06d}.fil'.format(db, j)
dat_fn = output_dir + '{:02}db_{:06d}.dat'.format(db, j)
frame = stg.Frame(fchans=1024,
tchans=32,
df=1.3969838619232178 * u.Hz,
dt=1.4316557653333333 * u.s,
fch1=6095.214842353016 * u.MHz,
data=np.load(npy_fn))
frame.save_fil(fil_fn)
try:
os.remove(dat_fn)
except FileNotFoundError:
pass
start_time = time.time()
find_seti_event = FindDoppler(fil_fn,
max_drift=31,
snr=10,
out_dir=output_dir)
find_seti_event.search()
end_time = time.time()
timing_array.append(end_time - start_time)
with open(dat_fn, 'r') as f:
data = [
line.split() for line in f.readlines() if line[0] != '#'
]
# Count number of times turboseti predicts wrong number
if len(data) != 1:
wrong_num_signals += 1
estimates = []
snrs = []
for signal in data:
snr = float(signal[2])
drift_rate = float(signal[1])
start_index = 1024 - int(signal[5])
end_index = frame.get_index(
frame.get_frequency(start_index) +
drift_rate * frame.tchans * frame.dt)
estimates.append([start_index, end_index])
snrs.append(snr)
if len(estimates) != 0:
# Get the ground truth positions from the saved dictionaries
true_indices = (
np.load(experiment_path +
'final_1sig_32bs_bright/test_predictions.npy',
allow_pickle=True).item()[fn] * 1024)[0]
# If turboseti found signals, choose the highest SNR one
turbo_rmse_dict[fn] = rmse(true_indices,
estimates[np.argsort(snrs)[-1]])
timing_array = np.array(timing_array)
print('Wrong: {} frames'.format(wrong_num_signals))
print('Total search: {:.2f} seconds'.format(time.time() - turbo_start))
np.save(output_dir + 'timing_array.npy', timing_array)
np.save(output_dir + 'test_predictions.npy', turbo_rmse_dict)
def frame_setup_from_h5():
my_path = os.path.abspath(os.path.dirname(__file__))
path = os.path.join(my_path, 'assets/sample.fil')
frame = stg.Frame(waterfall=path)
return frame
