Merge pull request #104 from jni/optimise-grain-finding-3

Optimise grain finding

.github/workflows/test.yml

@@ -9,7 +9,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.7", "3.8", "3.9", "3.10", "3.11"]
+        python-version: ["3.8", "3.9", "3.10", "3.11"]
 
     steps:
     - uses: actions/checkout@v2

defdap/_accelerated.py

@@ -0,0 +1,148 @@
+from numba import njit
+import numpy as np
+
+
+@njit
+def find_first(arr):
+    for i in range(len(arr)):
+        if arr[i]:
+            return i
+
+
+@njit
+def flood_fill(seed, index, points_remaining, grains, boundary_x, boundary_y,
+               added_coords):
+    """Flood fill algorithm that uses the x and y boundary arrays to
+    fill a connected area around the seed point. The points are inserted
+    into a grain object and the grain map array is updated.
+
+    Parameters
+    ----------
+    seed : tuple of 2 int
+        Seed point x for flood fill
+    index : int
+        Value to fill in grain map
+    points_remaining : numpy.ndarray
+        Boolean map of the points that have not been assigned a grain yet
+
+    Returns
+    -------
+    grain : defdap.ebsd.Grain
+        New grain object with points added
+    """
+    x, y = seed
+    grains[y, x] = index
+    points_remaining[y, x] = False
+    edge = [seed]
+    added_coords[0] = seed
+    npoints = 1
+
+    while edge:
+        x, y = edge.pop()
+        moves = [(x+1, y), (x-1, y), (x, y+1), (x, y-1)]
+        # get rid of any that go out of the map area
+        if x <= 0:
+            moves.pop(1)
+        elif x > grains.shape[1] - 2:
+            moves.pop(0)
+        if y <= 0:
+            moves.pop(-1)
+        elif y > grains.shape[0] - 2:
+            moves.pop(-2)
+
+        for (s, t) in moves:
+            if grains[t, s] > 0:
+                continue
+
+            add_point = False
+
+            if t == y:
+                # moving horizontally
+                if s > x:
+                    # moving right
+                    add_point = not boundary_x[y, x]
+                else:
+                    # moving left
+                    add_point = not boundary_x[t, s]
+            else:
+                # moving vertically
+                if t > y:
+                    # moving down
+                    add_point = not boundary_y[y, x]
+                else:
+                    # moving up
+                    add_point = not boundary_y[t, s]
+
+            if add_point:
+                added_coords[npoints] = s, t
+                grains[t, s] = index
+                points_remaining[t, s] = False
+                npoints += 1
+                edge.append((s, t))
+
+    return added_coords[:npoints]
+
+
+@njit
+def flood_fill_dic(seed, index, points_remaining, grains, added_coords):
+    """Flood fill algorithm that uses the combined x and y boundary array
+    to fill a connected area around the seed point. The points are returned and
+    the grain map array is updated.
+
+    Parameters
+    ----------
+    seed : tuple of 2 int
+        Seed point x for flood fill
+    index : int
+        Value to fill in grain map
+    points_remaining : numpy.ndarray
+        Boolean map of the points remaining to assign a grain yet
+    grains : numpy.ndarray
+    added_coords : numpy.ndarray
+        Buffer for points in the grain
+
+    Returns
+    -------
+    numpy.ndarray
+        Flooded points (n, 2)
+
+    """
+    # add first point to the grain
+    x, y = seed
+    grains[y, x] = index
+    points_remaining[y, x] = False
+    edge = [seed]
+    added_coords[0] = seed
+    npoints = 1
+
+    while edge:
+        x, y = edge.pop()
+
+        moves = [(x+1, y), (x-1, y), (x, y+1), (x, y-1)]
+        # get rid of any that go out of the map area
+        if x <= 0:
+            moves.pop(1)
+        elif x >= grains.shape[1] - 1:
+            moves.pop(0)
+        if y <= 0:
+            moves.pop(-1)
+        elif y >= grains.shape[0] - 1:
+            moves.pop(-2)
+
+        for (s, t) in moves:
+            add_point = False
+
+            if grains[t, s] == 0:
+                add_point = True
+                edge.append((s, t))
+
+            elif grains[t, s] == -1 and (s > x or t > y):
+                add_point = True
+
+            if add_point:
+                added_coords[npoints] = (s, t)
+                grains[t, s] = index
+                points_remaining[t, s] = False
+                npoints += 1
+
+    return added_coords[:npoints]

defdap/base.py

@@ -793,17 +793,6 @@ def __len__(self):
     def __str__(self):
         return f"Grain(ID={self.grain_id})"
 
-    def add_point(self, point):
-        """Append a coordinate and a quat to a grain.
-
-        Parameters
-        ----------
-        point : tuple
-            (x,y) coordinate to append
-
-        """
-        self.data.point.append(point)
-
     @property
     def extreme_coords(self):
         """Coordinates of the bounding box for a grain.
@@ -814,12 +803,7 @@ def extreme_coords(self):
             minimum x, minimum y, maximum x, maximum y.
 
         """
-        points = np.array(self.data.point, dtype=int)
-
-        x0, y0 = points.min(axis=0)
-        xmax, ymax = points.max(axis=0)
-
-        return x0, y0, xmax, ymax
+        return *self.data.point.min(axis=0), *self.data.point.max(axis=0)
 
     def centre_coords(self, centre_type="box", grain_coords=True):
         """
@@ -846,7 +830,7 @@ def centre_coords(self, centre_type="box", grain_coords=True):
             x_centre = round((xmax + x0) / 2)
             y_centre = round((ymax + y0) / 2)
         elif centre_type == "com":
-            x_centre, y_centre = np.array(self.data.point).mean(axis=0).round()
+            x_centre, y_centre = self.data.point.mean(axis=0).round()
         else:
             raise ValueError("centreType must be box or com")
 

defdap/ebsd.py

@@ -25,6 +25,7 @@
 from defdap.file_writers import EBSDDataWriter
 from defdap.quat import Quat
 from defdap import base
+from defdap._accelerated import flood_fill
 
 from defdap import defaults
 from defdap.plotting import MapPlot
@@ -881,13 +882,17 @@ def find_grains(self, min_grain_size=10):
         # Loop until all points (except boundaries) have been assigned
         # to a grain or ignored
         i = 0
+        coords_buffer = np.zeros((boundary_im_y.size, 2), dtype=np.intp)
         while found_point >= 0:
             # Flood fill first unknown point and return grain object
             seed = np.unravel_index(next_point, self.shape)
-            grain = self.flood_fill(
+
+            grain = Grain(grain_index - 1, self, group_id)
+            grain.data.point = flood_fill(
                 (seed[1], seed[0]), grain_index, points_left, grains,
-                boundary_im_x, boundary_im_y, group_id
+                boundary_im_x, boundary_im_y, coords_buffer
             )
+            coords_buffer = coords_buffer[len(grain.data.point):]
 
             if len(grain) < min_grain_size:
                 # if grain size less than minimum, ignore grain and set
@@ -962,81 +967,6 @@ def plot_grain_map(self, **kwargs):
 
         return plot
 
-    def flood_fill(self, seed, index, points_left, grains, boundary_im_x,
-                   boundary_im_y, group_id):
-        """Flood fill algorithm that uses the x and y boundary arrays to
-        fill a connected area around the seed point. The points are inserted
-        into a grain object and the grain map array is updated.
-
-        Parameters
-        ----------
-        seed : tuple of 2 int
-            Seed point x for flood fill
-        index : int
-            Value to fill in grain map
-        points_left : numpy.ndarray
-            Boolean map of the points that have not been assigned a grain yet
-
-        Returns
-        -------
-        grain : defdap.ebsd.Grain
-            New grain object with points added
-        """
-        # create new grain
-        grain = Grain(index - 1, self, group_id)
-
-        # add first point to the grain
-        x, y = seed
-        grain.add_point(seed)
-        grains[y, x] = index
-        points_left[y, x] = False
-        edge = [seed]
-
-        while edge:
-            x, y = edge.pop(0)
-
-            moves = [(x+1, y), (x-1, y), (x, y+1), (x, y-1)]
-            # get rid of any that go out of the map area
-            if x <= 0:
-                moves.pop(1)
-            elif x >= self.shape[1] - 1:
-                moves.pop(0)
-            if y <= 0:
-                moves.pop(-1)
-            elif y >= self.shape[0] - 1:
-                moves.pop(-2)
-
-            for (s, t) in moves:
-                if grains[t, s] > 0:
-                    continue
-
-                add_point = False
-
-                if t == y:
-                    # moving horizontally
-                    if s > x:
-                        # moving right
-                        add_point = not boundary_im_x[y, x]
-                    else:
-                        # moving left
-                        add_point = not boundary_im_x[t, s]
-                else:
-                    # moving vertically
-                    if t > y:
-                        # moving down
-                        add_point = not boundary_im_y[y, x]
-                    else:
-                        # moving up
-                        add_point = not boundary_im_y[t, s]
-
-                if add_point:
-                    grain.add_point((s, t))
-                    grains[t, s] = index
-                    points_left[t, s] = False
-                    edge.append((s, t))
-
-        return grain
-
     @report_progress("calculating grain mean orientations")
     def calc_grain_av_oris(self):
         """Calculate the average orientation of grains.