proxoperators.py 9.22 KB
Newer Older
s.ziehe's avatar
s.ziehe committed
1
2
3
4
5
# -*- coding: utf-8 -*-
"""
Created on Mon Sep 22 13:23:26 2014

@author: stefan
6
7

The "ProxOperators"-module contains various specific operators that do the actual calculations within the ProxToolbox.
s.ziehe's avatar
s.ziehe committed
8
9
"""

10
11
12
import numpy as np
from numpy import conj, dot, empty, ones, sqrt, sum, zeros, exp, nonzero, log
from numpy.fft import fft2
s.ziehe's avatar
s.ziehe committed
13

14
__all__ = ["P_diag","P_parallel","magproj", "Approx_P_JWST_Poisson", "P_amp"]
s.ziehe's avatar
s.ziehe committed
15

s.ziehe's avatar
s.ziehe committed
16
17
18
19
20
21
22


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

23
    def __init__(self, config):
rnahme's avatar
rnahme committed
24
25
26
        """
        Initialization method for a concrete instance
        
rnahme's avatar
rnahme committed
27
28
29
        Parameters
        ----------
        config : dict - Parameters to configure the algorithm
rnahme's avatar
rnahme committed
30
        """
31
        raise NotImplementedError("This is just an abstract interface")
s.ziehe's avatar
s.ziehe committed
32
33
34
35
36
    
    def work(self,u):
        """
        Applies a prox operator to some input data
        
rnahme's avatar
rnahme committed
37
38
39
        Parameters
        ----------
        u : array_like - Input data for the operator
s.ziehe's avatar
s.ziehe committed
40
        
rnahme's avatar
rnahme committed
41
42
43
        Returns
        -------
        array_like - Result of the operation
s.ziehe's avatar
s.ziehe committed
44
        """
45
        raise NotImplementedError("This is just an abstract interface")
s.ziehe's avatar
s.ziehe committed
46
47
48



rnahme's avatar
rnahme committed
49
50
51
52
def magproj(u, constr):
    """
    Projection operator onto a magnitude constraint
    
rnahme's avatar
rnahme committed
53
54
55
56
    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
rnahme's avatar
rnahme committed
57
    
rnahme's avatar
rnahme committed
58
59
60
    Returns
    -------
    array_like - The projection
rnahme's avatar
rnahme committed
61
    """
62
    modsq_u = conj(u) * u
rnahme's avatar
rnahme committed
63
    denom = modsq_u+3e-30
64
    denom2 = sqrt(denom)
rnahme's avatar
rnahme committed
65
66
    r_eps = (modsq_u/denom2) - constr
    dr_eps = (denom+3e-30)/(denom*denom2)
rnahme's avatar
rnahme committed
67
    
rnahme's avatar
rnahme committed
68
69
70
    return (1 - (dr_eps*r_eps)) * u


s.ziehe's avatar
s.ziehe committed
71
72
73
74
75
76
77

class P_diag(ProxOperator):
    """
    Projection onto the diagonal of a product space
    
    """
    
s.ziehe's avatar
s.ziehe committed
78
    def __init__(self,config):
s.ziehe's avatar
s.ziehe committed
79
        """
rnahme's avatar
rnahme committed
80
        Initialization
81
        
rnahme's avatar
rnahme committed
82
83
        Parameters
        ----------
84
        config : dict - Dictionary containing the problem configuration. It must contain the following mappings:
85
            
rnahme's avatar
rnahme committed
86
                'Nx': int
s.ziehe's avatar
s.ziehe committed
87
                    Row-dim of the product space elements
rnahme's avatar
rnahme committed
88
                'Ny': int
s.ziehe's avatar
s.ziehe committed
89
                    Column-dim of the product space elements
rnahme's avatar
rnahme committed
90
                'Nz': int
s.ziehe's avatar
s.ziehe committed
91
                    Depth-dim of the product space elements
rnahme's avatar
rnahme committed
92
                'dim': int
s.ziehe's avatar
s.ziehe committed
93
                    Size of the product space
s.ziehe's avatar
s.ziehe committed
94
        """
s.ziehe's avatar
s.ziehe committed
95
96
        self.n = config['Nx']; self.m = config['Ny']; self.p = config['Nz'];
        self.K = config['dim'];
s.ziehe's avatar
s.ziehe committed
97
98
99
    
    def work(self,u):
        """
rnahme's avatar
rnahme committed
100
101
        Projects the input data onto the diagonal of the product space given by its dimensions.
        
s.ziehe's avatar
s.ziehe committed
102
103
104
105
        """
        n = self.n; m = self.m; p = self.p; K = self.K;        
        
        if m == 1:
106
            tmp = sum(u,axis=1,dtype=u.dtype)
s.ziehe's avatar
s.ziehe committed
107
        elif n == 1:
108
            tmp = sum(u,axis=0,dtype=u.dtype)
s.ziehe's avatar
s.ziehe committed
109
        elif p == 1:
110
            tmp = zeros((n,m),dtype=u.dtype)
111
            for k in range(K):
rnahme's avatar
rnahme committed
112
                tmp += u[:,:,k]
s.ziehe's avatar
s.ziehe committed
113
        else:
114
            tmp = zeros((n,m,p),dtype=u.dtype)
115
            for k in range(K):
rnahme's avatar
rnahme committed
116
                tmp += u[:,:,:,k]
s.ziehe's avatar
s.ziehe committed
117
        
rnahme's avatar
rnahme committed
118
        tmp /= K
s.ziehe's avatar
s.ziehe committed
119
120
        
        if m == 1:
121
            return dot(tmp,ones((1,K),dtype=u.dtype))
s.ziehe's avatar
s.ziehe committed
122
        elif n == 1:
123
            return dot(ones((K,1),dtype=u.dtype),tmp)
s.ziehe's avatar
s.ziehe committed
124
        elif p == 1:
125
            u_diag = empty((n,m,K),dtype=u.dtype)
126
            for k in range(K):
rnahme's avatar
rnahme committed
127
128
                u_diag[:,:,k] = tmp
            return u_diag
s.ziehe's avatar
s.ziehe committed
129
        else:
130
            u_diag = empty((n,m,p,K),dtype=u.dtype)
131
            for k in range(K):
rnahme's avatar
rnahme committed
132
133
                u_diag[:,:,:,k] = tmp
            return u_diag
s.ziehe's avatar
s.ziehe committed
134
135
136
137
138
139
140
141
142
143
144





class P_parallel(ProxOperator):
    """
    Projection onto the diagonal of a product space
    
    """
    
s.ziehe's avatar
s.ziehe committed
145
    def __init__(self,config):
s.ziehe's avatar
s.ziehe committed
146
        """
rnahme's avatar
rnahme committed
147
        Initialization
148
        
rnahme's avatar
rnahme committed
149
150
        Parameters
        ----------
151
        config : dict - Dictionary containing the problem configuration. It must contain the following mappings:
152
            
rnahme's avatar
rnahme committed
153
                'Nx': int
s.ziehe's avatar
s.ziehe committed
154
                    Row-dim of the product space elements
rnahme's avatar
rnahme committed
155
                'Ny': int
s.ziehe's avatar
s.ziehe committed
156
                    Column-dim of the product space elements
rnahme's avatar
rnahme committed
157
                'Nz': int
s.ziehe's avatar
s.ziehe committed
158
                    Depth-dim of the product space elements
rnahme's avatar
rnahme committed
159
                'dim': int
s.ziehe's avatar
s.ziehe committed
160
                    Size of the product space
rnahme's avatar
rnahme committed
161
                'projs': sequence of ProxOperator
s.ziehe's avatar
s.ziehe committed
162
                    Sequence of prox operators to be used. (The classes, no instances)
s.ziehe's avatar
s.ziehe committed
163
        """
rnahme's avatar
rnahme committed
164
165
166
167
168
        self.n = config['Nx']
        self.m = config['Ny']
        self.p = config['Nz']
        self.K = config['dim']
        self.proj = []
s.ziehe's avatar
s.ziehe committed
169
        for p in config['projectors']:
rnahme's avatar
rnahme committed
170
            self.proj.append(p(config))
s.ziehe's avatar
s.ziehe committed
171
        
172
173
174
175
176
177
#        for p in config['projectors']:
#            self.proj.append(globals()[p](config))
#            pp = globals()[p]
#            self.proj.append(pp(config))

        
s.ziehe's avatar
s.ziehe committed
178
179
    def work(self,u):
        """
rnahme's avatar
rnahme committed
180
181
        Sequentially applies the projections of 'self.proj' to the input data.
        
rnahme's avatar
rnahme committed
182
183
184
        Parameters
        ----------
        u : array_like - Input
rnahme's avatar
rnahme committed
185
        
rnahme's avatar
rnahme committed
186
187
188
        Returns
        -------
        u_parallel : array_like - Projection
s.ziehe's avatar
s.ziehe committed
189
        """
rnahme's avatar
rnahme committed
190
        n = self.n; m = self.m; p = self.p; K = self.K; proj = self.proj
s.ziehe's avatar
s.ziehe committed
191
192

        if m == 1:
193
            u_parallel = empty((K,n),dtype=u.dtype)
194
            for k in range(K):
rnahme's avatar
rnahme committed
195
                u_parallel[k,:] = proj[k].work(u[k,:])
s.ziehe's avatar
s.ziehe committed
196
        elif n == 1:
197
            u_parallel = empty((m,K),dtype=u.dtype)
198
            for k in range(K):
rnahme's avatar
rnahme committed
199
                u_parallel[:,k] = proj[k].work(u[:,k])
s.ziehe's avatar
s.ziehe committed
200
        elif p == 1:
201
            u_parallel = empty((n,m,K),dtype=u.dtype)
202
            for k in range(K):
rnahme's avatar
rnahme committed
203
                u_parallel[:,:,k] = proj[k].work(u[:,:,k])
s.ziehe's avatar
s.ziehe committed
204
        else:
205
            u_parallel = empty((n,m,p,K),dtype=u.dtype)
206
            for k in range(K):
rnahme's avatar
rnahme committed
207
                u_parallel[:,:,:,k] = proj[k].work(u[:,:,:,k])
s.ziehe's avatar
s.ziehe committed
208
        
rnahme's avatar
rnahme committed
209
        return u_parallel
s.ziehe's avatar
s.ziehe committed
210
211


212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
#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'];
238
239
240
241
        self.indicator_ampl = config['indicator_ampl'];
        self.illumination_phase = config['illumination_phase'];
        self.data_zeros = config['data_zeros'];
        self.data_sq = config['data_sq'];
242
        self.product_space_dimension = config['product_space_dimension'];
243
244

    def work(self,u):
245
246
247
248
249
250
        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;
251
        product_space_dimension = self.product_space_dimension;
252
253

        for j in range(product_space_dimension-1):
254
            U = fft2( indicator_ampl *exp(1j*illumination_phase [:,:,j]) * u[:,:,j]);
255
            U_sq = U * conj(U);
256
            tmp = U_sq/data_sq[:,:,j];
257
258
259
260
261
262
263
            #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;
264
265
266
            IU= tmp==0;
            tmp[IU]=1;
            tmp=log(tmp);
267
            print(tmp);
268
            hU = sum(sum(U_sq *tmp + data_sq[:,:,j] - U_sq));
269
            if hU>=epsilon+TOL2:
270
                U0 = magproj(U,input.data[:,:,j]); #argument order changed compared to matlab implementation!!!
271
                u[:,:,j] = indicator_ampl *exp(-1j*illumination_phase[:,:,j]) *ifft2(U0);
272
273
274
275
276
277
            #else:
            # no change

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

283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
#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):
308
        return magproj(u,self.amplitude); #argument order changed compared to matlab implementation!!!
309