Separate out 3d and 4d combine functions

datashader/reductions.py

@@ -29,9 +29,9 @@
     cudf = cp = None
 
 from .utils import (
-    Expr, ngjit, nansum_missing, nanmax_in_place, nansum_in_place,
-    nanmax_n_in_place, nanmin_n_in_place, row_min_in_place,
-    row_max_n_in_place, row_min_n_in_place,
+    Expr, ngjit, nansum_missing, nanmax_in_place, nansum_in_place, row_min_in_place,
+    nanmax_n_in_place_4d, nanmax_n_in_place_3d, nanmin_n_in_place_4d, nanmin_n_in_place_3d,
+    row_max_n_in_place_4d, row_max_n_in_place_3d, row_min_n_in_place_4d, row_min_n_in_place_3d,
 )
 
 
@@ -1521,11 +1521,11 @@ def _build_combine(self, dshape, antialias, cuda, partitioned):
     @staticmethod
     def _combine(aggs):
         ret = aggs[0]
-        if ret.ndim == 3:  # ndim is either 3 (ny, nx, n) or 4 (ny, nx, ncat, n)
-            # 4d view of each agg
-            aggs = [np.expand_dims(agg, 2) for agg in aggs]
         for i in range(1, len(aggs)):
-            nanmax_n_in_place(aggs[0], aggs[i])
+            if ret.ndim == 3:  # ndim is either 3 (ny, nx, n) or 4 (ny, nx, ncat, n)
+                nanmax_n_in_place_3d(aggs[0], aggs[i])
+            else:
+                nanmax_n_in_place_4d(aggs[0], aggs[i])
         return ret
 
     @staticmethod
@@ -1593,11 +1593,11 @@ def _build_combine(self, dshape, antialias, cuda, partitioned):
     @staticmethod
     def _combine(aggs):
         ret = aggs[0]
-        if ret.ndim == 3:  # ndim is either 3 (ny, nx, n) or 4 (ny, nx, ncat, n)
-            # 4d view of each agg
-            aggs = [np.expand_dims(agg, 2) for agg in aggs]
         for i in range(1, len(aggs)):
-            nanmin_n_in_place(aggs[0], aggs[i])
+            if ret.ndim == 3:  # ndim is either 3 (ny, nx, n) or 4 (ny, nx, ncat, n)
+                nanmin_n_in_place_3d(aggs[0], aggs[i])
+            else:
+                nanmin_n_in_place_4d(aggs[0], aggs[i])
         return ret
 
     @staticmethod
@@ -2213,9 +2213,9 @@ def _combine(aggs):
         ret = aggs[0]
         if len(aggs) > 1:
             if ret.ndim == 3:  # ndim is either 3 (ny, nx, n) or 4 (ny, nx, ncat, n)
-                # 4d view of each agg
-                aggs = [np.expand_dims(agg, 2) for agg in aggs]
-            row_max_n_in_place(aggs[0], aggs[1])
+                row_max_n_in_place_3d(aggs[0], aggs[1])
+            else:
+                row_max_n_in_place_4d(aggs[0], aggs[1])
         return ret
 
     @staticmethod
@@ -2278,9 +2278,9 @@ def _combine(aggs):
         ret = aggs[0]
         if len(aggs) > 1:
             if ret.ndim == 3:  # ndim is either 3 (ny, nx, n) or 4 (ny, nx, ncat, n)
-                # 4d view of each agg
-                aggs = [np.expand_dims(agg, 2) for agg in aggs]
-            row_min_n_in_place(aggs[0], aggs[1])
+                row_min_n_in_place_3d(aggs[0], aggs[1])
+            else:
+                row_min_n_in_place_4d(aggs[0], aggs[1])
         return ret
 
     @staticmethod

datashader/utils.py

@@ -648,72 +648,108 @@ def nanmin_in_place(ret, other):
             ret[i] = other[i]
 
 
+@ngjit
+def _nanmax_n_impl(ret_pixel, other_pixel):
+    """Single pixel implementation of nanmax_n_in_place.
+    ret_pixel and other_pixel are both 1D arrays of the same length.
+
+    Walk along other_pixel a value at a time, find insertion index in
+    ret_pixel and shift values along to insert.  Next other_pixel value is
+    inserted at a higher index, so this walks the two pixel arrays just once
+    each.
+    """
+    n = len(ret_pixel)
+    istart = 0
+    for other_value in other_pixel:
+        if isnull(other_value):
+            break
+        else:
+            for i in range(istart, n):
+                if isnull(ret_pixel[i]) or other_value > ret_pixel[i]:
+                    # Shift values along then insert.
+                    for j in range(n-1, i, -1):
+                        ret_pixel[j] = ret_pixel[j-1]
+                    ret_pixel[i] = other_value
+                    istart = i+1
+                    break
+
+
 @ngjit_parallel
-def nanmax_n_in_place(ret, other):
+def nanmax_n_in_place_4d(ret, other):
     """Combine two max-n arrays, taking nans into account. Max-n arrays are 4D
     with shape (ny, nx, ncat, n) where ny and nx are the number of pixels,
     ncat the number of categories (will be 1 if not using a categorical
     reduction) and the last axis containing n values in descending order.
     If there are fewer than n values it is padded with nans.
     Return the first array.
     """
-    ny, nx, ncat, n = ret.shape
+    ny, nx, ncat, _n = ret.shape
     for y in nb.prange(ny):
         for x in range(nx):
             for cat in range(ncat):
-                ret_pixel = ret[y, x, cat]      # 1D array of n values for single pixel
-                other_pixel = other[y, x, cat]  # ditto
-                # Walk along other_pixel array a value at a time, find insertion
-                # index in ret_pixel and bump values along to insert.  Next
-                # other_pixel value is inserted at a higher index, so this walks
-                # the two pixel arrays just once each.
-                istart = 0
-                for other_value in other_pixel:
-                    if isnull(other_value):
-                        break
-                    else:
-                        for i in range(istart, n):
-                            if isnull(ret_pixel[i]) or other_value > ret_pixel[i]:
-                                # Bump values along then insert.
-                                for j in range(n-1, i, -1):
-                                    ret_pixel[j] = ret_pixel[j-1]
-                                ret_pixel[i] = other_value
-                                istart = i+1
-                                break
+                _nanmax_n_impl(ret[y, x, cat], other[y, x, cat])
 
 
 @ngjit_parallel
-def nanmin_n_in_place(ret, other):
+def nanmax_n_in_place_3d(ret, other):
+    """3d version of nanmax_n_in_place_4d, taking arrays of shape (ny, nx, n).
+    """
+    ny, nx, _n = ret.shape
+    for y in nb.prange(ny):
+        for x in range(nx):
+            _nanmax_n_impl(ret[y, x], other[y, x])
+
+
+@ngjit
+def _nanmin_n_impl(ret_pixel, other_pixel):
+    """Single pixel implementation of nanmin_n_in_place.
+    ret_pixel and other_pixel are both 1D arrays of the same length.
+
+    Walk along other_pixel a value at a time, find insertion index in
+    ret_pixel and shift values along to insert.  Next other_pixel value is
+    inserted at a higher index, so this walks the two pixel arrays just once
+    each.
+    """
+    n = len(ret_pixel)
+    istart = 0
+    for other_value in other_pixel:
+        if isnull(other_value):
+            break
+        else:
+            for i in range(istart, n):
+                if isnull(ret_pixel[i]) or other_value < ret_pixel[i]:
+                    # Shift values along then insert.
+                    for j in range(n-1, i, -1):
+                        ret_pixel[j] = ret_pixel[j-1]
+                    ret_pixel[i] = other_value
+                    istart = i+1
+                    break
+
+
+@ngjit_parallel
+def nanmin_n_in_place_4d(ret, other):
     """Combine two min-n arrays, taking nans into account. Min-n arrays are 4D
     with shape (ny, nx, ncat, n) where ny and nx are the number of pixels,
     ncat the number of categories (will be 1 if not using a categorical
     reduction) and the last axis containing n values in ascending order.
     If there are fewer than n values it is padded with nans.
     Return the first array.
     """
-    ny, nx, ncat, n = ret.shape
+    ny, nx, ncat, _n = ret.shape
     for y in nb.prange(ny):
         for x in range(nx):
             for cat in range(ncat):
-                ret_pixel = ret[y, x, cat]      # 1D array of n values for single pixel
-                other_pixel = other[y, x, cat]  # ditto
-                # Walk along other_pixel array a value at a time, find insertion
-                # index in ret_pixel and bump values along to insert.  Next
-                # other_pixel value is inserted at a higher index, so this walks
-                # the two pixel arrays just once each.
-                istart = 0
-                for other_value in other_pixel:
-                    if isnull(other_value):
-                        break
-                    else:
-                        for i in range(istart, n):
-                            if isnull(ret_pixel[i]) or other_value < ret_pixel[i]:
-                                # Bump values along then insert.
-                                for j in range(n-1, i, -1):
-                                    ret_pixel[j] = ret_pixel[j-1]
-                                ret_pixel[i] = other_value
-                                istart = i+1
-                                break
+                _nanmin_n_impl(ret[y, x, cat], other[y, x, cat])
+
+
+@ngjit_parallel
+def nanmin_n_in_place_3d(ret, other):
+    """3d version of nanmin_n_in_place_4d, taking arrays of shape (ny, nx, n).
+    """
+    ny, nx, _n = ret.shape
+    for y in nb.prange(ny):
+        for x in range(nx):
+            _nanmin_n_impl(ret[y, x], other[y, x])
 
 
 @ngjit_parallel
@@ -754,62 +790,94 @@ def row_min_in_place(ret, other):
 
 
 @ngjit
-def row_max_n_in_place(ret, other):
+def _row_max_n_impl(ret_pixel, other_pixel):
+    """Single pixel implementation of row_max_n_in_place.
+    ret_pixel and other_pixel are both 1D arrays of the same length.
+
+    Walk along other_pixel a value at a time, find insertion index in
+    ret_pixel and shift values along to insert.  Next other_pixel value is
+    inserted at a higher index, so this walks the two pixel arrays just once
+    each.
+    """
+    n = len(ret_pixel)
+    istart = 0
+    for other_value in other_pixel:
+        if other_value == -1:
+            break
+        else:
+            for i in range(istart, n):
+                if ret_pixel[i] == -1 or other_value > ret_pixel[i]:
+                    # Shift values along then insert.
+                    for j in range(n-1, i, -1):
+                        ret_pixel[j] = ret_pixel[j-1]
+                    ret_pixel[i] = other_value
+                    istart = i+1
+                    break
+
+
+@ngjit
+def row_max_n_in_place_4d(ret, other):
     """Combine two row_max_n signed integer arrays.
     Equivalent to nanmax_n_in_place with -1 replacing NaN for missing data.
     Return the first array.
     """
-    ny, nx, ncat, n = ret.shape
+    ny, nx, ncat, _n = ret.shape
     for y in range(ny):
         for x in range(nx):
             for cat in range(ncat):
-                ret_pixel = ret[y, x, cat]      # 1D array of n values for single pixel
-                other_pixel = other[y, x, cat]  # ditto
-                # Walk along other_pixel array a value at a time, find insertion
-                # index in ret_pixel and bump values along to insert.  Next
-                # other_pixel value is inserted at a higher index, so this walks
-                # the two pixel arrays just once each.
-                istart = 0
-                for other_value in other_pixel:
-                    if other_value == -1:
-                        break
-                    else:
-                        for i in range(istart, n):
-                            if ret_pixel[i] == -1 or other_value > ret_pixel[i]:
-                                # Bump values along then insert.
-                                for j in range(n-1, i, -1):
-                                    ret_pixel[j] = ret_pixel[j-1]
-                                ret_pixel[i] = other_value
-                                istart = i+1
-                                break
+                _row_max_n_impl(ret[y, x, cat], other[y, x, cat])
+
+
+@ngjit
+def row_max_n_in_place_3d(ret, other):
+    ny, nx, _n = ret.shape
+    for y in range(ny):
+        for x in range(nx):
+            _row_max_n_impl(ret[y, x], other[y, x])
 
 
 @ngjit
-def row_min_n_in_place(ret, other):
+def _row_min_n_impl(ret_pixel, other_pixel):
+    """Single pixel implementation of row_min_n_in_place.
+    ret_pixel and other_pixel are both 1D arrays of the same length.
+
+    Walk along other_pixel a value at a time, find insertion index in
+    ret_pixel and shift values along to insert.  Next other_pixel value is
+    inserted at a higher index, so this walks the two pixel arrays just once
+    each.
+    """
+    n = len(ret_pixel)
+    istart = 0
+    for other_value in other_pixel:
+        if other_value == -1:
+            break
+        else:
+            for i in range(istart, n):
+                if ret_pixel[i] == -1 or other_value < ret_pixel[i]:
+                    # Shift values along then insert.
+                    for j in range(n-1, i, -1):
+                        ret_pixel[j] = ret_pixel[j-1]
+                    ret_pixel[i] = other_value
+                    istart = i+1
+                    break
+
+
+@ngjit
+def row_min_n_in_place_4d(ret, other):
     """Combine two row_min_n signed integer arrays.
     Equivalent to nanmin_n_in_place with -1 replacing NaN for missing data.
     Return the first array.
     """
-    ny, nx, ncat, n = ret.shape
+    ny, nx, ncat, _n = ret.shape
     for y in range(ny):
         for x in range(nx):
             for cat in range(ncat):
-                ret_pixel = ret[y, x, cat]      # 1D array of n values for single pixel
-                other_pixel = other[y, x, cat]  # ditto
-                # Walk along other_pixel array a value at a time, find insertion
-                # index in ret_pixel and bump values along to insert.  Next
-                # other_pixel value is inserted at a higher index, so this walks
-                # the two pixel arrays just once each.
-                istart = 0
-                for other_value in other_pixel:
-                    if other_value == -1:
-                        break
-                    else:
-                        for i in range(istart, n):
-                            if ret_pixel[i] == -1 or other_value < ret_pixel[i]:
-                                # Bump values along then insert.
-                                for j in range(n-1, i, -1):
-                                    ret_pixel[j] = ret_pixel[j-1]
-                                ret_pixel[i] = other_value
-                                istart = i+1
-                                break
+                _row_min_n_impl(ret[y, x, cat], other[y, x, cat])
+
+
+@ngjit
+def row_min_n_in_place_3d(ret, other):
+    ny, nx, _n = ret.shape
+    for y in range(ny):
+        for x in range(nx):
+            _row_min_n_impl(ret[y, x], other[y, x])