def get_data_to_write(self, start_sample, nsamp): """ Read data to self.data, selects channels Optionally perform RFI filtering and zero-DM subtraction Args: start_sample (int): Start sample number to read from nsamp (int): Number of samples to read """ data = self.your_object.get_data(start_sample, nsamp) data = data[:, self.chan_min:self.chan_max] if self.flag_rfi: mask = sk_sg_filter( data, self.your_object, self.nchans, self.sk_sig, self.sg_fw, self.sg_sig, ) if self.your_object.your_header.dtype == np.uint8: data[:, mask] = np.around(np.mean(data[:, ~mask])) else: data[:, mask] = np.mean(data[:, ~mask]) if self.zero_dm_subt: logger.debug("Subtracting 0-DM time series from the data") data = data - data.mean(1)[:, None] data = data.astype(self.your_object.your_header.dtype) self.data = data
def get_data_to_write(self, start_sample, nsamp): """ Read data to self.data, selects channels Optionally perform RFI filtering and zero-DM subtraction Args: start_sample (int): Start sample number to read from nsamp (int): Number of samples to read """ data = self.your_object.get_data(start_sample, nsamp, npoln=self.npoln) if self.npoln == 1: data = np.expand_dims(data, 1) # shape of data is (nt, npoln, nf) data = data[:, :, self.chan_min : self.chan_max] if self.flag_rfi: for i in range(data.shape[1]): data_to_flag = data[:, i, :] mask = sk_sg_filter( data_to_flag, self.your_object, self.sk_sig, self.sg_fw, self.sg_sig, ) if self.replacement_policy == "mean": fill_value = np.mean(data_to_flag[:, ~mask]) elif self.replacement_policy == "median": fill_value = np.median(data_to_flag[:, ~mask]) else: fill_value = 0 if self.your_object.your_header.nbits < 32: fill_value = np.around(fill_value).astype( self.your_object.your_header.dtype ) data[:, i, mask] = fill_value if self.zero_dm_subt: if self.npoln > 1: raise NotImplementedError( "0-DM subtraction is implemented only for 1 output pol." ) logger.debug("Subtracting 0-DM time series from the data.") min_value = np.iinfo(self.your_object.your_header.dtype).min max_value = np.iinfo(self.your_object.your_header.dtype).max nt, npoln, nf = data.shape ts = data.mean(-1) bp = data.mean(0).squeeze() for channel in range(nf): data[:, :, channel] = np.clip( data[:, :, channel].astype("float32") - ts + bp[channel], min_value, max_value, ) data = data.astype(self.your_object.your_header.dtype) # shape of data is (nt, npoln, nf) if self.highest_frequency_first and self.your_object.your_header.foff > 0: data = data[:, :, ::-1] self.data = data
def get_chunk(self, tstart=None, tstop=None, for_preprocessing=True): """ Get a chunk of data. The data is saved in `self.data`. Args: tstart (float): start time of the chunk in seconds tstop (float): stop time of the chunk in seconds for_preprocessing (bool): if the data is to be preprocessed later. This will modify the number of samples read based on the width of the candidate """ if tstart is None: tstart = (self.tcand - self.dispersion_delay() - self.width * self.native_tsamp) if tstop is None: tstop = (self.tcand + self.dispersion_delay() + self.width * self.native_tsamp) logger.debug(f"tstart is {tstart}") logger.debug(f"tstop is {tstop}") nstart = int(tstart / self.native_tsamp) nsamp = int((tstop - tstart) / self.native_tsamp) nsamp_read = nsamp if for_preprocessing: if self.width > 2 and nsamp_read // (self.width // 2) < self.min_samp: nsamp_read = self.min_samp * self.width // 2 nstart_read = nstart - (nsamp_read - nsamp) // 2 if nsamp_read < self.min_samp: nsamp_read = self.min_samp nstart_read = nstart - (nsamp_read - nsamp) // 2 else: nstart_read = nstart else: nstart_read = nstart logging.debug( f"nstart_read is {nstart_read}, nsamp_read is {nsamp_read}," f"nstart is {nstart}, nsamp is {nsamp}") nspectra = int(self.your_header.nspectra) if nstart_read >= 0 and nstart_read + nsamp_read <= nspectra: logging.debug( f"All the data available in the file, no need to pad. \n" f"nstart_read({nstart_read})>=0 and \n" f"nstart_read({nstart_read})+nsamp_read({nsamp_read})<=nspectra({nspectra})" ) data = self.get_data(nstart=nstart_read, nsamp=nsamp_read) elif nstart_read < 0: if nstart_read + nsamp_read <= nspectra: logging.debug( f"nstart_read({nstart_read})<0 and nstart_read({nstart_read})\n" f"+nsamp_read({nsamp_read})<=nspectra({nspectra})") logging.info("Padding with median in the beginning") d = self.get_data(nstart=0, nsamp=nsamp_read + nstart_read) dmedian = np.median(d, axis=0) data = (np.ones( (nsamp_read, self.your_header.nchans), dtype=self.your_header.dtype, ) * dmedian[None, :]) data[-nstart_read:, :] = d else: logging.debug( f"nstart_read({nstart_read})<0 and nstart_read({nstart_read})" f"+nsamp_read({nsamp_read})>nspectra({nspectra})") logging.info( "Padding with median in the beginning and the end") d = self.get_data(nstart=0, nsamp=nspectra) dmedian = np.median(d, axis=0) data = (np.ones( (nsamp_read, self.your_header.nchans), dtype=self.your_header.dtype, ) * dmedian[None, :]) data[-nstart_read:-nstart_read + nspectra, :] = d else: logging.debug( f"nstart_read({nstart_read})>=0 and nstart_read({nstart_read})" f"+nsamp_read({nsamp_read})>nspectra({nspectra})") logging.info("Padding with median in the end") d = self.get_data(nstart=nstart_read, nsamp=nspectra - nstart_read) dmedian = np.median(d, axis=0) data = (np.ones((nsamp_read, self.your_header.nchans), dtype=self.your_header.dtype) * dmedian[None, :]) data[:nspectra - nstart_read, :] = d self.data = data.astype(self.your_header.dtype) if self.kill_mask.any(): logger.info("Applying the kill mask") assert len(self.kill_mask) == self.data.shape[1] data_copy = self.data.copy() data_copy[:, self.kill_mask] = 0 self.data = data_copy del data_copy if self.flag_rfi: data_copy = self.data.copy() mask = sk_sg_filter( data=data_copy, your_object=self, spectral_kurtosis_sigma=self.spectral_kurtosis_sigma, savgol_frequency_window=self.savgol_frequency_window, savgol_sigma=self.savgol_sigma, ) self.rfi_mask = mask data_copy[:, self.rfi_mask] = 0 self.data = data_copy del data_copy return self
def to_fil(self, nstart=None, nsamp=None, c=None, outdir=None, outname=None, flag_rfi=False, progress=None, sk_sig=4, sg_fw=15, sg_sig=4, zero_dm_subt=False): """ Writes out a Filterbank File. Args: nstart: Start sample to read from nsamp: Number of samples to write c: Required frequency channel range [min_chan, max_chan] (excludes the higher channel number) outdir: Output directory for Filterbank file outname: Name of the Filterbank file to write to progress: Turn on/off progress bar flag_rfi: To turn on RFI flagging sk_sig: Sigma for spectral kurtosis filter sg_fw: Filter window for savgol filter sg_sig: Sigma for savgol filter zero_dm_subt: Enable zero-DM RFI excision """ if c: min_c = int(np.min(c)) max_c = int(np.max(c)) else: min_c = 0 max_c = len(self.your_obj.chan_freqs) chan_freq = self.your_obj.chan_freqs[min_c:max_c] nchans = len(chan_freq) # Calculate loop of spectra if not nstart: nstart = 0 if not nsamp: nsamp = self.your_obj.your_header.native_nspectra interval = 4096 * 24 if nsamp < interval: interval = nsamp if nsamp > interval: nloops = 1 + nsamp // interval else: nloops = 1 nstarts = np.arange(nstart, interval * nloops, interval, dtype=int) nsamps = np.full(nloops, interval) if nsamp % interval != 0: nsamps = np.append(nsamps, [nsamp % interval]) original_dir, orig_basename = os.path.split( self.your_obj.your_header.filename) if not outname: name, ext = os.path.splitext(orig_basename) if ext == '.fits': temp = name.split('_') if len(temp) > 1: outname = '_'.join(temp[:-1]) + '_converted.fil' else: outname = name + '_converted.fil' else: outname = name + '_converted.fil' if not outdir: outdir = original_dir # Read data for st, samp in tqdm.tqdm(zip(nstarts, nsamps), total=len(nstarts), disable=progress): logger.debug( f'Reading spectra {st}-{st + samp} in file {self.your_obj.your_header.filename}' ) data = self.your_obj.get_data(st, samp) data = data[:, min_c:max_c] if flag_rfi: mask = sk_sg_filter(data, self.your_obj, sk_sig, nchans, sg_fw, sg_sig) if self.your_obj.your_header.dtype == np.uint8: data[:, mask] = np.around(np.mean(data[:, ~mask])) else: data[:, mask] = np.mean(data[:, ~mask]) if zero_dm_subt: logger.debug('Subtracting 0-DM time series from the data') data = data - data.mean(1)[:, None] data = data.astype(self.your_obj.your_header.dtype) logger.info( f'Writing data from spectra {st}-{st + samp} in the frequency channel range {min_c}-{max_c} ' f'to filterbank') write_fil(data, self.your_obj, nchans=nchans, chan_freq=chan_freq, outdir=outdir, filename=outname, nstart=nstart) logger.debug( f'Successfully written data from spectra {st}-{st + samp} to filterbank' ) logging.debug(f'Read all the necessary spectra')
def to_fits(self, nstart=None, c=None, nsamp=None, npsub=-1, outdir=None, outname=None, progress=None, flag_rfi=False, sk_sig=4, sg_fw=15, sg_sig=4, zero_dm_subt=False): """ Writes out a PSRFITS file Args: nstart: Start sample number to read from nsamp: Number of samples to read/write c: Required frequency channel range [min_chan, max_chan] (excludes the higher channel number) npsub: Number of spectra per subint outdir: Output directory for Filterbank file outname: Name of the PSRFITS file to write to progress: Turn on/off progress bar flag_rfi: To turn on RFI flagging sk_sig: Sigma for spectral kurtosis filter sg_fw: Filter window for savgol filter sg_sig: Sigma for savgol filter zero_dm_subt: Enable zero-DM RFI excision """ tsamp = self.your_obj.your_header.tsamp if npsub == -1: npsub = int(1.0 / tsamp) else: pass if nsamp: if nsamp < npsub: npsub = nsamp if not outname: original_dir, orig_basename = os.path.split( self.your_obj.your_header.filename) name, ext = os.path.splitext(orig_basename) if ext == '.fits': temp = name.split('_') if len(temp) > 1: outname = '_'.join(temp[:-1]) + '_converted.fits' else: outname = name + '_converted.fits' else: outname = name + '_converted.fits' if not outdir: outdir = os.getcwd() outfile = outdir + '/' + outname if c: min_c = int(np.min(c)) max_c = int(np.max(c)) else: min_c = 0 max_c = len(self.your_obj.chan_freqs) chan_freqs = self.your_obj.chan_freqs[min_c:max_c] nchans = len(chan_freqs) initialize_psrfits(outfile=outfile, y=self.your_obj, npsub=npsub, nstart=nstart, nsamp=nsamp, chan_freqs=chan_freqs) nifs = self.your_obj.your_header.npol logger.info("Filling PSRFITS file with data") # Open PSRFITS file hdulist = fits.open(outfile, mode='update') hdu = hdulist[1] nsubints = len(hdu.data[:]['data']) # Loop through chunks of data to write to PSRFITS n_read_subints = 10 if not nstart: nstart = 0 logger.info(f'Number of subints to write {nsubints}') for istart in tqdm.tqdm(np.arange(0, nsubints, n_read_subints), disable=progress): istop = istart + n_read_subints if istop > nsubints: istop = nsubints else: pass isub = istop - istart logger.info( f"Writing data to {outfile} from subint = {istart} to {istop}." ) # Read in nread samples from filfile nread = isub * npsub data = self.your_obj.get_data(nstart=nstart, nsamp=nread) data = data[:, min_c:max_c] if flag_rfi: mask = sk_sg_filter(data, self.your_obj, sk_sig, nchans, sg_fw, sg_sig) if self.your_obj.your_header.dtype == np.uint8: data[:, mask] = np.around(np.mean(data[:, ~mask])) else: data[:, mask] = np.mean(data[:, ~mask]) if zero_dm_subt: logger.debug('Subtracting 0-DM time series from the data') data = data - data.mean(1)[:, None] logger.debug(f'Shape of data array after get_data is {data.shape}') nstart += nread nvals = isub * npsub * nifs if data.shape[0] < nvals: logger.debug( f'nspectra in this chunk ({data.shape[0]}) < nsubints * npsub * nifs ({nvals})' ) logger.debug(f'Appending zeros at the end to fill the subint') pad_back = np.zeros((nvals - data.shape[0], data.shape[1])) data = np.vstack((data, pad_back)) else: pass data = np.reshape(data, (isub, npsub, nifs, nchans)) # If foff is negative, we need to flip the freq axis # if foff < 0: # logger.debug(f"Flipping band as {foff} < 0") # data = data[:, :, :, ::-1] # else: # pass # Put data in hdu data array logger.debug(f'Writing data of shape {data.shape} to {outfile}.') hdu.data[istart:istop]['data'][:, :, :, :] = data[:].astype( self.your_obj.your_header.dtype) # Write to file hdulist.flush() logger.info(f'All spectra written to {outfile}') # Close open FITS file hdulist.close()