in admm2.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 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 entries (prior to the inverse Fourier transform)
mask : numpy.ndarray or torch.LongTensor
indicators of the positions of the entries 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):
# 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):
# Pad x with zeros, obtaining y.
y = ctorch.ComplexTensor(x.real.new(m, n), x.imag.new(m, n)).zero_()
mask_flat = mask[:, 0] * n + mask[:, 1]
y.view(-1)[mask_flat] = 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.')