minor refactor and typo fix
@@ -146,6 +146,29 @@ def _get_coord_affine(ndim):
return ops
+def _unravel_loop_index(shape, uint_t='unsigned int'):
+ """
+ declare a multi-index array in_coord and unravel the 1D index, i into it.
+ This code assumes that the array is a C-ordered array.
+ """
+ ndim = len(shape)
+ code = [
+ """
+ {uint_t} in_coord[{ndim}];
+ {uint_t} s, t, idx = i;
+ """.format(uint_t=uint_t, ndim=ndim)]
+ for j in range(ndim - 1, 0, -1):
+ code.append("""
+ s = {size};
+ t = idx / s;
+ in_coord[{j}] = idx - t * s;
+ idx = t;
+ """.format(j=j, size=shape[j]))
+ code.append("""
+ in_coord[0] = idx;""")
+ return '\n'.join(code)
+
+
def _generate_interp_custom(coord_func, ndim, large_int, yshape, mode, cval,
order, name='', integer_output=False):
"""
@@ -175,7 +198,6 @@ def _generate_interp_custom(coord_func, ndim, large_int, yshape, mode, cval,
else:
uint_t = 'unsigned int'
int_t = 'int'
- ops.append('{uint_t} in_coord[{ndim}];'.format(uint_t=uint_t, ndim=ndim))
# determine strides for x along each axis
for j in range(ndim):
@@ -190,20 +212,8 @@ def _generate_interp_custom(coord_func, ndim, large_int, yshape, mode, cval,
)
)
- # determine nd coordinate in x corresponding to a given raveled coordinate,
- # i, in y.
- ops.append("""
- {uint_t} idx = i;
- {uint_t} s, t;
- """.format(uint_t=uint_t))
- for j in range(ndim - 1, 0, -1):
- ops.append("""
- s = {ysize_j};
- t = idx / s;
- in_coord[{j}] = idx - t * s;
- idx = t;
- """.format(j=j, ysize_j=yshape[j]))
- ops.append("in_coord[0] = idx;")
+ # create in_coords array to store the unraveled indices
+ ops.append(_unravel_loop_index(yshape, uint_t))
# compute the transformed (target) coordinates, c_j
ops = ops + coord_func(ndim)
@@ -303,7 +313,7 @@ def _generate_interp_custom(coord_func, ndim, large_int, yshape, mode, cval,
name = 'interpolate_{}_order{}_{}_{}d_y{}'.format(
name, order, mode, ndim, "_".join(["{}".format(j) for j in yshape]),
)
- if int_t == 'size_t':
+ if uint_t == 'size_t':
name += '_i64'
return operation, name
