matrix triple product DFT

mtp.py

def matrix_dft(f, alpha, npix, shift=None, unitary=False):
    if np.isscalar(alpha):
        ax = ay = alpha
    else:
        ax = ay = np.asarray(alpha)

    f = np.asarray(f)
    m, n = f.shape

    if np.isscalar(npix):
        M = N = npix
    else:
        M = N = np.asarray(npix)

    if shift is None:
        sx = sy = 0
    else:
        sx = sy = np.asarray(shift)

    # Y and X are (r,c) coordinates in the (m x n) input plane, f
    # V and U are (r,c) coordinates in the (M x N) output plane, F
    X = np.arange(n) - np.floor(n/2) - sx
    Y = np.arange(m) - np.floor(m/2) - sy
    U = np.arange(N) - np.floor(N/2) - sx
    V = np.arange(M) - np.floor(M/2) - sy

    E1 = np.exp(1j * -2 * np.pi * (ay/m) * np.outer(Y, V).T)
    E2 = np.exp(1j * -2 * np.pi * (ax/m) * np.outer(X, U))

    F = E1.dot(f).dot(E2)

    if unitary is True:
        norm_coef = np.sqrt((ay * ax)/(m * n * M * N))
        return F * norm_coef
    else:
        return F