
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import ticker, cm


L = 512
t = np.linspace(-1, 1, L)
X, Y = np.meshgrid(t, t)
coord = np.transpose(np.meshgrid(t, t))
extent = -1, 1, -1, 1

tumor_shape = .1

# rho, x0, y0, a, b, phi
ellipse_params = np.array([[5, 0, 0, .6, .9, 0],  # skull exterior
                           [-5, 0, 0, .09*6, .09*9, 0],  # skull interior
                           [2, -.15, -.2, tumor_shape, 1.2*tumor_shape, 1 * np.pi / 6], # tumor
                           [1, -.2, -.2, .23, .25, - np.pi / 20],
                           [1, -.2, +.0, .2, .6, - np.pi / 20],
                           [1, +.25, +.05, .2, .6, + np.pi / 20]
                          ])
K = len(ellipse_params)

def rotation_2D(theta):
    return np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]])

def ellipse_mask(coord, x_c, y_c, a, b, theta):
    return np.sum((np.diag([1/a, 1/b]) @ rotation_2D(theta).T @ (coord.T.reshape(2, -1) - np.array([x_c, y_c]).reshape(2, 1))).reshape(2, L, L)**2, axis=0) <= 1

Z = np.zeros((L, L))
for (rho, x_c, y_c, a, b, theta) in ellipse_params:
    Z = Z + rho * ellipse_mask(coord, x_c, y_c, a, b, theta)

fig = plt.figure(figsize=(10, 10))
plt.set_cmap("gray")
plt.contourf(X, Y, Z)
plt.show()


from scipy.special import j1
def ellipse_fourier(omega, x_c, y_c, a, b, theta):
    D = np.diag([a, b])
    R = rotation_2D(theta)
    r_c = np.array([x_c, y_c])
    return 2 * np.pi * a * b * complex(np.cos(-np.dot(omega, r_c)), np.sin(-np.dot(omega, r_c))) * j1(np.linalg.norm(D @ R.T @ omega)) / np.linalg.norm(D @ R.T @ omega)


t = np.linspace(-50, 50, 500)
X, Y = np.meshgrid(t, t)
FT = np.zeros((len(t), len(t), K), dtype = "complex_")
for k in range(K):
    for i in range(len(FT)):
        for j in range(len(FT)):
                FT[i, j, k] = ellipse_params[k, 0] * ellipse_fourier(np.array([X[i, j], Y[i, j]]), * ellipse_params[k, 1:]) # catchup nan issue, or remove 0 from meshgrid

fig = plt.figure(figsize=(10, 10))
plt.imshow(np.log(1+np.abs(FT[:, :, 0])), extent=(-50, 50, -50, 50), origin='lower')
plt.show()


fig = plt.figure(figsize=(10, 10))
plt.imshow(np.angle(FT[:, :, k]), extent=(-50, 50, -50, 50))
plt.show()


from scipy import ndimage, interpolate


S = 512
N = 25
image = Z
step = image.shape[0] / (np.max(X) - np.min(X))


sinogram = np.zeros((N, S))
for i in range(N):
    angle = np.degrees(np.pi * i / N)
    image_rot = ndimage.rotate(image, angle, reshape=False)
    sinogram[i] = np.sum(image_rot, axis=0)

fig = plt.figure(figsize=(10, 10))
plt.imshow(sinogram, aspect='auto')
plt.axis('off')
plt.show()


def radon_transform_ellipse(theta, p, a, b):
    if (a**2 * np.cos(theta)**2 + b**2 * np.sin(theta)**2 - p**2) > 0:
        return (2 * a * b * np.sqrt((a**2 * np.cos(theta)**2 + b**2 * np.sin(theta)**2 - p**2)))/(a**2 * np.cos(theta)**2 + b**2 * np.sin(theta)**2)
    else:
        return 0
    
def radon_transform_rotated_ellipse(theta, p, a, b, phi):
    return radon_transform_ellipse(theta - phi, p, a, b)

def radon_transform_rotated_translated_ellipse(theta, p, a, b, x0, y0, phi):
    return radon_transform_rotated_ellipse(theta, p - x0 * np.cos(theta) - y0 * np.sin(theta), a, b, phi)


sinogram_continuous = np.zeros((N, S))
for i in range(N):
    theta = np.pi * i / N
    P = np.linspace(-1, 1, S)
    for (rho, x0, y0, a, b, phi) in ellipse_params:
        for k, p in enumerate(P):
            sinogram_continuous[i, k] += rho * radon_transform_rotated_translated_ellipse(theta, p, a, b, x0, y0, phi)
fig = plt.figure(figsize=(10, 10))
plt.imshow(sinogram_continuous, aspect='auto')
plt.axis('off')
plt.show()


def normalize(M):
    return (M - np.mean(M)) / np.std(M)

fig, axes = plt.subplots(1, 3)
fig.set_size_inches(20, 6)
axes[0].imshow(normalize(sinogram_continuous), aspect='auto')
axes[0].set_title('(a) True sinogram')
axes[1].imshow(normalize(sinogram), aspect='auto')
axes[1].set_title('(b) Inverse crime sinogram')
axes[2].imshow(normalize(sinogram_continuous) - normalize(sinogram), aspect='auto')
axes[2].set_title('(c) Inverse crime - True sinogram')


for a in axes.flatten():
    a.axis('off')
plt.show()


sinogram = np.copy(sinogram_continuous)


sinogram = np.clip(sinogram, 0, np.max(sinogram))


from scipy.fftpack import fftshift, ifftshift, fft, ifft, fft2, ifft2

def shifted_fft(data):
    return fftshift(fft(ifftshift(data)))

def shifted_fft2(data):
    return fftshift(fft2(ifftshift(data)))

def shifted_ifft2(data):
    return fftshift(ifft2(ifftshift(data)))


sinogram_fft_rows = np.zeros_like(sinogram, dtype=complex)
for i in range(N):
    sinogram_fft_rows[i] = shifted_fft(sinogram[i])

fig, ax = plt.subplots(nrows=1, sharex=True, sharey=True)
ax.imshow(np.log(1 + np.abs(sinogram_fft_rows)), aspect='auto')
ax.axis('off')
plt.tight_layout()


t = [np.pi * i / N for i in range(N)]
w = np.arange(S) - S / 2
ww, tt = np.meshgrid(w, t)

uu = (S / 2) + ww * np.cos(tt)
vv = (S / 2) + ww * np.sin(tt)
u = uu.flatten()
v = vv.flatten()

plt.figure(figsize=(10, 10))
plt.scatter(u, v, c=np.log(1 + np.abs(sinogram_fft_rows.flatten())),
    marker='.', edgecolor='none')
plt.axis('off')
plt.show()


xx, yy = np.meshgrid(np.arange(S), np.arange(S))
x = xx.flatten()
y = yy.flatten()

Fuv_gridrec = interpolate.griddata((v, u), sinogram_fft_rows.flatten(), 
                                   (y, x), method='linear', fill_value=0.0).reshape((S, S))

plt.figure(figsize=(10, 10))
plt.imshow(np.log(1 + np.abs(Fuv_gridrec)), origin='lower')
plt.axis('off')
plt.show()


rec = np.real(shifted_ifft2(Fuv_gridrec))
plt.figure(figsize=(10, 10))
plt.imshow(np.clip(rec, 0, np.max(rec)), origin='lower')
plt.axis('off')
plt.show()

A. Phantom¶

A.1 Design¶

A.2 Fourier Transform¶

B. Tomography reconstruction¶

B.1 Sinogram computation / Inverse crime¶

B.2 Radial FFT¶

B.3 Fourier interpolation¶

B.4 Inverse Fourier Transform¶