def inject(infile, outfn, prof, period, dm, nbitsout=None,
block_size=BLOCKSIZE, pulsar_only=False, inplace=False):
if isinstance(infile, filterbank.FilterbankFile):
fil = infile
elif inplace:
fil = filterbank.FilterbankFile(infile, 'readwrite')
else:
fil = filterbank.FilterbankFile(infile, 'read')
print("Injecting pulsar signal into: %s" % fil.filename)
if False:
delays = psr_utils.delay_from_DM(dm, fil.frequencies)
delays -= delays[np.argmax(fil.frequencies)]
get_phases = lambda times: (times-delays)/period % 1
else:
get_phases = lambda times: times/period % 1
# Create the output filterbank file
if nbitsout is None:
nbitsout = fil.nbits
if inplace:
warnings.warn("Injecting pulsar signal *in-place*")
outfil = fil
else:
# Start an output file
print("Creating out file: %s" % outfn)
outfil = filterbank.create_filterbank_file(outfn, fil.header, \
nbits=nbitsout, mode='append')
if outfil.nbits == 8:
raise NotImplementedError("This code is out of date. 'delays' is not " \
"done in this way anymore..")
# Read the first second of data to get the global scaling to use
onesec = fil.get_timeslice(0, 1).copy()
onesec_nspec = onesec.shape[0]
times = np.atleast_2d(np.arange(onesec_nspec)*fil.tsamp).T+delays
phases = times/period % 1
onesec += prof(phases)
minimum = np.min(onesec)
median = np.median(onesec)
# Set median to 1/3 of dynamic range
global_scale = (256.0/3.0) / median
del onesec
else:
# No scaling to be performed
# These values will cause scaling to keep data unchanged
minimum = 0
global_scale = 1
sys.stdout.write(" %3.0f %%\r" % 0)
sys.stdout.flush()
oldprogress = -1
# Loop over data
lobin = 0
spectra = fil.get_spectra(0, block_size)
numread = spectra.shape[0]
while numread:
if pulsar_only:
# Do not write out data from input file
# zero it out
spectra *= 0
hibin = lobin+numread
# Sample at middle of time bin
times = (np.arange(lobin, hibin, 1.0/NINTEG_PER_BIN)+0.5/NINTEG_PER_BIN)*fil.dt
#times = (np.arange(lobin, hibin)+0.5)*fil.dt
phases = get_phases(times)
profvals = prof(phases)
shape = list(profvals.shape)
shape[1:1] = [NINTEG_PER_BIN]
shape[0] /= NINTEG_PER_BIN
profvals.shape = shape
toinject = profvals.mean(axis=1)
#toinject = profvals
if np.ndim(toinject) > 1:
injected = spectra+toinject
else:
injected = spectra+toinject[:,np.newaxis]
scaled = (injected-minimum)*global_scale
if inplace:
outfil.write_spectra(scaled, lobin)
else:
outfil.append_spectra(scaled)
# Print progress to screen
progress = int(100.0*hibin/fil.nspec)
if progress > oldprogress:
sys.stdout.write(" %3.0f %%\r" % progress)
sys.stdout.flush()
oldprogress = progress
# Prepare for next iteration
lobin = hibin
spectra = fil.get_spectra(lobin, block_size)
numread = spectra.shape[0]
sys.stdout.write("Done \n")
sys.stdout.flush()
def main(fits_fn, outfn, nbits, \
apply_weights, apply_scales, apply_offsets,
output_subints, skip):
start = time.time()
firstfits = pyfits.open(fits_fn[0],memmap=True)
nchan = firstfits['SUBINT'].header['NCHAN']
nsamps = firstfits['SUBINT'].header['NSBLK']
nsubints = firstfits['SUBINT'].header['NAXIS2']
fil_header = translate_header(firstfits,skip,output_subints)
fil_header['nbits'] = nbits
outfil = filterbank.create_filterbank_file(outfn, fil_header, \
nbits=nbits)
# if frequency channels are in ascending order
# band will need to be flipped
if firstfits['SUBINT'].header['CHAN_BW'] > 0:
flip_band=True
print("\nFits file frequencies in ascending order.")
print("\tFlipping frequency band.\n")
else:
flip_band=False
# check nbits for input
input_nbits = firstfits['SUBINT'].header['NBITS']
if input_nbits < 4:
raise ValueError('Does not support %d-bit data' % input_nbits)
if nbits != 32:
print("\nCalculating statistics on first subintegration...")
subint0 = read_subint(firstfits,0,nchan,nsamps, \
apply_weights, apply_scales, apply_offsets, \
input_nbits=input_nbits)
#new_max = np.mean(subint0) + 3*np.std(subint0)
new_max = 3 * np.median(subint0)
print("\t3*median =",new_max)
if new_max > 2.0**nbits:
scale = True
scale_fac = new_max / ( 2.0**nbits )
print("\tScaling data by",1/scale_fac)
print("\tValues larger than",new_max,"(pre-scaling) "\
"will be set to",2.0**nbits - 1,"\n")
else:
scale = False
scale_fac = 1
print("\tNo scaling necessary")
print("\tValues larger than",2.0**nbits-1,"(2^nbits) will "\
"be set to ",2.0**nbits-1,"\n")
else:
scale_fac = 1
print("\nNo scaling necessary for 32-bit float output file.")
firstfits.close()
fits = [pyfits.open(filename,memmap=True) for filename in fits_fn]
print("Writing data...")
sys.stdout.flush()
oldpcnt = ""
for i in range(skip+1,output_subints+skip+1):
index = (i-1)/320
subint_in_file = i-1-(index * 320)
subint = read_subint(fits[index],subint_in_file,nchan,nsamps, \
apply_weights, apply_scales, apply_offsets, \
input_nbits=input_nbits)
if flip_band:
subint = np.fliplr(subint)
subint /= scale_fac
outfil.append_spectra(subint)
pcnt = "%d" % (i*100.0/output_subints)
if pcnt != oldpcnt:
sys.stdout.write("% 4s%% complete\r" % pcnt)
sys.stdout.flush()
print("Done ")
outfil.close()
print("Runtime:",time.time() - start)
def main(fits_fn, outfn, nbits, \
apply_weights, apply_scales, apply_offsets):
start = time.time()
psrfits_file = psrfits.PsrfitsFile(fits_fn)
fil_header = translate_header(psrfits_file)
fil_header['nbits'] = nbits
outfil = filterbank.create_filterbank_file(outfn, fil_header, \
nbits=nbits)
# if frequency channels are in ascending order
# band will need to be flipped
if psrfits_file.fits['SUBINT'].header['CHAN_BW'] > 0:
flip_band=True
print("\nFits file frequencies in ascending order.")
print("\tFlipping frequency band.\n")
else:
flip_band=False
# check nbits for input
if psrfits_file.nbits < 4:
raise ValueError('Does not support %d-bit data' % \
psrfits_file.nbits)
if nbits != 32:
print("\nCalculating statistics on first subintegration...")
subint0 = psrfits_file.read_subint(0, \
apply_weights, apply_scales, apply_offsets)
#new_max = np.mean(subint0) + 3*np.std(subint0)
new_max = 3 * np.median(subint0)
print("\t3*median =",new_max)
if new_max > 2.0**nbits:
scale = True
scale_fac = new_max / ( 2.0**nbits )
print("\tScaling data by",1/scale_fac)
print("\tValues larger than",new_max,"(pre-scaling) "\
"will be set to",2.0**nbits - 1,"\n")
else:
scale = False
scale_fac = 1
print("\tNo scaling necessary")
print("\tValues larger than",2.0**nbits-1,"(2^nbits) will "\
"be set to ",2.0**nbits-1,"\n")
else:
scale_fac = 1
print("\nNo scaling necessary for 32-bit float output file.")
print("Writing data...")
sys.stdout.flush()
oldpcnt = ""
for isub in range(int(psrfits_file.nsubints)):
subint = psrfits_file.read_subint(isub, apply_weights, apply_scales, apply_offsets)
if flip_band:
subint = np.fliplr(subint)
subint /= scale_fac
outfil.append_spectra(subint)
pcnt = "%d" % (isub*100.0/psrfits_file.nsubints)
if pcnt != oldpcnt:
sys.stdout.write("% 4s%% complete\r" % pcnt)
sys.stdout.flush()
print("Done ")
outfil.close()
print("Runtime:",time.time() - start)
def main(args):
infn = args[0]
print("Reading filterbank file (%s)" % infn)
fil = filterbank.FilterbankFile(infn)
if options.start_time is None:
startbin = 0
else:
startbin = int(np.round(options.start_time/fil.tsamp))
if options.end_time is None:
endbin = fil.nspec
else:
endbin = int(np.round(options.end_time/fil.tsamp))+1
new_nspec = endbin-startbin
if new_nspec <= 0:
raise ValueError("Bad number of spectra to be written (%d). " \
"Check start/end times." % new_nspec)
# Determine lo/hi channels to write to file
# If high frequencies come first in spectra 'hichan' refers to
# the lo-freq cutoff and 'lochan' refers to the hi-freq cutoff.
if options.lo_freq is None:
if fil.foff > 0:
lochan = 0
else:
hichan = fil.nchans
else:
ichan = int(np.round((options.lo_freq-fil.fch1)/fil.foff))
if fil.foff > 0:
lochan = ichan
else:
hichan = ichan+1
if options.hi_freq is None:
if fil.foff > 0:
hichan = fil.nchans
else:
lochan = 0
else:
ichan = int(np.round((options.hi_freq-fil.fch1)/fil.foff))
if fil.foff > 0:
hichan = ichan+1
else:
lochan = ichan
new_nchans = hichan-lochan
if new_nchans <= 0:
raise ValueError("Bad number of channels to be written (%d). " \
"Check lo/hi frequencies." % new_nchans)
print("Will extract")
print(" %d bins (%d to %d incl.)" % (new_nspec, startbin, endbin-1))
print(" (Original num bins: %d)" % fil.nspec)
print(" %d channels (%d to %d incl.)" % (new_nchans, lochan, hichan-1))
print(" (Original num chans: %d)" % fil.nchans)
# Create output file
outfn = options.outname % fil.header
print("Creating out file: %s" % outfn)
outhdr = copy.deepcopy(fil.header)
outhdr['nchans'] = new_nchans
outhdr['fch1'] = fil.frequencies[lochan]
filterbank.create_filterbank_file(outfn, outhdr, nbits=fil.nbits)
outfil = filterbank.FilterbankFile(outfn, mode='write')
# Write data
sys.stdout.write(" %3.0f %%\r" % 0)
sys.stdout.flush()
nblocks = int(new_nspec/options.block_size)
remainder = new_nspec % options.block_size
oldprogress = -1
for iblock in np.arange(nblocks):
lobin = iblock*options.block_size + startbin
hibin = lobin+options.block_size
spectra = fil.get_spectra(lobin, hibin)
spectra = spectra[:,lochan:hichan] # restrict channels
outfil.append_spectra(spectra)
progress = int(100.0*((hibin-startbin)/new_nspec))
if progress > oldprogress:
sys.stdout.write(" %3.0f %%\r" % progress)
sys.stdout.flush()
oldprogress = progress
# Read all remaining spectra
if remainder:
spectra = fil.get_spectra(endbin-remainder, endbin)
spectra = spectra[:,lochan:hichan] # restrict channels
outfil.append_spectra(spectra)
sys.stdout.write("Done \n")
sys.stdout.flush()