Skip to content

Commit

Permalink
Separate out 3d and 4d combine functions
Browse files Browse the repository at this point in the history
  • Loading branch information
@ianthomas23
ianthomas23 committed Jun 30, 2023
1 parent 6e9b412 commit 049e036
Show file tree
Hide file tree
Showing 2 changed files with 169 additions and 101 deletions.
34 changes: 17 additions & 17 deletions datashader/reductions.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)


Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
236 changes: 152 additions & 84 deletions datashader/utils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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])

0 comments on commit 049e036

Please sign in to comment.