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.')