def test_adjoint_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 = tf.Variable( kbnufft_forward(nufft_ob._extract_nufft_interpob())(signal, ktraj_ori)) Idata = tf.Variable( kbnufft_adjoint(nufft_ob._extract_nufft_interpob())(kdata, 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: I_nufft = kbnufft_adjoint(nufft_ob._extract_nufft_interpob())(kdata, ktraj) A = get_fourier_matrix(ktraj, im_size, im_rank, do_ifft=True) I_ndft = tf.reshape(tf.transpose(tf.matmul(kdata, A), [0, 1, 2]), (batch_size, 1, *im_size)) loss_nufft = tf.math.reduce_mean(tf.abs(Idata - I_nufft)**2) loss_ndft = tf.math.reduce_mean(tf.abs(Idata - I_ndft)**2) tf_test = tf.test.TestCase() # Test if the NUFFT and NDFT operation is same tf_test.assertAllClose(I_nufft, I_ndft, atol=2e-3) # Test gradients with respect to kdata gradient_ndft_kdata = g.gradient(I_ndft, kdata)[0] gradient_nufft_kdata = g.gradient(I_nufft, kdata)[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(I_ndft, ktraj)[0] gradient_nufft_traj = g.gradient(I_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 __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 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_adjoint_gradient(): traj = ktraj_function() kspace = tf.zeros([1, 1, kspace_shape], dtype=tf.complex64) nufft_ob = KbNufftModule( im_size=(640, 400), grid_size=None, norm='ortho', ) backward_op = kbnufft_adjoint(nufft_ob._extract_nufft_interpob()) with tf.GradientTape() as tape: tape.watch(kspace) res = backward_op(kspace, traj) grad = tape.gradient(res, kspace) tf_test = tf.test.TestCase() tf_test.assertEqual(grad.shape, kspace.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