예제 #1
0
파일: pfb.py 프로젝트: jrs65/pfb-inverse
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
예제 #2
0
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)


# =========================
#
예제 #3
0
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)
예제 #4
0
파일: pfb.py 프로젝트: CHIMEFRB/pfb-inverse
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