def inverse_pfb_parallel(ts_pfb, ntap, nblock, window=sinc_hamming,
no_nyquist=False, skip_initial_blocks=True):
"""Invert the CHIME PFB timestream.
Parameters
----------
ts_pfb : np.ndarray[nsamp, nfreq]
The PFB timestream.
ntap : integer
The number of number of blocks combined into the final timestream.
window : function (ntap, lblock) -> np.ndarray[lblock * ntap]
The window function to apply to each block.
no_nyquist : boolean, optional
If True, we are missing the Nyquist frequency (i.e. CHIME PFB), and we
should add it back in (with zero amplitude).
skip_initial_blocks : boolean, optional
If True (default), throw away the initial, heavily unconstrained
blocks of samples.
"""
import mpiutil
# If we are missing the Nyquist freq (default for CHIME), add it back in
if no_nyquist:
new_shape = ts_pfb.shape[:-1] + (ts_pfb.shape[-1] + 1,)
pts2 = np.zeros(new_shape, dtype=np.float64)
pts2[..., :-1] = ts_pfb
ts_pfb = pts2
# Inverse fourier transform to get the pseudo-timestream
pseudo_ts = np.fft.irfft(ts_pfb, axis=-1)
# Pull out the number of blocks and their length
lblock = pseudo_ts.shape[-1]
ntsblock = nblock + ntap - 1
# Transpose timestream
pseudo_ts = mpiutil.transpose_blocks(pseudo_ts, (nblock, lblock))
pseudo_ts = pseudo_ts.T.copy()
# Get local offset range
local_off, start_off, end_off = mpiutil.split_local(lblock)
# Coefficients for the P matrix
# Create the window array
coeff_P = window(ntap, lblock).reshape(ntap, lblock)[:, start_off:end_off]
# Coefficients for the PP^T matrix
coeff_PPT = np.array([(coeff_P[:, np.newaxis, :] *
coeff_P[np.newaxis, :, :])
.diagonal(offset=k).sum(axis=-1)
for k in range(ntap)])
ax2size = nblock if skip_initial_blocks else ntsblock
rec_ts = np.zeros((local_off, ax2size), dtype=np.float64)
# print mpiutil.rank, local_off, pseudo_ts.shape, coeff_P.shape
for i_off in range(local_off):
# Create band matrix representation of P
band_P = np.zeros((ntap, ntsblock), dtype=np.float64)
band_P[:] = coeff_P[::-1, i_off, np.newaxis]
# Create band matrix representation of PP^T (symmetric)
band_PPT = np.zeros((ntap, nblock), dtype=np.float64)
band_PPT[:] = coeff_PPT[::-1, i_off, np.newaxis]
# Solve for intermediate vector
yh = la.solveh_banded(band_PPT, pseudo_ts[i_off])
# Project into timestream estimate
tt = band_mv(band_P, 0, 3, ntsblock, nblock, yh, trans=True)
# Trim off initial blocks
if skip_initial_blocks:
tt = tt[(ntap-1):]
rec_ts[i_off] = tt
# Transpose timestream back
rec_ts = mpiutil.transpose_blocks(rec_ts, (lblock, ax2size))
rec_ts = rec_ts.T.copy()
return rec_ts
print args # ========================= # # This is where we must load in the data. At the moment we have a very stupid method. # # At the end of it, each process must end up with a section of the file from # start_block to end_block, and these must be as doubles. # # ========================= pfb_data = np.load(args.file_in).reshape(-1, args.nfreq) # Load in whole file! nblock = pfb_data.shape[0] # Find out the file length in blocks local_block, start_block, end_block = mpiutil.split_local(nblock) # Work out how to split up the file into sections pfb_local = pfb_data[start_block:end_block] # Pull out the correct section of the file. # ========================= # Apply inverse PFB rects = pfb.inverse_pfb_parallel(pfb_local, args.ntap, nblock, no_nyquist=args.no_nyquist) # Calculate the range of local timestream blocks local_tsblock, start_tsblock, end_tsblock = mpiutil.split_local(nblock) # ========================= #
print args
#=========================
#
# This is where we must load in the data. At the moment we have a very stupid method.
#
# At the end of it, each process must end up with a section of the file from
# start_block to end_block, and these must be as doubles.
#
#=========================
pfb_data = np.load(args.file_in).reshape(-1, args.nfreq) # Load in whole file!
nblock = pfb_data.shape[0] # Find out the file length in blocks
local_block, start_block, end_block = mpiutil.split_local(
nblock) # Work out how to split up the file into sections
pfb_local = pfb_data[start_block:
end_block] # Pull out the correct section of the file.
#=========================
# Apply inverse PFB
rects = pfb.inverse_pfb_parallel(pfb_local,
args.ntap,
nblock,
no_nyquist=args.no_nyquist)
# Calculate the range of local timestream blocks
local_tsblock, start_tsblock, end_tsblock = mpiutil.split_local(nblock)
def inverse_pfb_parallel(
ts_pfb,
ntap,
nblock,
window=sinc_hamming,
no_nyquist=False,
skip_initial_blocks=True,
):
"""Invert the CHIME PFB timestream.
Parameters
----------
ts_pfb : np.ndarray[nsamp, nfreq]
The PFB timestream.
ntap : integer
The number of number of blocks combined into the final timestream.
window : function (ntap, lblock) -> np.ndarray[lblock * ntap]
The window function to apply to each block.
no_nyquist : boolean, optional
If True, we are missing the Nyquist frequency (i.e. CHIME PFB), and we
should add it back in (with zero amplitude).
skip_initial_blocks : boolean, optional
If True (default), throw away the initial, heavily unconstrained
blocks of samples.
"""
import mpiutil
# If we are missing the Nyquist freq (default for CHIME), add it back in
if no_nyquist:
new_shape = ts_pfb.shape[:-1] + (ts_pfb.shape[-1] + 1, )
pts2 = np.zeros(new_shape, dtype=np.float64)
pts2[..., :-1] = ts_pfb
ts_pfb = pts2
# Inverse fourier transform to get the pseudo-timestream
pseudo_ts = np.fft.irfft(ts_pfb, axis=-1)
# Pull out the number of blocks and their length
lblock = pseudo_ts.shape[-1]
ntsblock = nblock + ntap - 1
# Transpose timestream
pseudo_ts = mpiutil.transpose_blocks(pseudo_ts, (nblock, lblock))
pseudo_ts = pseudo_ts.T.copy()
# Get local offset range
local_off, start_off, end_off = mpiutil.split_local(lblock)
# Coefficients for the P matrix
# Create the window array
coeff_P = window(ntap, lblock).reshape(ntap, lblock)[:, start_off:end_off]
# Coefficients for the PP^T matrix
coeff_PPT = np.array([
(coeff_P[:, np.newaxis, :] *
coeff_P[np.newaxis, :, :]).diagonal(offset=k).sum(axis=-1)
for k in range(ntap)
])
ax2size = nblock if skip_initial_blocks else ntsblock
rec_ts = np.zeros((local_off, ax2size), dtype=np.float64)
# print mpiutil.rank, local_off, pseudo_ts.shape, coeff_P.shape
for i_off in range(local_off):
# Create band matrix representation of P
band_P = np.zeros((ntap, ntsblock), dtype=np.float64)
band_P[:] = coeff_P[::-1, i_off, np.newaxis]
# Create band matrix representation of PP^T (symmetric)
band_PPT = np.zeros((ntap, nblock), dtype=np.float64)
band_PPT[:] = coeff_PPT[::-1, i_off, np.newaxis]
# Solve for intermediate vector
yh = la.solveh_banded(band_PPT, pseudo_ts[i_off])
# Project into timestream estimate
tt = band_mv(band_P, 0, 3, ntsblock, nblock, yh, trans=True)
# Trim off initial blocks
if skip_initial_blocks:
tt = tt[(ntap - 1):]
rec_ts[i_off] = tt
# Transpose timestream back
rec_ts = mpiutil.transpose_blocks(rec_ts, (lblock, ax2size))
rec_ts = rec_ts.T.copy()
return rec_ts