proxoperators.py

# -*- coding: utf-8 -*-
"""
Created on Mon Sep 22 13:23:26 2014

@author: stefan

The "ProxOperators"-module contains various specific operators that do the actual calculations within the ProxToolbox.
"""

import numpy as np
from numpy import conj, dot, empty, ones, sqrt, sum, zeros, exp, nonzero, log
from numpy.fft import fft2

__all__ = ["P_diag","P_parallel","magproj", "Approx_P_JWST_Poisson", "P_amp"]



class ProxOperator:
    """
    Generic interface for prox operators
    """

    def __init__(self, config):
        """
        Initialization method for a concrete instance
        
        Parameters
        ----------
        config : dict - Parameters to configure the algorithm
        """
        raise NotImplementedError("This is just an abstract interface")
    
    def work(self,u):
        """
        Applies a prox operator to some input data
        
        Parameters
        ----------
        u : array_like - Input data for the operator
        
        Returns
        -------
        array_like - Result of the operation
        """
        raise NotImplementedError("This is just an abstract interface")



def magproj(u, constr):
    """
    Projection operator onto a magnitude constraint
    
    Parameters
    ----------
    u : array_like - The function to be projected onto constr (can be complex)
    constr : array_like - A nonnegative array that is the magnitude constraint
    
    Returns
    -------
    array_like - The projection
    """
    modsq_u = conj(u) * u
    denom = modsq_u+3e-30
    denom2 = sqrt(denom)
    r_eps = (modsq_u/denom2) - constr
    dr_eps = (denom+3e-30)/(denom*denom2)
    
    return (1 - (dr_eps*r_eps)) * u



class P_diag(ProxOperator):
    """
    Projection onto the diagonal of a product space
    
    """
    
    def __init__(self,config):
        """
        Initialization
        
        Parameters
        ----------
        config : dict - Dictionary containing the problem configuration. It must contain the following mappings:
            
                '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.n = config['Nx']; self.m = config['Ny']; self.p = config['Nz'];
        self.K = config['dim'];
    
    def work(self,u):
        """
        Projects the input data onto the diagonal of the product space given by its dimensions.
        
        """
        n = self.n; m = self.m; p = self.p; K = self.K;        
        
        if m == 1:
            tmp = sum(u,axis=1,dtype=u.dtype)
        elif n == 1:
            tmp = sum(u,axis=0,dtype=u.dtype)
        elif p == 1:
            tmp = zeros((n,m),dtype=u.dtype)
            for k in range(K):
                tmp += u[:,:,k]
        else:
            tmp = zeros((n,m,p),dtype=u.dtype)
            for k in range(K):
                tmp += u[:,:,:,k]
        
        tmp /= K
        
        if m == 1:
            return dot(tmp,ones((1,K),dtype=u.dtype))
        elif n == 1:
            return dot(ones((K,1),dtype=u.dtype),tmp)
        elif p == 1:
            u_diag = empty((n,m,K),dtype=u.dtype)
            for k in range(K):
                u_diag[:,:,k] = tmp
            return u_diag
        else:
            u_diag = empty((n,m,p,K),dtype=u.dtype)
            for k in range(K):
                u_diag[:,:,:,k] = tmp
            return u_diag





class P_parallel(ProxOperator):
    """
    Projection onto the diagonal of a product space
    
    """
    
    def __init__(self,config):
        """
        Initialization
        
        Parameters
        ----------
        config : dict - Dictionary containing the problem configuration. It must contain the following mappings:
            
                '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
                'projs': sequence of ProxOperator
                    Sequence of prox operators to be used. (The classes, no instances)
        """
        self.n = config['Nx']
        self.m = config['Ny']
        self.p = config['Nz']
        self.K = config['dim']
        self.proj = []
        for p in config['projectors']:
            self.proj.append(p(config))
        
#        for p in config['projectors']:
#            self.proj.append(globals()[p](config))
#            pp = globals()[p]
#            self.proj.append(pp(config))

        
    def work(self,u):
        """
        Sequentially applies the projections of 'self.proj' to the input data.
        
        Parameters
        ----------
        u : array_like - Input
        
        Returns
        -------
        u_parallel : array_like - Projection
        """
        n = self.n; m = self.m; p = self.p; K = self.K; proj = self.proj

        if m == 1:
            u_parallel = empty((K,n),dtype=u.dtype)
            for k in range(K):
                u_parallel[k,:] = proj[k].work(u[k,:])
        elif n == 1:
            u_parallel = empty((m,K),dtype=u.dtype)
            for k in range(K):
                u_parallel[:,k] = proj[k].work(u[:,k])
        elif p == 1:
            u_parallel = empty((n,m,K),dtype=u.dtype)
            for k in range(K):
                u_parallel[:,:,k] = proj[k].work(u[:,:,k])
        else:
            u_parallel = empty((n,m,p,K),dtype=u.dtype)
            for k in range(K):
                u_parallel[:,:,:,k] = proj[k].work(u[:,:,:,k])
        
        return u_parallel


#original matlab comment
#                      Approx_PM_Poisson.m
#             written on Feb. 18, 2011 by 
#                   Russell Luke
#   Inst. Fuer Numerische und Angewandte Mathematik
#                Universitaet Gottingen
#
# DESCRIPTION:  Projection subroutine for projecting onto Fourier
#               magnitude constraints
#
# INPUT:        input = data structure
#                       .data_ball is the regularization parameter described in 
#                        D. R. Luke, Nonlinear Analysis 75 (2012) 1531–1546.
#               u = function in the physical domain to be projected
#
# OUTPUT:       
#               
# USAGE: u_epsilon = P_M(input,u)
#
# Nonstandard Matlab function calls:  MagProj

class Approx_P_JWST_Poisson(ProxOperator):

    def __init__(self,config):
        self.TOL2 = config['TOL2'];
        self.epsilon = config['data_ball'];
        self.indicator_ampl = config['indicator_ampl'];
        self.illumination_phase = config['illumination_phase'];
        self.data_zeros = config['data_zeros'];
        self.data_sq = config['data_sq'];
        self.product_space_dimension = config['product_space_dimension'];

    def work(self,u):
        TOL2 = self.TOL2;
        epsilon = self.epsilon;
        indicator_ampl = self.indicator_ampl;
        illumination_phase = self.illumination_phase;
        data_zeros = self.data_zeros;
        data_sq = self.data_sq;
        product_space_dimension = self.product_space_dimension;

        for j in range(product_space_dimension-1):
            U = fft2( indicator_ampl *exp(1j*illumination_phase [:,:,j]) * u[:,:,j]);
            U_sq = U * conj(U);
            tmp = U_sq/data_sq[:,:,j];
            #tmp2= nonzero(data_zeros[:,j]);
            #tmpI = data_zeros[data_zeros[:,j] != 0, j];
            mask = ones(tmp.shape, np.bool);
            print(mask.shape);
            mask[data_zeros[j]] = 0;
            tmp[mask]=1;
            U_sq[mask]=0;
            IU= tmp==0;
            tmp[IU]=1;
            tmp=log(tmp);
            print(tmp);
            hU = sum(sum(U_sq *tmp + data_sq[:,:,j] - U_sq));
            if hU>=epsilon+TOL2:
                U0 = magproj(U,input.data[:,:,j]); #argument order changed compared to matlab implementation!!!
                u[:,:,j] = indicator_ampl *exp(-1j*illumination_phase[:,:,j]) *ifft2(U0);
            #else:
            # no change

        # now project onto the pupil constraint.
        # this is a qualitative constraint, so no 
        # noise is taken into account.
        j=product_space_dimension;
        u[:,:,j] = magproj(u[:,:,j],abs_illumination);#argument order changed compared to matlab implementation!!!
        
        return u;

#original matlab comment
#                      P_amp.m
#             written on Feb 3, 2012 by 
#                   Russell Luke
#   Inst. Fuer Numerische und Angewandte Mathematik
#                Universitaet Gottingen
#
# DESCRIPTION:  Projection subroutine for projecting onto
#               amplitude constraints

# INPUT:        input.amplitude =  nonnegative array
#               u = complex-valued array to be projected onto the ampltude constraints
#
# OUTPUT:       p_amp    = the projection 
#               
# USAGE: p_amp = P_amp(S,u)
#
# Nonstandard Matlab function calls:  MagProj

class P_amp(ProxOperator):

    def __init__(self,config):
        self.amplitude = config['amplitude'];

    def work(self,u):
        return magproj(u,self.amplitude); #argument order changed compared to matlab implementation!!!