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 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_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 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