def expand_results(result, nresults):
if len(result[0]) <= nresults:
result = SearchResult(expand_array_if_needed(result[0], nresults),
expand_array_if_needed(result[1], nresults),
expand_array_if_needed(result[2], nresults))
result.indices[nresults] = result.indices[nresults - 1]
result.positions[nresults] = result.positions[nresults - 1]
return result
def test_numba_expand_array_if_needed_int32():
ary0 = ary = np.arange(7).astype(np.int32)
for i in range(7):
ary = util.expand_array_if_needed(ary, i)
assert ary0 is ary
for i in range(7, 100):
ary = util.expand_array_if_needed(ary, i)
assert len(ary.shape) == 1
assert ary.shape[0] >= 100
assert np.all(ary[:len(ary0)] == ary0)
assert np.all(ary[len(ary0):] == -1)
def test_numba_expand_array_if_needed_2d1():
ary0 = ary = np.arange(7).reshape((7, 1))
for i in range(7):
ary = util.expand_array_if_needed(ary, i)
assert ary0 is ary
for i in range(7, 100):
ary = util.expand_array_if_needed(ary, i)
assert len(ary.shape) == 2
assert ary.shape[0] >= 100
assert ary.shape[1] == 1
assert np.all(ary[:len(ary0)] == ary0)
assert np.all(ary[len(ary0):] == -1)
def test_numba_expand_array_if_needed_7d():
ary0 = ary = np.stack([np.arange(7)] * 7, axis=1)
for i in range(7):
ary = util.expand_array_if_needed(ary, i)
assert ary0 is ary
for i in range(7, 100):
ary = util.expand_array_if_needed(ary, i)
assert len(ary.shape) == 2
assert ary.shape[0] >= 100
assert ary.shape[1] == 7
assert np.all(ary[:len(ary0)] == ary0)
assert np.all(ary[len(ary0):] == -1)
def expand_results(result, nresults):
if len(result[0]) <= nresults:
result = ResultJIT(
expand_array_if_needed(result[0], nresults),
expand_array_if_needed(result[1], nresults),
expand_array_if_needed(result[2], nresults),
result.stats,
)
result.idx[nresults] = result.idx[nresults - 1]
result.pos[nresults] = result.pos[nresults - 1]
result.err[nresults] = result.err[nresults - 1]
return result
