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