LQC finds neighbours based on radius or number(useful when ddic field is passed)

tools/DIC/kinematics.py

 """
-Library of SPAM functions for defining a regular grid in a reproducible way.
+Library of SPAM functions for post processing a deformation field.
 Copyright (C) 2020 SPAM Contributors
 
 This program is free software: you can redistribute it and/or modify it
@@ -22,14 +22,14 @@ import scipy.spatial
 import multiprocessing
 import progressbar
 
-
 nProcessesDefault = multiprocessing.cpu_count()
 
-def estimateLocalQuadraticCoherency(points, displacements, neighbourRadius=150, epsilon = 0.1, nProcesses=nProcessesDefault, verbose=True):
+def estimateLocalQuadraticCoherency(points, displacements, neighbourRadius=150, nNeighbours=None, epsilon=0.1, nProcesses=nProcessesDefault, verbose=True):
     '''
     This function computes the local quadratic coherency (LQC) of a set of displacement vectors as per Masullo and Theunissen 2016.
     LQC is the average residual between the point's displacement and a second-order (parabolic) surface Phi.
     The quadratic surface Phi is fitted to the point's closest N neighbours and evaluated at the point's position.
+    Neighbours are selected based on: radius (default option) or number (activated if nNeighbours is not None).
     A point with a LQC value smaller than a threshold (0.1 in Masullo and Theunissen 2016) is classified as coherent
 
     Parameters
@@ -42,8 +42,14 @@ def estimateLocalQuadraticCoherency(points, displacements, neighbourRadius=150,
 
         neighbourRadius: float, optional
             Distance in pixels around the point to extract neighbours
+            If nNeighbours is given, this option is disactivated (neighbours will be selected based on number)
             Default = 150
 
+        nNeighbours : int, optional
+            Number of the nearest neighbours to consider
+            If not None, then neighbourRadius is disactivated
+            Default = None
+
         epsilon: float, optional
             Background error as per (Westerweel and Scarano 2005)
             Default = 0.1
@@ -72,18 +78,29 @@ def estimateLocalQuadraticCoherency(points, displacements, neighbourRadius=150,
     # build KD-tree for quick neighbour identification
     treeCoord = scipy.spatial.KDTree(points)
 
+    # check if neighbours are selected based on radius or based on number, default is radius
+    ball = True
+    if nNeighbours is not None: ball = False
+
     # calculate coherency for each point
     global coherencyOnePoint
     def coherencyOnePoint(point):
-        radius  = neighbourRadius
-        indices = numpy.array(treeCoord.query_ball_point(points[point], radius))
-        while len(indices) <= 27: #TODO this could vary (even a treeCoord.query could be used instead)
-            radius *= 2
+        # select neighbours based on radius
+        if ball:
+            radius  = neighbourRadius
             indices = numpy.array(treeCoord.query_ball_point(points[point], radius))
+            # make sure that at least 27 neighbours are selected
+            while len(indices) <= 27:
+                radius *= 2
+                indices = numpy.array(treeCoord.query_ball_point(points[point], radius))
+            N = len(indices)
+        # select neighbours based on number
+        else:
+            _, indices = treeCoord.query(points[point], k=nNeighbours)
+            N = nNeighbours
 
-        N = len(indices)
         # fill in point+neighbours positions for the parabolic surface coefficients
-        X     = numpy.zeros((N, 10), dtype=float)
+        X = numpy.zeros((N, 10), dtype=float)
         for i, neighbour in enumerate(indices):
             pos     = points[neighbour]
             X[i, 0] = 1
@@ -158,11 +175,12 @@ def estimateLocalQuadraticCoherency(points, displacements, neighbourRadius=150,
 
     return LQC
 
-def estimateDisplacementFromQuadraticFit(points, displacements, coherency, coherencyThreshold=0.1, neighbourRadius=150, epsilon=0.1, nProcesses=nProcessesDefault, verbose=True):
+def estimateDisplacementFromQuadraticFit(points, displacements, coherency, coherencyThreshold=0.1, neighbourRadius=150, nNeighbours=None, epsilon=0.1, nProcesses=nProcessesDefault, verbose=True):
     '''
     This function estimates the displacement of an incoherent point based on a local quadratic fit
     of the displacements of N coherent neighbours, as per Masullo and Theunissen 2016.
-    A quadratic surface Phi is fitted to the point's closest coherent neighbours
+    A quadratic surface Phi is fitted to the point's closest coherent neighbours.
+    Neighbours are selected based on: radius (default option) or number (activated if nNeighbours is not None)
 
     Parameters
     ----------
@@ -184,6 +202,11 @@ def estimateDisplacementFromQuadraticFit(points, displacements, coherency, coher
             Distance around the point to extract neighbours
             Default = 150
 
+        nNeighbours : int, optional
+            Number of the nearest neighbours to consider
+            If not None, then neighbourRadius is disactivated
+            Default = None
+
         epsilon: float, optional
             Background error as per (Westerweel and Scarano 2005)
             Default = 0.1
@@ -218,18 +241,29 @@ def estimateDisplacementFromQuadraticFit(points, displacements, coherency, coher
     # build KD-tree of coherent points for quick neighbour identification
     treeCoord = scipy.spatial.KDTree(points[goodPoints])
 
+    # check if neighbours are selected based on radius or based on number, default is radius
+    ball = True
+    if nNeighbours is not None: ball = False
+
     # estimate disp for each incoherent point
     global dispOnePoint
     def dispOnePoint(badPoint):
-        radius = neighbourRadius
-        indices = numpy.array(treeCoord.query_ball_point(points[badPoints][badPoint], radius))
-        while len(indices) <= 27: #TODO this could vary (even a treeCoord.query could be used instead)
-            radius *= 2
+        # select neighbours based on radius
+        if ball:
+            radius = neighbourRadius
             indices = numpy.array(treeCoord.query_ball_point(points[badPoints][badPoint], radius))
+            # make sure that at least 27 neighbours are selected
+            while len(indices) <= 27:
+                radius *= 2
+                indices = numpy.array(treeCoord.query_ball_point(points[badPoints][badPoint], radius))
+            N = len(indices)
+        # select neighbours based on number
+        else:
+            _, indices = treeCoord.query(points[badPoints][badPoint], k=nNeighbours)
+            N = nNeighbours
 
-        N = len(indices)
         # fill in neighbours positions for the parabolic surface coefficients
-        X     = numpy.zeros((N, 10), dtype=float)
+        X = numpy.zeros((N, 10), dtype=float)
         for i, neighbour in enumerate(indices):
             pos     = points[goodPoints][neighbour]
             X[i, 0] = 1