Python numpy.fft.rfftn() Examples

def run_test_r2c_dtype(self, shape, axes, dtype=np.float32, scale=1., misalign=0):
        known_data = np.random.normal(size=shape).astype(np.float32)
        known_data = (known_data * scale).astype(dtype)

        # Force misaligned data
        padded_shape = shape[:-1] + (shape[-1] + misalign,)
        known_data = np.resize(known_data, padded_shape)
        idata = bf.ndarray(known_data, space='cuda')
        known_data = known_data[..., misalign:]
        idata = idata[..., misalign:]

        oshape = list(shape)
        oshape[axes[-1]] = shape[axes[-1]] // 2 + 1
        odata = bf.ndarray(shape=oshape, dtype='cf32', space='cuda')
        fft = Fft()
        fft.init(idata, odata, axes=axes)
        fft.execute(idata, odata)
        known_result = gold_rfftn(known_data.astype(np.float32) / scale, axes=axes)
        compare(odata.copy('system'), known_result) 

def rfftn_empty_aligned(shape, axes, dtype, order='C', n=None):
    """Construct an empty byte-aligned array for real FFTs.

    Construct an empty byte-aligned array for efficient use by :mod:`pyfftw`
    functions :func:`pyfftw.interfaces.numpy_fft.rfftn` and
    :func:`pyfftw.interfaces.numpy_fft.irfftn`. The shape of the
    empty array is appropriate for the output of
    :func:`pyfftw.interfaces.numpy_fft.rfftn` applied
    to an array of the shape specified by parameter `shape`, and for the
    input of the corresponding :func:`pyfftw.interfaces.numpy_fft.irfftn`
    call that reverses this operation.

    Parameters
    ----------
    shape : sequence of ints
      Output array shape
    axes : sequence of ints
      Axes on which the FFT will be computed
    dtype : dtype
      Real dtype from which the complex dtype of the output array is derived
    order : {'C', 'F'}, optional (default 'C')
      Specify whether arrays should be stored in row-major (C-style) or
      column-major (Fortran-style) order
    n : int, optional (default None)
      Output array should be aligned to n-byte boundary

    Returns
    -------
    a :  ndarray
      Empty array with required byte-alignment
    """

    ashp = list(shape)
    raxis = axes[-1]
    ashp[raxis] = ashp[raxis] // 2 + 1
    cdtype = complex_dtype(dtype)
    return pyfftw.empty_aligned(ashp, cdtype, order, n) 

def rfftn(a, s=None, axes=None):
    """Multi-dimensional discrete Fourier transform for real input.

    Compute the multi-dimensional discrete Fourier transform for real input.
    This function is a wrapper for :func:`pyfftw.interfaces.numpy_fft.rfftn`,
    with an interface similar to that of :func:`numpy.fft.rfftn`.

    Parameters
    ----------
    a : array_like
      Input array (taken to be real)
    s : sequence of ints, optional (default None)
      Shape of the output along each transformed axis (input is cropped
      or zero-padded to match).
    axes : sequence of ints, optional (default None)
      Axes over which to compute the DFT.

    Returns
    -------
    af : complex ndarray
      DFT of input array
    """

    return pyfftw.interfaces.numpy_fft.rfftn(
        a, s=s, axes=axes, overwrite_input=False,
        planner_effort=pyfftw_planner_effort, threads=pyfftw_threads) 

def fftconv(a, b, axes=(0, 1), origin=None):
    """Multi-dimensional convolution via the Discrete Fourier Transform.

    Compute a multi-dimensional convolution via the Discrete Fourier
    Transform. Note that the output has a phase shift relative to the
    output of :func:`scipy.ndimage.convolve` with the default `origin`
    parameter.

    Parameters
    ----------
    a : array_like
      Input array
    b : array_like
      Input array
    axes : sequence of ints, optional (default (0, 1))
      Axes on which to perform convolution
    origin : sequence of ints or None optional (default None)
      Indices of centre of `a` filter. The default of None corresponds
      to a centre at 0 on all axes of `a`

    Returns
    -------
    ab : ndarray
      Convolution of input arrays, `a` and `b`, along specified `axes`
    """

    if np.isrealobj(a) and np.isrealobj(b):
        fft = rfftn
        ifft = irfftn
    else:
        fft = fftn
        ifft = ifftn
    dims = np.maximum([a.shape[i] for i in axes], [b.shape[i] for i in axes])
    af = fft(a, dims, axes)
    bf = fft(b, dims, axes)
    ab = ifft(af * bf, dims, axes)
    if origin is not None:
        ab = np.roll(ab, -np.array(origin), axis=axes)
    return ab 

def rfl2norm2(xf, xs, axis=(0, 1)):
    r"""Compute the squared :math:`\ell_2` norm in the real DFT domain.

    Compute the squared :math:`\ell_2` norm in the DFT domain, taking
    into account the unnormalised DFT scaling, i.e. given the DFT of a
    multi-dimensional array computed via :func:`rfftn`, return the
    squared :math:`\ell_2` norm of the original array.

    Parameters
    ----------
    xf : array_like
      Input array
    xs : sequence of ints
      Shape of original array to which :func:`rfftn` was applied to
      obtain the input array
    axis : sequence of ints, optional (default (0,1))
      Axes on which the input is in the frequency domain

    Returns
    -------
    x : float
      :math:`\|\mathbf{x}\|_2^2` where the input array is the result of
      applying :func:`rfftn` to the specified axes of multi-dimensional
      array :math:`\mathbf{x}`
    """

    scl = 1.0 / np.prod(np.array([xs[k] for k in axis]))
    slc0 = (slice(None),) * axis[-1]
    nrm0 = np.linalg.norm(xf[slc0 + (0,)])
    idx1 = (xs[axis[-1]] + 1) // 2
    nrm1 = np.linalg.norm(xf[slc0 + (slice(1, idx1),)])
    if xs[axis[-1]] % 2 == 0:
        nrm2 = np.linalg.norm(xf[slc0 + (slice(-1, None),)])
    else:
        nrm2 = 0.0
    return scl*(nrm0**2 + 2.0*nrm1**2 + nrm2**2) 

def _rfftn(a, s=None, axes=None):
        return  npfft.rfftn(a, s, axes).astype(complex_dtype(a.dtype))