reponse_percussionelle.py 53.6 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
myron's avatar
myron committed
9
from scipy.ndimage import gaussian_filter 
10
import sys
11
import pickle
12
import os
13
14

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

20
from . import minimiser
21
22
23

import h5py

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

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
33
    print( "MPI LOADED , nprocs = ", nprocs)
34
35
36
except:
    nprocs=1
    myrank = 0
Alessandro Mirone's avatar
Alessandro Mirone committed
37
    print( "MPI NOT LOADED ")
38
39


40
41
global indent 
indent = ""
42

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


47

48
49
50
51
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
52
    for i in range(na):   ## qui sotto si pensa che lo zero e' all' inizio del pixel
53
54
55
56
57
58
59
60
        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:
61
62
                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
63
64
              
                res.append([i,j,(x1-x0)/nref ,  targetx0  ,targetx1 ])
65
66
    return res
            
Alessandro Mirone's avatar
Alessandro Mirone committed
67
def get_product(lut_1,lut_2, na2, nb2 , reponse_pixel):
68
    res=[]
Alessandro Mirone's avatar
Alessandro Mirone committed
69
70
71
72
73
74
75
76
    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
77
78
79
80
81
82
83
84
85

        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
86
            
87
        facts[:] /= (Y1-Y0) # media. Il fattore area c' e' gia in f1, f2
Alessandro Mirone's avatar
Alessandro Mirone committed
88
89
90
        
       
        for i2,j2,f2, x0,x1 in lut_2:
91
            # print j2
92
93
            X0 = dim2*x0
            X1 = dim2*x1
Alessandro Mirone's avatar
Alessandro Mirone committed
94
95
96
97
            iX0  = int(math. ceil(X0))
            iX1  = int(math.floor(X1))
            
            Fatt    = (facts[ iX0:iX1]).sum(axis=0)
Alessandro Mirone's avatar
Alessandro Mirone committed
98
99
100
101
102

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

Alessandro Mirone's avatar
Alessandro Mirone committed
103
            
Alessandro Mirone's avatar
Alessandro Mirone committed
104
105
106
107
108
            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
109
                
110
            Fatt /= (X1-X0) # media. Il fattore area c' e' gia in f1, f2
Alessandro Mirone's avatar
Alessandro Mirone committed
111
            if Fatt>1.1:
Alessandro Mirone's avatar
Alessandro Mirone committed
112
113
114
115
116
117
118
                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
119
120
121

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


Alessandro Mirone's avatar
Alessandro Mirone committed
127
def get_product4reponse(lut_1,lut_2, na2, nb2 , reponse_pixel, solution):
Alessandro Mirone's avatar
Alessandro Mirone committed
128
129
130
131
132
133
134
135
136
137
138
139
140
141
    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
142
143
144

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

Alessandro Mirone's avatar
Alessandro Mirone committed
152
153
154
155
        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
156
157
            
            
Alessandro Mirone's avatar
Alessandro Mirone committed
158
        ### facts[:]   /= (Y1-Y0) # media. Il fattore area c' e' gia in f1, f2
159
        factrec[:] /= (Y1-Y0) # *dim1
Alessandro Mirone's avatar
Alessandro Mirone committed
160
161
162
163
164

        # print " =========== "
        # TERM = factrec.sum(axis=0)
        # print factrec.sum(axis=0)
        
Alessandro Mirone's avatar
Alessandro Mirone committed
165
166
167
168
169
170
171
        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
172
173
            recolteX = np.array(recolte  [ :, iX0:iX1] )
            factrecX = np.array(factrec  [ :, iX0:iX1] ) 
174

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

            
Alessandro Mirone's avatar
Alessandro Mirone committed
178
179
180
181
182
183
184
185
186
187
188
189
190
191
            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
192
193
                
                
194
            factrecX     /= (X1-X0) # *dim2  # media. Il fattore area c' e' gia in f1, f2
Alessandro Mirone's avatar
Alessandro Mirone committed
195
196


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

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


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

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

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

227
228
229
    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
230
        if doproduct ==1:
231
232
233
            # LUT = get_product(lut_1,lut_2, na2, nb2 , reponse_pixel)
            
            # print " cy_product ", np.array(LUT).sum()
234
235
            #print np.array(lut_1,"f").shape
            #print np.array(lut_2,"f").shape
236
237
238
239
240
241
            # 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
242
            tmp = luts_cy.get_product( np.array(lut_1,"f"),
Alessandro Mirone's avatar
compila    
Alessandro Mirone committed
243
244
                                                np.array(lut_2,"f"),
                                                na2, nb2 ,
245
                                                np.array(reponse_pixel,"f"), ROI)
246
            LUT = np.array(tmp)
247
248
            
            
249
            #print LUT.sum()
Alessandro Mirone's avatar
compila    
Alessandro Mirone committed
250
            
Alessandro Mirone's avatar
Alessandro Mirone committed
251
        else :
252
253
254
255
256
257
            #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
258
259
260
261
262
263
264
            if 0:
                LUT= get_product4reponse(  lut_1,
                                           lut_2,
                                           na2, nb2 ,
                                           reponse_pixel,
                                           soluzione)
            else:
265
                global SYMM_RESPO
266
                tmp = luts_cy.get_product4reponse(np.array(lut_1,"f"),np.array(lut_2,"f"), na1,
267
                                                           na2, nb2 , np.array(reponse_pixel,"f"), np.array(soluzione,"f"),
268
                                                           SYMM_RESPO, ROI)
Alessandro Mirone's avatar
Alessandro Mirone committed
269
                LUT = np.array(tmp)
270
                
Alessandro Mirone's avatar
compila    
Alessandro Mirone committed
271
272

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

277
def calculate_grad(grad,data ,solution , s2d, d2s, solution_shape , parallel = 0 , beta=0.1 ) :
278
279
280
281
282
283
284
285
286
287
    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)

288
289
290
291
292
293
294
295
296
297
298
299
300
    #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]        
        
301
302
303
304
    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()
305
306


307
308
    # if myrank==0:
    #    print " fid ----<" , result, parallel
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
353
    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

354
        if myrank==0:
355
356
357
358
            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
359
        print( " ")
360

361

362
def    Fista( data , solution   ,   s2d, d2s, solution_shape , parallel = 0 , niter=500, beta=0.1 ):
363
    global indent
364
365
366
367
368
369
370
371
372
373
374
    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
375
    err=calculate_grad(grad,data ,solution , s2d, d2s, solution_shape , parallel = parallel , beta=beta) 
376
377
    # calculate_grad(grad, Volume, donnees, dim, n_row,n_col,  nnz, Bp, Bj, Bx, beta,auxs+3 ) ; 

378
379

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

392
393
394
    t=1.0
    y[:] = solution
    x_old[:] = solution
395
396
    for iter in range(abs(niter)):
        err = calculate_grad(grad,data ,y , s2d, d2s, solution_shape , parallel = parallel, beta=beta )
397
398
399
400
401
        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
402
403
        if niter<0:
            t=1
404
        x_old[:] = solution
405
        if myrank==0:
406
            if iter%1 ==0:
407
408
                # 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()) ))
409
                sys.stdout.flush()
410

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


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

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

    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]


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

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

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

442
443
444
    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         ]
445
        # lut1, lut2 =  get_LUT(data  ,  opticalPSF , center_pic, nref, reponse_pixel, doproduct = 0)
446
        LUT =  get_LUT(data  ,  opticalPSF , center_pic, nref, reponse_pixel, doproduct = 1, ROI=ROI)
447
448
449
450
451
452
        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)
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
        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())
474

475
476
    if retrieve_spots:
        return sim
477

478
479
480
481
    dd = (data*data).sum()
    ss=(sim*sim).sum()
    ds=(data*sim).sum()
    
482
    error = dd - ds*ds/ss
483
484
485
486
487
    
    ### 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:
488
489
        sys.stdout.write( " "*60+"\r"+indent+"trajectory error %e" % error )
        sys.stdout.flush()
490
491
    return error

492
493
494
495
496
497
498
# 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
499
        
500
501
502
503
504
#     def error(self):
#         X1 = self.variables[0].value
#         X2 = self.variables[1].value
#         Y1 = self.variables[2].value
#         Y2 = self.variables[3].value
505
        
506
507
508
509
510
511
512
#         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):
513
514
    global indent

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



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

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

    
527
528
529
530
531
532
533
    # 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  )
534
        
535
536
    #     miniob=minimiser.minimiser(tominimise,variables)
    #     miniob.amoeba(0.0001)  
537
        
538
539
540
541
    #     X1 = variables[0].value
    #     X2 = variables[1].value
    #     Y1 = variables[2].value
    #     Y2 = variables[3].value
542

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

    if suggerimento is not None:

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

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

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

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


    return trajectory    
585

586

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

    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  ):
596
597
        # if myrank==0:
        #    print n
598
        center_pic = [   trajectory.Y.intercept + n* trajectory.Y.slope ,      trajectory.X.intercept + n*  trajectory.X.slope         ]
599
        LUT =  get_LUT(data  ,solution , center_pic, nref, reponse_pixel, ROI=ROI)
600
601
602
603
604
        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)
605
606
        total_data.append( data.flatten()  )

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

615
616
617
618
619
    # if myrank==0:
    #     print " CREO MATRICE "
    #     print data.shape
    #     print I.shape
    #     print J.shape
620
        
621
622
623
    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  ]    )

624

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

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

Alessandro Mirone's avatar
Alessandro Mirone committed
642
643


644
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
645

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

    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
652
653
654
655
656
    total_I = []
    total_J = []
    total_F = []
    total_data = []

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

660
661
662
        if trajectory is None:
            continue

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


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

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

    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
695

696

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

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

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



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

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

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

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


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

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

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

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

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

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

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

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):
799
        self.N = N
800
801
802
803
        self.X.intercept = X1
        self.Y.intercept = Y1
        self.X.slope     = (X2-X1)/(N-1)
        self.Y.slope     = (Y2-Y1)/(N-1)
804
        
805
    def set_extrema_suggestion( self, X1, X2,   Nspots  , suggerimento):
806
807
808
809
810
811
        Nspots = self.N
        Y1    = suggerimento[0] + X1* suggerimento[1]
        Y2    = suggerimento[0] + X2* suggerimento[1]            
        self. set_extrema( X1, Y1, X2, Y2,Nspots)
        

812
    def  projectOnHint( self, suggerimento  ):
813
814
815
816
817
818
819
820
        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)
821
822
823
824
825
826
    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
827
828
        s=s+" Nspots     = %d " % self.N
        
829
        return s
830
831


alex's avatar
alex committed
832
833
834
835
836
837
class NoSignal(Exception):
    def __init__(self, value):
        self.value = value
    def __str__(self):
        return repr(self.value)
    
838
839
840
def get_Trajectory_byregression( O_spots    ):
    Xs=[]
    Ys=[]
841
    Ws=[]
842
843
844
845
846
847
848
    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
849
850
851
852
853
854
855
856
857
        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  ")
    
858
    ret_val = Trajectory()
859
860
861
862
863
864
865

    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 , )
866
    ret_val.N = len(O_spots)
867
868
869
870
    
    return ret_val


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

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

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

884
885
    if do_refine_trajectory ==2:

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

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

905
906
907
908
909
910
911
912
913
        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():
914
                    comm.Allreduce(  [np.array(integrated[n]), MPI.DOUBLE], [integrated[n], MPI.DOUBLE],  op=MPI.SUM)
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929

        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
930
931
932
933
934
935
            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]
936

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

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

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

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

            if trajectories_from_file is not None:
965
                " REUSING TRAJECTORIES "
966
                trajectory = trajectories_from_file[iterm]
967
                trajectory.N = O_spots.shape[0]
968
            else:
alex's avatar
alex committed
969
970
971
972
973
974
975
                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
976

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

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


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

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

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

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