def zero_pad()

in admm.py [0:0]


def zero_pad(m, n, x, mask):
    """
    Supplements the Fourier-domain input with zeros to fill out a full image.

    Pads x with rows of zeros to form a full m x n image, with the possibly
    nonzero locations specified by mask. Then applies the (two-dimensional)
    unitary inverse discrete Fourier transform.

    Parameters
    ----------
    m : int
        number of rows in the image being formed
    n : int
        number of columns in the image being formed
    x : numpy.ndarray or ctorch.ComplexTensor
        potentially nonzero rows (prior to the inverse Fourier transform)
    mask : numpy.ndarray or torch.LongTensor
        indicators of the positions of the rows of x in the full m x n array

    Returns
    -------
    numpy.ndarray or ctorch.ComplexTensor
        inverse Fourier transform of x padded with zeros
    """
    if isinstance(x, np.ndarray):
        assert x.shape[1] == n
        # Pad x with zeros, obtaining y.
        y = np.zeros((m, n), dtype=np.complex128)
        y[mask] = x
        # Scale the FFTs to make them unitary.
        return np.fft.ifft2(y) * np.sqrt(m * n)
    elif isinstance(x, ctorch.ComplexTensor):
        assert x.shape[1] == n
        # Create a tensor of the same type as x and zero its entries.
        y = x.new(m, n).zero_()
        # Copy the entries of x into y.
        y[mask] = x
        # Scale the FFTs to make them unitary.
        return ctorch.ifft2(y) * np.sqrt(m * n)
    else:
        raise TypeError('Input must be a numpy.ndarray ' +
                        'or a ctorch.ComplexTensor.')