Functions for experimental TorchIO notebooks

functions.py

import numpy as np
import matplotlib.pyplot as plt

def fourier_transform(array: np.ndarray):
    transformed = np.fft.fftn(array)
    fshift = np.fft.fftshift(transformed)
    return fshift

def inv_fourier_transform(fshift: np.ndarray):
    f_ishift = np.fft.ifftshift(fshift)
    img_back = np.fft.ifftn(f_ishift)
    return img_back

def get_1d_space(N, dt):
    """
    Return a shifted linspace with N points and step dt.
    The samples are shifted dt/2."""
    n = N/2
    t_min = -n * dt
    t_max = (n-1) * dt
    t = np.linspace(t_min, t_max, num=N)
    return t, t_min, t_max

def plotx(x, f):
    plt.plot(x, np.real(f))
    plt.plot(x, np.imag(f))
    plt.xlabel('$x~[mm]$')

def plotk(k, F):
    plt.plot(k, np.real(F))
    plt.plot(k, np.imag(F))
    plt.xlabel('$k~[mm^{-1}]$')

def imshow(image, fourier=False, axis=None, extent=None,
           labels=None, title=None, v=None):
    if fourier:
        np.seterr(divide='ignore')
        log_image = np.log(np.abs(image))
        log_image[log_image == -np.inf] = 0
        image = log_image
    if axis is None:
        fig = plt.figure()
        axis = fig.gca()
    if labels is not None:
        axis.set_xlabel(labels[0])
        axis.set_ylabel(labels[1])
    if title is not None:
        axis.set_title(title)
    return axis.imshow(image, cmap='gray', extent=extent, clim=v)