def main():
""" Command line tool for plotting dynamic spectrum. """
parser = ArgumentParser(
description=
"Process and plot dynamic spectrum chunks from one or more data sets.")
optional = parser._action_groups.pop()
required = parser.add_argument_group('required arguments')
required.add_argument('-i',
action='store',
required=True,
dest='inputs_cfg',
type=str,
help="Configuration script of inputs")
parser._action_groups.append(optional)
if len(sys.argv) == 1:
parser.print_help()
sys.exit(1)
parse_args = parser.parse_args()
# Initialize parameter values
inputs_cfg = parse_args.inputs_cfg
hotpotato = set_defaults(read_config(inputs_cfg))
# Profile code execution.
prog_start_time = time.time()
# Run data processing.
data, freqs_GHz, times = myexecute(hotpotato)
# Calculate total run time for the code.
prog_end_time = time.time()
run_time = (prog_end_time - prog_start_time) / 60.0
print('Code run time = %.5f minutes' % (run_time))
return data, freqs_GHz, times, hotpotato
from psrdynspec.modules.filters1d import blockavg1d, savgol_lowpass
from psrdynspec.modules.fft import fft1d_mask
from psrdynspec.modules.fold import fold_ts, execute_plan, fold_rotations_ts
from psrdynspec.modules.folding_plan import ProcessingPlan
from psrdynspec.modules.normalize_Nd import normalize_stdnormal
from psrdynspec.plotting.config import *
from psrdynspec.plotting.fft_plot import fft_gridplot
from psrdynspec.plotting.fold_plot import subplots_metric_profile, plot_foldedprofile_rotations
from riptide.running_median import fast_running_median
from astropy.stats import sigma_clip
import os, time
import numpy as np
############################################################################
# INPUTS
dict = read_config('search_presto_dat.cfg')
# Data
if (dict['DM']==''):
if 'DM' in dict['dat_file']:
dict['DM'] = float(dict['dat_file'].split('DM')[1].split('.dat')[0])
else:
dict['DM'] = 0.0
if (dict['low_freq_data']==''):
dict['low_freq_data'] = None
if (dict['high_freq_data']==''):
dict['high_freq_data'] = None
# Plotting
if (dict['show_plot']==''):
dict['show_plot'] = False
if (dict['plot_format']==''):
def __MPI_MAIN__(parser):
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
stat = MPI.Status()
nproc = comm.Get_size()
# Parent processor
if rank == 0:
print('STARTING RANK 0')
# Profile code execution.
prog_start_time = time.time()
parse_args = parser.parse_args()
# Initialize parameter values
inputs_cfg = parse_args.inputs_cfg
# Construct list of calls to run from shell.
hotpotato = set_defaults(read_config(inputs_cfg))
# Create output directory if non-existent.
if not os.path.isdir(hotpotato['OUTPUT_DIR']):
os.makedirs(hotpotato['OUTPUT_DIR'])
# Load information on single pulse candidates.
metadata, cand_DMs, cand_sigma, cand_dedisp_times, cand_dedisp_samples, cand_downfact, select_indices = filter_spcands(
hotpotato)
if len(select_indices) == 0:
print(
'No single pulse candidates fit the user-supplied selection criteria. Quitting program'
)
sys.exit(1)
# Read header of filterbank file.
hdr = Header(hotpotato['DATA_DIR'] + '/' + hotpotato['fil_file'],
file_type='filterbank') # Returns a Header object
tot_time_samples = hdr.ntsamples # Total no. of time samples in entire dynamic spectrum.
t_samp = hdr.t_samp # Sampling time (s)
chan_bw = hdr.chan_bw # Channel bandwidth (MHz)
nchans = hdr.nchans # No. of channels
npol = hdr.npol # No. of polarizations
n_bytes = hdr.primary['nbits'] / 8.0 # No. of bytes per data sample
hdr_size = hdr.primary['hdr_size'] # Header size (bytes)
# Set up frequency array. Frequencies in GHz.
freqs_GHz = (hdr.fch1 + np.arange(nchans) * chan_bw) * 1e-3
print(hdr)
# Flip frequency axis if chan_bw<0.
if (chan_bw < 0):
print('Channel bandwidth is negative.')
freqs_GHz = np.flip(freqs_GHz)
print('Frequencies rearranged in ascending order.')
# Chop bandpass edges.
hotpotato['ind_band_low'] = np.where(
freqs_GHz >= hotpotato['freq_band_low'])[0][0]
hotpotato['ind_band_high'] = np.where(
freqs_GHz <= hotpotato['freq_band_high'])[0][-1] + 1
# Clip bandpass edges.
freqs_GHz = freqs_GHz[
hotpotato['ind_band_low']:hotpotato['ind_band_high']]
# Load median bandpass, if pre-computed.
if hotpotato['bandpass_method'] == 'file':
print('Loading median bandpass from %s' %
(hotpotato['bandpass_npz']))
hotpotato['median_bp'] = np.load(
hotpotato['BANDPASS_DIR'] + '/' + hotpotato['bandpass_npz'],
allow_pickle=True)['Median bandpass']
hotpotato['median_bp'] = hotpotato['median_bp'][
hotpotato['ind_band_low']:hotpotato['ind_band_high']]
print('Median bandpass loaded.')
elif hotpotato['bandpass_method'] not in ['file', 'compute']:
print('Unrecognized bandpass computation method.')
sys.exit(1)
# Load rfifind mask.
print('Reading rfifind mask: %s' % (hotpotato['rfimask']))
nint, int_times, ptsperint, mask_zap_chans, mask_zap_ints, mask_zap_chans_per_int = read_rfimask(
hotpotato['RFIMASK_DIR'] + '/' + hotpotato['rfimask'])
mask_zap_chans, mask_zap_chans_per_int = modify_zapchans_bandpass(
mask_zap_chans, mask_zap_chans_per_int, hotpotato['ind_band_low'],
hotpotato['ind_band_high'])
if nproc == 1:
f = open(hotpotato['DATA_DIR'] + '/' + hotpotato['fil_file'], 'rb')
for i in range(len(select_indices)):
cand_index = select_indices[i]
downfact = cand_downfact[cand_index]
myexecute(cand_index, cand_DMs, cand_sigma, cand_dedisp_times,
downfact, metadata, int_times, mask_zap_chans,
mask_zap_chans_per_int, freqs_GHz, tot_time_samples,
t_samp, chan_bw, npol, nchans, n_bytes, hdr_size,
hotpotato, f, rank)
f.close()
else:
# Distribute candidates evenly among child processors.
indices_dist_list = np.array_split(select_indices, nproc - 1)
# Send data to child processors.
for indx in range(1, nproc):
comm.send((indices_dist_list[indx - 1], cand_DMs, cand_sigma,
cand_dedisp_times, cand_downfact, metadata,
int_times, mask_zap_chans, mask_zap_chans_per_int,
freqs_GHz, tot_time_samples, t_samp, chan_bw, npol,
nchans, n_bytes, hdr_size, hotpotato),
dest=indx,
tag=indx)
comm.Barrier(
) # Wait for all child processors to receive sent call.
# Receive Data from child processors after execution.
comm.Barrier()
# Calculate total run time for the code.
prog_end_time = time.time()
run_time = (prog_end_time - prog_start_time) / 60.0
print('Code run time = %.5f minutes' % (run_time))
print('FINISHING RANK 0')
else:
# Recieve data from parent processor.
indx_vals, cand_DMs, cand_sigma, cand_dedisp_times, cand_downfact, metadata, int_times, mask_zap_chans, mask_zap_chans_per_int, freqs_GHz, tot_time_samples, t_samp, chan_bw, npol, nchans, n_bytes, hdr_size, hotpotato = comm.recv(
source=0, tag=rank)
comm.Barrier()
print('STARTING RANK: ', rank)
f = open(hotpotato['DATA_DIR'] + '/' + hotpotato['fil_file'], 'rb')
for i in range(len(indx_vals)):
cand_index = indx_vals[i]
downfact = cand_downfact[cand_index]
myexecute(cand_index, cand_DMs, cand_sigma, cand_dedisp_times,
downfact, metadata, int_times, mask_zap_chans,
mask_zap_chans_per_int, freqs_GHz, tot_time_samples,
t_samp, chan_bw, npol, nchans, n_bytes, hdr_size,
hotpotato, f, rank)
f.close()
print('FINISHING RANK: ', rank)
comm.Barrier()
from psrdynspec import read_config, Header, load_psrfits_data
from psrdynspec.modules.ds_systematics import remove_additive_time_noise, correct_bandpass
from psrdynspec.modules.filters1d import blockavg1d
from psrdynspec.modules.filters2d import smooth_master
from psrdynspec.modules.dedisperse import dedisperse_ds
from psrdynspec.plotting.config import *
from psrdynspec.plotting.dedisperse_plot import plot_dedispersed_ds
from psrdynspec.plotting.bandpass_plot import plot_bandpass
from psrdynspec.plotting.ds_plot import plot_ds
import os, time
import numpy as np
############################################################################
# INPUTS
dict = read_config('build_dedispts_from_psrfits.cfg')
# Set default values for empty dictionary items.
if (dict['BANDPASS_DIR'] == ''):
dict['BANDPASS_DIR'] = dict['OUTPUT_DIR']
if (dict['DS_SAVE_DIR'] == ''):
dict['DS_SAVE_DIR'] = dict['OUTPUT_DIR']
if (dict['show_plot'] == ''):
dict['show_plot'] = False
if (dict['log_colorbar_ds'] == ''):
dict['log_colorbar_ds'] = False
if (dict['vmin_percentile'] == ''):
dict['vmin_percentile'] = None
if (dict['vmax_percentile'] == ''):
dict['vmax_percentile'] = None
if (dict['DM'] == ''):
def calc_bandpass(inputs_cfg):
# Profile code execution.
prog_start_time = time.time()
# Read inputs from config file and set default parameter values, if applicable.
hotpotato = read_config(inputs_cfg)
hotpotato = set_defaults(hotpotato)
# Read header.
print('Reading header of file %s' % (hotpotato['data_file']))
hdr = Header(hotpotato['DATA_DIR'] + '/' + hotpotato['data_file'],
file_type=hotpotato['file_type'])
print(hdr)
tot_time_samples = hdr.ntsamples # Total no. of time samples in entire dynamic spectrum.
t_samp = hdr.t_samp # Sampling time (s)
chan_bw = hdr.chan_bw # Channel bandwidth (MHz)
nchans = hdr.nchans # No. of channels
npol = hdr.npol # No. of polarizations
freqs_GHz = (hdr.fch1 + np.arange(nchans) *
chan_bw) * 1e-3 # 1D array of radio frequencies (GHz)
# Start and stop indices along the time axis.
t_start = int(hotpotato['start_time'] // t_samp)
if (hotpotato['duration'] == None):
t_stop = tot_time_samples
else:
t_stop = int(
(hotpotato['start_time'] + hotpotato['duration']) // t_samp)
times = np.arange(t_start, t_stop) * t_samp
# Extracting the data.
print('Reading %s:' % (hotpotato['data_file']))
if (hotpotato['file_type'] == 'psrfits'):
raw_ds, hdr, acc_start_time = extract_psrfits_datachunk(
hotpotato['DATA_DIR'] + '/' + hotpotato['data_file'], times[0],
times[-1], hotpotato['pol'])
elif (hotpotato['file_type'] == 'filterbank'):
f = open(hotpotato['DATA_DIR'] + '/' + hotpotato['data_file'], 'rb')
raw_ds = load_fil_data(f, t_start, t_stop, npol, nchans,
hdr.primary['nbits'] / 8.0,
hdr.primary['hdr_size'], hotpotato['pol'],
f.tell())
f.close()
else:
sys.exit('Unsupported file type: %s' % (hotpotato['file_type']))
print('Data extracted. Data shape = ', raw_ds.shape)
# Flip along the frequency axis if channel bandwidth is negative.
if (chan_bw < 0):
raw_ds = np.flip(raw_ds, axis=0)
freqs_GHz = np.flip(freqs_GHz)
print('Flipped frequency axis of dynamic spectrum.')
# Calculate median bandpass shape.
print("Calculating median bandpass...")
median_bp = calc_median_bandpass(raw_ds)
print('Median bandpass computed.')
# Create output directory if non-existent.
if not os.path.isdir(hotpotato['OUTPUT_DIR']):
os.makedirs(hotpotato['OUTPUT_DIR'])
# Section the raw DS into subbands as specified.
if (hotpotato['freqs_extract'] is not None):
freqs_subband = [
] # List to store frequency information for each subband.
bandpass_fit_subband = [
] # List to store calculated bandpass shape for each subband.
print('Sectioning raw DS into specified subbands..')
for j in range(len(hotpotato['freqs_extract'])):
ind_low = np.where(
freqs_GHz >= hotpotato['freqs_extract'][j, 0])[0][0]
ind_high = np.where(
freqs_GHz > hotpotato['freqs_extract'][j, 1])[0][0]
f_subband = freqs_GHz[ind_low:
ind_high] # Frequency array for this subband
# Smooth the median bandpass with a low-pass filter.
print('Smoothing median bandpass over the subband: %s - %s GHz' %
(hotpotato['freqs_extract'][j, 0],
hotpotato['freqs_extract'][j, 1]))
smooth_bp = savgol_lowpass(median_bp[ind_low:ind_high],
hotpotato['window_length'],
hotpotato['poly_degree'])
# Update lists to reflect extracted subband information.
freqs_subband.append(f_subband)
bandpass_fit_subband.append(smooth_bp)
# Covert lists to arrays.
freqs_subband = np.array(freqs_subband)
bandpass_fit_subband = np.array(bandpass_fit_subband)
# Plot bandpass with color-coded subbands.
plot_bandpass_subbands(
freqs_GHz, median_bp, freqs_subband, bandpass_fit_subband, 'GHz',
'arb. units',
hotpotato['OUTPUT_DIR'] + '/' + hotpotato['basename'],
hotpotato['show_plot'])
else:
freqs_subband = None
bandpass_fit_subband = None
plot_bandpass(freqs_GHz, median_bp, 'GHz', 'arb. units',
hotpotato['OUTPUT_DIR'] + '/' + hotpotato['basename'],
hotpotato['show_plot'])
# Save bandpass information to npz file
file_name = hotpotato['basename'] + '_bandpass'
save_array = [
freqs_GHz, median_bp, hotpotato['freqs_extract'], freqs_subband,
bandpass_fit_subband, times, hotpotato['window_length'],
hotpotato['poly_degree']
]
save_keywords = [
'Radio frequency (GHz)', 'Median bandpass', 'Band edges (GHz)',
'Subband Frequency (GHz)', 'Smooth subband bandpass', 'Time (s)',
'Savgol window size', 'Savgol polynomial deg'
]
np.savez(hotpotato['OUTPUT_DIR'] + '/' + file_name,
**{name: value
for name, value in zip(save_keywords, save_array)})
# Calculate total run time for the code.
prog_end_time = time.time()
run_time = (prog_end_time - prog_start_time) / 60.0
print('Code run time = %.2f minutes' % (run_time))