def add_to_data(self, delta_t, freq, data, scintillate=True, bandwidth=300.):
""" Method to add already-dedispersed pulse
to background noise data. Includes frequency-dependent
width (smearing, scattering, etc.) and amplitude
(scintillation, spectral index).
"""
NFREQ = data.shape[0]
NTIME = data.shape[1]
tmid = NTIME//2
freq_mid = freq[len(freq)//2]
scint_amp = self.scintillation(freq)
self._fluence /= np.sqrt(NFREQ)
# stds_perchan = np.std(data)#*np.sqrt(NFREQ)
SNRTools = tools.SNR_Tools()
stds_perchan, med = SNRTools.sigma_from_mad(data.flatten())
for ii, f in enumerate(freq):
# Do not add FRB to missing channels
if data[ii].sum()==0:
continue
width_ = self.calc_width(self._dm, self._f_ref*1e-3,
bw=bandwidth, NFREQ=NFREQ,
ti=self._width, tsamp=delta_t, tau=0)
index_width = max(1, (np.round((width_/ delta_t))).astype(int))
tpix = int(self.arrival_time(f) / delta_t)
if abs(tpix) >= tmid:
# ensure that edges of data are not crossed
continue
pp = self.pulse_profile(NTIME, index_width, f,
tau=self._scat_factor, t0=tpix)
val = pp.copy()
val /= val.max()
val *= self._fluence
# val *= (100.0 / stds_perchan) # hack
val /= (width_ / delta_t)
val = val * (f / self._f_ref) ** self._spec_ind
if scintillate is True:
val = (0.1 + scint_amp[ii]) * val
data[ii] += val
# data[ii, tpix] += 5*np.std(data[ii])
if f == freq_mid:
width_eff = width_
return width_eff
def proc_trigger(fn_fil, dm0, t0, sig_cut,
ndm=50, mk_plot=False, downsamp=1,
beamno='', fn_mask=None, nfreq_plot=32,
ntime_plot=250,
cmap='RdBu', cand_no=1, multiproc=False,
rficlean=False, snr_comparison=-1,
outdir='./', sig_thresh_local=5.0,
subtract_zerodm=False,
threshold_time=3.25, threshold_frequency=2.75,
bin_size=32, n_iter_time=3,
n_iter_frequency=3, clean_type='time',
freq=1370.0, sb_generator=None, sb=None,
dumb_mask=True, save_sb_fil=False):
""" Locate data within filterbank file (fn_fi)
at some time t0, and dedisperse to dm0, generating
plots
Parameters:
----------
fn_fil : str
name of filterbank file
dm0 : float
trigger dm found by single pulse search software
t0 : float
time in seconds where trigger was found
sig_cut : np.float
sigma of detected trigger at (t0, dm0)
ndm : int
number of DMs to use in DM transform
mk_plot : bool
make three-panel plots
downsamp : int
factor by which to downsample in time. comes from searchsoft.
beamno : str
beam number, for fig names
nfreq_plot : int
number of frequencies channels to plot
freq : int
central frequency used to find zapped channels file
sb_generator: SBGenerator object
synthesized beam mapper from DARC (None for TAB/IAB)
sb : int
synthesized beam to generate (None for TAB/IAB)
Returns:
-------
full_dm_arr_downsamp : np.array
data array with downsampled dm-transformed intensities
full_freq_arr_downsamp : np.array
data array with downsampled freq-time intensities
"""
if dumb_mask:
try:
fndmask='/home/arts/.controller/amber_conf/zapped_channels_{:.0f}.conf'.format(int(freq))
rfimask = np.loadtxt(fndmask)
rfimask = rfimask.astype(int)
except:
rfimask = np.array([])
logging.warning("Could not load dumb RFIMask")
else:
rfimask = np.array([])
SNRtools = tools.SNR_Tools()
downsamp = min(4096, downsamp)
if downsamp >= 100:
wideclean = int(downsamp)
else:
wideclean = None
# store path to filterbanks
if sb is not None:
prefix_fil = fn_fil
# get first file
fn_fil = prefix_fil + '_00.fil'
if not os.path.exists(fn_fil):
fn_fil = prefix_fil + '00.fil'
rawdatafile = filterbank.filterbank(fn_fil)
dfreq_MHz = rawdatafile.header['foff']
dt = rawdatafile.header['tsamp']
freq_up = rawdatafile.header['fch1']
nfreq = rawdatafile.header['nchans']
# fix RFI mask order
rfimask = nfreq - rfimask
freq_low = freq_up + nfreq*rawdatafile.header['foff']
ntime_fil = (os.path.getsize(fn_fil) - rawdatafile.header_size)/nfreq
tdm = np.abs(8.3*1e-6*dm0*dfreq_MHz*(freq_low/1000.)**-3)
dm_min = max(0, dm0-40)
dm_max = dm0 + 40
dms = np.linspace(dm_min, dm_max, ndm, endpoint=True)
# make sure dm0 is in the array
dm_max_jj = np.argmin(abs(dms-dm0))
dms += (dm0-dms[dm_max_jj])
dms[0] = max(0, dms[0])
global t_min, t_max
# if smearing timescale is < 4*pulse width,
# downsample before dedispersion for speed
downsamp_smear = max(1, int(downsamp*dt/tdm/2.))
# ensure that it's not larger than pulse width
downsamp_smear = int(min(downsamp, downsamp_smear))
downsamp_res = int(downsamp//downsamp_smear)
downsamp = int(downsamp_res*downsamp_smear)
time_res = dt * downsamp
logging.info("Width_full:%d Width_smear:%d Width_res: %d" %
(downsamp, downsamp_smear, downsamp_res))
start_bin = int(t0/dt - ntime_plot*downsamp//2)
width = abs(4.148e3 * dm0 * (freq_up**-2 - freq_low**-2))
chunksize = int(width/dt + ntime_plot*downsamp)
t_min, t_max = 0, ntime_plot*downsamp
if start_bin < 0:
extra = start_bin
start_bin = 0
t_min += extra
t_max += extra
t_min, t_max = int(t_min), int(t_max)
snr_max = 0
# Account for the pre-downsampling to speed up dedispersion
t_min /= downsamp_smear
t_max /= downsamp_smear
ntime = t_max-t_min
if ntime_fil < (start_bin+chunksize):
logging.info("Trigger at end of file, skipping")
return [], [], [], []
# get data of all files (SB) or one file (TAB/IAB)
if sb is not None:
ntab = 12
data = np.zeros((ntab, nfreq, chunksize))
# get list of unique TABs in required SB
sb_map = list(set(sb_generator.get_map(sb)))
logging.info("SB {} consists of TABs {}".format(sb, sb_map))
threads = []
for tab in range(ntab):
# skip if we do not need this TAB
if not tab in sb_map:
continue
fname = prefix_fil + '_{:02d}.fil'.format(tab)
if not os.path.exists(fname):
fname = prefix_fil + '{:02d}.fil'.format(tab)
load_tab_data(fname, start_bin, chunksize, out=data, tab=tab)
for thread in threads:
logging.info("Waiting for loading of {}".format(thread.name))
thread.join()
# generate sb
logging.info("Synthesizing beam {}".format(sb))
data = sb_generator.synthesize_beam(data, sb=sb)
# convert to a spectra object, mimicking filterbank.get_spectra
data = spectra.Spectra(rawdatafile.frequencies,
rawdatafile.tsamp, data,
starttime=start_bin*rawdatafile.tsamp,
dm=0)
else:
data = rawdatafile.get_spectra(start_bin, chunksize)
if save_sb_fil:
fn_sbfil_out = '%s/data/CB%s_snr%d_dm%d_t0%d_sb%d.fil' % \
(outdir, beamno, sig_cut, dm0, t0, sb)
reader.create_new_filterbank(fn_sbfil_out, telescope='Apertif')
reader.write_to_fil(data.data.transpose(), rawdatafile.header, fn_sbfil_out)
# fil_obj = reader.filterbank_.FilterbankFile(fn_sbfil_out, mode='readwrite')
# fil_obj.append_spectra(data.data.transpose())
# np.save(fn_fig_out, data.data)
rawdatafile.close()
# apply dumb mask
if len(rfimask)>0:data.data[rfimask] = 0.
if rficlean is True:
data.data = cleandata(data.data, threshold_time,
threshold_frequency, bin_size,
n_iter_time, n_iter_frequency,
clean_type, wideclean=wideclean)
if subtract_zerodm:
data.data -= np.mean(data.data, axis=0, keepdims=True)
freq_ref = 0.5*(freq_up+freq_low)
# Downsample before dedispersion up to 1/4th
# DM smearing limit
data.downsample(downsamp_smear)
data.data -= np.median(data.data, axis=-1)[:, None]
# full_arr = np.empty([int(ndm), int(ntime)])
if not fn_mask is None:
pass
# rfimask = rfifind.rfifind(fn_mask)
# mask = get_mask(rfimask, start_bin, chunksize)
# data = data.masked(mask, maskval='median-mid80')
if multiproc is True:
pass
else:
logging.info("\nDedispersing Serially\n")
data_copy = copy.deepcopy(data)
data_copy.dedisperse(dm0)
data_dm_max = data_copy.data[:, max(0, t_min):t_max]
snr_max = SNRtools.calc_snr_matchedfilter(data_dm_max.mean(0), widths=[downsamp_res])[0]
if t_min<0:
Z = np.zeros([nfreq, np.abs(t_min)])
data_dm_max = np.concatenate([Z, data_dm_max], axis=1)
# scale max DM by pulse width, 5 units for each ms
dm_max_trans = 10. + 5*time_res/0.001 + 10*dm0/1000.
dm_min_trans = -10. - 5*time_res/0.001 - 10*dm0/1000.
if dm0+dm_min_trans<=0:
dm_min_trans = 0.
dm_max_trans = 2*dm0
dm_center = dm0
else:
dm_center = 0.
full_arr, dms, times = RTproc.dm_transform(data_dm_max,
(freq_up, freq_low),
dt=dt*downsamp_smear,
dm_max=dm_max_trans,
dm_min=dm_min_trans,
freq_ref=freq_ref,ndm=ndm,
dm0=dm_center)
dms += dm0
# bin down to nfreq_plot freq channels
full_freq_arr_downsamp = data_dm_max[:nfreq//nfreq_plot*nfreq_plot, :]
full_freq_arr_downsamp = full_freq_arr_downsamp.reshape(\
nfreq_plot, -1, ntime).mean(1)
# bin down in time by factor of downsamp
full_freq_arr_downsamp = full_freq_arr_downsamp[:,:ntime//downsamp_res*downsamp_res]
full_freq_arr_downsamp = full_freq_arr_downsamp.reshape(-1, ntime//downsamp_res,
downsamp_res).mean(-1)
if snr_max < sig_thresh_local:
logging.info("\nSkipping trigger below local threshold %.2f:" % sig_thresh_local)
logging.info("snr_local=%.2f snr_trigger=%.2f\n" % (snr_max, sig_cut))
return [], [], [], []
times = np.linspace(0,ntime_plot*downsamp*dt,len(full_freq_arr_downsamp[0]))
full_dm_arr_downsamp = full_arr[:, :ntime//downsamp_res*downsamp_res]
full_dm_arr_downsamp = full_dm_arr_downsamp.reshape(-1,
ntime//downsamp_res, downsamp_res).mean(-1)
full_freq_arr_downsamp /= np.std(full_freq_arr_downsamp)
full_dm_arr_downsamp /= np.std(full_dm_arr_downsamp)
suptitle = " CB:%s S/N$_{pipe}$:%.1f S/N$_{presto}$:%.1f\
S/N$_{compare}$:%.1f \nDM:%d t:%.1fs width:%d" %\
(beamno, sig_cut, snr_max, snr_comparison,
dm0, t0, downsamp)
if not os.path.isdir('%s/plots' % outdir):
os.system('mkdir -p %s/plots' % outdir)
if sb is None:
sbname = -1
else:
sbname = sb
fn_fig_out = '%s/plots/CB%s_snr%d_dm%d_t0%d_sb%d.pdf' % \
(outdir, beamno, sig_cut, dm0, t0, sbname)
params = sig_cut, dms[dm_max_jj], downsamp, t0, dt
tmed = np.median(full_freq_arr_downsamp, axis=-1, keepdims=True)
full_freq_arr_downsamp -= tmed
if mk_plot is True:
logging.info(fn_fig_out)
if ndm == 1:
plotter.plot_two_panel(full_freq_arr_downsamp, params,
prob=None,
freq_low=freq_low, freq_up=freq_up,
cand_no=cand_no, times=times,
suptitle=suptitle,
fnout=fn_fig_out)
else:
plotter.plot_three_panel(full_freq_arr_downsamp,
full_dm_arr_downsamp, params, dms,
times=times, freq_low=freq_low,
freq_up=freq_up,
suptitle=suptitle, fnout=fn_fig_out,
cand_no=cand_no)
return full_dm_arr_downsamp, full_freq_arr_downsamp, time_res, params
# TAB: filterbank filename is specified
# SB: filterbank prefix is specified (i.e. without _<TABno>.fil)
fn_fil = args[0]
fn_sp = args[1]
if options.save_data == 'concat':
data_dm_time_full = []
data_freq_time_full = []
params_full = []
if options.sb:
data_sb_full = []
if options.multiproc is True:
import multiprocessing
SNRTools = tools.SNR_Tools()
if options.compare_trig is not None:
res = SNRTools.compare_snr(fn_sp, options.compare_trig,
dm_min=options.dm_min,
dm_max=options.dm_max,
save_data=False,
sig_thresh=options.sig_thresh,
max_rows=None,
t_window=0.25)
res = par_1, par_2, par_match_arr, ind_missed, ind_matched
snr_1, snr_2 = par_1[0], par_2[0]
snr_comparison_arr = np.zeros_like(snr_1)
ind_missed = np.array(ind_missed)
def proc_trigger(fn_fil,
dm0,
t0,
sig_cut,
ndm=50,
mk_plot=False,
downsamp=1,
beamno='',
fn_mask=None,
nfreq_plot=32,
ntime_plot=250,
cmap='RdBu',
cand_no=1,
multiproc=False):
""" Locate data within filterbank file (fn_fi)
at some time t0, and dedisperse to dm0, generating
plots
Parameters:
----------
fn_fil : str
name of filterbank file
dm0 : float
trigger dm found by single pulse search software
t0 : float
time in seconds where trigger was found
sig_cut : np.float
sigma of detected trigger at (t0, dm0)
ndm : int
number of DMs to use in DM transform
mk_plot : bool
make three-panel plots
downsamp : int
factor by which to downsample in time. comes from searchsoft.
beamno : str
beam number, for fig names
nfreq_plot : int
number of frequencies channels to plot
Returns:
-------
full_dm_arr_downsamp : np.array
data array with downsampled dm-transformed intensities
full_freq_arr_downsamp : np.array
data array with downsampled freq-time intensities
"""
SNRtools = tools.SNR_Tools()
rawdatafile = filterbank.filterbank(fn_fil)
mask = [
] #np.array([ 5, 6, 9, 32, 35, 49, 75, 76, 78, 82, 83, 87, 92,
# 93, 97, 98, 108, 110, 111, 112, 114, 118, 122, 123, 124, 157,
# 160, 299, 300, 301, 302, 303, 304, 305, 306, 307, 308, 309, 660,
# 661])
dt = rawdatafile.header['tsamp']
freq_up = rawdatafile.header['fch1']
nfreq = rawdatafile.header['nchans']
freq_low = freq_up + nfreq * rawdatafile.header['foff']
time_res = dt * downsamp
ntime_fil = (os.path.getsize(fn_fil) - 467.) / nfreq
dm_min = max(0, dm0 - 20)
dm_max = dm0 + 20
dms = np.linspace(dm_min, dm_max, ndm, endpoint=True)
# make sure dm0 is in the array
dm_max_jj = np.argmin(abs(dms - dm0))
dms += (dm0 - dms[dm_max_jj])
dms[0] = max(0, dms[0])
# Read in 5 disp delays
width = 2 * abs(4.14e3 * dm0 * (freq_up**-2 - freq_low**-2))
tdisp = width / dt
tplot = ntime_plot * downsamp
global t_min, t_max
if tdisp > tplot:
# Need to read in more data than you'll plot
# because of large dispersion time
chunksize = int(tdisp)
t_min = chunksize // 2 - (ntime_plot * downsamp) // 2
t_max = chunksize // 2 + (ntime_plot * downsamp) // 2
else:
# Only need ot read in enough to plot
chunksize = int(tplot)
t_min, t_max = 0, chunksize
start_bin = int(t0 / dt - chunksize / 2.)
if start_bin < 0:
extra = start_bin
start_bin = 0
t_min += extra
t_max += extra
t_min, t_max = int(t_min), int(t_max)
ntime = t_max - t_min
snr_max = 0
if ntime_fil < (start_bin + chunksize):
print("Trigger at end of file, skipping")
return 0, 0, 0
data = rawdatafile.get_spectra(start_bin, chunksize)
data.data -= np.median(data.data, axis=-1)[:, None]
data.data[mask] = 0.
full_arr = np.empty([int(ndm), int(ntime)])
if not fn_mask is None:
rfimask = rfifind.rfifind(fn_mask)
mask = get_mask(rfimask, start_bin, chunksize)
data = data.masked(mask, maskval='median-mid80')
if multiproc is True:
print("\nDedispersing in Parallel\n")
t0 = time.time()
global datacopy
ndm_ = min(10, ndm)
for kk in range(ndm // ndm_):
dms_ = dms[10 * kk:10 * (kk + 1)]
datacopy = copy.deepcopy(data)
pool = multiprocessing.Pool(processes=ndm_)
data_tuple = pool.map(multiproc_dedisp, [i for i in dms_])
pool.close()
data_tuple = np.concatenate(data_tuple)
ddm = np.concatenate(data_tuple[0::2]).reshape(ndm_, -1)
df = np.concatenate(data_tuple[1::2]).reshape(ndm_, nfreq, -1)
print(time.time() - t0)
full_arr[10 * kk:10 * (kk + 1)] = ddm[:, t_min:t_max]
ind_kk = range(10 * kk, 10 * (kk + 1))
if dm_max_jj in ind_kk:
print(dms_[ind_kk.index(dm_max_jj)])
data_dm_max = df[ind_kk.index(dm_max_jj)] #dm_max_jj]hack
del ddm, df
else:
print("\nDedispersing Serially\n")
for jj, dm_ in enumerate(dms):
print("Dedispersing to dm=%0.1f at t=%0.1f sec with width=%.2f" %
(dm_, start_bin * dt, downsamp))
data_copy = copy.deepcopy(data)
data_copy.dedisperse(dm_)
dm_arr = data_copy.data[:, t_min:t_max].mean(0)
snr_ = SNRtools.calc_snr(dm_arr)
full_arr[jj] = copy.copy(dm_arr)
if jj == dm_max_jj:
data_dm_max = data_copy.data[:, t_min:t_max]
downsamp = int(downsamp)
# bin down to nfreq_plot freq channels
full_freq_arr_downsamp = data_dm_max[:nfreq//nfreq_plot*nfreq_plot, :].reshape(\
nfreq_plot, -1, ntime).mean(1)
# bin down in time by factor of downsamp
full_freq_arr_downsamp = full_freq_arr_downsamp[:, :ntime//downsamp*downsamp\
].reshape(-1, ntime//downsamp, downsamp).mean(-1)
times = np.linspace(0, ntime * dt, len(full_freq_arr_downsamp[0]))
full_dm_arr_downsamp = full_arr[:, :ntime // downsamp * downsamp]
full_dm_arr_downsamp = full_dm_arr_downsamp.reshape(
-1, ntime // downsamp, downsamp).mean(-1)
full_freq_arr_downsamp /= np.std(full_freq_arr_downsamp)
full_dm_arr_downsamp /= np.std(full_dm_arr_downsamp)
suptitle = "beam%s snr%d dm%d t0%d width%d" %\
(beamno, sig_cut, dms[dm_max_jj], t0, downsamp)
fn_fig_out = './plots/train_data_beam%s_snr%d_dm%d_t0%d.pdf' % \
(beamno, sig_cut, dms[dm_max_jj], t0)
if mk_plot is True:
print(fn_fig_out)
if ndm == 1:
params = snr_, dm_, downsamp, t0
plotter.plot_two_panel(full_freq_arr_downsamp,
params,
prob=None,
freq_low=1250.09765625,
freq_up=1549.90234375,
cand_no=cand_no)
else:
print(fn_fig_out)
plotter.plot_three_panel(full_freq_arr_downsamp,
full_dm_arr_downsamp,
times,
dms,
freq_low=freq_low,
freq_up=freq_up,
suptitle=suptitle,
fnout=fn_fig_out,
cand_no=cand_no)
return full_dm_arr_downsamp, full_freq_arr_downsamp, time_res
def proc_trigger(fn_fil, dm0, t0, sig_cut,
ndm=50, mk_plot=False, downsamp=1,
beamno='', fn_mask=None, nfreq_plot=32,
ntime_plot=250,
cmap='RdBu', cand_no=1, multiproc=False,
rficlean=False, snr_comparison=-1,
outdir='./', sig_thresh_local=5.0,
subtract_zerodm=False,
threshold_time=3.25, threshold_frequency=2.75, bin_size=32,
n_iter_time=3, n_iter_frequency=3, clean_type='time'):
""" Locate data within filterbank file (fn_fi)
at some time t0, and dedisperse to dm0, generating
plots
Parameters:
----------
fn_fil : str
name of filterbank file
dm0 : float
trigger dm found by single pulse search software
t0 : float
time in seconds where trigger was found
sig_cut : np.float
sigma of detected trigger at (t0, dm0)
ndm : int
number of DMs to use in DM transform
mk_plot : bool
make three-panel plots
downsamp : int
factor by which to downsample in time. comes from searchsoft.
beamno : str
beam number, for fig names
nfreq_plot : int
number of frequencies channels to plot
Returns:
-------
full_dm_arr_downsamp : np.array
data array with downsampled dm-transformed intensities
full_freq_arr_downsamp : np.array
data array with downsampled freq-time intensities
"""
try:
rfimask = np.loadtxt('/home/arts/ARTS-obs/amber_conf/zapped_channels_1400.conf')
rfimask = rfimask.astype(int)
except:
rfimask = []
logging.warning("Could not load dumb RFIMask")
SNRtools = tools.SNR_Tools()
downsamp = min(4096, downsamp)
rawdatafile = filterbank.filterbank(fn_fil)
dfreq_MHz = rawdatafile.header['foff']
mask = []
dt = rawdatafile.header['tsamp']
freq_up = rawdatafile.header['fch1']
nfreq = rawdatafile.header['nchans']
# fix RFI mask order
rfimask = nfreq - rfimask
freq_low = freq_up + nfreq*rawdatafile.header['foff']
ntime_fil = (os.path.getsize(fn_fil) - 467.)/nfreq
tdm = np.abs(8.3*1e-6*dm0*dfreq_MHz*(freq_low/1000.)**-3)
dm_min = max(0, dm0-40)
dm_max = dm0 + 40
dms = np.linspace(dm_min, dm_max, ndm, endpoint=True)
# make sure dm0 is in the array
dm_max_jj = np.argmin(abs(dms-dm0))
dms += (dm0-dms[dm_max_jj])
dms[0] = max(0, dms[0])
global t_min, t_max
# if smearing timescale is < 4*pulse width,
# downsample before dedispersion for speed
downsamp_smear = max(1, int(downsamp*dt/tdm/2.))
# ensure that it's not larger than pulse width
downsamp_smear = int(min(downsamp, downsamp_smear))
downsamp_res = int(downsamp//downsamp_smear)
downsamp = int(downsamp_res*downsamp_smear)
time_res = dt * downsamp
tplot = ntime_plot * downsamp
logging.info("Width_full:%d Width_smear:%d Width_res: %d" %
(downsamp, downsamp_smear, downsamp_res))
# print("Width_full:%d Width_smear:%d Width_res: %d" %
# (downsamp, downsamp_smear, downsamp_res))
start_bin = int(t0/dt - ntime_plot*downsamp//2)
width = abs(4.148e3 * dm0 * (freq_up**-2 - freq_low**-2))
chunksize = int(width/dt + ntime_plot*downsamp)
t_min, t_max = 0, ntime_plot*downsamp
if start_bin < 0:
extra = start_bin
start_bin = 0
t_min += extra
t_max += extra
t_min, t_max = int(t_min), int(t_max)
snr_max = 0
# Account for the pre-downsampling to speed up dedispersion
t_min /= downsamp_smear
t_max /= downsamp_smear
ntime = t_max-t_min
if ntime_fil < (start_bin+chunksize):
logging.info("Trigger at end of file, skipping")
# print("Trigger at end of file, skipping")
return [],[],[],[]
data = rawdatafile.get_spectra(start_bin, chunksize)
# apply dumb mask
data.data[rfimask] = 0.
if rficlean is True:
data = cleandata(data, threshold_time, threshold_frequency, bin_size, \
n_iter_time, n_iter_frequency, clean_type)
if subtract_zerodm:
data.data -= np.mean(data.data, axis=0)[None]
# Downsample before dedispersion up to 1/4th
# DM smearing limit
data.downsample(downsamp_smear)
data.data -= np.median(data.data, axis=-1)[:, None]
full_arr = np.empty([int(ndm), int(ntime)])
if not fn_mask is None:
pass
# rfimask = rfifind.rfifind(fn_mask)
# mask = get_mask(rfimask, start_bin, chunksize)
# data = data.masked(mask, maskval='median-mid80')
if multiproc is True:
tbeg=time.time()
global datacopy
size_arr = sys.getsizeof(data.data)
nproc = int(32.0e9/size_arr)
ndm_ = min(min(nproc, ndm), 10)
for kk in range(ndm//ndm_):
dms_ = dms[ndm_*kk:ndm_*(kk+1)]
datacopy = copy.deepcopy(data)
pool = multiprocessing.Pool(processes=ndm_)
data_tuple = pool.map(multiproc_dedisp, [i for i in dms_])
pool.close()
data_tuple = np.concatenate(data_tuple)
ddm = np.concatenate(data_tuple[0::2]).reshape(ndm_, -1)
df = np.concatenate(data_tuple[1::2]).reshape(ndm_, nfreq, -1)
print(time.time()-tbeg)
full_arr[ndm_*kk:ndm_*(kk+1)] = ddm[:, t_min:t_max]
ind_kk = range(ndm_*kk, ndm_*(kk+1))
if dm_max_jj in ind_kk:
data_dm_max = df[ind_kk.index(dm_max_jj)]
del ddm, df
else:
logging.info("\nDedispersing Serially\n")
#print("\nDedispersing Serially\n")
for jj, dm_ in enumerate(dms):
tcopy = time.time()
data_copy = copy.deepcopy(data)
t0_dm = time.time()
data_copy.dedisperse(dm_)
dm_arr = data_copy.data[:, max(0, t_min):t_max].mean(0)
full_arr[jj, np.abs(min(0, t_min)):] = copy.copy(dm_arr)
logging.info("Dedispersing to dm=%0.1f at t=%0.1fsec with width=%.1f S/N=%.1f" %
(dm_, t0, downsamp, sig_cut))
# print("Dedispersing to dm=%0.1f at t=%0.1fsec with width=%.1f S/N=%.1f" %
# (dm_, t0, downsamp, sig_cut))
if jj==dm_max_jj:
data_dm_max = data_copy.data[:, max(0, t_min):t_max]
snr_max = SNRtools.calc_snr_matchedfilter(data_dm_max.mean(0), widths=[downsamp_res])[0]
if t_min<0:
Z = np.zeros([nfreq, np.abs(t_min)])
data_dm_max = np.concatenate([Z, data_dm_max], axis=1)
# bin down to nfreq_plot freq channels
full_freq_arr_downsamp = data_dm_max[:nfreq//nfreq_plot*nfreq_plot, :].reshape(\
nfreq_plot, -1, ntime).mean(1)
# bin down in time by factor of downsamp
full_freq_arr_downsamp = full_freq_arr_downsamp[:, :ntime//downsamp_res*downsamp_res\
].reshape(-1, ntime//downsamp_res, downsamp_res).mean(-1)
# snr_max = SNRtools.calc_snr_mad(full_freq_arr_downsamp.mean(0))
if snr_max < sig_thresh_local:
logging.info("\nSkipping trigger below local threshold %.2f:" % sig_thresh_local)
logging.info("snr_local=%.2f snr_trigger=%.2f\n" % (snr_max, sig_cut))
return [],[],[],[]
times = np.linspace(0, ntime_plot*downsamp*dt, len(full_freq_arr_downsamp[0]))
full_dm_arr_downsamp = full_arr[:, :ntime//downsamp_res*downsamp_res]
full_dm_arr_downsamp = full_dm_arr_downsamp.reshape(-1,
ntime//downsamp_res, downsamp_res).mean(-1)
full_freq_arr_downsamp /= np.std(full_freq_arr_downsamp)
full_dm_arr_downsamp /= np.std(full_dm_arr_downsamp)
suptitle = " CB:%s S/N$_{pipe}$:%.1f S/N$_{presto}$:%.1f\
S/N$_{compare}$:%.1f \nDM:%d t:%.1fs width:%d" %\
(beamno, sig_cut, snr_max, snr_comparison, \
dms[dm_max_jj], t0, downsamp)
if not os.path.isdir('%s/plots' % outdir):
os.system('mkdir -p %s/plots' % outdir)
fn_fig_out = '%s/plots/CB%s_snr%d_dm%d_t0%d.pdf' % \
(outdir, beamno, sig_cut, dms[dm_max_jj], t0)
params = sig_cut, dms[dm_max_jj], downsamp, t0, dt
tmed = np.median(full_freq_arr_downsamp, axis=-1, keepdims=True)
full_freq_arr_downsamp -= tmed
if mk_plot is True:
logging.info(fn_fig_out)
if ndm==1:
plotter.plot_two_panel(full_freq_arr_downsamp, params, prob=None,
freq_low=freq_low, freq_up=freq_up,
cand_no=cand_no, times=times, suptitle=suptitle,
fnout=fn_fig_out)
else:
plotter.plot_three_panel(full_freq_arr_downsamp,
full_dm_arr_downsamp, params, dms,
times=times, freq_low=freq_low,
freq_up=freq_up,
suptitle=suptitle, fnout=fn_fig_out,
cand_no=cand_no)
return full_dm_arr_downsamp, full_freq_arr_downsamp, time_res, params
def fil_trigger(fn_fil, dm0, t0, sig_cut,
ndm=50, mk_plot=False, downsamp=1,
beamno='', fn_mask=None, nfreq_plot=32,
ntime_plot=250,
cmap='RdBu', cand_no=1, multiproc=False,
rficlean=False, snr_comparison=-1,
outdir='./', sig_thresh_local=5.0,
threshold_time=3.25, threshold_frequency=2.75, bin_size=32,
n_iter_time=3, n_iter_frequency=3, clean_type='time'):
try:
rfimask = np.loadtxt('/home/arts/ARTS-obs/amber_conf/zapped_channels_1400.conf')
rfimask = rfimask.astype(int)
except:
rfimask = []
logging.info("Could not load dumb RFIMask")
SNRtools = tools.SNR_Tools()
downsamp = min(4096, downsamp)
rawdatafile = filterbank.filterbank(fn_fil)
dfreq_MHz = rawdatafile.header['foff']
mask = []
dt = rawdatafile.header['tsamp']
freq_up = rawdatafile.header['fch1']
nfreq = rawdatafile.header['nchans']
freq_low = freq_up + nfreq*rawdatafile.header['foff']
ntime_fil = (os.path.getsize(fn_fil) - 467.)/nfreq
tdm = np.abs(8.3*1e-6*dm0*dfreq_MHz*(freq_low/1000.)**-3)
dm_min = max(0, dm0-40)
dm_max = dm0 + 40
dms = np.linspace(dm_min, dm_max, ndm, endpoint=True)
# make sure dm0 is in the array
dm_max_jj = np.argmin(abs(dms-dm0))
dms += (dm0-dms[dm_max_jj])
dms[0] = max(0, dms[0])
global t_min, t_max
# if smearing timescale is < 4*pulse width,
# downsample before dedispersion for speed
downsamp_smear = max(1, int(downsamp*dt/tdm/2.))
# ensure that it's not larger than pulse width
downsamp_smear = int(min(downsamp, downsamp_smear))
downsamp_res = int(downsamp//downsamp_smear)
downsamp = int(downsamp_res*downsamp_smear)
time_res = dt * downsamp
tplot = ntime_plot * downsamp
# print("Width_full:%d Width_smear:%d Width_res: %d" %
# (downsamp, downsamp_smear, downsamp_res))
start_bin = int(t0/dt - ntime_plot*downsamp//2)
width = abs(4.148e3 * dm0 * (freq_up**-2 - freq_low**-2))
chunksize = int(width/dt + ntime_plot*downsamp)
t_min, t_max = 0, ntime_plot*downsamp
if start_bin < 0:
extra = start_bin
start_bin = 0
t_min += extra
t_max += extra
t_min, t_max = int(t_min), int(t_max)
snr_max = 0
# Account for the pre-downsampling to speed up dedispersion
t_min /= downsamp_smear
t_max /= downsamp_smear
ntime = t_max-t_min
data = rawdatafile.get_spectra(start_bin, chunksize)
if rficlean is True:
data = cleandata(data, threshold_time, threshold_frequency, \
bin_size, n_iter_time, n_iter_frequency, clean_type)
return data, downsamp, downsamp_smear
def inject_in_filterbank(fn_fil,
fn_out_dir,
N_FRB=1,
NFREQ=1536,
NTIME=2**15,
rfi_clean=False,
dm=250.0,
freq=(1550, 1250),
dt=0.00004096,
chunksize=5e4,
calc_snr=True,
start=0,
freq_ref=1400.,
subtract_zero=False,
clipping=None):
""" Inject an FRB in each chunk of data
at random times. Default params are for Apertif data.
Parameters:
-----------
fn_fil : str
name of filterbank file
fn_out_dir : str
directory for output files
N_FRB : int
number of FRBs to inject
NTIME : int
number of time samples per data chunk
rfi_clean : bool
apply rfi filters
dm : float / tuple
dispersion measure(s) to inject FRB with
freq : tuple
(freq_bottom, freq_top)
dt : float
time resolution
chunksize : int
size of data in samples to read in
calc_snr : bool
calculates S/N of injected pulse
start : int
start sample
freq_ref : float
reference frequency for injection code
subtract_zero : bool
subtract zero DM timestream from data
clipping :
zero out bright events in zero-DM timestream
Returns:
--------
None
"""
if type(dm) is not tuple:
max_dm = dm
else:
max_dm = max(dm)
t_delay_max = abs(4.14e3 * max_dm * (freq[0]**-2 - freq[1]**-2))
t_delay_max_pix = int(t_delay_max / dt)
# ensure that dispersion sweep is not too large
# for chunksize
f_edge = 0.3
while chunksize <= t_delay_max_pix / f_edge:
chunksize *= 2
NTIME *= 2
print('Increasing to NTIME:%d, chunksize:%d' % (NTIME, chunksize))
ii = 0
params_full_arr = []
ttot = int(N_FRB * chunksize * dt)
timestr = time.strftime("%Y%m%d-%H%M")
fn_fil_out = '%s/dm%s_nfrb%d_%s_sec_%s.fil' % (fn_out_dir, dm, N_FRB, ttot,
timestr)
fn_params_out = fn_fil_out.strip('.fil') + '.txt'
f_params_out = open(fn_params_out, 'w+')
f_params_out.write(
'# DM Sigma Time (s) Sample Downfact\n')
f_params_out.close()
for ii in xrange(N_FRB):
# drop FRB in random location in data chunk
offset = int(
np.random.uniform(0.1 * chunksize, (1 - f_edge) * chunksize))
data_filobj, freq_arr, delta_t, header = reader.read_fil_data(
fn_fil, start=start + chunksize * ii, stop=chunksize)
if ii == 0:
fn_rfi_clean = reader.write_to_fil(np.zeros([NFREQ, 0]), header,
fn_fil_out)
data = data_filobj.data
# injected pulse time in seconds since start of file
t0_ind = offset + NTIME // 2 + chunksize * ii
t0_ind = start + chunksize * ii + offset # hack because needs to agree with presto
t0 = t0_ind * delta_t
if len(data) == 0:
break
data_event = (data[:, offset:offset + NTIME]).astype(np.float)
data_event, params = simulate_frb.gen_simulated_frb(
NFREQ=NFREQ,
NTIME=NTIME,
sim=True,
# fluence=3000*(1+0.1*ii),
fluence=(200, 500),
spec_ind=0,
width=(delta_t, delta_t * 100),
dm=dm + 10 * ii,
scat_factor=(-4, -3.5),
background_noise=data_event,
delta_t=delta_t,
plot_burst=False,
freq=(freq_arr[0], freq_arr[-1]),
FREQ_REF=freq_ref,
scintillate=False)
dm_ = params[0]
params.append(offset)
print("%d/%d Injecting with DM:%d width: %.2f offset: %d" %
(ii, N_FRB, dm_, params[2], offset))
data[:, offset:offset + NTIME] = data_event
#params_full_arr.append(params)
width = params[2]
downsamp = max(1, int(width / delta_t))
t_delay_mid = 4.15e3 * dm_ * (freq_ref**-2 - freq_arr[0]**-2)
# this is an empirical hack. I do not know why
# the PRESTO arrival times are different from t0
# by the dispersion delay between the reference and
# upper frequency
t0 -= t_delay_mid
if rfi_clean is True:
data = rfi_test.apply_rfi_filters(data.astype(np.float32), delta_t)
if subtract_zero is True:
print("Subtracting zero DM")
data_ts_zerodm = data.mean(0)
data -= data_ts_zerodm[None]
if clipping is not None:
# Find tsamples > 8sigma and replace them with median
assert type(clipping) in (float,
int), 'clipping must be int or float'
data_ts_zerodm = data.mean(0)
stds, med = sigma_from_mad(data_ts_zerodm)
ind = np.where(np.absolute(data_ts_zerodm - med) > 8.0 * stds)[0]
data[:, ind] = np.median(data, axis=-1, keepdims=True)
if ii < 0:
fn_rfi_clean = reader.write_to_fil(data.transpose(), header,
fn_fil_out)
elif ii >= 0:
fil_obj = reader.filterbank.FilterbankFile(fn_fil_out,
mode='readwrite')
fil_obj.append_spectra(data.transpose())
if calc_snr is True:
data_filobj.data = data
data_filobj.dedisperse(dm_)
end_t = abs(4.15e3 * dm_ * (freq[0]**-2 - freq[1]**-2))
end_pix = int(end_t / dt)
end_pix_ds = int(end_t / dt / downsamp)
data_rb = (data_filobj.data).copy()
data_rb = data_rb[:, :-end_pix].mean(0)
data_rb -= np.median(data_rb)
SNRTools = tools.SNR_Tools()
snr_max, width_max = SNRTools.calc_snr_widths(data_rb,
widths=range(100))
# snr_max2, width_max2 = tools.calc_snr_widths(data_rb,
# )
print("S/N: %.2f width_used: %.3f width_tru: %.3f DM: %.1f" %
(snr_max, width_max, width / delta_t, dm_))
else:
snr_max = 10.0
width_max = int(width / dt)
f_params_out = open(fn_params_out, 'a+')
f_params_out.write('%2f %2f %5f %7d %d\n' %
(params[0], snr_max, t0, t0_ind, width_max))
f_params_out.close()
del data, data_event
params_full_arr = np.array(params_full_arr)
