#!/usr/bin/env python
# coding: utf-8

# In[1]:


get_ipython().run_cell_magic('file', 'matrixmodel.py', '\n# -*- coding: utf-8 -*-\n"""\nThis file defines a class for matrix forward operators that are useful for \niterative reconstruction algorithms.\n\nCreated on Tue Jul 26 2016\n\n@author: U. S. Kamilov, 2016.\n"""\n\nimport numpy as np\n\n\nclass MatrixOperator:\n    """\n    Class for all matrix forward operators\n    """\n\n    def __init__(self, H_matrix, sigSize):\n        """\n        Class constructor\n        """\n        \n        self.H_matrix = H_matrix\n        self.sigSize = sigSize\n        self.L = np.linalg.norm(H_matrix, 2) ** 2\n\n\n    def apply(self, f):\n        """\n        Apply the forward operator\n        """\n\n        fvec = np.reshape(f, (f.size, 1))\n        z = np.dot(self.H_matrix, fvec)\n        return z\n\n\n    def applyAdjoint(self, z):\n        """\n        Apply the adjoint of the forward operator\n        """\n\n        fvec = np.dot(np.transpose(self.H_matrix), z)\n        f = np.reshape(fvec, self.sigSize)\n        return f\n\n\n    def getLipschitzConstant(self):\n        """\n        Return the largest eigen value of HTH\n        """\n        \n        return (self.L)\n')


# In[2]:


get_ipython().run_cell_magic('file', 'algorithm.py', '\n# -*- coding: utf-8 -*-\n"""\nThis file implements fast iterative shrinkage/thresholding algorithm (FISTA)\n\nFor more details on this code see corresponding paper:\n\nU. S. Kamilov, "Minimizing Isotropic Total Variation without Subiterations," \nProc. 3rd International Traveling Workshop on Interactions between Sparse models \nand Technology (iTWIST 2016) (Aalborg, Denmark, August 24-26)\n\nCreated on Tue Jul 16 07:08:59 2016\n\n@author: U. S. Kamilov, 2016.\n"""\n\nimport numpy as np\n\n\ndef fistaEst(y, forwardObj, tau, numIter = 100, accelerate = True):\n    """\n    Iterative reconstruction with FISTA\n    """\n    \n    # Lipschitz constant is used for obtaining the step-size\n    stepSize = 1/forwardObj.getLipschitzConstant()\n    \n    # Size of the reconstructed image\n    sigSize = forwardObj.sigSize\n    \n    # Define iterates\n    fhat = np.zeros(sigSize)\n    s = fhat\n    q = 1\n    \n    # Tracks evolution of the energy functional\n    cost = np.zeros((numIter, 1))\n\n    \n    # Iterate\n    for iter in range(numIter):\n        \n        # Gradient step\n        fhatnext = s - stepSize*forwardObj.applyAdjoint(forwardObj.apply(s)-y)\n        \n        # Proximal step\n        fhatnext = softThresh(fhatnext, stepSize*tau)\n        \n        # Update the relaxation parameter\n        if accelerate:\n            qnext = 0.5*(1+np.sqrt(1+4*(q**2)))\n        else:\n            qnext = 1\n        \n        # Relaxation step\n        s = fhatnext + ((q-1)/qnext)*(fhatnext-fhat)\n        \n        # Compute cost\n        cost[iter] = evaluateCost(y, forwardObj, fhat, tau)\n        \n        # Update variables\n        q = qnext\n        fhat = fhatnext\n        \n    return fhat, cost\n\n\ndef softThresh(w, lam):\n    """\n    Soft-thresholding function\n    """\n\n    # norm along axis 3 of differences\n    abs_w = np.abs(w)\n\n    # compute shrinkage\n    shrinkFac = np.maximum(abs_w-lam, 0)\n    abs_w[abs_w <= 0] = 1  # to avoid division by zero\n    shrinkFac = shrinkFac/abs_w\n\n    # save output\n    what = shrinkFac * w\n\n    return what\n\ndef evaluateCost(y, forwardObj, f, tau):\n    """\n    Evaluate the energy functional\n    """\n    \n    cost = 0.5*np.sum((y-forwardObj.apply(f))**2) + tau*np.sum(np.abs(f))\n    cost = cost/(0.5*np.sum(y ** 2))\n    return cost\n')


# In[3]:


get_ipython().run_line_magic('reset', '-f')
get_ipython().run_line_magic('matplotlib', 'inline')


# In[4]:


import numpy as np
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
from scipy import misc
from matrixmodel import MatrixOperator
from algorithm import fistaEst


# In[5]:


# Size of the signal
sigSize = (32, 32)

# Read and resize the Shepp-Logan phantom
f = np.random.binomial(1, 0.1, sigSize) * np.random.randn(sigSize[0], sigSize[1])

# Total number of pixels
numPixels = np.size(f)

# Plot image
plt.figure()
imgplot = plt.imshow(f, interpolation='nearest')
imgplot.set_cmap('gray')
plt.colorbar()
plt.axis("off")
plt.title("Sparse Signal")
plt.show()


# In[6]:


# Define the measurement model and its transpose

# Number of measurements
numMeasurements = np.ceil(2*numPixels/3).astype(int)

# Measurement matrix
H_matrix = np.random.randn(numMeasurements, numPixels)

# Define a class object for easier manipulation
forwardObj = MatrixOperator(H_matrix, sigSize)


# In[7]:


# Noise standard deviation
stdDev = 0.01

# Generate noise
noise = stdDev*np.random.randn(numMeasurements, 1)

# Generate measurements
y = forwardObj.apply(f) + noise


# In[8]:


# Benchmark reconstruction

# Regularization parameter
tau = 1

# Total number of iterations
numIter = 100

# Reconstruct with ISTA
[fhatISTA, costISTA] = fistaEst(y, forwardObj, tau, numIter, accelerate=False)

# Reconstruct with FISTA 
[fhatFISTA, costFISTA] = fistaEst(y, forwardObj, tau, numIter, accelerate=True)


# In[10]:


t = np.linspace(1, numIter, numIter)

plt.figure()
plt.semilogy(t, costISTA, 'r-', t, costFISTA, 'b-')
plt.xlim(1, numIter)
plt.grid(True)
plt.legend(["ISTA", "FISTA"])
plt.title("Evolution of the Cost")
plt.xlabel("t")
plt.ylabel("Relative Cost")

plt.figure()

# Plot original
plt.subplot(1, 2, 1)
imgplot = plt.imshow(fhatISTA, interpolation='nearest')
imgplot.set_cmap('gray')
plt.axis("off")
plt.title("ISTA")

# Plot reconstruction
plt.subplot(1, 2, 2)
imgplot = plt.imshow(fhatFISTA, interpolation='nearest')
imgplot.set_cmap('gray')
plt.axis("off")
plt.title("FISTA")
plt.show()


# In[ ]: