def test_hitsearch(): """ Test the hit search routines """ n_timesteps = 32 n_chan = 4096 signal_bw = 16 # Create test data metadata = {'fch1': 1000*u.MHz, 'dt': 1.0*u.s, 'df': 1.0*u.Hz} frame = stg.Frame(fchans=n_chan*u.pixel, tchans=n_timesteps*u.pixel, df=metadata['df'], dt=metadata['dt'], fch1=metadata['fch1']) 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, metadata = dedoppler(darray.data, metadata, boxcar_size=16, max_dd=1.0) hits0 = hitsearch(dedopp, metadata, threshold=1000).sort_values('snr') 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 --- ") dedopp, metadata, hits = run_pipeline(darray.data, 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 --- ") dedopp, md, 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.data, metadata, 'channel', 'timestep') plt.subplot(1,2,2) imshow_dedopp(dedopp, metadata, 'channel', 'driftrate') plt.savefig('docs/figs/test_hitsearch.png') plt.show()
def test_dedoppler_boxcar(): """ Test that boxcar averaging works as expected """ def generate_drifting_tone(n_chan, n_timesteps, n_drift_per_step, n_beams=1, sigval=10): """ Simple tone generator to generate smeared tones """ bg = np.zeros((n_timesteps, n_beams, n_chan)) for ii in range(0, bg.shape[0]): for nd in range(n_drift_per_step): z = n_drift_per_step * ii + nd bg[ii, :, bg.shape[2]//2 + z] = sigval / n_drift_per_step return bg def maxhold_dedoppler(data): """ A simple and crappy dedoppler algorithm Finds the top value in each timestep and adds together. This method only works for single channel tones with high SNR """ manual_dd_tot = 0 for ii in range(bg.shape[0]): manual_dd_tot += np.max(bg[ii]) return manual_dd_tot # Drift rate of 2 channels / integration # To simulate channel smearing metadata = {'fch1': 1000*u.MHz, 'dt': 1.0*u.s, 'df': 1.0*u.Hz} bg = generate_drifting_tone(n_chan=256, n_timesteps=32, n_drift_per_step=2, sigval=10) print(f"Total power in frame: {np.sum(bg)}") # Compute dedoppler using basic maxhold method maxhold_res = maxhold_dedoppler(bg) print(f"MAXHOLD recovered power: {maxhold_res}") # With boxcar_size = 1 we should recover 160 dedopp, metadata = dedoppler(bg, metadata, boxcar_size=1, max_dd=2.0) maxpixel = np.argmax(dedopp) mdrift, mchan = (maxpixel // 1024, maxpixel % 1024) maxpixel_val = np.max(dedopp) print(f"dedopp recovered power (boxcar 1): {maxpixel_val}") assert maxpixel_val == maxhold_res # With boxcar_size = 2 we should recover 320 (full amount) metadata = {'fch1': 1000*u.MHz, 'dt': 1.0*u.s, 'df': 1.0*u.Hz} dedopp, metadata = dedoppler(bg, metadata, boxcar_size=2, max_dd=4.0) maxpixel = np.argmax(dedopp) mdrift, mchan = (maxpixel // 1024, maxpixel % 1024) # <----------- UNUSED maxpixel_val = np.max(dedopp) print(f"dedopp recovered power (boxcar 2): {maxpixel_val}") assert maxpixel_val == np.sum(bg) # plot plt.figure(figsize=(10, 4)) plt.subplot(1,2,1) imshow_waterfall(bg, metadata) plt.subplot(1,2,2) imshow_dedopp(dedopp, metadata, 'channel', 'driftrate') plt.savefig('docs/figs/test_dedoppler_boxcar.png') plt.show()
def test_dedoppler(): """ Basic tests of the dedoppler functionality """ # zero drift test, no normalization test_data = np.ones(shape=(32, 1, 1024)) test_data[:, :, 511] = 10 metadata = {'fch1': 1000*u.MHz, 'dt': 1.0*u.s, 'df': 1.0*u.Hz} dedopp, metadata = dedoppler(test_data, metadata, boxcar_size=1, max_dd=1.0) assert np.max(dedopp) == 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['df'], dt=metadata['dt'], fch1=metadata['fch1']) tone = {'f_start': frame.get_frequency(index=500), 'drift_rate': dr_test * u.Hz / u.s, 'snr': 500, 'width': metadata['df']} 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.data, metadata, boxcar_size=1, max_dd=1.0) # 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, metadata, show_colorbar=False) maxpixel = np.argmax(dedopp) mdrift, mchan = (maxpixel // 1024, maxpixel % 1024) optimal_drift = metadata['drift_trials'][mdrift] maxpixel_val = np.max(dedopp) 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 assert np.abs(mchan - 500) <= 1 # Drift rate should be detected +/- 1 drift resolution assert np.abs(optimal_drift - dr_test) <= 1.01*metadata['dd'].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('docs/figs/test_dedoppler.png') plt.show()