reponse_percussionelle.py 53.3 KB
Newer Older
1
2
3
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
4
5
6
7
8
import numpy as np
import math
import scipy
import scipy.stats
import scipy.ndimage.filters as filt
9
import sys
10
import pickle
11
import os
12
13

from scipy.optimize import minimize
14
15
16
from six.moves import filter
from six.moves import range
from six.moves import zip
17
scipymin=1
18

19
from . import minimiser
20
21
22

import h5py

myron's avatar
myron committed
23
if(sys.argv[0][-12:]!="sphinx-build"):
myron's avatar
myron committed
24
    from .XRStools_c import luts_cy
25
26
27
28
29
30
31

try:
    from mpi4py import MPI
    myrank = MPI.COMM_WORLD.Get_rank()
    nprocs = MPI.COMM_WORLD.Get_size()
    procnm = MPI.Get_processor_name()
    comm = MPI.COMM_WORLD
Alessandro Mirone's avatar
Alessandro Mirone committed
32
    print( "MPI LOADED , nprocs = ", nprocs)
33
34
35
except:
    nprocs=1
    myrank = 0
Alessandro Mirone's avatar
Alessandro Mirone committed
36
    print( "MPI NOT LOADED ")
37
38


39
40
global indent 
indent = ""
41

42
43
44
45
def filterRoiList(l):
    return [t for t in l if t not in [ "motorDict"] ]


46

47
48
49
50
def get_LUT_1d( na ,nb, cp, cb , nref):
    # ecco un caso dove bisogna considerare che il pixel va da -0.5 a 0.5
    res=[]
    # print na, nb, cp, cb , nref
Alessandro Mirone's avatar
Alessandro Mirone committed
51
    for i in range(na):   ## qui sotto si pensa che lo zero e' all' inizio del pixel
52
53
54
55
56
57
58
59
        X0 = (cb+0.5) +(i-(cp+0.5))*nref
        X1 = (cb+0.5) +(i+1-(cp+0.5))*nref
        I0 = int(math.floor(X0))
        I1 = int(math.ceil(X1))
        for j in range(I0,I1):
            x0 = float(max(j,X0))
            x1 = float(min(j+1,X1))
            if j>=0 and j<nb:
60
61
                targetx0   =  ( x0-(cb+0.5) )/nref +(cp+0.5) -i  ## la posizione esatta nel pixel su cui cade
                targetx1   =  ( x1-(cb+0.5) )/nref +(cp+0.5) -i  ##  per maschera, sono compresi fra i e i+1
Alessandro Mirone's avatar
Alessandro Mirone committed
62
63
              
                res.append([i,j,(x1-x0)/nref ,  targetx0  ,targetx1 ])
64
65
    return res
            
Alessandro Mirone's avatar
Alessandro Mirone committed
66
def get_product(lut_1,lut_2, na2, nb2 , reponse_pixel):
67
    res=[]
Alessandro Mirone's avatar
Alessandro Mirone committed
68
69
70
71
72
73
74
75
    dim1,dim2 = reponse_pixel.shape
    for i1,j1,f1, y0,y1 in lut_1:

        Y0 = dim1*y0
        Y1 = dim1*y1
        iY0  = int(math.ceil(Y0))
        iY1  = int(math.floor(Y1))
        facts = (reponse_pixel[ iY0:iY1]).sum(axis=0)
Alessandro Mirone's avatar
Alessandro Mirone committed
76
77
78
79
80
81
82
83
84

        fiY0 = min(  iY0, Y1 )
        fiY1 = max(  iY1, Y0 )
        
        if  fiY0-Y0>1.0e-10 :
            facts[:] += (fiY0-Y0)  *  reponse_pixel[int(math.floor(Y0))]
        if(iY1>=iY0):
            if(Y1-fiY1>1.0e-10):
                facts[:] += (Y1-fiY1)  *  reponse_pixel[iY1]
Alessandro Mirone's avatar
Alessandro Mirone committed
85
            
86
        facts[:] /= (Y1-Y0) # media. Il fattore area c' e' gia in f1, f2
Alessandro Mirone's avatar
Alessandro Mirone committed
87
88
89
        
       
        for i2,j2,f2, x0,x1 in lut_2:
90
            # print j2
91
92
            X0 = dim2*x0
            X1 = dim2*x1
Alessandro Mirone's avatar
Alessandro Mirone committed
93
94
95
96
            iX0  = int(math. ceil(X0))
            iX1  = int(math.floor(X1))
            
            Fatt    = (facts[ iX0:iX1]).sum(axis=0)
Alessandro Mirone's avatar
Alessandro Mirone committed
97
98
99
100
101

            
            fiX0 = min(  iX0, X1 )
            fiX1 = max(  iX1, X0 )

Alessandro Mirone's avatar
Alessandro Mirone committed
102
            
Alessandro Mirone's avatar
Alessandro Mirone committed
103
104
105
106
107
            if ( fiX0-X0>1.0e-10):
                Fatt   += (fiX0-X0)  *  facts[int(math.floor(X0))]
            if(iX1>=iX0):
                if(X1-fiX1>1.0e-10):
                    Fatt  += (X1-fiX1)  *  facts[iX1]
Alessandro Mirone's avatar
Alessandro Mirone committed
108
                
109
            Fatt /= (X1-X0) # media. Il fattore area c' e' gia in f1, f2
Alessandro Mirone's avatar
Alessandro Mirone committed
110
            if Fatt>1.1:
Alessandro Mirone's avatar
Alessandro Mirone committed
111
112
113
114
115
116
117
                print( Fatt, " Fatt ")
                print( facts, Y0, Y1, iY0, iY1)
                print( "  X0,X1 " , X0,X1)
                print( " iX0, iX1 " ,iX0, iX1)
                print( "(facts[ iX0:iX1]).sum(axis=0)  " , (facts[ iX0:iX1]).sum(axis=0))
                print( " (fiX0-X0)  *  facts[int(math.floor(X0))] " ,(fiX0-X0)  *  facts[int(math.floor(X0))])
                print( " (X1-fiX1)  *  facts[iX1] " ,(X1-fiX1)  *  facts[iX1])
Alessandro Mirone's avatar
Alessandro Mirone committed
118
119
120

                
                raise
121
            newel =    [  i1*na2+i2   ,  j1*nb2+j2  , f1*f2* Fatt ] 
122
123
            res.append(newel)
    return res
Alessandro Mirone's avatar
Alessandro Mirone committed
124
125


Alessandro Mirone's avatar
Alessandro Mirone committed
126
def get_product4reponse(lut_1,lut_2, na2, nb2 , reponse_pixel, solution):
Alessandro Mirone's avatar
Alessandro Mirone committed
127
128
129
130
131
132
133
134
135
136
137
138
139
140
    res=[]
    dim1,dim2 = reponse_pixel.shape
    repnseq = np.arange(reponse_pixel.size)
    repnseq.shape = reponse_pixel.shape
    for i1,j1,f1, y0,y1 in lut_1:

        Y0 = dim1*y0
        Y1 = dim1*y1
        iY0  = int(math.ceil(Y0))
        iY1  = int(math.floor(Y1))
        facts = (reponse_pixel[ iY0:iY1]).sum(axis=0)
        
        recolte = np.array(repnseq[ iY0:iY1] )
        factrec = np.ones( repnseq[ iY0:iY1].shape   )
Alessandro Mirone's avatar
Alessandro Mirone committed
141
142
143

        fiY0 = min(  iY0, Y1 )
        fiY1 = max(  iY1, Y0 )
144
        
Alessandro Mirone's avatar
Alessandro Mirone committed
145
146
        
        if fiY0-Y0>1.0e-10 :
Alessandro Mirone's avatar
Alessandro Mirone committed
147
            where = int(math.floor(Y0))
Alessandro Mirone's avatar
Alessandro Mirone committed
148
149
            recolte = np.concatenate(  [recolte ,   [repnseq[where]]   ]   )            
            factrec = np.concatenate(  [factrec , [ [ (fiY0-Y0) ]  *  repnseq.shape[1]]   ]  )
Alessandro Mirone's avatar
Alessandro Mirone committed
150

Alessandro Mirone's avatar
Alessandro Mirone committed
151
152
153
154
        if(iY1 > Y0):            
            if( Y1-fiY1 > 1.0e-10 ):
                recolte = np.concatenate( [ [repnseq[iY1]],  recolte ]    )
                factrec = np.concatenate(  [ [[(Y1-fiY1)]*repnseq.shape[1]],     factrec  ]  )
Alessandro Mirone's avatar
Alessandro Mirone committed
155
156
            
            
Alessandro Mirone's avatar
Alessandro Mirone committed
157
        ### facts[:]   /= (Y1-Y0) # media. Il fattore area c' e' gia in f1, f2
158
        factrec[:] /= (Y1-Y0) # *dim1
Alessandro Mirone's avatar
Alessandro Mirone committed
159
160
161
162
163

        # print " =========== "
        # TERM = factrec.sum(axis=0)
        # print factrec.sum(axis=0)
        
Alessandro Mirone's avatar
Alessandro Mirone committed
164
165
166
167
168
169
170
        for i2,j2,f2, x0,x1 in lut_2:

            X0 = dim2*x0
            X1 = dim2*x1
            iX0  = int(math. ceil(X0))
            iX1  = int(math.floor(X1))
            
Alessandro Mirone's avatar
Alessandro Mirone committed
171
172
            recolteX = np.array(recolte  [ :, iX0:iX1] )
            factrecX = np.array(factrec  [ :, iX0:iX1] ) 
173

Alessandro Mirone's avatar
Alessandro Mirone committed
174
            # print " PEZZZO " , TERM[ iX0:iX1]
Alessandro Mirone's avatar
Alessandro Mirone committed
175
176

            
Alessandro Mirone's avatar
Alessandro Mirone committed
177
178
179
180
181
182
183
184
185
186
187
188
189
190
            fiX0 = min(  iX0, X1 )
            fiX1 = max(  iX1, X0 )
        
            if(fiX0-X0>1.0e-10):
                where = int(math.floor(X0))
                recolteX  = np.concatenate(  [ recolte[:, where:where+1]               ,  recolteX]    , axis=1 )
                factrecX  = np.concatenate(  [   (fiX0-X0) * factrec [:, where:where+1] ,  factrecX]    , axis=1  )
                # print " PEZZZO " , (fiX0-X0) * TERM[ where:where+1]
           
            if(iX1>=iX0):
                if(X1-fiX1>1.0e-10):
                    recolteX = np.concatenate(  [  recolteX  , recolte[:, iX1:iX1+1]                 ] , axis=1   )
                    factrecX = np.concatenate(  [  factrecX  ,  (X1-fiX1) * factrec[ :, iX1:iX1+1 ]   ] , axis=1 )
                    # print " PEZZZO " ,  (X1-fiX1) * TERM[ iX1:iX1+1]
Alessandro Mirone's avatar
Alessandro Mirone committed
191
192
                
                
193
            factrecX     /= (X1-X0) # *dim2  # media. Il fattore area c' e' gia in f1, f2
Alessandro Mirone's avatar
Alessandro Mirone committed
194
195


Alessandro Mirone's avatar
Alessandro Mirone committed
196
197
            # print " ---> " , factrecX.sum()
            if  (   factrecX.sum()>1.1   ) :
Alessandro Mirone's avatar
Alessandro Mirone committed
198
199
200
                print(    X0,X1,Y0,Y1)
                print(    iX0,iX1,iY0,iY1)
                print(    fiX0,fiX1,fiY0,fiY1)
Alessandro Mirone's avatar
Alessandro Mirone committed
201
202
203
204
                raise
            ## factrecX[:] = factrecX/factrecX.sum()

            
Alessandro Mirone's avatar
Alessandro Mirone committed
205
            for J,F in zip( recolteX.flatten()   ,   factrecX.flatten()  ):
Alessandro Mirone's avatar
Alessandro Mirone committed
206
207
                # print j1, nb2, j2
                newel =    [  i1*na2+i2   ,  J  , f1*f2* solution[j1,j2]*F      ] 
Alessandro Mirone's avatar
Alessandro Mirone committed
208
            # newel =    [  i1*na2+i2   ,  j1*nb2+j2  , f1*f2, Fatt ] 
209
                res.append(newel)
Alessandro Mirone's avatar
Alessandro Mirone committed
210
211
212
213
214
    return res
            
                


215
def get_LUT( mat_a,mat_b, center_pic, nref, reponse_pixel, doproduct = 1, soluzione=None, ROI = None):
216
217
218
219

    na1, na2  = mat_a.shape
    nb1, nb2  = mat_b.shape

220
221
222
    if ROI is None:
        ROI = np.ones_like(mat_a)
    ROI = ROI.astype("f")
223
    
224
225
    center_b = np.array(    [ (nb1-1)/2.0, (nb2-1)/2.0 ]   )

226
227
228
    lut_1 = get_LUT_1d( na1 ,nb1 , center_pic[0] ,center_b[0]  , nref)
    lut_2 = get_LUT_1d( na2 ,nb2 , center_pic[1] ,center_b[1]  , nref)
    if doproduct:
Alessandro Mirone's avatar
Alessandro Mirone committed
229
        if doproduct ==1:
230
231
232
            # LUT = get_product(lut_1,lut_2, na2, nb2 , reponse_pixel)
            
            # print " cy_product ", np.array(LUT).sum()
233
234
            #print np.array(lut_1,"f").shape
            #print np.array(lut_2,"f").shape
235
236
237
238
239
240
            # print " SHAPE "
            # print np.array(reponse_pixel,"f").shape, np.array(lut_1,"f").shape, np.array(lut_2,"f").shape, na2, nb2 , len( np.array(lut_1,"f")  ), len( np.array(lut_2,"f") ) 


            if len(lut_1)==0 or len(lut_2)==0:
                return None
241
            tmp = luts_cy.get_product( np.array(lut_1,"f"),
Alessandro Mirone's avatar
compila    
Alessandro Mirone committed
242
243
                                                np.array(lut_2,"f"),
                                                na2, nb2 ,
244
                                                np.array(reponse_pixel,"f"), ROI)
245
            LUT = np.array(tmp)
246
247
            
            
248
            #print LUT.sum()
Alessandro Mirone's avatar
compila    
Alessandro Mirone committed
249
            
Alessandro Mirone's avatar
Alessandro Mirone committed
250
        else :
251
252
253
254
255
256
            #print soluzione.shape
            #print mat_a.shape
            #print " ------------------ " 
            #print np.array(lut_2)[:,1].max()
            #print np.array(lut_1)[:,1].max()
            #print soluzione.shape
Alessandro Mirone's avatar
Alessandro Mirone committed
257
258
259
260
261
262
263
            if 0:
                LUT= get_product4reponse(  lut_1,
                                           lut_2,
                                           na2, nb2 ,
                                           reponse_pixel,
                                           soluzione)
            else:
264
                global SYMM_RESPO
265
                tmp = luts_cy.get_product4reponse(np.array(lut_1,"f"),np.array(lut_2,"f"), na1,
266
                                                           na2, nb2 , np.array(reponse_pixel,"f"), np.array(soluzione,"f"),
267
                                                           SYMM_RESPO, ROI)
Alessandro Mirone's avatar
Alessandro Mirone committed
268
                LUT = np.array(tmp)
269
                
Alessandro Mirone's avatar
compila    
Alessandro Mirone committed
270
271

            
272
273
274
275
        return LUT
    else:
        return lut_1, lut_2

276
def calculate_grad(grad,data ,solution , s2d, d2s, solution_shape , parallel = 0 , beta=0.1 ) :
277
278
279
280
281
282
283
284
285
286
    nsol, ndata = d2s.shape
    
    proj =   s2d.dot(solution)
    if data is not None:
        err  =   data - proj
    else:
        err  =    -proj
    fid  =   np.dot(err,err)
    grad[:] =   d2s.dot(err)

287
288
289
290
291
292
293
294
295
296
297
298
299
    #print solution_shape

    if parallel and nprocs>1:
        
        grad_res = np.zeros_like(grad)
        comm.Allreduce([ grad , MPI.FLOAT  ], [grad_res, MPI.FLOAT  ], op=MPI.SUM)
        grad[:]= grad_res

        a    = np.array( [fid]    ,  "f"  )
        a_r  = np.array( [0]    ,  "f"  )
        comm.Allreduce([ a , MPI.FLOAT  ], [a_r, MPI.FLOAT  ], op=MPI.SUM)
        fid = a_r[0]        
        
300
301
302
303
    sol_on_surf = np.reshape( solution, solution_shape )
    laplacian    = filt.laplace(sol_on_surf , mode='reflect')
    result =  fid/2.0 - beta*        np.dot( laplacian.flatten(), sol_on_surf.flatten()  ) /2
    grad[:] +=   beta * laplacian.flatten()
304
305


306
307
    # if myrank==0:
    #    print " fid ----<" , result, parallel
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
    return result


def    cg( data , solution   ,   s2d, d2s, solution_shape  ):
    maxdim = max(  s2d.shape[0], s2d.shape[1]   )
    auxs = [   np.zeros( [maxdim]  ,"f")  for i in range(5)                            ]

    nsol, ndata = d2s.shape
    
    grad=auxs[0][:nsol]
    grad_old=auxs[1][:nsol]
    p=auxs[2][:nsol]

    err = 0.0
    err=calculate_grad(grad,data ,solution , s2d, d2s, solution_shape  ) 
    # calculate_grad(grad, Volume, donnees, dim, n_row,n_col,  nnz, Bp, Bj, Bx, beta,auxs+3 ) ; 
      
    p[:] = grad[:]
    
    rrold=0.0
    rr   = 0.0
    
    rrold +=  np.dot( grad, grad  )
    rr=rrold;

    for iter in range(500):
        grad_old[:] = grad[:]
        calculate_grad(grad,None ,p , s2d, d2s, solution_shape  )
        pap = np.dot(  p, grad )

        solution[:] -= p*(rr/pap)

        err=calculate_grad(grad , data ,solution , s2d, d2s, solution_shape  ) 

        rrold=rr;
        
        rr = np.dot(  grad,   (grad-grad_old)  )

        
        beta = rr/rrold
        if beta<0 :
            beta = 0

        p[:] =  grad[:]+  p[:]*beta

353
        if myrank==0:
354
355
356
357
            if iter%1 ==0:
                sys.stdout.write(" "*60+"\r"+indent+( "iter %d errore est %e  mod_grad est  %e" % ( iter,  err,rr )    )  )
                sys.stdout.flush()
    if myrank==0:
Alessandro Mirone's avatar
Alessandro Mirone committed
358
        print( " ")
359

360

361
def    Fista( data , solution   ,   s2d, d2s, solution_shape , parallel = 0 , niter=500, beta=0.1 ):
362
    global indent
363
364
365
366
367
368
369
370
371
372
373
    maxdim = max(  s2d.shape[0], s2d.shape[1]   )
    auxs = [   np.zeros( [maxdim]  ,"f")  for i in range(5)                            ]

    nsol, ndata = d2s.shape
    
    grad =auxs[0][:nsol]
    grad2=auxs[1][:nsol]
    x_old=auxs[2][:nsol]
    y    =auxs[3][:nsol]

    err = 0.0
374
    err=calculate_grad(grad,data ,solution , s2d, d2s, solution_shape , parallel = parallel , beta=beta) 
375
376
    # calculate_grad(grad, Volume, donnees, dim, n_row,n_col,  nnz, Bp, Bj, Bx, beta,auxs+3 ) ; 

377
378

    if myrank==0:
Alessandro Mirone's avatar
Alessandro Mirone committed
379
        print( indent+"CALCULATING LIPSCHITZ FACTOR ")
myron's avatar
myron committed
380
    for i in range(10):
381
382
        calculate_grad(grad2,None ,grad , s2d, d2s, solution_shape , parallel = parallel , beta=beta)
        Lip = math.sqrt( np.linalg.norm(grad2/100000)   )*100000
383
        grad[:]  = grad2/ Lip
384
        if myrank==0:
385
386
            sys.stdout.write( " "*60+"\r"+indent+"LIP  %e"% Lip    ) 
            sys.stdout.flush()
387
388
    #if myrank==0:
        #print ""
389
390
    

391
392
393
    t=1.0
    y[:] = solution
    x_old[:] = solution
394
395
    for iter in range(abs(niter)):
        err = calculate_grad(grad,data ,y , s2d, d2s, solution_shape , parallel = parallel, beta=beta )
396
397
398
399
400
        solution[:] =  y + grad/Lip
        solution[:] = np.maximum(solution, 0)
        tnew = ( 1+math.sqrt(1.0+4*t*t) )/2
        y[:] = solution +(t-1)/tnew *( solution - x_old )
        t = tnew
401
402
        if niter<0:
            t=1
403
        x_old[:] = solution
404
        if myrank==0:
405
            if iter%1 ==0:
406
407
                # sys.stdout.write(" "*60+"\r"+indent+("FISTA iter %d  errore est %e  mod_grad est  %e" % ( iter,  err, grad.std()) ))
                sys.stdout.write(indent+("FISTA iter %d  errore est %e  mod_grad est  %e\n" % ( iter,  err, grad.std()) ))
408
                sys.stdout.flush()
409

410
    if myrank==0:
Alessandro Mirone's avatar
Alessandro Mirone committed
411
        print( " " )
412
413
414


        
415
def   trajectory_error( XYXY    ,  O_spots, opticalPSF,  nref, reponse_pixel  , retrieve_spots, suggerimento=None, ROI = None):
Alessandro Mirone's avatar
Alessandro Mirone committed
416
    if myrank==0: print( XYXY)
417

418
    Nspots = O_spots.shape[0]
419
420
421
422
423
424
425
426
427

    if(len(XYXY)==4):
        X1,Y1,X2,Y2 = XYXY
    else:
        X1,X2 = XYXY
        Y1    = suggerimento[0] + X1* suggerimento[1]
        Y2    = suggerimento[0] + X2* suggerimento[1]


428
429
430
431
432
433
434
    trajectory = Trajectory()
    trajectory.set_extrema(X1,Y1,X2,Y2, Nspots  )

    err = 0.0
    dd = 0.0
    ss = 0.0
    ds = 0.0
435
436
437
438
439
440

    total_I = []
    total_J = []
    total_F = []
    total_data = []

441
442
443
    for n, data in enumerate( O_spots  ):
        # print n
        center_pic = [   trajectory.Y.intercept + n* trajectory.Y.slope  ,      trajectory.X.intercept + n*  trajectory.X.slope         ]
444
        # lut1, lut2 =  get_LUT(data  ,  opticalPSF , center_pic, nref, reponse_pixel, doproduct = 0)
445
        LUT =  get_LUT(data  ,  opticalPSF , center_pic, nref, reponse_pixel, doproduct = 1, ROI=ROI)
446
447
448
449
450
451
        if LUT is not  None:        
            I,J,F = np.array(LUT).T
            assert(J.max()<    opticalPSF.size )
            total_I.append( n*data.size + I  )
            total_J.append( J  )
            total_F.append(F)
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
        total_data.append( data.flatten()  )

        # tmp = np.zeros(  [ data.shape[0] , opticalPSF.shape[1]   ],"f"  )
        # tmp2= np.zeros(  [ data.shape[0] ,       data.shape[1]   ],"f"  )
        # for i,j,f, dum1, dum2 in lut1:
        #     tmp [i,:] += opticalPSF[j,: ] *f
        # for i,j,f , dum1, dum2 in lut2:
        #     tmp2[:,i] += tmp[:,j]  *f 
        # ss += (tmp2*tmp2).sum()
        # dd += (data*data).sum()
        # ds += (tmp2*data).sum()

    #print " CONCATENO " 
    I = np.concatenate(total_I)
    J = np.concatenate(total_J)
    F = np.concatenate(total_F)
    data = np.concatenate(total_data )
    
    s2d_coo = scipy.sparse.coo_matrix( (F,(I,J)) , shape = [   O_spots.size   ,  opticalPSF.size   ])
    s2d =  s2d_coo.tocsr()
    sim = s2d.dot( opticalPSF.flatten())
473

474
475
    if retrieve_spots:
        return sim
476

477
478
479
480
    dd = (data*data).sum()
    ss=(sim*sim).sum()
    ds=(data*sim).sum()
    
481
    error = dd - ds*ds/ss
482
483
484
485
486
    
    ### f2 * ss + dd -2*ds*f    ==>  2f*ss -2 ds *f = 0
    ###  f = ds/ss        error =  ds*ds/ss +dd -2ds*ds/ss = dd - ds*ds/ss
    
    if myrank==0:
487
488
        sys.stdout.write( " "*60+"\r"+indent+"trajectory error %e" % error )
        sys.stdout.flush()
489
490
    return error

491
492
493
494
495
496
497
# class Tobject:
#     def __init__(self,  variables ,  O_spots, opticalPSF,  nref , reponse_pixel   ) :
#         self.variables      =  variables
#         self.O_spots        =  O_spots
#         self.opticalPSF     =  opticalPSF
#         self.nref           =  nref
#         self.reponse_pixel  =  reponse_pixel
498
        
499
500
501
502
503
#     def error(self):
#         X1 = self.variables[0].value
#         X2 = self.variables[1].value
#         Y1 = self.variables[2].value
#         Y2 = self.variables[3].value
504
        
505
506
507
508
509
510
511
#         res = trajectory_error(np.array([X1,Y1,X2,Y2]),  self.O_spots, self.opticalPSF,  self.nref , self.reponse_pixel,0  )
#         if myrank==0:
#             print " error = " , res
#         return res


def refine_trajectory(O_spots, opticalPSF, trajectory   , nref  , reponse_pixel , retrieve_spots = 0 , suggerimento = None , ROI = None):
512
513
    global indent

514
    Nspots = O_spots.shape[0]
515
516
517



518
519
520
    X1,Y1 = trajectory.get_coords(0)
    X2,Y2 = trajectory.get_coords(Nspots-1)

521
    if retrieve_spots:
Alessandro Mirone's avatar
Alessandro Mirone committed
522
        res = trajectory_error(np.array([X1,Y1,X2,Y2]),  O_spots, opticalPSF,  nref , reponse_pixel ,   1 , None, ROI)
523
        return res
524
525

    
526
527
528
529
530
531
532
    # if not scipymin or 0:
    #     variables = [ minimiser.Variable( X1,X1-2, X1+2),
    #                   minimiser.Variable( X2,X2-2, X2+2),
    #                   minimiser.Variable( Y1,Y1-2, Y1+2),
    #                   minimiser.Variable( Y2,Y2-2, Y2+2)
    #     ]
    #     tominimise=Tobject(variables ,  O_spots, opticalPSF,  nref , reponse_pixel  )
533
        
534
535
    #     miniob=minimiser.minimiser(tominimise,variables)
    #     miniob.amoeba(0.0001)  
536
        
537
538
539
540
    #     X1 = variables[0].value
    #     X2 = variables[1].value
    #     Y1 = variables[2].value
    #     Y2 = variables[3].value
541

542
    
543
544
545
546
547
    #     trajectory.set_extrema( X1,Y1,X2,Y2  , Nspots)
    #     return trajectory    

    if suggerimento is not None:

548
        if myrank==0:
Alessandro Mirone's avatar
Alessandro Mirone committed
549
            print( indent +"IMPROVING TRAJECTORY ")
550

Alessandro Mirone's avatar
Alessandro Mirone committed
551
        res = minimize(trajectory_error,np.array([X1,X2]),( O_spots, opticalPSF,  nref , reponse_pixel ,0 ,suggerimento , None, ROI),     # method='Powell',
552
                       method='Nelder-Mead',
553
                       options={'disp': False, 'maxiter': 40,
554
555
556
                                'return_all': False,
                                'maxfev': None, 'xtol': 0.001,
                       'ftol': 0.001 }  )
557
        if myrank==0:
Alessandro Mirone's avatar
Alessandro Mirone committed
558
559
            print( " ...  MINIMO IN ",    res.x,)
            print( " INIZIALE  ",    [X1,X2])
560
           #  print " FINALE   ",    res.x[0], res.x[1]
561
562
563
564

        trajectory.set_extrema_suggestion( res.x[0], res.x[1],   Nspots  , suggerimento)
       

565
    else:
566
        dodebug = False
567
        if myrank==0:
Alessandro Mirone's avatar
Alessandro Mirone committed
568
            print( indent + "IMPROVING TRAJECTORY ")
569
            dodebug = True
570
        res = minimize(trajectory_error,np.array([X1,Y1,X2,Y2]),( O_spots, opticalPSF,  nref , reponse_pixel ,0, None, ROI ),     # method='Powell',
571
                       method='Nelder-Mead',
572
573
                       options={'disp': dodebug, 'maxiter': 40, 'return_all': False,  'maxfev': None,  'ftol': 0.001})
                       # options={'disp': dodebug, 'maxiter': 40, 'return_all': False,  'maxfev': None, 'xtol': None, 'ftol': 0.001})
574
575
        # res = minimize(trajectory_error,np.array([X1,Y1,X2,Y2]),( O_spots, opticalPSF,  nref ))
        if myrank==0:
Alessandro Mirone's avatar
Alessandro Mirone committed
576
577
            print( " ...  MINIMO IN ",    res.x,)
            print( " INIZIALE  ",    [X1,Y1,X2,Y2])
578
            # print " FINALE   ",    res.x[0], res.x[1], res.x[2],  res.x[3]
579
580
            
        trajectory.set_extrema( res.x[0], res.x[1], res.x[2],  res.x[3],   Nspots)
581
582
583


    return trajectory    
584

585

586
def fit_reponse(trajectory, O_spots , nref, reponse_pixel, beta=0.1, niter=500, ROI = None )    :
587
588
589
590
591
592
593
594

    solution = np.zeros( [O_spots[0].shape[0]*nref, O_spots[0].shape[1]*nref    ] , "f" )
    total_I = []
    total_J = []
    total_F = []
    total_data = []

    for n, data in enumerate( O_spots  ):
595
596
        # if myrank==0:
        #    print n
597
        center_pic = [   trajectory.Y.intercept + n* trajectory.Y.slope ,      trajectory.X.intercept + n*  trajectory.X.slope         ]
598
        LUT =  get_LUT(data  ,solution , center_pic, nref, reponse_pixel, ROI=ROI)
599
600
601
602
603
        if LUT is not None:
            I,J,F = np.array(LUT).T
            total_I.append( n*data.size + I  )
            total_J.append( J  )
            total_F.append(F)
604
605
        total_data.append( data.flatten()  )

606
607
608
    # if myrank==0:
    #     print " CONCATENO "
    
609
610
611
612
613
    I = np.concatenate(total_I)
    J = np.concatenate(total_J)
    F = np.concatenate(total_F)
    data = np.concatenate(total_data )

614
615
616
617
618
    # if myrank==0:
    #     print " CREO MATRICE "
    #     print data.shape
    #     print I.shape
    #     print J.shape
619
        
620
621
622
    s2d_coo = scipy.sparse.coo_matrix( (F,(I,J)) , shape = [   O_spots.size   ,  solution.size   ]    )
    d2s_coo = scipy.sparse.coo_matrix( (F,(J,I)) , shape = [   solution.size  ,  O_spots.size  ]    )

623

624
625
626
    # if myrank==0:
    #     print " CAMBIO FORMATO "
    
627
628
629
    s2d =  s2d_coo.tocsr()
    d2s =  d2s_coo.tocsr()

630
631
    # if myrank==0:
    #     print " FISTA  "
632
        
633
634
635
636
    b=solution[:]
    solution_shape = solution.shape
    solution=np.reshape(solution,[-1])
    # cg( data , solution    ,   s2d, d2s, solution_shape  )
637
    Fista( data , solution    ,   s2d, d2s, solution_shape, niter=niter, beta=beta )
638
639
640
    solution.shape = solution_shape
    return solution

Alessandro Mirone's avatar
Alessandro Mirone committed
641
642


643
def fit_reponse_pixel(trajectory_list, O_spots_list , nref, reponse_pixel, solution_list, niter=500, beta=100 , rois = None):
Alessandro Mirone's avatar
Alessandro Mirone committed
644

Alessandro Mirone's avatar
Alessandro Mirone committed
645
646
647
648
649
650

    if type(trajectory_list)!=type([]):
        trajectory_list = [trajectory_list]
        O_spots_list = [O_spots_list]
        solution_list = [solution_list]
        
Alessandro Mirone's avatar
Alessandro Mirone committed
651
652
653
654
655
    total_I = []
    total_J = []
    total_F = []
    total_data = []

656
    tot_size_data = 0
657
    for trajectory , O_spots, solution, ROI  in zip(trajectory_list,O_spots_list, solution_list, rois ) :
Alessandro Mirone's avatar
Alessandro Mirone committed
658

659
660
661
        if trajectory is None:
            continue

Alessandro Mirone's avatar
Alessandro Mirone committed
662
663
664
665
        assert( solution.shape == (O_spots[0].shape[0]*nref, O_spots[0].shape[1]*nref ) )


        for n, data in enumerate( O_spots  ):
666
            # if myrank==0:
667
        #     print n
Alessandro Mirone's avatar
Alessandro Mirone committed
668
            center_pic = [   trajectory.Y.intercept + n* trajectory.Y.slope ,      trajectory.X.intercept + n*  trajectory.X.slope         ]
669
            LUT =  get_LUT(data  ,solution , center_pic, nref, reponse_pixel, soluzione = solution, doproduct = 2, ROI = ROI )
670
671
672
673
674
            if LUT is not None:
                I,J,F = np.array(LUT).T
                total_I.append( tot_size_data+ n*data.size + I  )
                total_J.append( J  )
                total_F.append(F)
Alessandro Mirone's avatar
Alessandro Mirone committed
675
            total_data.append( data.flatten()  )
676
677
678
679
            
        tot_size_data +=O_spots .size
        
    if myrank==0:
Alessandro Mirone's avatar
Alessandro Mirone committed
680
        print( " CONCATENO ", )
Alessandro Mirone's avatar
Alessandro Mirone committed
681
682
683
684
685
    I = np.concatenate(total_I)
    J = np.concatenate(total_J)
    F = np.concatenate(total_F)
    data = np.concatenate(total_data )

686
687
688
689
690
    # if myrank==0:
    #     print " CREO MATRICE "
    #     print data.shape
    #     print I.shape
    #     print J.shape
691
692
693

    s2d_coo = scipy.sparse.coo_matrix( (F,(I,J)) , shape = [   tot_size_data   ,  reponse_pixel.size   ]    )
    d2s_coo = scipy.sparse.coo_matrix( (F,(J,I)) , shape = [   reponse_pixel.size  ,  tot_size_data  ]    )
Alessandro Mirone's avatar
Alessandro Mirone committed
694

695

696
697
    # if myrank==0:
    #     print " CAMBIO FORMATO "
698
        
Alessandro Mirone's avatar
Alessandro Mirone committed
699
700
701
702
703
    s2d =  s2d_coo.tocsr()
    d2s =  d2s_coo.tocsr()

    datasim = s2d.dot( reponse_pixel.flatten()   )
    datasim = datasim - data
704

705
    if myrank==0:
Alessandro Mirone's avatar
Alessandro Mirone committed
706
707
        print( " FIDELITY ", (datasim*datasim).sum(),)
        print( " FISTA  ",)
708
        
Alessandro Mirone's avatar
Alessandro Mirone committed
709
710
    reponse_pixel_shape = reponse_pixel.shape
    reponse_pixel=np.reshape(reponse_pixel,[-1])
711
    
712
    Fista( data , reponse_pixel    ,   s2d, d2s, reponse_pixel_shape , parallel = 1 , niter=niter, beta=beta )
Alessandro Mirone's avatar
Alessandro Mirone committed
713
714
715
716
717
    reponse_pixel.shape = reponse_pixel_shape
    return reponse_pixel



718
def get_spots_list(  filename  , groupname , filter_rois =1 ):
719
720
721
    res=[]
    nomi_scan = []
    stats   = []
722
    origini = {}
723
    
724
    h5 = h5py.File(filename,"r")
725

726
727
    xscales = {}
    enescan = None
myron's avatar
myron committed
728
729
730

    tmp_list = list(filterRoiList(list(h5[groupname].keys())))
    tmp_list = sorted(tmp_list, key = int )
731
    
myron's avatar
myron committed
732
    for sn in tmp_list:
733
        
Alessandro Mirone's avatar
Alessandro Mirone committed
734
        print( groupname+"/"+sn+"/matrix")
735
        m  = h5[groupname+"/"+sn+"/matrix"][:]
Alessandro Mirone's avatar
Alessandro Mirone committed
736
        print( m.shape)
737

myron's avatar
myron committed
738
739
740
741
742
743
744
        if groupname+"/"+sn+"/xscale" in h5:
            xscales[sn] = h5[groupname+"/"+sn+"/xscale"][:]
        else:
            xscales[sn] = np.arange( m.shape[0]).astype("f") 


        
745
        if groupname+"/motorDict/energy" in h5:
myron's avatar
myron committed
746
            enescan = h5[groupname+"/motorDict/energy"][()]
747
748

        
749
750
        if m.shape!=(0,):
            stat = m.sum(axis=-1).sum(axis=-1)
751
752
            
            if (not filter_rois) or (stat.min()>stat.max()/2.0 and stat.min() >250.0):
753
754
755
                res.append(m)
                nomi_scan.append(  sn   )
                stats.append(stat)
756
757
758
759
                origini[sn] =  h5[groupname+"/"+sn+"/cornerpos"][:]
                
    h5.close()
    rois = []
760
761
                                    ## METTERE D UFFICIO ANCHE LE ROI QUANDO SI SCRIVE UNO SCAN E CAMBIARE QUI IL ALLINDIETRO
    for i in range(1):
762
763
        pos = groupname.rfind("/")
        groupname=groupname[:pos]
764

765
    # groupname=groupname+"/rois_definition/rois_dict/"
766

Alessandro Mirone's avatar
Alessandro Mirone committed
767
    print( " NOMI SCAN ", nomi_scan)
768
769
770
771
    
    h5 = h5py.File(filename,"r")
    for sn in nomi_scan:

Alessandro Mirone's avatar
Alessandro Mirone committed
772
        print( " DATANAME ", groupname+"/ROI%02d/mask" %int(sn), filename)
773
774
        m  = h5[groupname+"/ROI%02d/mask" %int(sn)][:]
        rois.append(m)
775
    h5.close()
776

Alessandro Mirone's avatar
Alessandro Mirone committed
777
    print( res)
alex's avatar
alex committed
778
779
    for m in res:
        print( m.sum())
Alessandro Mirone's avatar
Alessandro Mirone committed
780
781
    print( nomi_scan)
    print( rois)
782
    return res, nomi_scan, stats, origini, rois, xscales, enescan
783
784

   
785
786
class pippo:
    pass
787
788
789
790
791
792
793
794
795
796
797

class Trajectory :
    def __init__(self):
        self.X=pippo()
        self.Y=pippo()

    def get_coords(self,n):
        X1 = self.X.intercept + self.X.slope*n
        Y1 = self.Y.intercept + self.Y.slope*n
        return X1, Y1
    def set_extrema(self, X1, Y1, X2, Y2,N):
798
        self.N = N
799
800
801
802
        self.X.intercept = X1
        self.Y.intercept = Y1
        self.X.slope     = (X2-X1)/(N-1)
        self.Y.slope     = (Y2-Y1)/(N-1)
803
        
804
    def set_extrema_suggestion( self, X1, X2,   Nspots  , suggerimento):
805
806
807
808
809
810
        Nspots = self.N
        Y1    = suggerimento[0] + X1* suggerimento[1]
        Y2    = suggerimento[0] + X2* suggerimento[1]            
        self. set_extrema( X1, Y1, X2, Y2,Nspots)
        

811
    def  projectOnHint( self, suggerimento  ):
812
813
814
815
816
817
818
819
        Nspots = self.N

        X1,Y1 = self.get_coords(0)
        X2,Y2 = self.get_coords(Nspots-1)
        Y1    = suggerimento[0] + X1* suggerimento[1]
        Y2    = suggerimento[0] + X2* suggerimento[1]
            
        self. set_extrema( X1, Y1, X2, Y2,Nspots)
820
821
822
823
824
825
    def __repr__(self):
        s=""
        s=s+" X.intercept = %e " % self.X.intercept
        s=s+" X.slope     = %e " % self.X.slope
        s=s+" Y.intercept = %e " % self.Y.intercept
        s=s+" Y.slope     = %e " % self.Y.slope
826
827
        s=s+" Nspots     = %d " % self.N
        
828
        return s
829
830


alex's avatar
alex committed
831
832
833
834
835
836
class NoSignal(Exception):
    def __init__(self, value):
        self.value = value
    def __str__(self):
        return repr(self.value)
    
837
838
839
def get_Trajectory_byregression( O_spots    ):
    Xs=[]
    Ys=[]
840
    Ws=[]
841
842
843
844
845
846
847
    for slice in O_spots:
        mmax  = slice.max()
        thres = mmax / 20
        slice = np.array(slice)
        slice[ slice<thres  ] = 0

        m0 = slice.sum()
alex's avatar
alex committed
848
849
850
851
852
853
854
855
856
        if m0>0:
            X  = (  slice * np.arange(slice.shape[1] )  ).sum()/m0
            Y  = (  slice.T * np.arange(slice.shape[0] )  ).sum()/m0
            Xs.append(X)
            Ys.append(Y)
            Ws.append(m0)
    if  len(Xs)<=1:
        raise NoSignal(" All images were zero  ")
    
857
    ret_val = Trajectory()
858
859
860
861
862
863
864

    ret_val.X.slope, ret_val.X.intercept = np.polyfit( np.arange(len(Xs)), Xs ,1,  w = Ws)
    ret_val.Y.slope, ret_val.Y.intercept = np.polyfit( np.arange(len(Ys)), Ys ,1,w= Ws )

    
    # ret_val.X.slope, ret_val.X.intercept, Xr_value, Xp_value, Xstd_err = scipy.stats.linregress(np.arange(len(Xs)),Xs , )
    # ret_val.Y.slope, ret_val.Y.intercept, Yr_value, Yp_value, Ystd_err = scipy.stats.linregress(np.arange(len(Ys)),Ys , )
865
    ret_val.N = len(O_spots)
866
867
868
869
    
    return ret_val


870
871
def DOFIT(filename=None, groupname=None, nref=5, niter_optical=500, beta_optical=0.1 ,
          beta_pixel=1000.0, niter_pixel = -20,
872
          niter_global  = 50, pixel_dim=6, simmetrizza=1, do_refine_trajectory=1, target_file="responses.h5",target_groupname="FIT",
873
          trajectory_reference_scansequence_filename = None, trajectory_reference_scansequence_groupname = None ,  trajectory_threshold = None,
874
          trajectory_file=None  , filter_rois=1 , fit_lines = False):
875

876
877
878
    ###### filename = "../nonregressions/demo_imaging.hdf5"
    ###### groupname = "ROI_B/foil_scanXX/scans/Scan273/"
    # print filename,  groupname
879

880
    O_spots_list, nomi_rois, stats, origini_foil, rois, xscales, energy = get_spots_list( filename,  groupname,  filter_rois= filter_rois )
881
    stats = np.array(stats)
myron's avatar
myron committed
882

883
884
    if do_refine_trajectory ==2:

885
886
887
        h5f = h5py.File(trajectory_reference_scansequence_filename,"r")
        h5  = h5f[trajectory_reference_scansequence_groupname]

888
        zscan_keys =  sorted(  list(filterRoiList(list(h5.keys())))     , key = lambda x:  int(list(filter(str.isdigit, str(x) ))) )
889
890
891
        ZDIM = len(zscan_keys)
        myrange = np.array_split( np.arange( ZDIM ), nprocs   )[myrank]
        myZDIM =   len(myrange)
892
893
894
895
        
        roi_keys  = nomi_rois
        integrated     = {}
        origini_ref    = {}
896
897
        for rk in roi_keys:
            zkey = zscan_keys[ 0 ]
898
            m  = h5[zkey][rk]["matrix"][:]
899
            
900
901
902
903
            mm = m.sum(axis=0)
            integrated[rk ] = np.zeros_like( mm ) .astype("d")
            origini_ref   [rk ] = h5[zkey][rk]["cornerpos"][:]

904
905
906
907
908
909
910
911
912
        for iz in range(myZDIM):        
            zkey = zscan_keys[myrange[iz] ]
            for rk in roi_keys:
                m = h5[ zkey ][ rk ]["matrix"][:]
                msum = m.sum(axis=0)
                integrated[rk] = integrated[rk]+msum
                
        if nprocs>1:    
            for n in  integrated.keys():
913
                    comm.Allreduce(  [np.array(integrated[n]), MPI.DOUBLE], [integrated[n], MPI.DOUBLE],  op=MPI.SUM)
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928

        trajectory_hints={}
        for n in  integrated.keys():
            C     = integrated[n]
            pesiC = C.sum(axis=0)
            medieC = (np.arange(C.shape[0])[:,None]*C).sum(axis=0)/pesiC
            coords  = np.arange(len( medieC  )) 

            maxP=pesiC.max()            
            whereBigger = np.where( pesiC > maxP*trajectory_threshold   )[0]
            
            inizio = whereBigger.min()
            fine   = whereBigger.max()
            pesifit = np.zeros( len(pesiC),"f")
            pesifit[inizio:fine+1]   =   ( pesiC[inizio:fine+1]> maxP*trajectory_threshold  )*1
929
930
931
932
933
934
            pfit=np.polynomial.polynomial.polyfit(coords, medieC, 1,  w=pesifit)
            tmp_intercept = pfit[0]
            tmp_slope     = pfit[1]
            
            DY = origini_foil[n][0] -origini_ref[n][0]
            DX = origini_foil[n][1] -origini_ref[n][1]
935

936
937
938
            hint_slope = tmp_slope
            hint_intercept  = tmp_intercept - DY  +  hint_slope*DX
            
939
940
941
942
            trajectory_hints[n] = hint_intercept,hint_slope
    else:
        trajectory_hints={}

943
944
    global SYMM_RESPO
    SYMM_RESPO = simmetrizza
Alessandro Mirone's avatar
Alessandro Mirone committed
945

946
    reponse_pixel = np.ones([ pixel_dim , pixel_dim  ],"f")
Alessandro Mirone's avatar
Alessandro Mirone committed
947
    
948
949
950
951
    #### nref = 5
    
    trajectory_list = []
    solution_list   = []
952
953
    trajectories_from_file = None
    if trajectory_file is not None:
954
        ### file = open(trajectory_file,"r")
Alessandro Mirone's avatar
Alessandro Mirone committed
955
        print( "RELOADING TRAJECTORIES ")
956
        trajectories_from_file = reload_trajectories(trajectory_file,  nomi_rois)
957
958

    
959
    for iterm,  (O_spots, name, ROI)  in enumerate(zip(O_spots_list, nomi_rois, rois)):
960
        # print "MYRANK %d fa "%myrank, iterm
961
        if (iterm)%nprocs == myrank:
962
963

            if trajectories_from_file is not None:
964
                " REUSING TRAJECTORIES "
965
                trajectory = trajectories_from_file[iterm]
966
                trajectory.N = O_spots.shape[0]
967
            else:
alex's avatar
alex committed
968
969
970
971
972
973
974
                try:
                    trajectory = get_Trajectory_byregression( O_spots    )
                except NoSignal as inst:
                    msg = " Was not able to determine trajactory because not enough signal in ROI %s"%name
                    print(msg)
                    inst.value = msg
                    raise
975

Alessandro Mirone's avatar
Alessandro Mirone committed
976
            print( " TRAJECTORY ", trajectory)
977
978
979
980

            if  do_refine_trajectory  ==2 :      
  
                PRIMA = trajectory.Y.intercept, trajectory.Y.slope
Alessandro Mirone's avatar
Alessandro Mirone committed
981
                print( "============= projectOnHint============= ", )
982
983


984
                trajectory.projectOnHint(  trajectory_hints[  name  ]  )
985
986
                
  
Alessandro Mirone's avatar
Alessandro Mirone committed
987
                print( "Intercept : ", PRIMA[0] , "===>", trajectory.Y.intercept , "Slope : ", PRIMA[1] , "===>", trajectory.Y.slope)
988

989
            if myrank ==0:
Alessandro Mirone's avatar
Alessandro Mirone committed
990
                print( "Process 0 FITTING OPTICAL RESPONSE for scan ", name)
991

992
            solution =  fit_reponse( trajectory, O_spots , nref  , reponse_pixel,niter=niter_optical, beta= beta_optical, ROI=ROI)   ## niter 500    beta_optical 0.1
993
994
995
996
997
998
999
1000

        else:
            solution = None
            trajectory = None
       
        solution_list.append(solution)
        trajectory_list.append(trajectory)