De-dispersion using cumulative sums of dynamical spectra

de_dispersion.py

import numpy as np

vint = np.vectorize(int)
vround = np.vectorize(round)


def de_disperse(image, nu_0, d_nu, d_t, dm_values):
    """
    De-disperse dynamical spestra with grid of user specifies values of DM.

    :param image:
        2D numpy array of dynamical spectra (#freq, #t).
    :param nu_0:
        Frequency of the highest frequency channel [MHz].
    :param d_nu:
        Width of spectral channel [MHz].
    :param d_t:
        Time step [s].
    :param dm_values:
        Array-like of DM values to dedisperse [cm^3 /pc].

    :return:
        2D numpy array (#DM, #t)

    """
    dm_values = np.array(dm_values)
    n_nu, n_t = image.shape
    nu = np.arange(n_nu, dtype=float)
    nu = (nu_0 - nu * d_nu)[::-1]
    print "Pre-calculating cumulative sums..."
    cumsums = np.cumsum(image[::-1, :], axis=0)

    ## MHz ** 2 * cm ** 3 * s / pc
    k = 1. / (2.410331 * 10 ** (-4))
    # Calculate shift of time caused by de-dispersion for all channels and all
    # values of DM
    dt_all = k * dm_values[:, np.newaxis] * (1. / nu ** 2. - 1. / nu_0 ** 2.)
    # Find what number of time bins corresponds to this shifts
    nt_all = vint(vround(dt_all / d_t))[:, ::-1]

    # Create array for TDM
    values = np.zeros((len(dm_values), n_t), dtype=float)
    # Fill DM=0 row
    values[0] = cumsums[-1]

    # Cycle over DM values and fill TDM array for others DM values
    for i, nt in enumerate(nt_all[1:]):
        # Find at wich frequency channels time shifts have occurred
        indx = np.where(nt[1:] - nt[:-1] == 1)[0]
        result = cumsums[indx[0]] + np.roll(cumsums[-1] - cumsums[indx[-1]],
                                            -nt[-1])
        diff = [np.roll(cumsums[indx[j + 1]] - cumsums[indx[j]],
                        -nt[indx[j + 1]]) for j in range(len(indx) - 1)]
        result += np.sum(diff, axis=0)
        values[i + 1] = result

    return values