def maskfile(maskfn, data, start_bin, nbinsextra):
rfimask = rfifind.rfifind(maskfn)
mask = get_mask(rfimask, start_bin, nbinsextra)[::-1]
masked_chans = mask.all(axis=1)
# Mask data
data = data.masked(mask, maskval='median-mid80')
#datacopy = copy.deepcopy(data)
return data, masked_chans
def waterfall(rawdatafile, start, duration, dm=None, nbins=None, nsub=None,\
subdm=None, zerodm=False, downsamp=1, scaleindep=False,\
width_bins=1, mask=False, maskfn=None, bandpass_corr=False,
ref_freq=None):
"""
Create a waterfall plot (i.e. dynamic specrum) from a raw data file.
Inputs:
rawdatafile - a PsrfitsData instance.
start - start time of the data to be read in for waterfalling.
duration - duration of data to be waterfalled.
Optional Inputs:
dm - DM to use when dedispersing data.
Default: Don't de-disperse
nbins - Number of time bins to plot. This option overrides
the duration argument.
Default: determine nbins from duration.
nsub - Number of subbands to use. Must be a factor of number of channels.
Default: Number of channels.
subdm - DM to use when subbanding. Default: same as dm argument.
zerodm - subtract mean of each time-sample from data before
de-dispersing.
downsamp - Factor to downsample in time by. Default: Don't downsample.
scaleindep - Scale each channel independently.
Default: Scale using global maximum.
width_bins - Smooth each channel/subband with a boxcar width_bins wide.
Default: Don't smooth.
maskfn - Filename of RFIFIND mask to use for masking data.
Default: Don't mask data.
bandpass_corr - Correct for the bandpass. Requires an rfifind
mask provided by maskfn keyword argument.
Default: Do not remove bandpass.
ref_freq - Reference frequency to de-disperse to.
If subbanding and de-dispersing the start time
will be corrected to account for change in
reference frequency.
Default: Frequency of top channel.
Outputs:
data - Spectra instance of waterfalled data cube.
nbinsextra - number of time bins read in from raw data.
nbins - number of bins in duration.
start - corrected start time.
"""
if subdm is None:
subdm = dm
# Read data
if ref_freq is None:
ref_freq = rawdatafile.freqs.max()
if nsub and dm:
df = rawdatafile.freqs[1] - rawdatafile.freqs[0]
nchan_per_sub = rawdatafile.nchan / nsub
top_ctrfreq = rawdatafile.freqs.max() - \
0.5*nchan_per_sub*df # center of top subband
start += 4.15e3 * np.abs(1. / ref_freq**2 - 1. / top_ctrfreq**2) * dm
start_bin = np.round(start / rawdatafile.tsamp).astype('int')
dmfac = 4.15e3 * np.abs(1. / rawdatafile.frequencies[0]**2 -
1. / rawdatafile.frequencies[-1]**2)
if nbins is None:
nbins = np.round(duration / rawdatafile.tsamp).astype('int')
if dm:
nbinsextra = np.round(
(duration + dmfac * dm) / rawdatafile.tsamp).astype('int')
else:
nbinsextra = nbins
# If at end of observation
if (start_bin + nbinsextra) > rawdatafile.nspec - 1:
nbinsextra = rawdatafile.nspec - 1 - start_bin
data = rawdatafile.get_spectra(start_bin, nbinsextra)
# Masking
if mask and maskfn:
data, masked_chans = maskfile(maskfn, data, start_bin, nbinsextra)
else:
masked_chans = np.zeros(rawdatafile.nchan, dtype=bool)
# Bandpass correction
if maskfn and bandpass_corr:
bandpass = rfifind.rfifind(maskfn).bandpass_avg[::-1]
#bandpass[bandpass == 0] = np.min(bandpass[np.nonzero(bandpass)])
masked_chans[bandpass == 0] = True
# ignore top and bottom 1% of band
ignore_chans = np.ceil(0.01 * rawdatafile.nchan)
masked_chans[:ignore_chans] = True
masked_chans[-ignore_chans:] = True
data_masked = np.ma.masked_array(data.data)
data_masked[masked_chans] = np.ma.masked
data.data = data_masked
if bandpass_corr:
data.data /= bandpass[:, None]
# Zerodm filtering
if (zerodm == True):
data.data -= data.data.mean(axis=0)
# Subband data
if (nsub is not None) and (subdm is not None):
data.subband(nsub, subdm, padval='mean')
# Dedisperse
if dm:
data.dedisperse(dm, padval='mean')
# Downsample
data.downsample(downsamp)
# scale data
data = data.scaled(scaleindep)
# Smooth
if width_bins > 1:
data.smooth(width_bins, padval='mean')
return data, nbinsextra, nbins, start
def waterfall(rawdatafile, start, duration, dm=None, nbins=None, nsub=None,\
subdm=None, zerodm=False, downsamp=1, scaleindep=False,\
width_bins=1, mask=False, maskfn=None, csv_file=None, bandpass_corr=False, \
ref_freq=None):
"""
Create a waterfall plot (i.e. dynamic specrum) from a raw data file.
Inputs:
rawdatafile - a PsrfitsData instance.
start - start time of the data to be read in for waterfalling.
duration - duration of data to be waterfalled.
Optional Inputs:
dm - DM to use when dedispersing data.
Default: Don't de-disperse
nbins - Number of time bins to plot. This option overrides
the duration argument.
Default: determine nbins from duration.
nsub - Number of subbands to use. Must be a factor of number of channels.
Default: Number of channels.
subdm - DM to use when subbanding. Default: same as dm argument.
zerodm - subtract mean of each time-sample from data before
de-dispersing.
downsamp - Factor to downsample in time by. Default: Don't downsample.
scaleindep - Scale each channel independently.
Default: Scale using global maximum.
width_bins - Smooth each channel/subband with a boxcar width_bins wide.
Default: Don't smooth.
maskfn - Filename of RFIFIND mask to use for masking data.
Default: Don't mask data.
bandpass_corr - Correct for the bandpass. Requires an rfifind
mask provided by maskfn keyword argument.
Default: Do not remove bandpass.
ref_freq - Reference frequency to de-disperse to.
If subbanding and de-dispersing the start time
will be corrected to account for change in
reference frequency.
Default: Frequency of top channel.
Outputs:
data - Spectra instance of waterfalled data cube.
nbinsextra - number of time bins read in from raw data.
nbins - number of bins in duration.
start - corrected start time.
"""
if subdm is None:
subdm = dm
# Read data
if ref_freq is None:
ref_freq = rawdatafile.freqs.max()
if nsub and dm:
df = rawdatafile.freqs[1] - rawdatafile.freqs[0]
nchan_per_sub = rawdatafile.nchan / nsub
top_ctrfreq = rawdatafile.freqs.max() - \
0.5*nchan_per_sub*df # center of top subband
start += 4.15e3 * np.abs(1. / ref_freq**2 - 1. / top_ctrfreq**2) * dm
try:
source_name = rawdatafile.header['source_name']
except:
source_name = "Unknown"
start_bin = np.round(start / rawdatafile.tsamp).astype('int')
dmfac = 4.15e3 * np.abs(1. / rawdatafile.frequencies[0]**2 -
1. / rawdatafile.frequencies[-1]**2)
if nbins is None:
nbins = np.round(duration / rawdatafile.tsamp).astype('int')
if dm:
nbinsextra = np.round(
(duration + dmfac * dm) / rawdatafile.tsamp).astype('int')
else:
nbinsextra = nbins
# If at end of observation
if (start_bin + nbinsextra) > rawdatafile.nspec - 1:
nbinsextra = rawdatafile.nspec - 1 - start_bin
data = rawdatafile.get_spectra(start_bin, nbinsextra)
# Masking
if mask and maskfn:
data, masked_chans = maskfile(maskfn, data, start_bin, nbinsextra)
else:
masked_chans = np.zeros(rawdatafile.nchan, dtype=bool)
# Bandpass correction
if maskfn and bandpass_corr:
bandpass = rfifind.rfifind(maskfn).bandpass_avg[::-1]
#bandpass[bandpass == 0] = np.min(bandpass[np.nonzero(bandpass)])
masked_chans[bandpass == 0] = True
# ignore top and bottom 5% of band
ignore_chans = int(np.ceil(0.05 * rawdatafile.nchan))
masked_chans[:ignore_chans] = True
masked_chans[-ignore_chans:] = True
data_masked = np.ma.masked_array(data.data)
data_masked[masked_chans] = np.ma.masked
data.data = data_masked
if bandpass_corr:
data.data /= bandpass[:, None]
# Zerodm filtering
if (zerodm == True):
data.data -= data.data.mean(axis=0)
# Subband data
if (nsub is not None) and (subdm is not None):
data.subband(nsub, subdm, padval='mean')
# Dedisperse
if dm:
data.dedisperse(dm, padval='rotate')
# Downsample
# Moving it after the DM-vs-time plot
#data.downsample(downsamp)
#Orignal scale has error if chan is flagged so writting this here
if not scaleindep:
std = data.data.std()
for ii in range(data.numchans):
chan = data.get_chan(ii)
median = np.median(chan)
if scaleindep:
std = chan.std()
if std:
chan[:] = (chan - median) / std
else:
chan[:] = chan
# scale data
#data = data.scaled(scaleindep)
#data = data.scaled(False)
# Smooth
if width_bins > 1:
for ii in range(data.numchans):
chan = data.get_chan(ii)
nbin = len(chan)
x = list(range(nbin))
if (nbin > 4000): deg = 10
if (nbin < 4000 and nbin > 2000): deg = 8
if (nbin < 2000 and nbin > 1000): deg = 6
if (nbin < 1000 and nbin > 150): deg = 5
if (nbin < 150): deg = 2
base = np.polyfit(x, chan, deg)
p = np.poly1d(base)
chan[:] = chan[:] - p(x)
return data, nbinsextra, nbins, start, source_name