def test_fft2(shape, named):
shape = shape + [2]
data = create_input(shape, named=named)
if named:
dim = ("height", "width")
else:
dim = (-2, -1)
out_torch = transforms.fft2(data, dim=dim).numpy()
out_torch = out_torch[..., 0] + 1j * out_torch[..., 1]
data_numpy = tensor_to_complex_numpy(data)
data_numpy = np.fft.ifftshift(data_numpy, (-2, -1))
out_numpy = np.fft.fft2(data_numpy, norm="ortho")
out_numpy = np.fft.fftshift(out_numpy, (-2, -1))
z = out_torch - out_numpy
assert np.allclose(out_torch, out_numpy)
def test_fft2(shape, named):
shape = shape + [2]
data = create_input(shape, named=named)
if named:
dim = ('height', 'width')
else:
dim = (-3, -2)
out_torch = transforms.fft2(data, dim=dim).numpy()
out_torch = out_torch[..., 0] + 1j * out_torch[..., 1]
data_numpy = tensor_to_complex_numpy(data)
data_numpy = np.fft.ifftshift(data_numpy, (-2, -1))
out_numpy = np.fft.fft2(data_numpy, norm='ortho')
out_numpy = np.fft.fftshift(out_numpy, (-2, -1))
z = out_torch - out_numpy
print(z.real.max(), z.real.min(), z.imag.max(), z.imag.min())
assert np.allclose(out_torch, out_numpy)
# Torch
input_image = input_image.align_to("batch", "height", "width", "complex")
sensitivity_map = sensitivity_map.align_to("batch", "coil", "height", "width",
"complex")
masked_kspace = masked_kspace.align_to("batch", "coil", "height", "width",
"complex")
mul = transforms.complex_multiplication(sensitivity_map,
input_image.align_as(sensitivity_map))
mul_names = mul.names
mr_forward = torch.where(
sampling_mask.rename(None) == 0,
torch.tensor([0.0], dtype=masked_kspace.dtype).to(masked_kspace.device),
transforms.fft2(mul).rename(None),
)
error = mr_forward - torch.where(
sampling_mask.rename(None) == 0,
torch.tensor([0.0], dtype=masked_kspace.dtype).to(masked_kspace.device),
masked_kspace.rename(None),
)
error = error.refine_names(*mul_names)
mr_backward = transforms.ifft2(error)
out = transforms.complex_multiplication(transforms.conjugate(sensitivity_map),
mr_backward).sum("coil")
# numpy