grid.py 31.9 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
"""
Library of SPAM functions for defining a regular grid in a reproducible way.
Copyright (C) 2020 SPAM Contributors

This program is free software: you can redistribute it and/or modify it
under the terms of the GNU General Public License as published by the Free
Software Foundation, either version 3 of the License, or (at your option)
any later version.

This program is distributed in the hope that it will be useful, but WITHOUT
ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for
more details.

You should have received a copy of the GNU General Public License along with
this program.  If not, see <http://www.gnu.org/licenses/>.
"""

from __future__ import print_function
20

21
import numpy
Edward Andò's avatar
Edward Andò committed
22
import spam.DIC.correlate
Edward Andò's avatar
Edward Andò committed
23
import sys
Emmanuel Roubin's avatar
Emmanuel Roubin committed
24
import progressbar
Edward Andò's avatar
Edward Andò committed
25
26
27
28
if sys.version[0] == '2':
    import Queue as queue
else:
    import queue as queue
29

30

31
def makeGrid(imageSize, nodeSpacing):
32
33
34
35
36
37
38
39
40
41
42
43
44
    """
        Define a grid of correlation points.

        Parameters
        ----------
        imageSize : 3-item list
            Size of volume to spread the grid inside

        nodeSpacing : 3-item list or int
            Spacing between nodes

        Returns
        -------
45
        nodePositions : nPointsx3 numpy.array
46
47
            Array containing Z, Y, X positions of each point in the grid
    """
48
49
50

    if len(imageSize) != 3:
        print("\tgrid.makeGrid(): imageSize doesn't have three dimensions, exiting")
51
52
        return

53
    if type(nodeSpacing) == int or type(nodeSpacing) == float:
Emmanuel Roubin's avatar
Emmanuel Roubin committed
54
55
        nodeSpacing = [nodeSpacing] * 3
    elif len(nodeSpacing) != 3:
56
        print("\tgrid.makeGrid(): nodeSpacing is not an int or float and doesn't have three dimensions, exiting")
57
58
        return

Emmanuel Roubin's avatar
Emmanuel Roubin committed
59
60
61
62
    if imageSize[0] == 1:
        twoD = True
    else:
        twoD = False
63

64
65
    # Note: in this cheap node spacing, the first node is always at a distance of --nodeSpacing-- from the origin
    # The following could just be done once in principle...
66
67
68
    nodesMgrid = numpy.mgrid[nodeSpacing[0]:imageSize[0]:nodeSpacing[0],
                             nodeSpacing[1]:imageSize[1]:nodeSpacing[1],
                             nodeSpacing[2]:imageSize[2]:nodeSpacing[2]]
69
70
71

    # If twoD then overwrite nodesMgrid
    if twoD:
Emmanuel Roubin's avatar
Emmanuel Roubin committed
72
        nodesMgrid = numpy.mgrid[0: 1: 1,
73
74
                                 nodeSpacing[1]:imageSize[1]:nodeSpacing[1],
                                 nodeSpacing[2]:imageSize[2]:nodeSpacing[2]]
75

Emmanuel Roubin's avatar
Emmanuel Roubin committed
76
    nodesDim = (nodesMgrid.shape[1], nodesMgrid.shape[2], nodesMgrid.shape[3])
77

78
    numberOfNodes = int(nodesMgrid.shape[1] * nodesMgrid.shape[2] * nodesMgrid.shape[3])
79

80
    nodePositions = numpy.zeros((numberOfNodes, 3))
81

Emmanuel Roubin's avatar
Emmanuel Roubin committed
82
83
84
    nodePositions[:, 0] = nodesMgrid[0].ravel()
    nodePositions[:, 1] = nodesMgrid[1].ravel()
    nodePositions[:, 2] = nodesMgrid[2].ravel()
85

Edward Andò's avatar
Edward Andò committed
86
    return nodePositions, nodesDim
87
88


89
def pixelSearchOnGrid(im1, im2, nodePositions, halfWindowSize, searchRange, PhiField=None, minMaskCoverage=0.5, im1mask=None, greyThreshold=[-numpy.inf, numpy.inf], mpi=False):
90
    """
91
92
    This function handles grid-based local correlation, offering an initial rough dispalcement-only guess.
    At the moment matching of windows is done with a Normalised-Correlation-Coefficient approach.
93

94
95
    Parameters
    ----------
96
97
98
99
100
101
102
103
        im1 : 3D numpy array
            A 3D image of greylevels defining a reference configuration for the pixel search

        im2 : 3D numpy array
            A deformed 3D image of greylevels

        nodePositions : nPoints*3 numpy array
            Array containing Z, Y, X positions of each point in the grid, as returned by ``makeGrid`` for example
104
            defined at im1
105
106
107
108
109
110
111
112
113
114

        halfWindowSize : 3-item list or int
            Size of subvolumes to perform the image correlation on, as a data range taken either side of the voxel on which the node is placed.
            The subvolume will be 2*halfWindowSize + 1 pixels on each side.
            A general recommendation is to make this half the node spacing

        searchRange : dictionary
            Search range as a dictionary containing 3 keys: 'zRange', 'yRange', and 'xRange',
            Each of which contains a list with two items

Edward Andò's avatar
Edward Andò committed
115
        PhiField : nPoints*4*4 numpy array, optional
116
            Optional field of ``F`` transformation operators defined for each node.
117
118
            Currently, only the translational components of F will be taken into account.
            Default = No displacement
119

120
        minMaskCoverage : float, optional
121
            Minimum number of pixels in a subvolume for it to be correlated (only considered in the case of im1mask).
122
            Default = 0.5
123
124

        im1mask : 3D boolean numpy array, optional
125
126
            A mask for im1 which is true in the zones to correlate.
            Default = None
127

128
        greyThreshold : list of two floats, optional
129
            Threshold for the mean greylevel in each im1 subvolume.
130
131
            If the mean is below the first value or above the second value, the grid point is not correlated.
            Default = [ -inf, inf ]
132

133
134
135
        mpi : bool, optional (default = False)
            Are we being called by an MPI run?

136
137
    Returns
    -------
138
139
        Dictionary containing:

140
        Keys
Edward Andò's avatar
Edward Andò committed
141
            PhiField : nNodes*4*4 numpy array of floats
142
143
144
145
146
                For each node, the measured transformatio operator (displacement only)

            pixelSearchCC : nNodes numpy array of floats
                For each node, the NCC score obtained
    """
147

148
    def getImagettes(nodeNumber, nodePositions, PhiField, searchRange, halfWindowSize, im1, im2, minMaskVolume, greyThreshold):
149
150
        returnStatus = 2

Emmanuel Roubin's avatar
Emmanuel Roubin committed
151
        initialDisplacement = PhiField[nodeNumber, 0:3, 3].astype(int)
152
153

        # Add initial displacement guess to search range
Emmanuel Roubin's avatar
Emmanuel Roubin committed
154
155
156
        searchRangeForThisNode = {'zRange': [searchRange['zRange'][0] + initialDisplacement[0], searchRange['zRange'][1] + initialDisplacement[0]],
                                  'yRange': [searchRange['yRange'][0] + initialDisplacement[1], searchRange['yRange'][1] + initialDisplacement[1]],
                                  'xRange': [searchRange['xRange'][0] + initialDisplacement[2], searchRange['xRange'][1] + initialDisplacement[2]]}
157
158

        # point in im2 that we are searching around
Emmanuel Roubin's avatar
Emmanuel Roubin committed
159
160
161
        searchCentre = [halfWindowSize[0] - searchRangeForThisNode['zRange'][0],
                        halfWindowSize[1] - searchRangeForThisNode['yRange'][0],
                        halfWindowSize[2] - searchRangeForThisNode['xRange'][0]]
162

163
164
165
166
167
168
169
170
171
172
173
        # 2020-09-25 OS and EA: Prepare startStop array for imagette 1 to be extracted with new slicePadded
        ## Prepare slice for im1 -- this is always in the same place
        #subVolSlice1 = (slice(int(nodePositions[nodeNumber, 0] - halfWindowSize[0]), int(nodePositions[nodeNumber, 0] + halfWindowSize[0] + 1)),
                        #slice(int(nodePositions[nodeNumber, 1] - halfWindowSize[1]), int(nodePositions[nodeNumber, 1] + halfWindowSize[1] + 1)),
                        #slice(int(nodePositions[nodeNumber, 2] - halfWindowSize[2]), int(nodePositions[nodeNumber, 2] + halfWindowSize[2] + 1)))
        ## Extract it...
        #imagette1 = im1[subVolSlice1]
        startStopIm1 = [int(nodePositions[nodeNumber, 0] - halfWindowSize[0]), int(nodePositions[nodeNumber, 0] + halfWindowSize[0] + 1),
                        int(nodePositions[nodeNumber, 1] - halfWindowSize[1]), int(nodePositions[nodeNumber, 1] + halfWindowSize[1] + 1),
                        int(nodePositions[nodeNumber, 2] - halfWindowSize[2]), int(nodePositions[nodeNumber, 2] + halfWindowSize[2] + 1)]
        imagette1 = spam.helpers.slicePadded(im1, startStopIm1)
174
175
176

        # Make sure imagette is not 0-dimensional in any dimension
        if numpy.all(numpy.array(imagette1.shape) > 0):
177
            if numpy.nanmean(imagette1) > greyThreshold[0] and numpy.nanmean(imagette1) < greyThreshold[1] and len(imagette1.ravel()) > minMaskVolume:
Emmanuel Roubin's avatar
Emmanuel Roubin committed
178
                # Slice for image 2
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
                ## 2020-09-25 OS and EA: Prepare startStop array for imagette 2 to be extracted with new slicePadded
                #subVolSlice2 = (slice(int(nodePositions[nodeNumber, 0] - halfWindowSize[0] + searchRangeForThisNode['zRange'][0]),
                                      #int(nodePositions[nodeNumber, 0] + halfWindowSize[0] + searchRangeForThisNode['zRange'][1] + 1)),
                                #slice(int(nodePositions[nodeNumber, 1] - halfWindowSize[1] + searchRangeForThisNode['yRange'][0]),
                                      #int(nodePositions[nodeNumber, 1] + halfWindowSize[1] + searchRangeForThisNode['yRange'][1] + 1)),
                                #slice(int(nodePositions[nodeNumber, 2] - halfWindowSize[2] + searchRangeForThisNode['xRange'][0]),
                                      #int(nodePositions[nodeNumber, 2] + halfWindowSize[2] + searchRangeForThisNode['xRange'][1] + 1)))
                ## Extract it...
                #imagette2 = im2[subVolSlice2]
                startStopIm2 = [int(nodePositions[nodeNumber, 0] - halfWindowSize[0] + searchRangeForThisNode['zRange'][0]),
                                int(nodePositions[nodeNumber, 0] + halfWindowSize[0] + searchRangeForThisNode['zRange'][1] + 1),
                                int(nodePositions[nodeNumber, 1] - halfWindowSize[1] + searchRangeForThisNode['yRange'][0]),
                                int(nodePositions[nodeNumber, 1] + halfWindowSize[1] + searchRangeForThisNode['yRange'][1] + 1),
                                int(nodePositions[nodeNumber, 2] - halfWindowSize[2] + searchRangeForThisNode['xRange'][0]),
                                int(nodePositions[nodeNumber, 2] + halfWindowSize[2] + searchRangeForThisNode['xRange'][1] + 1)]
                imagette2 = spam.helpers.slicePadded(im2, startStopIm2)
195

196
            # Failed minMaskVolume or greylevel condition
197
198
199
200
201
202
203
204
205
206
207
            else:
                returnStatus = -5
                imagette1 = None
                imagette2 = None

        # Failed 0-dimensional imagette test
        else:
            returnStatus = -5
            imagette1 = None
            imagette2 = None

208
209
        return {'imagette1': imagette1,
                'imagette2': imagette2,
210
211
212
213
214
                'returnStatus': returnStatus,
                'initialDisplacement': initialDisplacement,
                'searchRangeForThisNode': searchRangeForThisNode,
                'searchCentre': searchCentre}

215
216
217
218
219
220
221
222
    if mpi:
        import mpi4py.MPI

        mpiComm = mpi4py.MPI.COMM_WORLD
        mpiSize = mpiComm.Get_size()
        mpiRank = mpiComm.Get_rank()
        mpiStatus = mpi4py.MPI.Status()

Emmanuel Roubin's avatar
Emmanuel Roubin committed
223
        boss = mpiSize - 1
224
225

        numberOfWorkers = mpiSize - 1
Emmanuel Roubin's avatar
Emmanuel Roubin committed
226
        workersActive = numpy.zeros(numberOfWorkers)
227
228
229
230
231
232
    else:
        # not mpi
        numberOfWorkers = 1
        workersActive = numpy.array([0])

    # start setting up
233
234
235
    numberOfNodes = nodePositions.shape[0]

    # Check input sanity
236
    if type(halfWindowSize) == int or type(halfWindowSize) == float:
Emmanuel Roubin's avatar
Emmanuel Roubin committed
237
        halfWindowSize = [halfWindowSize] * 3
238

239
240
241
242
        # Check minMaskVolume
    minMaskVolume = int(minMaskCoverage * (1+halfWindowSize[0]*2)*
                                          (1+halfWindowSize[1]*2)*
                                          (1+halfWindowSize[2]*2))
243
244

    # Check F field
Edward Andò's avatar
Edward Andò committed
245
246
    if PhiField is None:
        PhiField = numpy.zeros((numberOfNodes, 4, 4))
Emmanuel Roubin's avatar
Emmanuel Roubin committed
247
248
        for nodeNumber in range(numberOfNodes):
            PhiField[nodeNumber] = numpy.eye(4)
249
250

    # Create pixelSearchCC vector
251
    pixelSearchCC = numpy.zeros((numberOfNodes))
252

253
    # Add nodes to a queue -- mostly useful for MPI
Edward Andò's avatar
Edward Andò committed
254
    q = queue.Queue()
Emmanuel Roubin's avatar
Emmanuel Roubin committed
255
256
257
    for node in range(numberOfNodes):
        q.put(node)
    finishedNodes = 0
258

259
260
261
    writeReturns = False

    print("\n\tStarting Pixel search")
262
263
264
    widgets = [progressbar.FormatLabel(''), ' ', progressbar.Bar(), ' ', progressbar.AdaptiveETA()]
    pbar = progressbar.ProgressBar(widgets=widgets, maxval=numberOfNodes)
    pbar.start()
265
266
267
268
    while finishedNodes != numberOfNodes:
        # If there are workers not working, satify their requests...
        #   Note: this condition is alyas true if we are not in MPI and there are jobs to do
        if workersActive.sum() < numberOfWorkers and not q.empty():
Emmanuel Roubin's avatar
Emmanuel Roubin committed
269
            worker = numpy.where(workersActive == False)[0][0]
270
            # Get the next node off the queue
Emmanuel Roubin's avatar
Emmanuel Roubin committed
271
            nodeNumber = q.get()
272

273
            imagetteReturns = getImagettes(nodeNumber, nodePositions, PhiField, searchRange, halfWindowSize, im1, im2, minMaskVolume, greyThreshold)
274

275
276
277
            if imagetteReturns['returnStatus'] == 2:
                if mpi:
                    # build message for pixel search worker
Emmanuel Roubin's avatar
Emmanuel Roubin committed
278
279
280
                    m = {'nodeNumber': nodeNumber,
                         'im1': imagetteReturns['imagette1'],
                         'im2': imagetteReturns['imagette2'],
281
                         'searchRangeForThisNode': imagetteReturns['searchRangeForThisNode'],
Emmanuel Roubin's avatar
Emmanuel Roubin committed
282
283
284
                         'searchCentre': imagetteReturns['searchCentre'],
                         'initialDisplacement': imagetteReturns['initialDisplacement']
                         }
285

Emmanuel Roubin's avatar
Emmanuel Roubin committed
286
                    # print "\tBoss: sending node {} to worker {}".format( nodeNumber, worker )
287
                    mpiComm.send(m, dest=worker, tag=1)
288

289
290
                    # Mark this worker as working
                    workersActive[worker] = True
291

292
293
                    # NOTE: writeReturns is defined later when receiving messages

Emmanuel Roubin's avatar
Emmanuel Roubin committed
294
                else:  # Not MPI
295
296
                    returns = spam.DIC.correlate.pixelSearch(imagetteReturns['imagette1'],
                                                             imagetteReturns['imagette2'],
Emmanuel Roubin's avatar
Emmanuel Roubin committed
297
298
                                                             searchRange=imagetteReturns['searchRangeForThisNode'],
                                                             searchCentre=imagetteReturns['searchCentre'])
299
300
301
                    initialDisplacement = imagetteReturns['initialDisplacement']
                    writeReturns = True

Emmanuel Roubin's avatar
Emmanuel Roubin committed
302
            else:  # Regardless of MPI or single proc
303
                # Failed to extract imagettes or something
Emmanuel Roubin's avatar
Emmanuel Roubin committed
304
305
306
307
                pixelSearchCC[nodeNumber] = 0.0
                finishedNodes += 1
                PhiField[nodeNumber, 0:3, 0:3] = numpy.eye(3)
                PhiField[nodeNumber, 0:3, 3] = numpy.nan
308

309
310
311
        # Otherwise spend time looking waiting for replies from workers
        elif mpi:
            message = mpiComm.recv(source=mpi4py.MPI.ANY_SOURCE, tag=2, status=mpiStatus)
Emmanuel Roubin's avatar
Emmanuel Roubin committed
312
            tag = mpiStatus.Get_tag()
313
            if tag == 2:
Emmanuel Roubin's avatar
Emmanuel Roubin committed
314
315
316
317
318
319
320
                worker = message[0]
                nodeNumber = message[1]
                returns = message[2]
                initialDisplacement = message[3]
                # print "\tBoss: received node {} from worker {}".format( nodeNumber, worker )
                workersActive[worker] = False
                writeReturns = True
321
            else:
Emmanuel Roubin's avatar
Emmanuel Roubin committed
322
                print("\tBoss: Don't recognise tag ", tag)
323
324
325

        # If we have new DVC returns, save them in our output matrices
        if writeReturns:
Emmanuel Roubin's avatar
Emmanuel Roubin committed
326
327
328
329
            finishedNodes += 1
            writeReturns = False
            # set translation for this node
            PhiField[nodeNumber, 0:3, 3] = numpy.array(returns['transformation']['t'])
330

Emmanuel Roubin's avatar
Emmanuel Roubin committed
331
332
333
            pixelSearchCC[nodeNumber] = returns['cc']
            widgets[0] = progressbar.FormatLabel("  CC={:0>7.5f} ".format(pixelSearchCC[nodeNumber]))
            pbar.update(finishedNodes)
334
335

    pbar.finish()
336
    print("\n")
337

Edward Andò's avatar
Edward Andò committed
338
    return {'PhiField': PhiField,
339
            'pixelSearchCC': pixelSearchCC}
340
341


342
def registerOnGrid(im1, im2, nodePositions, halfWindowSize, PhiField=None, margin=None, maxIterations=None, deltaPhiMin=None, interpolationOrder=None, minMaskCoverage=None, im1mask=None, greyThreshold=[-numpy.inf, numpy.inf], mpi=False, killWorkersWhenDone=True):
343
    """
Edward Andò's avatar
Edward Andò committed
344
    This function handles grid-based local correlation, performing a "register" subpixel refinement.
345
    Here we minimise a residual which is the difference between im1 and im2.
346

347
348
    Parameters
    ----------
349
350
351
352
353
354
355
356
357
358
359
360
361
362
        im1 : 3D numpy array
            A 3D image of greylevels defining a reference configuration for the pixel search

        im2 : 3D numpy array
            A deformed 3D image of greylevels

        nodePositions : nPoints*3 numpy array
            Array containing Z, Y, X positions of each point in the grid, as returned by ``makeGrid`` for example

        halfWindowSize : 3-item list or int
            Size of subvolumes to perform the image correlation on, as a data range taken either side of the voxel on which the node is placed.
            The subvolume will be 2*halfWindowSize + 1 pixels on each side.
            A general recommendation is to make this half the node spacing

363
364
365
        PhiField : nPoints*4*4 numpy array, optional
            Optional field of Phi deformation functions defined for each node
            Default = None
366

367
        margin : int or list, optional
368
            Margin to extract for subpixel interpolation
369
            Default = None (use ``register``'s default)
370

Edward Andò's avatar
Edward Andò committed
371
        maxIterations : int, optional (default = None (use ``register``'s default))
372
            Number of iterations for subpixel refinement
373

374
        deltaPhiMin : float, optional
375
            Stop iterating when norm of F gets below this value
376
            Default = None (use ``register``'s default)
377

378
        interpolationOrder : int, optional
379
            Greyscale interpolation order
380
            Default = None (use ``register``'s default)
381

382
        minMaskCoverage : float, optional
383
            Minimum number of pixels in a subvolume for it to be correlated (only considered in the case of im1mask).
384
            Default = 0.5
385

386
        im1mask : 3D boolean numpy array, optional
387
388
            A mask for the whole of im1 which is true in the zones to correlate.
            This is NOT used to mask the image, but to detect subvolumes that
389
390
            Fall inside the mask (see minMaskVolume for proportion)
            Default = None
391

392
        greyThreshold : list of two floats, optional
393
            Threshold for the mean greylevel in each im1 subvolume.
394
395
396
            If the mean is below the first value or above the second value, the grid point is not correlated.
            Default = [ -inf, inf ]

397
398
        mpi : bool, optional (default = False)
            Are we being called by an MPI run?
399
400
401
402
403
404
405
406
407

    Returns
    --------
        Dictionary containing:
            PhiField
            error
            iterations
            returnStatus
            deltaPhiNorm
408
    """
409

410
    # Check input sanity
411
    if type(margin) == int or type(margin) == float:
Emmanuel Roubin's avatar
Emmanuel Roubin committed
412
        margin = [margin] * 3
413

414
    # Detect unpadded 2D image first:
415
    if len(im1.shape) == 2:
416
417
418
419
420
        # pad them
        im1 = im1[numpy.newaxis, ...]
        im2 = im2[numpy.newaxis, ...]

    # Detect 2D images
Emmanuel Roubin's avatar
Emmanuel Roubin committed
421
422
423
424
    if im1.shape[0] == 1:
        twoD = True
    else:
        twoD = False
425

Emmanuel Roubin's avatar
Emmanuel Roubin committed
426
427
428
429
    if interpolationOrder == 1:
        interpolator = 'C'
    else:
        interpolator = 'python'
430
    # Override interpolator for python in 2D
Emmanuel Roubin's avatar
Emmanuel Roubin committed
431
432
    if twoD:
        interpolator = 'python'
433

434
    numberOfNodes = nodePositions.shape[0]
435

436
    def getImagettes(nodeNumber, im1, im2, PhiField, nodePositions, margin, halfWindowSize, greyThreshold, im1mask, minMaskVolume):
437
        # Initialise defaults
438
        imagette1mask = None
Edward Andò's avatar
Edward Andò committed
439
        PhiInit = None
440
        nodeDisplacement = None
441

Edward Andò's avatar
Edward Andò committed
442
        # Make sure all components of F are real numbers
Edward Andò's avatar
Edward Andò committed
443
        if numpy.isfinite(PhiField[nodeNumber]).sum() == 16:
Emmanuel Roubin's avatar
Emmanuel Roubin committed
444
            nodeDisplacement = numpy.round(PhiField[nodeNumber][0:3, -1])
Edward Andò's avatar
Edward Andò committed
445

446
            # prepare slice for imagette 1 -- this is fixed +1 for the margin which is always 1
447
448
449
450
451
452
453
454
455
456
            # 2020-09-25 OS and EA: Prepare startStop array for imagette 1 to be extracted with new slicePadded
            #subVolSlice1 = (slice(int(nodePositions[nodeNumber, 0] - halfWindowSize[0] - 1), int(nodePositions[nodeNumber, 0] + halfWindowSize[0] + 1 + 1)),
                            #slice(int(nodePositions[nodeNumber, 1] - halfWindowSize[1] - 1), int(nodePositions[nodeNumber, 1] + halfWindowSize[1] + 1 + 1)),
                            #slice(int(nodePositions[nodeNumber, 2] - halfWindowSize[2] - 1), int(nodePositions[nodeNumber, 2] + halfWindowSize[2] + 1 + 1)))
            #imagette1 = im1[subVolSlice1].copy()
            startStopIm1 = [int(nodePositions[nodeNumber, 0] - halfWindowSize[0] - 1), int(nodePositions[nodeNumber, 0] + halfWindowSize[0] + 1 + 1),
                            int(nodePositions[nodeNumber, 1] - halfWindowSize[1] - 1), int(nodePositions[nodeNumber, 1] + halfWindowSize[1] + 1 + 1),
                            int(nodePositions[nodeNumber, 2] - halfWindowSize[2] - 1), int(nodePositions[nodeNumber, 2] + halfWindowSize[2] + 1 + 1)]
            imagette1, imagette1mask = spam.helpers.slicePadded(im1, startStopIm1, createMask=True)

Edward Andò's avatar
Edward Andò committed
457

458
            # If there is a mask defined for im1, compute the coverage of this correlation window
Edward Andò's avatar
Edward Andò committed
459
460
            # Check Mask volume condition
            if im1mask is not None:
461
462
                # AND masks together, the one coming from the im1mask -- i.e., allowed correlation zones
                # and the imagette1sliceMask indicating whether we've gone outside the volume
463
                imagette1mask = numpy.logical_and(imagette1mask, spam.helpers.slicePadded(im1mask, startStopIm1))
464
            maskVolCondition = imagette1mask.sum() > minMaskVolume
465

466
467
468
469
470
            # Make sure imagette is not 0-dimensional in any dimension
            if numpy.all(numpy.array(imagette1.shape) > 0):
                # Check if out extracted imagette 1 is above grey threshold, if there is enough under the mask if defined, and that dispalcements are not NaN
                if numpy.nanmean(imagette1) > greyThreshold[0] and numpy.nanmean(imagette1) < greyThreshold[1] and maskVolCondition:

471
472
473
                    # Prepare im2 imagette
                    # start from nodePosition in im1 and move it according to the nodeDisplacement
                    # and move it + 1 for the margin which is always 1
474
475
476
477
478
479
480
481
482
483
                    # 2020-09-25 OS and EA: Prepare startStop array for imagette 2 to be extracted with new slicePadded
                    #subVolSlice2 = (slice(int(nodePositions[nodeNumber, 0] - halfWindowSize[0] + nodeDisplacement[0] - margin[0] - 1), int(nodePositions[nodeNumber, 0] + halfWindowSize[0] + nodeDisplacement[0] + margin[0] + 1 + 1)),
                                    #slice(int(nodePositions[nodeNumber, 1] - halfWindowSize[1] + nodeDisplacement[1] - margin[1] - 1), int(nodePositions[nodeNumber, 1] + halfWindowSize[1] + nodeDisplacement[1] + margin[1] + 1 + 1)),
                                    #slice(int(nodePositions[nodeNumber, 2] - halfWindowSize[2] + nodeDisplacement[2] - margin[2] - 1), int(nodePositions[nodeNumber, 2] + halfWindowSize[2] + nodeDisplacement[2] + margin[2] + 1 + 1)))
                    ## Extract it
                    #imagette2 = im2[subVolSlice2].copy()
                    startStopIm2 = [int(nodePositions[nodeNumber, 0] - halfWindowSize[0] + nodeDisplacement[0] - margin[0] - 1), int(nodePositions[nodeNumber, 0] + halfWindowSize[0] + nodeDisplacement[0] + margin[0] + 1 + 1),
                                    int(nodePositions[nodeNumber, 1] - halfWindowSize[1] + nodeDisplacement[1] - margin[1] - 1), int(nodePositions[nodeNumber, 1] + halfWindowSize[1] + nodeDisplacement[1] + margin[1] + 1 + 1),
                                    int(nodePositions[nodeNumber, 2] - halfWindowSize[2] + nodeDisplacement[2] - margin[2] - 1), int(nodePositions[nodeNumber, 2] + halfWindowSize[2] + nodeDisplacement[2] + margin[2] + 1 + 1)]
                    imagette2 = spam.helpers.slicePadded(im2, startStopIm2)
484
485

                    # Extract initial F for correlation, remove int() part of displacement since it's already used to extract imagette2
Edward Andò's avatar
Edward Andò committed
486
                    PhiInit = PhiField[nodeNumber].copy()
Emmanuel Roubin's avatar
Emmanuel Roubin committed
487
488
                    PhiInit[0:3, -1] -= nodeDisplacement.copy()
                    # print "PhiInit:\n", PhiInit, '\n'
Edward Andò's avatar
Edward Andò committed
489
                    #PhiInit = numpy.eye(4)
490

Emmanuel Roubin's avatar
Emmanuel Roubin committed
491
                    if (numpy.array(imagette2.shape) - numpy.array(imagette1.shape) == numpy.array(margin) * 2).all():
492
493
494
495
496
497
498
499
500
                        returnStatus = 2

                    # Failed innermost condition -- im1 and im2 margin alignment -- this is a harsh condition
                    else:
                        returnStatus = -4
                        imagette1 = None
                        imagette2 = None

                # Failed mask or greylevel condition
501
                else:
502
                    returnStatus = -5
503
504
                    imagette1 = None
                    imagette2 = None
Edward Andò's avatar
Edward Andò committed
505

506
            # Failed 0-dimensional imagette test
507
508
509
510
            else:
                returnStatus = -5
                imagette1 = None
                imagette2 = None
511

512
513
514
515
516
        # Failed non-NaN components in F
        else:
            returnStatus = -6
            imagette1 = None
            imagette2 = None
517

518
        return {'imagette1': imagette1, 'imagette2': imagette2,
519
                'imagette1mask': imagette1mask,
520
                'returnStatus': returnStatus,
Edward Andò's avatar
Edward Andò committed
521
                'PhiInit': PhiInit,
522
                'nodeDisplacement': nodeDisplacement}
523

524
525
526
527
528
529
530
531
    if mpi:
        import mpi4py.MPI

        mpiComm = mpi4py.MPI.COMM_WORLD
        mpiSize = mpiComm.Get_size()
        mpiRank = mpiComm.Get_rank()
        mpiStatus = mpi4py.MPI.Status()

Emmanuel Roubin's avatar
Emmanuel Roubin committed
532
        boss = mpiSize - 1
533

534
        numberOfWorkers = mpiSize - 1
Emmanuel Roubin's avatar
Emmanuel Roubin committed
535
        workersActive = numpy.zeros(numberOfWorkers)
536
537
    else:
        numberOfWorkers = 1
538
        workersActive = numpy.array([0])
539
540

    # Check input sanity
541
    if type(halfWindowSize) == int or type(halfWindowSize) == float:
Emmanuel Roubin's avatar
Emmanuel Roubin committed
542
        halfWindowSize = [halfWindowSize] * 3
543

544
545
546
547
    # Check minMaskVolume
    minMaskVolume = int(minMaskCoverage * (1+halfWindowSize[0]*2)*
                                          (1+halfWindowSize[1]*2)*
                                          (1+halfWindowSize[2]*2))
548
549

    # Check F field
Edward Andò's avatar
Edward Andò committed
550
551
    if PhiField is None:
        PhiField = numpy.zeros((numberOfNodes, 4, 4))
Emmanuel Roubin's avatar
Emmanuel Roubin committed
552
553
        for nodeNumber in range(numberOfNodes):
            PhiField[nodeNumber] = numpy.eye(4)
554

555
    # Add nodes to a queue -- mostly useful for MPI
Edward Andò's avatar
Edward Andò committed
556
    q = queue.Queue()
Emmanuel Roubin's avatar
Emmanuel Roubin committed
557
558
559
    for node in range(numberOfNodes):
        q.put(node)
    finishedNodes = 0
560

Emmanuel Roubin's avatar
Emmanuel Roubin committed
561
562
    subpixelError = numpy.zeros((numberOfNodes))
    subpixelIterations = numpy.zeros((numberOfNodes))
563
    subpixelReturnStatus = numpy.zeros((numberOfNodes))
Emmanuel Roubin's avatar
Emmanuel Roubin committed
564
    subpixelDeltaFnorm = numpy.zeros((numberOfNodes))
565
566
567

    writeReturns = False

Edward Andò's avatar
Edward Andò committed
568
    print("\n\tStarting Correlation")
569
570
571
    widgets = [progressbar.FormatLabel(''), ' ', progressbar.Bar(), ' ', progressbar.AdaptiveETA()]
    pbar = progressbar.ProgressBar(widgets=widgets, maxval=numberOfNodes)
    pbar.start()
572
573
574
575
    while finishedNodes != numberOfNodes:
        # If there are workers not working, satify their requests...
        #   Note: this condition is alyas true if we are not in MPI and there are jobs to do
        if workersActive.sum() < numberOfWorkers and not q.empty():
Emmanuel Roubin's avatar
Emmanuel Roubin committed
576
            worker = numpy.where(workersActive == False)[0][0]
577
            # Get the next node off the queue
Emmanuel Roubin's avatar
Emmanuel Roubin committed
578
            nodeNumber = q.get()
579

580
            imagetteReturns = getImagettes(nodeNumber, im1, im2, PhiField, nodePositions, margin, halfWindowSize, greyThreshold, im1mask, minMaskVolume)
581

582
            if imagetteReturns['returnStatus'] == 2:
583
584
                if mpi:
                    # build message for lukas kanade worker
Emmanuel Roubin's avatar
Emmanuel Roubin committed
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
                    m = {'nodeNumber': nodeNumber,
                         'im1': imagetteReturns['imagette1'],
                         'im2': imagetteReturns['imagette2'],
                         'im1mask': imagetteReturns['imagette1mask'],
                         'PhiInit': imagetteReturns['PhiInit'],
                         # 'PhiInitBinRatio' : PhiInitBinRatio,
                         'margin': 1,  # see top of this file for compensation
                         'maxIterations': maxIterations,
                         'deltaPhiMin': deltaPhiMin,
                         'interpolationOrder': interpolationOrder,
                         'interpolator': interpolator,
                         'nodeDisplacement': imagetteReturns['nodeDisplacement']
                         }

                    # print "\tBoss: sending node {} to worker {}".format( nodeNumber, worker )
Edward Andò's avatar
Edward Andò committed
600
                    mpiComm.send(m, dest=worker, tag=3)
601

602
                    # Mark this worker as working
603
                    workersActive[worker] = True
604

605
                    # NOTE: writeReturns is defined later when receiving messages
606

Emmanuel Roubin's avatar
Emmanuel Roubin committed
607
                else:  # Not MPI
Edward Andò's avatar
Edward Andò committed
608
                    returns = spam.DIC.correlate.register(imagetteReturns['imagette1'],
609
                                                             imagetteReturns['imagette2'],
Emmanuel Roubin's avatar
Emmanuel Roubin committed
610
611
612
613
614
615
616
                                                             im1mask=imagetteReturns['imagette1mask'],
                                                             PhiInit=imagetteReturns['PhiInit'],
                                                             margin=1,  # see top of this file for compensation
                                                             maxIterations=maxIterations,
                                                             deltaPhiMin=deltaPhiMin,
                                                             interpolationOrder=interpolationOrder,
                                                             verbose=False,
617
                                                             imShowProgress=False)
618
619
                    nodeDisplacement = imagetteReturns['nodeDisplacement']
                    writeReturns = True
620

Emmanuel Roubin's avatar
Emmanuel Roubin committed
621
            else:  # Regardless of MPI or single proc
622
                # Failed to extract imagettes or something
Emmanuel Roubin's avatar
Emmanuel Roubin committed
623
624
625
626
627
628
629
                subpixelError[nodeNumber] = numpy.inf
                subpixelIterations[nodeNumber] = 0
                subpixelReturnStatus[nodeNumber] = imagetteReturns['returnStatus']
                subpixelDeltaFnorm[nodeNumber] = numpy.inf
                PhiField[nodeNumber, 0:3, 0:3] = numpy.eye(3)
                PhiField[nodeNumber, 0:3, 3] = numpy.nan
                finishedNodes += 1
630
631
632

        # Otherwise spend time looking waiting for replies from workers
        elif mpi:
Edward Andò's avatar
Edward Andò committed
633
            message = mpiComm.recv(source=mpi4py.MPI.ANY_SOURCE, tag=4, status=mpiStatus)
Emmanuel Roubin's avatar
Emmanuel Roubin committed
634
            tag = mpiStatus.Get_tag()
Edward Andò's avatar
Edward Andò committed
635
            if tag == 4:
Emmanuel Roubin's avatar
Emmanuel Roubin committed
636
637
638
639
640
641
642
                worker = message[0]
                nodeNumber = message[1]
                returns = message[2]
                nodeDisplacement = message[3]
                # print "\tBoss: received node {} from worker {}".format( nodeNumber, worker )
                workersActive[worker] = False
                writeReturns = True
643
            else:
Emmanuel Roubin's avatar
Emmanuel Roubin committed
644
                print("\tBoss: Don't recognise tag ", tag)
645

646
647
        # If we have new DVC returns, save them in our output matrices
        if writeReturns:
Emmanuel Roubin's avatar
Emmanuel Roubin committed
648
649
650
            finishedNodes += 1
            writeReturns = False
            # Overwrite transformation operator for this node
651
            PhiField[nodeNumber] = returns['Phi']
Emmanuel Roubin's avatar
Emmanuel Roubin committed
652
653
654
655
            # Add back in the translation from the initial guess
            PhiField[nodeNumber, 0:3, 3] += nodeDisplacement

            subpixelError[nodeNumber] = returns['error']
656
            subpixelIterations[nodeNumber] = returns['iterations']
Emmanuel Roubin's avatar
Emmanuel Roubin committed
657
658
659
660
            subpixelReturnStatus[nodeNumber] = returns['returnStatus']
            subpixelDeltaFnorm[nodeNumber] = returns['deltaPhiNorm']

            #widgets[0] = progressbar.FormatLabel("err={:0>7.03f}\tit={:0>3d}\tdFnorm={:0>5.03f}\trs={:+1d}".format( returns['error'], returns['iterationNumber'], returns['deltaFnorm'], returns['returnStatus'] ))
661
            widgets[0] = progressbar.FormatLabel("  it={:0>3d}  dPhiNorm={:0>6.4f}  rs={:+1d} ".format(returns['iterations'], returns['deltaPhiNorm'], returns['returnStatus']))
Emmanuel Roubin's avatar
Emmanuel Roubin committed
662
663
664
            pbar.update(finishedNodes)
            #print("\r\t\tCorrelating node {:04d} of {:04d}".format( nodeNumber+1, numberOfNodes ), end='' )
            #print("\terror={:05.0f}\titerations={:02d}\tdeltaFnorm={:0.5f}\treturnStat={:+1d}".format( returns['error'], returns['iterationNumber'], returns['deltaFnorm'], returns['returnStatus'] ), end='')
665

666
667
    pbar.finish()
    print("\n")
668

Edward Andò's avatar
Edward Andò committed
669
    return {"PhiField": PhiField,
670
671
672
673
            "error": subpixelError,
            "iterations": subpixelIterations,
            "returnStatus": subpixelReturnStatus,
            "deltaPhiNorm": subpixelDeltaFnorm,
674
            }