SimpleAlgortihm.py

# Simple Algorithm
# Analog to AlgorithmWrapper in ProxMatlab

from numpy import zeros, angle, trace, exp, sqrt, sum, ndarray
from numpy.linalg import norm


def nd_norm_difference(arr1: ndarray, arr2: ndarray, norm_factor: float = 1.):
    """
    Root-sum-square norm of the difference of 2 numpy arrays, equal to Frobenius norm for arr.ndim == 2
    """
    return sqrt(sum(abs(arr1 - arr2) ** 2)) / norm_factor


def phase_offset_compensated_norm(arr1: ndarray, arr2: ndarray, norm_factor: float = 1.,
                                  norm_type: str = 'fro') -> float:
    """
    Norm of the difference between 2 numpy arrays.

    :param arr1: numpy array
    :param arr2: numpy array like arr1
    :param norm_factor: normalization factor
    :param norm_type: for 1d and 2d arrays, pass norm_type to numpy.linalg.norm
    :return: float
    """
    if arr1.ndim < 3:
        phase_offset = angle(sum(arr1 * arr2.conj()))
        return norm(arr1 - arr2 * exp(1j * phase_offset), norm_type) / norm_factor
    else:
        phase_offset = angle(sum(arr1 * arr2.conj()))
        return nd_norm_difference(arr1, arr2 * exp(1j * phase_offset), norm_factor=norm_factor)


class SimpleAlgorithm:

    def __init__(self, config):
        """
        Parameters
        ----------
        config : dict        
         Dictionary containing the problem configuration.
         It must contain the following mappings:

        proxoperators: 2 ProxOperators
            Tuple of ProxOperators (the class, no instance)
        beta_0: number
            Starting relaxation parmater
        beta_max: number
            Maximum relaxation parameter
        beta_switch: int
            Iteration at which beta moves from beta_0 -> beta_max
        norm_data: number
            ?
        Nx: int
            Row-dim of the product space elements
        Ny: int
            Column-dim of the product space elements
        Nz: int
            Depth-dim of the product space elements
        dim: int
            Size of the product space
        """

        self.prox1 = config['proxoperators'][0](config)
        self.prox2 = config['proxoperators'][1](config)
        self.norm_data = config['norm_data']
        self.Nx = config['Nx']
        self.Ny = config['Ny']
        self.Nz = config['Nz']
        self.product_space_dimension = config['product_space_dimension']
        self.iter = 0
        self.config = config

        if 'truth' in config:
            self.truth = config['truth']
            self.truth_dim = config['truth_dim']
            self.norm_truth = config['norm_truth']

        if 'diagnostic' in config:
            self.diagnostic = config['diagnostic']
        else:
            self.diagnostic = False

    def evaluate(self, u):
        """
        Method to override in specific algorithm subclass
        :param u: old iterate
        :return: new iterate
        """
        return u

    def run(self, u, tol, maxiter):
        """
        Runs the algorithm for the specified input data

        :param u: iterate
        :param tol: tolerance for convergence, if change<tol, stop
        :param maxiter:
        :return:
        """
        ##### PREPROCESSING
        config = self.config
        norm_data = self.norm_data

        iter = self.iter

        # TODO: select the right case also for a 3d problem.
        #  This could involve using the parameter self.product_space_dimension
        if u.ndim < 3:  # temp note: as used in (e.g.) 2d CDI, or orbital tomography
            p = 1
            q = 1
        elif u.ndim == 3:
            p = u.shape[2]
            q = 1
        else:
            p = u.shape[2]
            q = u.shape[3]

        change = zeros(maxiter + 1)
        change[0] = 999

        ######################################
        #  set up diagnostic arrays
        ######################################
        if self.diagnostic:
            gap = change.copy()
            shadow_change = change.copy()
        if hasattr(self, 'truth'):
            Relerrs = change.copy()

        # Different algorithms will have different qualities that one should monitor.
        # Since we take a fixed point perspective, the main thing to monitor is the change in the iterates.
        shadow = self.prox1.work(u)

        ##### LOOP
        if 'tqdm_progressbar' in self.config and self.config['tqdm_progressbar'] is not None:
            pbar = self.config['tqdm_progressbar'](total=maxiter)

        while iter < maxiter and change[iter] >= tol:
            iter += 1
            config['iter'] = iter
            if 'tqdm_progressbar' in self.config and self.config['tqdm_progressbar'] is not None:
                pbar.update()
            # makes call to algorithm:
            config['u'] = u
            u_new = self.evaluate(u)

            if self.diagnostic:
                # the next prox operation only gets used in the computation of the size of the gap.
                # This extra step is not required in alternating projections, which makes RAAR a bit slower
                u2 = self.prox2.work(u_new)
                u1 = self.prox1.work(u2)

            # compute the normalized change in successive iterates
            tmp_change = 0
            tmp_gap = 0
            tmp_shadow = 0
            # the simple case, where everything can be calculated in one difference
            if self.product_space_dimension == 1 or (p == 1 and q == 1):
                tmp_change = phase_offset_compensated_norm(u, u_new, norm_type='fro', norm_factor=norm_data) ** 2
                if self.diagnostic:
                    # For Douglas-Rachford,in general it is appropriate to monitor the SHADOWS of the iterates,
                    # since in the convex case these converge even for beta=1. (see Bauschke-Combettes-Luke,
                    # J. Approx. Theory, 2004)
                    tmp_shadow = phase_offset_compensated_norm(u2, shadow, norm_factor=norm_data, norm_type='fro') ** 2
                    tmp_gap = phase_offset_compensated_norm(u1, u2, norm_factor=norm_data, norm_type='fro') ** 2
                    if hasattr(self, 'truth'):
                        if self.truth_dim[0] == 1:
                            z = u1[0, :]
                        elif self.truth_dim[1] == 1:
                            z = u1[:, 0]
                        else:
                            z = u1
                        Relerrs[iter] = phase_offset_compensated_norm(self.truth, z, norm_factor=self.norm_truth,
                                                                      norm_type='fro')

            elif q == 1:  # more complex: loop over product space dimension
                for j in range(self.product_space_dimension):  # this loop might be avoided using the 3d norm function
                    tmp_change = tmp_change + (norm(u[:, :, j] - u_new[:, :, j], 'fro') / norm_data) ** 2
                    if self.diagnostic:
                        # compute (||P_SP_Mx-P_Mx||/norm_data)^2:
                        tmp_gap = tmp_gap + (norm(u1[:, :, j] - u2[:, :, j], 'fro') / norm_data) ** 2
                        tmp_shadow = tmp_shadow + (norm(u2[:, :, j] - shadow[:, :, j], 'fro') / norm_data) ** 2
                if self.diagnostic:
                    if hasattr(self, 'truth'):
                        z = u1[:, :, 0]
                        Relerrs[iter] = norm((self.truth - exp(-1j * angle(trace(self.truth.T.transpose() * z))) * z),
                                             'fro') / self.norm_truth

            else:  # loop over product space dimension as well as additional z dimension
                if self.diagnostic:
                    if hasattr(self, 'truth'):
                        Relerrs[iter] = 0
                for j in range(self.product_space_dimension):  # this loop might be avoided using a 4d norm function
                    for k in range(self.Nz):  # this loop might be avoided using the 3d norm function
                        tmp_change = tmp_change + (norm(u[:, :, k, j] - u_new[:, :, k, j], 'fro') / norm_data) ** 2
                        if self.diagnostic:
                            # compute (||P_Sx-P_Mx||/norm_data)^2:
                            tmp_gap = tmp_gap + (norm(u1[:, :, k, j] - u2[:, :, k, j], 'fro') / (norm_data)) ** 2
                            tmp_shadow = tmp_shadow + (
                                    norm(u2[:, :, k, j] - shadow[:, :, k, j], 'fro') / (norm_data)) ** 2
                if hasattr(self, 'truth') and (j == 0):
                    Relerrs[iter] = Relerrs[iter] + norm(
                        self.truth - exp(-1j * angle(trace(self.truth.T * u1[:, :, k, 1]))) * u1[:, :, k, 1],
                        'fro') / self.norm_Truth

            # Add values to the list of change, gap and shadow_change
            change[iter] = sqrt(tmp_change)
            if self.diagnostic:
                gap[iter] = sqrt(tmp_gap)
                shadow_change[iter] = sqrt(tmp_shadow)  # this is the Euclidean norm of the gap to
                # the unregularized set. To monitor the Euclidean norm of the gap to the regularized set is
                # expensive to calculate, so we use this surrogate. Since the stopping criteria is on the change
                # in the iterates, this does not matter.

            # update iterate
            u = u_new
            if self.diagnostic:
                # For Douglas-Rachford, in general it is appropriate to monitor the SHADOWS of the iterates, since in
                # the convex case these converge even for beta=1. (see Bauschke-Combettes-Luke, J. Approx. Theory, 2004)
                shadow = u2

        ##### POSTPROCESSING
        u2 = self.prox2.work(u)
        u1 = self.prox1.work(u2)
        if self.Nx == 1:
            u1 = u1[:, 0]
            u2 = u2[:, 0]
        elif self.Ny == 1:
            u1 = u1[0, :]
            u2 = u2[0, :]
        elif self.Nz == 1 and u1.ndim > 2:
            u1 = u1[:, :, 0]
            u2 = u2[:, :, 0]

        change = change[1:iter + 1]
        output = {'u': u, 'u1': u1, 'u2': u2, 'iter': iter, 'change': change}
        if self.diagnostic:
            gap = gap[1:iter + 1]
            shadow_change = shadow_change[1:iter + 1]
            output['gap'] = gap
            output['shadow_change'] = shadow_change
            if hasattr(self, 'truth'):
                output['Relerrs'] = Relerrs
        return output