def gen_simpulse(self, ntime=10000):
""" Generate pulse dynamic spectrum
with simpulse
"""
undispersed_arrival_time = 0.5 * ntime * self.dt
# undispersed_arrival_time -= 4148*self.dm*(self.freq_hi_MHz**-2)
sp = simpulse.single_pulse(
ntime,
self._nfreq,
self._freq_lo_MHz,
self._freq_hi_MHz,
self._dm,
self._scat_tau_ref,
self._width_i,
self._fluence,
self.spec_ind,
undispersed_arrival_time,
)
data_simpulse = np.zeros([self._nfreq, ntime])
sp.add_to_timestream(data_simpulse, 0.0, ntime * self._dt)
data_simpulse = data_simpulse[::-1]
return data_simpulse
def fit_sim_frb(dm=1000.0, nfreq=1536, freq=(1219.70092773, 1519.50561523),
scat_tau_ref=0.001, spec_ind=0.0, dt=8.192e-5, width=0.0001, save_data=False):
ntime = np.int(2*4148*dm*np.abs(freq[0]**-2-freq[-1]**-2)/dt)
print(nfreq, ntime)
undispersed_arrival_time = 0.5*ntime*dt
undispersed_arrival_time -= 4148*dm*(max(freq)**-2)
sp = simpulse.single_pulse(ntime, nfreq, min(freq), max(freq),
dm, scat_tau_ref, width, 10.0,
spec_ind, undispersed_arrival_time)
data_simpulse = np.zeros([nfreq, ntime])
sp.add_to_timestream(data_simpulse, 0.0, ntime*dt)
data_event = data_simpulse[::-1]
data_dedisp = tools.dedisperse(data_event, dm, freq=(freq[1], freq[0]))
# plt.imshow(data_event, aspect='auto')
# plt.show()
if save_data:
np.save('./herial.npy', data_dedisp)
f,s = fit_width_freq(data_dedisp, freq=(freq[1], freq[0]), plot=True)
s = np.array(s)
plt.plot(f, s/s[len(s)//2], color='k', lw=3)
plt.plot(f, (f/np.median(f))**-2.0)
# plt.plot(f, (f/np.median(f))**-3.0)
plt.plot(f, (f/np.median(f))**-4.0)
plt.legend(['data','-2','-4'])
plt.loglog()
plt.show()
def plot1():
"""
Some Gaussians with a little bit of dispersion and scattering.
Visual checks:
- arrival times of pulses should be 141.5, 118.4, 110.4 ms
- fluences should be { 13.3, 30, 53.3 }
- low frequency pulse should show scattering
- no dispersion broadening visible
"""
pulse_nt = 1024
nfreq = 512
freq_lo_MHz = 1000.0
freq_hi_MHz = 2000.0
dm = 10.0
sm = 4.0
intrinsic_width = 0.005
fluence = 30.0
spectral_index = 2.0
undispersed_arrival_time = 0.100
sp = simpulse.single_pulse(pulse_nt, nfreq, freq_lo_MHz, freq_hi_MHz, dm,
sm, intrinsic_width, fluence, spectral_index,
undispersed_arrival_time)
make_plot(sp,
ifreq_list=[0, 256, 511],
color_list=['r', 'b', 'm'],
label_list=['1 GHz', '1.5 GHz', '2.0 GHz'],
filename='plot1.png')
def plot3():
"""
Same as plot1, but with intrisic_width set to zero, so that pulse width is dominated by scattering,
and spectral_inex set to zero.
"""
pulse_nt = 1024
nfreq = 512
freq_lo_MHz = 1000.0
freq_hi_MHz = 2000.0
dm = 10.0
sm = 4.0
intrinsic_width = 0.0
fluence = 30.0
spectral_index = 0.0
undispersed_arrival_time = 0.100
sp = simpulse.single_pulse(pulse_nt, nfreq, freq_lo_MHz, freq_hi_MHz, dm,
sm, intrinsic_width, fluence, spectral_index,
undispersed_arrival_time)
make_plot(sp,
ifreq_list=[0, 256, 511],
color_list=['r', 'b', 'm'],
label_list=['1 GHz', '1.5 GHz', '2.0 GHz'],
filename='plot3.png')
def plot2():
"""Boxcars labeled by time intervals"""
pulse_nt = 1024
nfreq = 7
freq_lo_MHz = 1000.0
freq_hi_MHz = 2000.0
dm = 10.0
sm = 0.0
intrinsic_width = 0.0
fluence = 1.0
spectral_index = 0.0
undispersed_arrival_time = 0.100
sp = simpulse.single_pulse(pulse_nt, nfreq, freq_lo_MHz, freq_hi_MHz, dm,
sm, intrinsic_width, fluence, spectral_index,
undispersed_arrival_time)
make_plot(sp,
ifreq_list=[0, 2, 4, 6],
color_list=['r', 'b', 'm', 'k'],
label_list=[
'131.8-141.5', '116.8-120.3', '114.1-116.8', '110.4-112.0'
],
filename='plot2.png')
def _make_single_pulse_object(self, nupfreq=1):
return simpulse.single_pulse(
nt=1024,
nfreq=self.nfreq * nupfreq,
freq_lo_MHz=self.freq_lo_MHz,
freq_hi_MHz=self.freq_hi_MHz,
dm=self.dm,
sm=self.sm,
intrinsic_width=self.intrinsic_width,
fluence=self.fluence,
spectral_index=self.spectral_index,
undispersed_arrival_time=self.undispersed_arrival_time)
def _make_single_pulse_object(self, nupfreq=1):
return simpulse.single_pulse(
nt = 1024,
nfreq = self.nfreq * nupfreq,
freq_lo_MHz = self.freq_lo_MHz,
freq_hi_MHz = self.freq_hi_MHz,
dm = self.dm,
sm = self.sm,
intrinsic_width = self.intrinsic_width,
fluence = self.fluence,
spectral_index = self.spectral_index,
undispersed_arrival_time = self.undispersed_arrival_time
)
def _bind_transform(self, json_attrs):
try:
import simpulse
except ImportError:
raise RuntimeError("rf_pipelines: couldn't import the 'simpulse' module. You may need to clone https://github.com/kmsmith137/simpulse and install.")
self.freq_lo_MHz = json_attrs['freq_lo_MHz']
self.freq_hi_MHz = json_attrs['freq_hi_MHz']
self.dt_sample = json_attrs['dt_sample']
self.t_initial = json_attrs.get('t_initial', 0)
self.t_final = self.t_initial
# Construct a simpulse.single_pulse object, which represents a single FRB signal, to be added to the timestream in chunks.
#
# The mechanism implemented in simpulse for obtaining a pulse with specified signal-to-noise ratio is a little awkward!
# First we construct the pulse with an arbitrary normalization, then we compute the signal-to-noise of the pulse given
# its arbitrary normalization, then we adjust the normalization accordingly.
self.pulse = simpulse.single_pulse(1024, # number of samples used "under the hood" in the simpulse library
self.nfreq,
self.freq_lo_MHz,
self.freq_hi_MHz,
self.dm,
self.sm,
self.intrinsic_width,
1.0, # fluence, placeholder value to be changed below based on S/N
self.spectral_index,
self.undispersed_arrival_time)
# Compute the signal-to-noise ratio of the pulse, for the arbitrary normalization specified in the constructor.
#
# Currently, we take sample_rms to be a scalar, which simpulse interprets as meaning that all frequency
# channels have the same rms. More generally, if sample_rms is a 1D numpy array of length nfreq, then
# simpulse interprets it as a per-frequency rms.
#
# If the rms is frequency-dependent, then the weighting of the channels is also important when determining
# the signal-to-noise. This is specified via the 'channel_weights' argument, also an array of length nfreq,
# to simpulse.single_pulse.get_signal_to_noise(). The best thing to use here would be a per-frequency estimate
# of the average value of the pipeline 'weights' array, but I propose just taking the weights to be 0.0 in
# bad frequency channels (defined as channels with many v2 failures during variance estimation), and 1.0 in
# non-bad channels. Note that leaving channel_weights unspecified (or equivalently setting channel_weights=None)
# defaults to 1/sigma_rms^2 weighting, which is inconsistent with what we're currently doing in the pipeline.
#
# (BTW the general issue of optimal frequency channel weighting is a nontrivial one that we'll want to revisit,
# so the channel-weighting logic in the previous paragraph will probably be changed later!)
snr0 = self.pulse.get_signal_to_noise(self.dt_sample, self.sample_rms, self.sample_weights)
# The pulse has been simulated with nominal fluence 1, and its computed signal-to-noise is 'snr0'.
# We want to end up with signal-to-noise 'snr'. Therefore, the following fluence gives the signal-to-noise we want.
self.pulse.fluence = self.snr / snr0
def run(self, nmc, noisy=True):
"""
Runs 'nmc' Monte Carlo simulations, and appends the results to the 'class ensemble' internal state.
In principle the interface is general enough to make either the simulation or the dedisperser the
outer loop. For example:
# Run 1000 Monte Carlo simulations, with outer loop over simulation and inner loop over dedisperser.
ensemble.run(1000)
# Run 1000 Monte Carlo simulations, with outer loop over dedipserser and inner loop over simulation.
for i in xrange(1000):
ensemble.run(1)
However, I ended up deciding to use the latter form (outer loop over dedipserser and inner loop over simulation)
exclusively, so the interface could be simplified by replacing run(nmc) by run_one().
"""
nfreq = self.search_params.nfreq
nsamples = self.search_params.nsamples
freq_lo_MHz = self.search_params.freq_lo_MHz
freq_hi_MHz = self.search_params.freq_hi_MHz
dm_max = self.search_params.dm_max
nmc_in = len(self.sim_json)
intensity = np.zeros((nfreq, nsamples), dtype=np.float32)
timestream_length = nsamples * self.search_params.dt_sample
max_dispersion_delay = dispersion_delay(
dm_max, freq_lo_MHz) - dispersion_delay(dm_max, freq_hi_MHz)
noise_seeds = []
# Outer loop over dedispersers.
for (id, d) in enumerate(self.dedisperser_list):
d.allocate()
# Inner loop over simulations.
for imc in xrange(nmc):
if noisy:
print 'frb_olympics: dedisperser %d/%d (%s), simulation %d' % (
id + 1, len(self.dedisperser_list), d.tex_label,
nmc_in + imc + 1)
if id == 0:
# Simulate random FRB params.
true_params = {
'dm':
np.random.uniform(self.search_params.dm_min,
self.search_params.dm_max),
'sm':
np.random.uniform(self.search_params.sm_min,
self.search_params.sm_max),
'spectral_index':
np.random.uniform(
self.search_params.spectral_index_min,
self.search_params.spectral_index_max),
'intrinsic_width':
np.random.uniform(
self.search_params.intrinsic_width_min,
self.search_params.intrinsic_width_max),
'snr':
np.random.uniform(self.search_params.snr_min,
self.search_params.snr_max),
'tmid':
np.random.uniform(
0.51 * max_dispersion_delay,
timestream_length - 0.51 * max_dispersion_delay)
}
this_sim = {
'true_params': true_params,
'recovered_params':
[] # this list will be populated by the dedispersers
}
# Save FRB params and RNG state
self.sim_json.append(this_sim)
noise_seeds.append(np.random.get_state())
else:
# Use same random FRB params and RNG state as previous simulation.
true_params = self.sim_json[nmc_in + imc]['true_params']
np.random.set_state(noise_seeds[imc])
# Simulate Gaussian random noise. There is no float32 gaussian random number generator in numpy,
# so we simulate in float64 and down-convert. We do this in slices to save memory!
if self.add_noise:
for i in xrange(nsamples):
intensity[:, i] = np.random.standard_normal(size=nfreq)
else:
intensity[:, :] = 0.
# Add simulated FRB. Note that we pay the computational cost of simulating the pulse
# "from scratch" for every dedisperser. This sometimes (if nmc is large) saves memory,
# and the cost of simulating the pulse is small.
# The simpulse library uses the undispersed arrival time 'tu' of the pulse,
# whereas frb_olympics uses the central arrival time 'tmid', so we need to translate.
dt_i = dispersion_delay(true_params['dm'], freq_hi_MHz)
dt_f = dispersion_delay(true_params['dm'], freq_lo_MHz)
true_tu = true_params['tmid'] - dt_i - (dt_f - dt_i) / 2.0
# Note that 'nt' is the number of samples used internally by simpulse to represent
# the pulse. (FIXME: tune this parameter.)
p = simpulse.single_pulse(
nt=1024,
nfreq=nfreq,
freq_lo_MHz=freq_lo_MHz,
freq_hi_MHz=freq_hi_MHz,
dm=true_params['dm'],
sm=true_params['sm'],
intrinsic_width=true_params['intrinsic_width'],
fluence=1.0,
spectral_index=true_params['spectral_index'],
undispersed_arrival_time=true_tu)
# We simulate the pulse with an nominal normalization (fluence=1.0), and
# rescale to the target SNR.
nominal_snr = p.get_signal_to_noise(
self.search_params.dt_sample)
rescaling_factor = true_params['snr'] / nominal_snr
p.fluence *= rescaling_factor
# Add FRB here!
# Note that we order frequency channels from lowest to highest (the intuitive ordering,
# but some dedispersers assume the opposite and will need to reverse it, for example
# bonsai and heimdall).
p.add_to_timestream(intensity, 0.0, timestream_length)
# Run dedisperser here!
dedisperser_output = d.dedisperse(intensity)
# The return value of d.dedisperser() should be a dictionary containing 'snr', 'dm',
# and precisely one of { 'tmid', 'tini', or 'tfin' }.
if not isinstance(dedisperser_output, dict):
raise RuntimeError(
'expected return value of %s.dedisperse() to be a dict, got %s'
% (d.__class__.__name__,
dedisperser_output.__class__.__name__))
required_keys = set(['snr', 'dm'])
optional_keys = set(['tini', 'tmid', 'tfin'])
missing_keys = required_keys.difference(
dedisperser_output.keys())
unrecognized_keys = set(dedisperser_output.keys()).difference(
required_keys).difference(optional_keys)
if len(missing_keys) > 0:
raise RuntimeError(
'return value of %s.dedisperse() does not contain key(s) %s'
% (d.__class__.__name__, sorted(missing_keys)))
if len(unrecognized_keys) > 0:
raise RuntimeError(
'return value of %s.dedisperse() contains unrecognized key(s) %s'
% (d.__class__.__name__, sorted(unrecognized_keys)))
if len(dedisperser_output.keys()) != 3:
raise RuntimeError(
'return value of %s.dedisperse() must contain precisely one of %s'
% (d.__class__.__name__, sorted(optional_keys)))
# If tini or tfin were specified, translate to a value of tmid.
dm = dedisperser_output['dm']
dt = dispersion_delay(dm, freq_lo_MHz) - dispersion_delay(
dm, freq_hi_MHz)
if not dedisperser_output.has_key(
'tmid') and dedisperser_output.has_key('tini'):
dedisperser_output[
'tmid'] = dedisperser_output['tini'] + dt / 2.0
if not dedisperser_output.has_key(
'tmid') and dedisperser_output.has_key('tfin'):
dedisperser_output[
'tmid'] = dedisperser_output['tfin'] - dt / 2.0
assert dedisperser_output.has_key('tmid')
# Done with this (simulation, dedisperser) pair.
# Record the results in the 'sim_json' data structure.
self.sim_json[nmc_in + imc]['recovered_params'].append(
dedisperser_output)
# End of inner loop over simulations; back to outer loop over dedispersers.
d.deallocate()
freq_lo_MHz = 1000.0
freq_hi_MHz = 2000.0
dm = 1000.0
sm = 0.0
intrinsic_width = 0.0005
fluence = 1.0
spectral_index = 0.
undispersed_arrival_time = 0.10
dedisp_delay = 4148 * dm * (freq_lo_MHz**-2. - freq_hi_MHz**-2.)
pulse_nt = int(2 * dedisp_delay / dt + 1000)
nfreq = 1024
freq_ref = 1500.
freq_arr = np.linspace(freq_lo_MHz, freq_hi_MHz, nfreq)
sp = simpulse.single_pulse(pulse_nt, nfreq, freq_lo_MHz, freq_hi_MHz, dm, sm,
intrinsic_width, fluence, spectral_index,
undispersed_arrival_time)
data_simpulse = np.zeros([nfreq, pulse_nt])
sp.add_to_timestream(data_simpulse, 0.0, pulse_nt * dt)
data_simpulse = data_simpulse[::-1]
data_injfrb, p = simulate_frb.gen_simulated_frb(
NFREQ=nfreq,
NTIME=pulse_nt,
sim=True,
fluence=fluence,
spec_ind=0.0,
width=intrinsic_width,
dm=dm,
background_noise=np.zeros([nfreq, pulse_nt]),
def inject_in_filterbank(fn_fil,
fn_out_dir,
N_FRB=1,
NFREQ=1536,
NTIME=2**15,
rfi_clean=False,
dm=1000.0,
chunksize=2,
calc_snr=True,
start=0,
freq_ref=None,
clipping=None,
gaussian=False,
gaussian_noise=True,
upchan_factor=2,
upsamp_factor=2,
simulator='injectfrb',
paramslist=None,
noise_std=18.,
nbit=8):
""" 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
dt : float
time resolution
chunksize : int
size of data in samples to read in
calc_snr : bool
calculates S/N of injected pulse using pulse profile
as matched filter
start : int
start sample
freq_ref : float
reference frequency for injection code
clipping :
zero out bright events in zero-DM timestream
nbit : int
number of bits in filterbank data
Returns:
--------
None
"""
assert simulator in ['simpulse',
'injectfrb'], "Do not recognize simulator backend"
if simulator == 'simpulse':
import simpulse
if paramslist is not None:
params_arr = np.loadtxt(paramslist)
params_arr = params_arr.transpose()
dm_max = params_arr[0].max()
dm_min = params_arr[0].min()
if len(params_arr.shape) == 1:
params_arr = params_arr[:, None]
else:
params_arr = None
dm_max = 1000.
dm_min = 10.
SNRTools = tools.SNR_Tools()
data_fil_obj_skel, freq_arr, dt, header = reader.read_fil_data(fn_fil,
start=0,
stop=1)
NFREQ = header['nchans']
BW = np.abs(header['nchans'] * header['foff'])
if freq_ref is None:
freq_ref = 0.5 * (freq_arr[0] + freq_arr[-1])
if type(dm) is not tuple:
max_dm = dm
else:
max_dm = max(dm)
max_dm = max(max_dm, dm_max)
t_delay_max = abs(k_DM * max_dm * (freq_arr[0]**-2 - freq_arr[-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
ii = 0
params_full_arr = []
ttot = int(N_FRB * chunksize * dt)
timestr = time.strftime("%Y%m%d-%H%M")
fn_fil_out = '%s/%s_nfrb%d_DM%d-%d_%ssec_%s.fil' % (
fn_out_dir, simulator, N_FRB, dm_min, dm_max, 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 Width_int With_obs Spec_ind \
Scat_tau_ref Tsamp (s) BW_MHz Freq_hi Nchan Freq_ref\n'
)
f_params_out.close()
fmt_out = '%8.3f %5.2f %8.4f %9d %5d %1.6f %5f %5.2f %1.4f %1.4f %8.2f %8.2f %d %8.2f\n'
if gaussian == True:
fn_fil_out = fn_fil_out.strip('.fil') + '_gaussian.fil'
inject_in_filterbank_gaussian(data_fil_obj_skel, header, fn_fil_out,
N_FRB)
exit()
print("============ HEADER INFORMATION ============")
reader.print_filheader(header)
kk = 0
samplecounter = 0
for ii in xrange(N_FRB):
if params_arr is not None:
dm = params_arr[0, ii]
fluence = params_arr[1, ii]
width_sec = params_arr[2, ii]
spec_ind = params_arr[3, ii]
disp_ind = params_arr[4, ii]
scat_tau_ref = 0.
# tdm = 8.3e-6 * dm * np.diff(freq_arr)[0] / np.mean(freq_arr*1e-3)**3
# fluence *= ((width_sec**2 + dt**2 + tdm**2)**0.5 / width_sec)
else:
np.random.seed(np.random.randint(12312312))
fluence = np.random.uniform(0, 1)**(-2 / 3.)
dm = np.random.uniform(10., 2000.)
scat_tau_ref = 0.
spec_ind = 0.
width_sec = 2 * dt
if gaussian_noise is True:
t_delay_max = abs(k_DM * dm * (freq_arr[0]**-2 - freq_arr[-1]**-2))
t_delay_max_pix = np.int(3 * t_delay_max / dt)
t_chunksize_min = 3.0
chunksize = 2**(np.ceil(np.log2(t_delay_max_pix)))
chunksize = np.int(max(t_chunksize_min / dt, chunksize))
NTIME = chunksize
offset = 0
data_filobj, freq_arr, dt, header = reader.read_fil_data(fn_fil,
start=0,
stop=1)
data = np.empty([NFREQ, NTIME])
else:
fluence = 1000.
# drop FRB in random location in data chunk
#offset = random.randint(np.int(0.1*chunksize), np.int((1-f_edge)*chunksize))
offset = 0
data_filobj, freq_arr, dt, header = reader.read_fil_data(
fn_fil, start=start + chunksize * (ii - kk), stop=chunksize)
data = data_filobj.data
if ii == 0:
fn_rfi_clean = reader.write_to_fil(np.zeros([NFREQ, 0]), header,
fn_fil_out)
# injected pulse time in seconds since start of file
t0_ind = offset + NTIME // 2 + chunksize * ii
t0 = t0_ind * dt
if len(data) == 0:
break
if gaussian_noise is True:
if simulator == 'injectfrb':
data_event = np.zeros(
[upchan_factor * NFREQ, upsamp_factor * NTIME])
noise_event = np.random.normal(128, noise_std,
NFREQ * NTIME).reshape(
NFREQ, NTIME)
elif simulator == 'simpulse':
data_event = np.zeros([NFREQ, NTIME])
noise_event = np.random.normal(128, noise_std,
NFREQ * NTIME).reshape(
NFREQ, NTIME)
else:
print(
"Do not recognize simulator, neither (injectfrb, simpulse)"
)
exit()
else:
NTIME = np.int(2 * t_delay_max / dt)
if data.shape[1] < NTIME:
print("Not enough data in the filterbank file. Not injecting.")
return
data_event = (data[:, offset:offset + NTIME]).astype(np.float)
if simulator == 'injectfrb':
data_event, params = simulate_frb.gen_simulated_frb(
NFREQ=upchan_factor * NFREQ,
NTIME=upsamp_factor * NTIME,
sim=True,
fluence=fluence,
spec_ind=spec_ind,
width=width_sec,
dm=dm,
scat_tau_ref=scat_tau_ref,
background_noise=data_event,
delta_t=dt / upsamp_factor,
plot_burst=False,
freq=(freq_arr[0], freq_arr[-1]),
FREQ_REF=freq_ref,
scintillate=False)
print(data_event.shape)
data_event = data_event.reshape(NFREQ, upchan_factor, NTIME,
upsamp_factor).mean(-1).mean(1)
data_event *= (20. * noise_std / np.sqrt(NFREQ))
elif simulator == 'simpulse':
# Scaling to match fluence vals with injectfrb
fluence *= 5e-4
sp = simpulse.single_pulse(NTIME, NFREQ, freq_arr.min(),
freq_arr.max(), dm, scat_tau_ref,
width_sec, fluence, spec_ind, 0.)
sp.add_to_timestream(data_event, 0.0, NTIME * dt)
data_event = data_event[::-1]
data_event *= (10. * noise_std / np.sqrt(NFREQ))
# [dm, fluence, width, spec_ind, disp_ind, scat_tau_ref]
params = [dm, fluence, width_sec, spec_ind, 2., scat_tau_ref]
dm_ = params[0]
params.append(offset)
width_obs = np.sqrt(width_sec**2 + dt**2 + scat_tau_ref**2)
params[2] = width_obs
print("width_intrinsic: %0.2f\nwidth_obs: %0.2f" %
(1e3 * width_sec, 1e3 * width_obs))
print(
"%d/%d Injecting with DM:%d width_samp: %.1f offset: %d using %s" %
(ii + 1, N_FRB, dm_, width_obs / dt, offset, simulator))
data[:, offset:offset + NTIME] = data_event
width = width_obs
downsamp = max(1, int(width / dt))
t_delay_mid = k_DM * 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 #hack
#data_filobj.data = copy.copy(data)
data_filobj.data = data
# Note presto dedisperse assumes
if calc_snr is True:
print("Calculating true filter")
prof_true_filobj = copy.deepcopy(data_filobj)
prof_true_filobj.dedisperse(dm_, ref_freq=freq_ref)
prof_true = np.mean(prof_true_filobj.data, 0)
prof_true = prof_true[np.where(prof_true > prof_true.max() * 0.01)]
sig_total = np.sqrt((prof_true**2).sum())
else:
print("not calculating")
prof_true = None
if gaussian_noise is True:
data[:, offset:offset + NTIME] += noise_event
data[data > (2**nbit - 1)] = 2**nbit - 1
data_filobj.data = copy.copy(data)
data_filobj.dedisperse(dm_, ref_freq=freq_ref)
start_t = abs(k_DM * dm_ * (freq_arr[0]**-2 - freq_ref**-2))
start_pix = int(start_t / dt)
end_t = abs(k_DM * dm_ * (freq_arr[-1]**-2 - freq_ref**-2))
end_pix = int(end_t / dt)
data_rb = data_filobj.data
data_rb = data_rb[:, start_pix:-end_pix].mean(0)
widths_snr = range(int(max(downsamp / 2., 1)),
int(min(downsamp * 2, 2500)))
snr_max, width_max = SNRTools.calc_snr_matchedfilter(
data_rb, widths=widths_snr, true_filter=prof_true)
#plt.plot(data_filobj.data.mean(0))
#plt.legend(['%f %f' % (snr_max, width_max)])
#plt.axvline(np.argmax(data_filobj.data.mean(0)), color='red')
#plt.show()
local_thresh = 0.
if snr_max <= local_thresh:
print("S/N <= %d: Not writing to file" % local_thresh)
kk += 1
continue
print("S/N: %.2f width_used: %.1f width_tru: %.1f DM: %.1f" %
(snr_max, width_max, width / dt, dm_))
# Presto dedisperses to top of band, so this is at fmax
# t0_ind = np.argmax(data_filobj.data.mean(0)) + chunksize*ii
t0_ind = np.argmax(data_filobj.data.mean(0)) + samplecounter
t0 = t0_ind * dt #huge hack
if rfi_clean is True:
data = rfi_test.apply_rfi_filters(data.astype(np.float32), dt)
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'
print("Clipping data")
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())
samplecounter += data.shape[1]
f_params_out = open(fn_params_out, 'a+')
f_params_out.write(
fmt_out %
(params[0], snr_max, t0, t0_ind, downsamp, width_sec, width_obs,
spec_ind, scat_tau_ref, dt, BW, header['fch1'], NFREQ, freq_ref))
f_params_out.close()
del data, data_event