<jupyter_start><jupyter_text>Fourier Compressed SensingIn this notebook, we consider the problem of Fourier compressed sensing. In particular, we are interested in recovering an image acquired by a randomly subsampled Fourier sensing measurment system.   SetupLet us import relevant packages and load a brain dataset.<jupyter_code>%matplotlib notebook
import numpy as np
import matplotlib.pyplot as plt
img = np.load('data/brain_img.npy')

fig, ax = plt.subplots()
ax.imshow(np.abs(img), cmap='gray')
ax.set_title('Ground Truth')
ksp = np.load('data/brain_ksp.npy')

fig, ax = plt.subplots()
ax.imshow(np.abs(ksp)**0.1, cmap='gray')
ax.set_title('Log k-space magnitude')<jupyter_output><empty_output><jupyter_text>Native ReconstructionSimply do an inverse FFT.<jupyter_code>img_naive = np.fft.ifftshift(ksp)
img_naive = np.fft.ifftn(img_naive, norm='ortho')
img_naive = np.fft.fftshift(img_naive)

fig, ax = plt.subplots()
ax.imshow(np.abs(img_naive), cmap='gray')
ax.set_title('Naive Reconstruction')<jupyter_output><empty_output><jupyter_text>L1 Wavelet ReconstructionThe L1 wavelet regularization reconstruction solves the following problem:$$\min_x \frac{1}{2} \| \underbrace{S F W}_{A} x - y \|_2^2 + \lambda \| x \|_1$$where $S$ is the sampling operator, $F$ is the Fourier transform operator, $W$ is the inverse wavelet transform operator, $x$ is the wavelet coefficients and $y$ is the acquired k-space measurements.The go-to method to solve this problem is the iterative soft-thresholding algorithm (ISTA), which performs a gradient step w.r.t. the data consistency followed by a soft-thresholding step.The reconstruction image can be generated from $W x$.We will create an L1 wavelet reconstruction by doing:- Create linear operators $S, F, W$- Create the soft-threshold function  Linear operators (Linop)In the following, we will create the neccessary linear operators, $F$, $S$, and $W$.First we will create a generic Linop class such that we can perform forward, adjoint and compose.    A(x)    A.H(x)    C = A * B<jupyter_code>class Linop(object):
    def __call__(self, x):
        return self._forward(x)
    
    def H(self, x):
        return self._adjoint(x)
    
    def __mul__(self, B):
        return Compose(self, B)
    
    def _forward(self, x):
        raise NotImplementedError

    def _adjoint(self, x):
        raise NotImplementedError
class Compose(Linop):
    def __init__(self, A, B):
        self.A = A
        self.B = B
        
    def _forward(self, x):
        return self.A(self.B(x))
    
    def _adjoint(self, x):
        return self.B.H(self.A.H(x))<jupyter_output><empty_output><jupyter_text>$F$ LinopTo start, we will create an FFT linear operator class.  It takes the input array shape and application axes as arguments. We will also apply its adjoint, which is the inverse FFT, to the k-space array `ksp` to verify if it's working.<jupyter_code>class FFTLinop(Linop):
    
    def _forward(self, x):
        y = np.fft.ifftshift(x)
        y = np.fft.fftn(y, norm='ortho')
        y = np.fft.fftshift(y)
        return y

    def _adjoint(self, x):
        y = np.fft.ifftshift(x)
        y = np.fft.ifftn(y, norm='ortho')
        y = np.fft.fftshift(y)
        return y

F = FFTLinop()
fig, ax = plt.subplots()
ax.imshow(np.abs(F.H(ksp)), cmap='gray')
ax.set_title('Naive Reconstruction using Linop')<jupyter_output><empty_output><jupyter_text>$S$ LinopGiven the sampling mask, this operator simply multiplies the input array with the mask. We will estimate the sampling mask from the non-zero entries of the k-space arrays.<jupyter_code>mask = np.abs(ksp) > 0

fig, ax = plt.subplots()
ax.imshow(mask, cmap='gray')
ax.set_title('Sampling Mask')<jupyter_output><empty_output><jupyter_text>This can also be made into a linear operator<jupyter_code>class SamplingLinop(Linop):

    def _forward(self, x):
        return mask * x

    _adjoint = _forward
    
S = SamplingLinop()<jupyter_output><empty_output><jupyter_text>$W$ LinopWe will perform wavelet transforms using the `pywt` library. By default, the wavelet decomposition function `wavdecn` outputs a dictionary with wavelet coefficients from each subband. We will concatenate that into an array using the function `coeffs_to_array`.<jupyter_code>import pywt

coeff = pywt.wavedecn(img, 'db4', mode='periodic', level=3)
coeff, coeff_slices = pywt.coeffs_to_array(coeff)

fig, ax = plt.subplots()
ax.imshow(np.abs(coeff)**0.3, cmap='gray')
ax.set_title('Wavelet Coefficients')
class InverseWaveletLinop(Linop):

    def _forward(self, x):
        y = pywt.array_to_coeffs(x, coeff_slices)
        y = pywt.waverecn(y, 'db4', mode='periodic')
        return y

    def _adjoint(self, x):
        coeff = pywt.wavedecn(x, 'db4', mode='periodic', level=3)
        coeff, coeff_slices = pywt.coeffs_to_array(coeff)
        return coeff
    
W = InverseWaveletLinop()
A = S * F * W<jupyter_output><empty_output><jupyter_text>Soft-thresholdGiven a function $g(x)$, a proximal operator is called on a scalar ($\alpha$) and an array $x$ to compute:$$\text{prox}_{\alpha g} (y) = \text{argmin}_x \frac{1}{2} || x - y ||_2^2 + \alpha g(x)$$Here, our function $g(x) = \lambda \| x \|_1$ is a scaled L1-norm function. The proximal operator becomes the soft-thresholding function:$$\text{prox}_{\alpha g} (y) = (|y| - \alpha \lambda)_+ \frac{y}{|y|}$$<jupyter_code>def soft_thresh(y, lamda):
    mag = np.abs(y) - lamda
    mag *= mag > 0
    
    sign = np.divide(y, np.abs(y), where=np.abs(y) != 0)
    return mag * sign<jupyter_output><empty_output><jupyter_text>With the soft-thresholding function, we can then soft-thresh the wavelet coefficients.<jupyter_code>lamda = 0.05
coeff_thresh = soft_thresh(coeff, lamda)

fig, ax = plt.subplots()
ax.imshow(np.abs(coeff_thresh)**0.3, cmap='gray')
ax.set_title('Soft-thresholded Wavelet Coefficients')
img_thresh = W(coeff_thresh)

fig, ax = plt.subplots()
ax.imshow(np.abs(img_thresh), cmap='gray')
ax.set_title('WavThresh Denoised Image')<jupyter_output><empty_output><jupyter_text>Gradient DescentWe will do gradient descent. Let us first define the parameters and gradient function<jupyter_code>max_iter = 100
alpha = 1

def gradf(x):
    return A.H(A(x) - ksp)<jupyter_output><empty_output><jupyter_text>Then we do the following iterative reconstruction<jupyter_code>coeff_hat = np.zeros_like(coeff)

for it in range(max_iter):
    coeff_hat -= alpha * gradf(coeff_hat)
    coeff_hat = soft_thresh(coeff_hat, lamda * alpha)
    
img_hat = W(coeff_hat)<jupyter_output><empty_output><jupyter_text>I'd like to highlight a very convenient package `tqdm`. `tqdm` adds a progress bar very easily. You can add it in the for loop `tqdm(range(max_iter))`.<jupyter_code>from tqdm.auto import tqdm

coeff_hat = np.zeros_like(coeff)

for it in tqdm(range(max_iter)):
    coeff_hat -= alpha * gradf(coeff_hat)
    coeff_hat = soft_thresh(coeff_hat, lamda * alpha)
    
img_hat = W(coeff_hat)
fig, ax = plt.subplots()
ax.imshow(np.abs(img), cmap='gray', vmax=1)
ax.set_title('Ground Truth')

fig, ax = plt.subplots()
ax.imshow(np.abs(img_naive), cmap='gray', vmax=1)
ax.set_title('Naive Reconstruction')

fig, ax = plt.subplots()
ax.imshow(np.abs(img_hat), cmap='gray', vmax=1)
ax.set_title('L1 Wavelet Regularized Reconstruction')<jupyter_output><empty_output><jupyter_text>WrapupI'd like to wrap up with a package that I've been developing, SigPy. SigPy is a package for signal processing with iterative methods. It is built to operate directly on NumPy arrays on CPU and CuPy arrays on GPU. SigPy provides three supporting classes for building iterative reconstruction [App](https://sigpy.readthedocs.io/en/latest/generated/sigpy.app.App.htmlsigpy.app.App)'s: - [Linop](https://sigpy.readthedocs.io/en/latest/generated/sigpy.linop.Linop.htmlsigpy.linop.Linop) for linear operator- [Prox](https://sigpy.readthedocs.io/en/latest/generated/sigpy.prox.Prox.htmlsigpy.prox.Prox) for proximal operator- [Alg](https://sigpy.readthedocs.io/en/latest/generated/sigpy.alg.Alg.htmlsigpy.alg.Alg) for iterative algorithms![architecture](https://sigpy.readthedocs.io/en/latest/_images/architecture.pdf)<jupyter_code>import sigpy as sp

class L1WaveletRecon(sp.app.App):
    def __init__(self, ksp, lamda, max_iter):
        mask = ksp != 0
        
        F = sp.linop.FFT(ksp.shape)
        S = sp.linop.Multiply(ksp.shape, mask)
        self.W = sp.linop.InverseWavelet(ksp.shape)
        A = S * F * self.W
        
        proxg = sp.prox.L1Reg(A.ishape, lamda)
        
        self.coeff = np.zeros(A.ishape, np.complex)
        alpha = 1
        def gradf(x):
            return A.H * (A * x - ksp)

        alg = sp.alg.GradientMethod(gradf, self.coeff, alpha, proxg=proxg, 
                                    max_iter=max_iter)
        super().__init__(alg)
        
    def _output(self):
        return self.W(self.coeff)
img_hat = L1WaveletRecon(ksp, lamda, max_iter).run()
fig, ax = plt.subplots()
ax.imshow(np.abs(img_hat), cmap='gray', vmax=1)
ax.set_title('L1 Wavelet Regularized Reconstruction using SigPy')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12336160370084812, 'fraction_numerical': 0.0050666372948562615, 'mean_word_length': 4.150311968235961, 'pattern_counts': {'":': 0, '<': 53, '<?xml version=': 0, '>': 55, 'https://': 5, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '13', 'dir': 'jupyter-structured-clean-dedup', 'id': '445766', 'license': 'BSD-3-Clause', 'n_tokens_mistral': 3324, 'n_tokens_neox': 3222, 'n_words': 827, 'path': '02_fourier_compressed_sensing.ipynb', 'repo_name': 'frankong/ima2019_python_tutorial', 'size': '9079'}