def test_forward_and_gradients(im_size, batch_size):
tf.random.set_seed(0)
grid_size = tuple(np.array(im_size) * 2)
im_rank = len(im_size)
M = im_size[0] * 2**im_rank
nufft_ob = KbNufftModule(im_size=im_size,
grid_size=grid_size,
norm='ortho',
grad_traj=True)
# Generate Trajectory
ktraj_ori = tf.Variable(
tf.random.uniform(
(batch_size, im_rank, M), minval=-1 / 2, maxval=1 / 2) * 2 * np.pi)
# Have a random signal
signal = tf.Variable(
tf.cast(tf.random.uniform((batch_size, 1, *im_size)), tf.complex64))
kdata = kbnufft_forward(nufft_ob._extract_nufft_interpob())(signal,
ktraj_ori)
ktraj_noise = np.copy(ktraj_ori)
ktraj_noise += 0.01 * tf.Variable(
tf.random.uniform(
(batch_size, im_rank, M), minval=-1 / 2, maxval=1 / 2) * 2 * np.pi)
ktraj = tf.Variable(ktraj_noise)
with tf.GradientTape(persistent=True) as g:
kdata_nufft = kbnufft_forward(nufft_ob._extract_nufft_interpob())(
signal, ktraj)
A = get_fourier_matrix(ktraj, im_size, im_rank, do_ifft=False)
kdata_ndft = tf.matmul(
tf.reshape(signal, (batch_size, 1, tf.reduce_prod(im_size))),
tf.transpose(A, [0, 2, 1]))
loss_nufft = tf.math.reduce_mean(tf.abs(kdata - kdata_nufft)**2)
loss_ndft = tf.math.reduce_mean(tf.abs(kdata - kdata_ndft)**2)
tf_test = tf.test.TestCase()
# Test if the NUFFT and NDFT operation is same
tf_test.assertAllClose(kdata_nufft, kdata_ndft, atol=2e-3)
# Test gradients with respect to kdata
gradient_ndft_kdata = g.gradient(kdata_ndft, signal)[0]
gradient_nufft_kdata = g.gradient(kdata_nufft, signal)[0]
tf_test.assertAllClose(gradient_ndft_kdata,
gradient_nufft_kdata,
atol=6e-3)
# Test gradients with respect to trajectory location
gradient_ndft_traj = g.gradient(kdata_ndft, ktraj)[0]
gradient_nufft_traj = g.gradient(kdata_nufft, ktraj)[0]
tf_test.assertAllClose(gradient_ndft_traj, gradient_nufft_traj, atol=6e-3)
# Test gradients in chain rule with respect to ktraj
gradient_ndft_loss = g.gradient(loss_ndft, ktraj)[0]
gradient_nufft_loss = g.gradient(loss_nufft, ktraj)[0]
tf_test.assertAllClose(gradient_ndft_loss, gradient_nufft_loss, atol=5e-4)
def from_volume_to_nc_kspace_and_traj(volume):
if acq_type == 'radial_stacks':
traj = get_stacks_of_radial_trajectory(volume_size, **acq_kwargs)
elif acq_type == 'spiral_stacks':
traj = get_stacks_of_spiral_trajectory(volume_size, **acq_kwargs)
else:
raise NotImplementedError(
f'{acq_type} dataset not implemented yet.')
if compute_dcomp:
interpob = nfft_ob._extract_nufft_interpob()
nufftob_forw = kbnufft_forward(interpob)
nufftob_back = kbnufft_adjoint(interpob)
dcomp = calculate_radial_dcomp_tf(
interpob,
nufftob_forw,
nufftob_back,
traj[0],
stacks=True,
)
# need to add batch and coil dimension to the volume
nc_kspace = nufft(nfft_ob, volume[None, None, ...], traj, volume_size)
nc_kspace_scaled = nc_kspace * scale_factor
volume_scaled = tf.abs(volume * scale_factor)
volume_channeled = volume_scaled[None, ..., None]
nc_kspaces_channeled = nc_kspace_scaled[..., None]
orig_shape = tf.shape(volume)[None, ...]
extra_args = (orig_shape, )
if compute_dcomp:
dcomp = tf.ones([1, tf.shape(dcomp)[0]],
dtype=dcomp.dtype) * dcomp[None, :]
extra_args += (dcomp, )
return (nc_kspaces_channeled, traj, extra_args), volume_channeled
def from_kspace_to_nc_kspace_and_traj(images, kspaces):
if acq_type == 'radial':
traj = get_radial_trajectory(image_size, **acq_kwargs)
elif acq_type == 'cartesian':
traj = get_debugging_cartesian_trajectory()
elif acq_type == 'spiral':
traj = get_spiral_trajectory(image_size, **acq_kwargs)
else:
raise NotImplementedError(f'{acq_type} dataset not implemented yet.')
interpob = nfft_ob._extract_nufft_interpob()
nufftob_back = kbnufft_adjoint(interpob, multiprocessing=multiprocessing)
nufftob_forw = kbnufft_forward(interpob, multiprocessing=multiprocessing)
dcomp = calculate_density_compensator(
interpob,
nufftob_forw,
nufftob_back,
traj[0],
)
traj = tf.repeat(traj, tf.shape(images)[0], axis=0)
orig_image_channels = tf_ortho_ifft2d(kspaces)
nc_kspace = nufft(nfft_ob, orig_image_channels, traj, image_size, multiprocessing=multiprocessing)
nc_kspace_scaled = nc_kspace * scale_factor
images_scaled = images * scale_factor
images_channeled = images_scaled[..., None]
nc_kspaces_channeled = nc_kspace_scaled[..., None]
orig_shape = tf.ones([tf.shape(kspaces)[0]], dtype=tf.int32) * tf.shape(kspaces)[-1]
dcomp = tf.ones([tf.shape(kspaces)[0], tf.shape(dcomp)[0]], dtype=dcomp.dtype) * dcomp[None, :]
extra_args = (orig_shape, dcomp)
smaps = non_cartesian_extract_smaps(nc_kspace, traj, dcomp, nufftob_back, orig_shape)
return (nc_kspaces_channeled, traj, smaps, extra_args), images_channeled
def test_density_compensators_tf():
# This is a simple test to ensure that the code works only!
# We still dont have a method to test if the results are correct
ktraj, nufft_ob, torch_forward, torch_backward = setup()
interpob = nufft_ob._extract_nufft_interpob()
tf_ktraj = tf.convert_to_tensor(ktraj)
nufftob_back = kbnufft_adjoint(interpob)
nufftob_forw = kbnufft_forward(interpob)
tf_dcomp = calculate_density_compensator(interpob, nufftob_forw,
nufftob_back, tf_ktraj)
def nufft(nufft_ob, image, ktraj, image_size=None, multicoil=True):
if image_size is not None:
image = adjust_image_size(
image,
image_size,
multicoil=multicoil,
)
forward_op = kbnufft_forward(nufft_ob._extract_nufft_interpob(),
multiprocessing=True)
shape = tf.shape(image)[-1]
kspace = forward_op(image, ktraj)
return kspace
def __init__(self, multicoil=False, im_size=(640, 472), density_compensation=False, **kwargs):
super(NFFTBase, self).__init__(**kwargs)
self.multicoil = multicoil
self.im_size = im_size
self.nufft_ob = KbNufftModule(
im_size=im_size,
grid_size=None,
norm='ortho',
)
self.density_compensation = density_compensation
self.forward_op = kbnufft_forward(self.nufft_ob._extract_nufft_interpob())
self.backward_op = kbnufft_adjoint(self.nufft_ob._extract_nufft_interpob())
def nufft(nufft_ob, image, ktraj, image_size=None):
forward_op = kbnufft_forward(nufft_ob._extract_nufft_interpob())
shape = tf.shape(image)[2:]
if image_size is not None:
image_adapted = tf.cond(
tf.reduce_any(tf.math.greater(shape, image_size)),
lambda: _crop_for_nufft(image, image_size),
lambda: _pad_for_nufft(image, image_size),
)
else:
image_adapted = image
kspace = forward_op(image_adapted, ktraj)
return kspace
def test_calculate_radial_dcomp_tf():
ktraj, nufft_ob, torch_forward, torch_backward = setup()
interpob = nufft_ob._extract_nufft_interpob()
tf_nufftob_forw = kbnufft_forward(interpob)
tf_nufftob_back = kbnufft_adjoint(interpob)
tf_ktraj = tf.convert_to_tensor(ktraj)
torch_ktraj = torch.tensor(ktraj).unsqueeze(0)
tf_dcomp = calculate_radial_dcomp_tf(interpob, tf_nufftob_forw, tf_nufftob_back, tf_ktraj)
torch_dcomp = calculate_radial_dcomp_pytorch(torch_forward, torch_backward, torch_ktraj)
np.testing.assert_allclose(
tf_dcomp.numpy(),
torch_dcomp[0].numpy(),
rtol=1e-5,
atol=1e-5,
)
def test_forward_gradient():
traj = ktraj_function()
image = tf.zeros([1, 1, *image_shape], dtype=tf.complex64)
nufft_ob = KbNufftModule(
im_size=(640, 400),
grid_size=None,
norm='ortho',
)
forward_op = kbnufft_forward(nufft_ob._extract_nufft_interpob())
with tf.GradientTape() as tape:
tape.watch(image)
res = forward_op(image, traj)
grad = tape.gradient(res, image)
tf_test = tf.test.TestCase()
tf_test.assertEqual(grad.shape, image.shape)
def test_ncpdnet_init_and_call_3d(dcomp, volume_shape):
model = NCPDNet(
n_iter=1,
n_primal=2,
n_filters=2,
multicoil=False,
im_size=volume_shape,
three_d=True,
dcomp=dcomp,
fastmri=False,
)
af = 16
traj = get_stacks_of_radial_trajectory(volume_shape, af=af)
spokelength = volume_shape[-2]
nspokes = volume_shape[-1] // af
nstacks = volume_shape[0]
kspace_shape = nspokes * spokelength * nstacks
extra_args = (tf.constant([volume_shape]), )
if dcomp:
nufft_ob = KbNufftModule(
im_size=volume_shape,
grid_size=None,
norm='ortho',
)
interpob = nufft_ob._extract_nufft_interpob()
nufftob_forw = kbnufft_forward(interpob)
nufftob_back = kbnufft_adjoint(interpob)
dcomp = calculate_radial_dcomp_tf(
interpob,
nufftob_forw,
nufftob_back,
traj[0],
stacks=True,
)
dcomp = tf.ones([1, tf.shape(dcomp)[0]],
dtype=dcomp.dtype) * dcomp[None, :]
extra_args += (dcomp, )
res = model([
tf.zeros([1, 1, kspace_shape, 1], dtype=tf.complex64),
traj,
extra_args,
])
assert res.shape[1:4] == volume_shape
def profile_tfkbnufft(
image,
ktraj,
im_size,
device,
):
if device == 'CPU':
num_nuffts = 20
else:
num_nuffts = 50
print(f'Using {device}')
device_name = f'/{device}:0'
with tf.device(device_name):
image = tf.constant(image)
if device == 'GPU':
image = tf.cast(image, tf.complex64)
ktraj = tf.constant(ktraj)
nufft_ob = KbNufftModule(im_size=im_size, grid_size=None, norm='ortho')
forward_op = kbnufft_forward(nufft_ob._extract_nufft_interpob())
adjoint_op = kbnufft_adjoint(nufft_ob._extract_nufft_interpob())
# warm-up computation
for _ in range(2):
y = forward_op(image, ktraj)
start_time = time.perf_counter()
for _ in range(num_nuffts):
y = forward_op(image, ktraj)
end_time = time.perf_counter()
avg_time = (end_time - start_time) / num_nuffts
print('forward average time: {}'.format(avg_time))
# warm-up computation
for _ in range(2):
x = adjoint_op(y, ktraj)
# run the adjoint speed tests
start_time = time.perf_counter()
for _ in range(num_nuffts):
x = adjoint_op(y, ktraj)
end_time = time.perf_counter()
avg_time = (end_time - start_time) / num_nuffts
print('backward average time: {}'.format(avg_time))
def __init__(self, **kwargs):
super().__init__(**kwargs)
interpob = nufft_ob._extract_nufft_interpob()
self.nufftob_back = kbnufft_adjoint(interpob,
multiprocessing=False)
self.nufftob_forw = kbnufft_forward(interpob,
multiprocessing=False)
if acq_type == 'radial':
self.traj = get_radial_trajectory(image_size, af=af)
elif acq_type == 'cartesian':
self.traj = get_debugging_cartesian_trajectory()
elif acq_type == 'spiral':
self.traj = get_spiral_trajectory(image_size, af=af)
else:
raise NotImplementedError(
f'{acq_type} dataset not implemented yet.')
self.dcomp = calculate_density_compensator(
interpob,
self.nufftob_forw,
self.nufftob_back,
self.traj[0],
)
def preprocessing(self, image, kspace):
traj = self.generate_trajectory()
interpob = self.nufft_obj._extract_nufft_interpob()
nufftob_forw = kbnufft_forward(interpob, multiprocessing=True)
nufftob_back = kbnufft_adjoint(interpob, multiprocessing=True)
if self.dcomp:
dcomp = calculate_density_compensator(
interpob,
nufftob_forw,
nufftob_back,
traj[0],
)
traj = tf.repeat(traj, tf.shape(image)[0], axis=0)
orig_image_channels = ortho_ifft2d(kspace)
if self.crop_image_data:
image = adjust_image_size(image, self.image_size)
nc_kspace = nufft(self.nufft_obj, orig_image_channels, traj, self.image_size, multicoil=self.multicoil)
nc_kspace, image = scale_tensors(nc_kspace, image, scale_factor=self.scale_factor)
image = image[..., None]
nc_kspaces_channeled = nc_kspace[..., None]
orig_shape = tf.ones([tf.shape(kspace)[0]], dtype=tf.int32) * self.image_size[-1]
if not self.crop_image_data:
output_shape = tf.shape(image)[1:][None, :]
output_shape = tf.tile(output_shape, [tf.shape(image)[0], 1])
extra_args = (orig_shape,)
if self.dcomp:
dcomp = tf.ones(
[tf.shape(kspace)[0], tf.shape(dcomp)[0]],
dtype=dcomp.dtype,
) * dcomp[None, :]
extra_args += (dcomp,)
model_inputs = (nc_kspaces_channeled, traj)
if self.multicoil:
smaps = non_cartesian_extract_smaps(nc_kspace, traj, dcomp, nufftob_back, self.image_size)
model_inputs += (smaps,)
if not self.crop_image_data:
model_inputs += (output_shape,)
model_inputs += (extra_args,)
return model_inputs, image
def from_kspace_to_nc_kspace_and_traj(images, kspaces):
if acq_type == 'radial':
traj = get_radial_trajectory(image_size, **acq_kwargs)
elif acq_type == 'cartesian':
traj = get_debugging_cartesian_trajectory()
elif acq_type == 'spiral':
traj = get_spiral_trajectory(image_size, **acq_kwargs)
else:
raise NotImplementedError(
f'{acq_type} dataset not implemented yet.')
if compute_dcomp:
interpob = nfft_ob._extract_nufft_interpob()
nufftob_forw = kbnufft_forward(interpob)
nufftob_back = kbnufft_adjoint(interpob)
dcomp = calculate_radial_dcomp_tf(
interpob,
nufftob_forw,
nufftob_back,
traj[0],
)
traj = tf.repeat(traj, tf.shape(images)[0], axis=0)
orig_image = tf_unmasked_adj_op(kspaces[..., None])
nc_kspace = nufft(nfft_ob, orig_image[:, None, ..., 0], traj,
image_size)
nc_kspace_scaled = nc_kspace * scale_factor
images_scaled = images * scale_factor
images_channeled = images_scaled[..., None]
nc_kspaces_channeled = nc_kspace_scaled[..., None]
orig_shape = tf.ones([tf.shape(kspaces)[0]],
dtype=tf.int32) * tf.shape(kspaces)[-1]
extra_args = (orig_shape, )
if compute_dcomp:
dcomp = tf.ones([tf.shape(kspaces)[0],
tf.shape(dcomp)[0]],
dtype=dcomp.dtype) * dcomp[None, :]
extra_args += (dcomp, )
return (nc_kspaces_channeled, traj, extra_args), images_channeled
def from_kspace_to_nc_kspace_and_traj(images, kspaces):
if acq_type == 'radial':
traj = get_radial_trajectory(image_size, **acq_kwargs)
elif acq_type == 'cartesian':
traj = get_debugging_cartesian_trajectory()
elif acq_type == 'spiral':
traj = get_spiral_trajectory(image_size, **acq_kwargs)
else:
raise NotImplementedError(f'{acq_type} dataset not implemented yet.')
if compute_dcomp:
interpob = nfft_ob._extract_nufft_interpob()
nufftob_back = kbnufft_adjoint(interpob, multiprocessing=True)
nufftob_forw = kbnufft_forward(interpob, multiprocessing=True)
dcomp = calculate_density_compensator(
interpob,
nufftob_forw,
nufftob_back,
traj[0],
)
traj = tf.repeat(traj, tf.shape(images)[0], axis=0)
orig_image = tf_unmasked_adj_op(kspaces[..., None])
nc_kspace = nufft(nfft_ob, orig_image[:, None, ..., 0], traj, image_size)
nc_kspace_scaled = nc_kspace * scale_factor
images_scaled = images * scale_factor
# Here implement gridding
if gridding:
pi = tf.constant(math.pi)
def tf_grid_nc(nc_kspace_traj):
nc_kspace, traj = nc_kspace_traj
X_grid, Y_grid = tf.meshgrid(
tf.range(-pi, pi, 2*pi / image_size[0]),
tf.range(-pi, pi, 2*pi / image_size[1]),
)
nc_kspace = tf.numpy_function(
grid_non_cartesian,
[traj, nc_kspace, X_grid, Y_grid],
tf.complex64,
)
return nc_kspace
nc_kspace_scaled = tf.nest.map_structure(tf.stop_gradient, tf.map_fn(
tf_grid_nc,
(nc_kspace_scaled, traj),
fn_output_signature=tf.complex64,
parallel_iterations=multiprocessing.cpu_count(),
))
nc_kspace_scaled.set_shape([
None,
*image_size,
])
traj = tf.ones_like(nc_kspace_scaled)
images_channeled = images_scaled[..., None]
nc_kspaces_channeled = nc_kspace_scaled[..., None]
orig_shape = tf.ones([tf.shape(kspaces)[0]], dtype=tf.int32) * tf.shape(kspaces)[-1]
extra_args = (orig_shape,)
if compute_dcomp:
dcomp = tf.ones([tf.shape(kspaces)[0], tf.shape(dcomp)[0]], dtype=dcomp.dtype) * dcomp[None, :]
extra_args += (dcomp,)
if gridding:
return (nc_kspaces_channeled, traj), images_channeled
else:
return (nc_kspaces_channeled, traj, extra_args), images_channeled
