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

# # Optimisation Framework
# 
# The $\textbf{cil.optimisation}$ framework contains three structures, namely $\texttt{Function}$\, $\texttt{Operator}$ and $\texttt{Algorithm}$ that formalise a generic optimisation problem for imaging applications as
# $$
# u^{*} =\argmin_{u\in\mathbb{X}} f(Ku) + g(u) \equiv \argmin_{u\in \mathbb{X}} \sum_{i=0}^{n-1}f_{i}K_{i}(u) + g(u).
# \tag{1}
# $$
# 
# We let $\mathbb{X}, \mathbb{Y}$ denote finite-dimensional vector spaces, $K:\mathbb{X}\rightarrow \mathbb{Y}$ a linear operator with operator norm $\|K\| = \max \{ \|Ku\|_{Y}: \|u\|_{\mathbb{X}}\leq 1 \}$ and proper, convex functions 
# - $f: \mathbb{Y}\rightarrow\overline{\mathbb{R}}$
# - $g: \mathbb{X}\rightarrow\overline{\mathbb{R}}$
# 
# ---
# <a id='BlockFunction'></a>
# **Note:** In certain cases, it is convenient to decompose $\mathbb{Y}=Y_{0}\times\dots \times Y_{n-1}$, $n\geq1$ and consider a separable function, e.g., $\texttt{BlockFunction}$, 
# 
# - $f(y) := f(y_{0},\dots,y_{n-1}) = \sum_{i=0}^{n-1}f_{i}(y_{i})$
# 
# which results to the right-side formulation in $eqn.(1)$.
# 
# With the above form, $f(y)$ is a __separable sum__ of decoupled functions, which is useful when we need to compute the proximal operator of $f$, i.e., 
# 
# $$\mathrm{prox}_{\tau f}(z) = 
# \begin{bmatrix}
# \mathrm{prox}_{\tau f_0}(z_0)\\
# \mathrm{prox}_{\tau f_1}(z_1)\\
# \vdots\\
# \mathrm{prox}_{\tau f_{n-1}}(z_{n-1})
# \end{bmatrix} = 
# \begin{bmatrix}
# \underset{y_{0}\in \mathbb{Y_{0}}}{\operatorname{argmin}} \tau f_{0}(y_{0}) + \frac{1}{2}\|y_{0} - z_{0}\|^{2}\\
# \underset{y_{1}\in \mathbb{Y_{1}}}{\operatorname{argmin}} \tau f_{1}(y_{1}) + \frac{1}{2}\|y_{1} - z_{1}\|^{2}\\
# \vdots\\
# \underset{y_{{n-1}}\in \mathbb{Y_{n-1}}}{\operatorname{argmin}} \tau f_{{n-1}}(y_{{n-1}}) + \frac{1}{2}\|y_{{n-1}} - z_{{n-1}}\|^{2}\\
# \end{bmatrix}
# $$
# 
# ---
# 
# 
# # Gallery Algorithms
# 
# 1. [Simultaneous Iterative Reconstruction Technique](#SIRT)
# 2. [A Fast Iterative Shrinkage-Thresholding Algorithm](#FISTA)
# 3. [Primal Dual Hybrid Gradient](#PDHG)
# 4. [Stochastic Primal Dual Hybrid Gradient](#SPDHG)
# 
# <hr style="border:2px solid gray"> </hr>
# 
# ## SIRT
# 
# **Initialize:** $x_{0}\in \mathbb{X}$\
# **Update:** $x_{k}$
# 
# **for $k\geq0$**
# 
#    $x_{k+1} = \mathcal{P}_{C}(x_{k} + D \mathcal{A}^{T} M ( g - \mathcal{A} x_{k}))$
# 
# where,
#  $$
#     \begin{cases}
#     M = \frac{\mathbb{1}}{\mathcal{A}\mathbb{1}}, \quad m_{ii} = \frac{1}{\sum_{j} a_{ij}}, \quad\text{ sum over columns }\\
#     D = \frac{\mathbb{1}}{\mathcal{A}^{T}\mathbb{1}}, \quad d_{jj} = \frac{1}{\sum_{i} a_{ij}},\quad\text{ sum over rows }\\
#     \end{cases}
#     $$
#     
# <hr style="border:2px solid gray"> </hr>    
# 
# ## FISTA
# 
# **Inputs:** $\tau = \frac{1}{L}$, $L$ Lipschitz constant of $\nabla \mathcal{F}$, $t_{0}=1$
# 
# 
# **Initialize:** $x_{0}\in \mathbb{X}$\
# **Update:** $x_{k}, t_{k}, y_{k}$
# 
# **for $k\geq0$**
# 
# 1. $x_{k} = \mathrm{prox}_{\tau \mathcal{G}}(y_{k} - \tau\nabla\mathcal{F}(y^{k}))$
# 
# 2. $t_{k+1} = 1 + \frac{\sqrt{1+4 t_{k}^{2}}}{2}$
# 
# 3. $y_{k+1} = x_{k+1} + \frac{t_{k}-1}{t_{k+1}}(x_{k}-x_{k-1})$
# 
# 
#     
# <hr style="border:2px solid gray"> </hr>    
# 
# 
# ## PDHG
# 
# **Inputs:** $\tau,\sigma>0$, $\tau\sigma\|K\|^{2}<1$, $\theta\in [0,1]$
# 
# 
# **Initialize:** $x_{0}\in \mathbb{X}$, $y_{0}\in \mathbb{Y}$,$\overline{x}_{0}=x_{0}$
# **Update:** $x_{k}, y_{k}, \overline{x_{k}}$
# 
# **for $k\geq0$**
# 
# 1. $y_{k+1} = \mathrm{prox}_{\sigma \mathcal{F}^{*}}(y_{k} + \sigma K \overline{x}_{k})$
# 
# 2. $x_{k+1} = \mathrm{prox}_{\tau \mathcal{G}}(x_{k} - \tau K y_{k+1})$
# 
# 3. $\overline{x}_{k+1} = \overline{x}_{k+1} + \theta (x_{k+1} - x_{k})$
#     
# <hr style="border:2px solid gray"> </hr>  
# 
# 
# # SPDHG
# 
# **Inputs:** $n = \# \mathrm{subsets}$, $\gamma=1.0$,
# $\rho=0.99$,$\mathrm{probability}$ $p_{i}$, $i=0,\dots,n-1$
# $\\\mathrm{Stepsizes}: S_{i} = \gamma\frac{\rho}{\|A_{i}\|},\, T = \min\limits_{i} \frac{1}{\gamma}\frac{\rho p_{i}}{\|A_{i}\|},
# i=0,\dots,n-1$
#     
#     
# **Initialize:** $x^{0}, z^{0}=\overline{z}^{0}, y^{0}$\
# **Update:** $x_{k}, y_{k}, \overline{z_{k}}$
# 
# **for $k\geq0$**
# 
# 1. $x^{k+1} = \mathrm{prox}^{T}_{\mathcal{G}}(x^{k} - T\overline{z}^{k})$ 
# 
# 2. Select $i\in [n]$ at random with probability $p_{i}$
# 
# 3. $y_{i}^{k+1} = \mathrm{prox}^{S_{i}}_{\mathcal{F}_{i}^{*}}(y_{i}^{k} + S_{i}A_{i}x^{k+1})$
# 
# 4. $\Delta z = A_{i}^{T}(y_{i}^{k+1} - y_{i}^{k})$
# 
# 5. $z^{k+1} = z^{k} + \Delta z$
# 
# 6. $\overline{z}^{k+1} = z^{k+1} + \frac{\Delta z}{p_{i}}$
#     
# <hr style="border:2px solid gray"> </hr>
# 
# 
#