6. Non-Cartesian spiral under-sampling

#DISPLAY T2* MR IMAGE
%matplotlib inline

import numpy as np
import os.path as op
import os
import math ; import cmath
import matplotlib.pyplot as plt
import sys

import brainweb_dl as bwdl
import mrinufft
from mrinufft.trajectories import display
from mrinufft import get_density
from mrinufft.density import voronoi
from mrinufft.trajectories import initialize_2D_spiral
from mrinufft.trajectories.display import display_2D_trajectory


from skimage import data, img_as_float, io, filters
from modopt.math.metrics import ssim

plt.rcParams["image.origin"]="lower"
plt.rcParams["image.cmap"]='Greys_r'

mri_img = bwdl.get_mri(4, "T1")[70, ...].astype(np.float32)
print(mri_img.shape)
img_size = mri_img.shape[0]
plt.figure()
plt.imshow(abs(mri_img))
plt.title("Original brain image")
plt.show()

# set up the first shot
rfactor = 4
num_shots = math.ceil(img_size/rfactor)
print("number of shots: {}".format(num_shots))

# define the regularly spaced samples on a single shot
#nsamples = (np.arange(0,img_size) - img_size//2)/(img_size)
num_samples = img_size
num_samples = (num_samples + 1) // 2
print("number of samples: {}".format(num_samples))
num_revolutions = 1

shot = np.arange(0, num_samples, dtype=np.complex_)
radius = shot / num_samples * 1 / (2 * np.pi) * (1 - np.finfo(float).eps)
angle = np.exp(2 * 1j * np.pi * shot / num_samples * num_revolutions)
# first half of the spiral
single_shot = np.multiply(radius, angle)
# add second half of the spiral
#single_shot = np.append(np.flip(single_shot, axis=0), -single_shot[1:])
single_shot = np.append(np.flip(single_shot, axis=0), -single_shot)
print("number of samples per shot: {}".format(np.size(single_shot)))

k_shots = np.array([], dtype = np.complex_)
#for i in vec_shots:
for i in np.arange(0, num_shots):
    shot_rotated = single_shot * np.exp(1j * 2 * np.pi * i / (num_shots * 2))
    k_shots = np.append(k_shots, shot_rotated)

print(k_shots.shape)
kspace_loc = np.zeros((len(k_shots),2))
kspace_loc[:,0] = k_shots.real
kspace_loc[:,1] = k_shots.imag

#Plot full initialization
kspace = plt.figure(figsize = (8,8))
#plot shots
plt.scatter(kspace_loc[::4,0], kspace_loc[::4,1], marker = '.')
plt.title("Spiral undersampling R = %d" %rfactor)

axes = plt.gca()
plt.grid()

print(np.arange(0, num_shots))

#data=convert_locations_to_mask(kspace_loc, mri_img.shape)

NufftOperator = mrinufft.get_operator("finufft")

# For better image quality we use a density compensation
density = "voronoi"
#density = "cell_count"
#ensity = None
# And create the associated operator.
if density == None:
    nufft = NufftOperator(
        kspace_loc.reshape(-1, 2), shape=mri_img.shape, density=None, n_coils=1
    )
elif density == "voronoi":
    voronoi_weights = get_density("voronoi", kspace_loc.reshape(-1, 2))
    nufft = NufftOperator(
        kspace_loc.reshape(-1, 2), shape=mri_img.shape, density=voronoi_weights, n_coils=1
    )
else:
    cell_count_weights = get_density("cell_count", kspace_loc.reshape(-1, 2), shape=mri_img.shape, osf=2.0)
    nufft = NufftOperator(
        kspace_loc.reshape(-1, 2), shape=mri_img.shape, density=cell_count_weights, n_coils=1, upsampfac=2.0
    )

kspace_obs = nufft.op(mri_img)  # Image -> Kspace
image2 = nufft.adj_op(kspace_obs)  # Kspace -> Image

grid_space = np.linspace(-0.5, 0.5, num=mri_img.shape[0])
grid2D = np.meshgrid(grid_space, grid_space)
#rid_soln = gridded_inverse_fourier_transform_nd(kspace_loc, kspace_obs,
#                                                tuple(grid2D), 'linear')
plt.imshow(np.abs(image2), cmap='gray')
# Calculate SSIM
base_ssim = ssim(image2, mri_img)
plt.title('Gridded Solution\nSSIM = ' + str(base_ssim))
plt.show()

from scipy.spatial import Voronoi, voronoi_plot_2d

vor = Voronoi(kspace_loc.reshape(-1, 2))
fig = voronoi_plot_2d(vor)

fig = voronoi_plot_2d(vor, show_vertices=False, line_colors='orange',
                      line_width=2, line_alpha=0.6, point_size=2)
plt.show()

mri_2D = bwdl.get_mri(4, "T2")[150, ...].astype(np.float32)
print(mri_2D.shape)
plt.figure()
plt.imshow(abs(mri_2D))
plt.show()

from mrinufft import get_density, get_operator, check_backend

traj = initialize_2D_spiral(256, 256).astype(np.float32)

nufft = get_operator("finufft")(traj, mri_2D.shape, density=False)
kspace = nufft.op(mri_2D)
adjoint = nufft.adj_op(kspace)

fig, axs = plt.subplots(1, 3, figsize=(15, 5))
axs[0].imshow(abs(mri_2D))
display_2D_trajectory(traj, subfigure=axs[1])
axs[2].imshow(abs(adjoint))

voronoi_weights = get_density("voronoi", traj)
nufft_voronoi = get_operator("finufft")(
    traj, shape=mri_2D.shape, density=voronoi_weights
)
adjoint_voronoi = nufft_voronoi.adj_op(kspace)
fig, axs = plt.subplots(1, 3, figsize=(15, 5))
axs[0].imshow(abs(mri_2D))
display_2D_trajectory(traj, subfigure=axs[1])
axs[2].imshow(abs(adjoint_voronoi))

flat_traj = traj.reshape(-1, 2)
weights = np.sqrt(np.sum(flat_traj**2, axis=1))
nufft = get_operator("finufft")(traj, shape=mri_2D.shape, density=weights)
adjoint_manual = nufft.adj_op(kspace)
fig, axs = plt.subplots(1, 3, figsize=(15, 5))
axs[0].imshow(abs(mri_2D))
axs[0].set_title("Ground Truth")
axs[1].imshow(abs(adjoint))
axs[1].set_title("no density compensation")
axs[2].imshow(abs(adjoint_manual))
axs[2].set_title("manual density compensation")