def test_compiled_for_loop_fn_0(self):
test_data = numpy.array([1, 2, 3], dtype = 'l')
compiled_for_loop_fn = numba_compile(
arg_types = [['l']])(for_loop_fn_0)
result = compiled_for_loop_fn(test_data)
self.assertEqual(result, 6)
self.assertEqual(result, for_loop_fn_0(testdata))
def test_vectorized_sum2d(self):
usum2d = numba_compile(arg_types=[d[:, :]],
ret_type=d)(sum2d)
image = numpy.random.random(10, 10)
plain_old_result = sum2d(image)
hot_new_result = usum2d(image)
self.assertTrue((abs(plain_old_result - hot_new_result) < 1e-9).all())
def test_getattr_shape_2(self):
test_data2 = numpy.array([[1., 2., 3.], [4., 5., 6.]])
compiled_fn2 = numba_compile(ret_type = 'i%d*' % (_plat_bits // 8),
arg_types = [[['d']]])(get_ndarray_shape)
result = compiled_fn2(test_data2)
self.assertEqual(result[0], 2)
self.assertEqual(result[1], 3)
def test_set_index_fn_1(self):
control_arr = numpy.zeros((50, 50, 2))
test_arr = numpy.zeros_like(control_arr)
set_index_fn_1(-1., 1., -1., control_arr)
compiled_fn = numba_compile(
arg_types=['d', 'd', 'd', ['d']])(set_index_fn_1)
compiled_fn(-1., 1., -1., test_arr)
self.assertTrue((numpy.abs(control_arr - test_arr) < 1e9).all())
def main (*args, **kws):
compiled_demo_function = numba_compile(
arg_types = ['d', 'd', 'd', [[['d']]]])(demo_function)
control_arr = numpy.zeros((5, 5, 2))
demo_function(-1., 1., -1., control_arr)
test_arr = numpy.zeros_like(control_arr)
compiled_demo_function(-1., 1., -1., test_arr)
assert (numpy.abs(control_arr - test_arr) < 1e9).all()
def test_get_imag_fn (self):
num0 = 0 - 2j
num1 = numpy.complex128(num0)
compiled_get_imag_fn = numba_compile(arg_types = ['D'])(get_imag_fn)
self.assertEqual(compiled_get_imag_fn(num0), -2.)
self.assertEqual(get_imag_fn(num0), compiled_get_imag_fn(num0))
self.assertEqual(compiled_get_imag_fn(num1), -2.)
self.assertEqual(get_imag_fn(num1), compiled_get_imag_fn(num1))
def test_set_index_fn_1 (self):
control_arr = numpy.zeros((50, 50, 2))
test_arr = numpy.zeros_like(control_arr)
set_index_fn_1(-1., 1., -1., control_arr)
compiled_fn = numba_compile(
arg_types = ['d', 'd', 'd', ['d']])(set_index_fn_1)
compiled_fn(-1., 1., -1., test_arr)
self.assertTrue((numpy.abs(control_arr - test_arr) < 1e9).all())
def test_getattr_data_1(self):
test_data = numpy.array([1., 2., 3.])
compiled_fn = numba_compile(ret_type = 'd*',
arg_types = [['d']])(get_ndarray_data)
result = compiled_fn(test_data)
self.assertEqual(result[0], 1.)
self.assertEqual(result[1], 2.)
self.assertEqual(result[2], 3.)
def test_get_imag_fn(self):
num0 = 0 - 2j
num1 = numpy.complex128(num0)
compiled_get_imag_fn = numba_compile(arg_types=['D'])(get_imag_fn)
self.assertEqual(compiled_get_imag_fn(num0), -2.)
self.assertEqual(get_imag_fn(num0), compiled_get_imag_fn(num0))
self.assertEqual(compiled_get_imag_fn(num1), -2.)
self.assertEqual(get_imag_fn(num1), compiled_get_imag_fn(num1))
def test_get_conj_fn(self):
num0 = 4 - 1.5j
num1 = numpy.complex128(num0)
compiled_get_conj_fn = numba_compile(arg_types=['D'],
ret_type='D')(get_conj_fn)
self.assertEqual(compiled_get_conj_fn(num0), 4 + 1.5j)
self.assertEqual(get_conj_fn(num0), compiled_get_conj_fn(num0))
self.assertEqual(compiled_get_conj_fn(num1), 4 + 1.5j)
self.assertEqual(get_conj_fn(num1), compiled_get_conj_fn(num1))
def test_vectorized_filter2d(self):
ufilter2d = numba_compile(arg_types=[d[:, :], d[:, :]],
ret_type=d[:, :])(filter2d)
image = numpy.random.random((50, 50))
filt = numpy.random.random((5, 5))
filt /= filt.sum()
plain_old_result = filter2d(image, filt)
hot_new_result = ufilter2d(image, filt)
self.assertTrue((abs(plain_old_result - hot_new_result) < 1e-9).all())
def test_getattr_shape_2_unpack(self):
compiler_fn = numba_compile(ret_type = 'i%d' % (_plat_bits // 8),
arg_types = [[['d']]])
dim0_fn, dim1_fn = (compiler_fn(fn)
for fn in (get_ndarray_2_shape_unpack_0,
get_ndarray_2_shape_unpack_1))
test_data2 = numpy.array([[1., 2., 3.], [4., 5., 6.]])
self.assertEqual(dim0_fn(test_data2), 2)
self.assertEqual(dim1_fn(test_data2), 3)
def test_get_conj_fn (self):
num0 = 4 - 1.5j
num1 = numpy.complex128(num0)
compiled_get_conj_fn = numba_compile(arg_types = ['D'],
ret_type = 'D')(get_conj_fn)
self.assertEqual(compiled_get_conj_fn(num0), 4 + 1.5j)
self.assertEqual(get_conj_fn(num0), compiled_get_conj_fn(num0))
self.assertEqual(compiled_get_conj_fn(num1), 4 + 1.5j)
self.assertEqual(get_conj_fn(num1), compiled_get_conj_fn(num1))
def test_vectorized_filter2d(self):
ufilter2d = numba_compile(arg_types=[d[:, :], d[:, :]],
ret_type=d[:, :])(filter2d)
image = numpy.random.random((50, 50))
filt = numpy.random.random((5, 5))
filt /= filt.sum()
plain_old_result = filter2d(image, filt)
hot_new_result = ufilter2d(image, filt)
self.assertTrue((abs(plain_old_result - hot_new_result) < 1e-9).all())
def test_prod_sum_fn (self):
compiled_prod_sum_fn = numba_compile(arg_types = ['D', 'D', 'D'],
ret_type = 'D')(prod_sum_fn)
rng = numpy.arange(-1., 1.1, 0.5)
for ar, ai, xr, xi, br, bi in itertools.product(rng, rng, rng, rng, rng,
rng):
a = numpy.complex128(ar + ai * 1j)
x = numpy.complex128(xr + xi * 1j)
b = numpy.complex128(br + bi * 1j)
self.assertEqual(prod_sum_fn(a, x, b),
compiled_prod_sum_fn(a, x, b))
def test_getattr_data_2(self):
test_data = numpy.array([[1., 2., 3.], [4., 5., 6.]])
compiled_fn = numba_compile(ret_type = 'd*',
arg_types = [[['d']]])(get_ndarray_data)
result = compiled_fn(test_data)
self.assertEqual(result[0], 1.)
self.assertEqual(result[1], 2.)
self.assertEqual(result[2], 3.)
self.assertEqual(result[3], 4.)
self.assertEqual(result[4], 5.)
self.assertEqual(result[5], 6.)
def test_prod_sum_fn(self):
compiled_prod_sum_fn = numba_compile(arg_types=['D', 'D', 'D'],
ret_type='D')(prod_sum_fn)
rng = numpy.arange(-1., 1.1, 0.5)
for ar, ai, xr, xi, br, bi in itertools.product(
rng, rng, rng, rng, rng, rng):
a = numpy.complex128(ar + ai * 1j)
x = numpy.complex128(xr + xi * 1j)
b = numpy.complex128(br + bi * 1j)
self.assertEqual(prod_sum_fn(a, x, b),
compiled_prod_sum_fn(a, x, b))
def test_getattr_ndim_1(self):
test_data1 = numpy.array([1., 2., 3.])
compiled_fn1 = numba_compile(ret_type = 'i',
arg_types = [['d']])(get_ndarray_ndim)
self.assertEqual(compiled_fn1(test_data1), 1)
def test_tuple_fn_0 (self):
test_arr = numpy.zeros((4,4,4))
compiled_fn = numba_compile(arg_types = [['d']])(tuple_fn_0)
self.assertEqual(compiled_fn(test_arr), 0.)
def test_if_fn_1(self):
if_fn_1c = numba_compile()(if_fn_1)
self.assertEqual(if_fn_1c(-1.), 42.)
self.assertEqual(if_fn_1c(1.), 22.)
def test_set_index_fn_0(self):
arr = numpy.ones((4, 4, 4))
compiled_fn = numba_compile(arg_types=[['d']])(set_index_fn_0)
self.assertEqual(arr[1, 2, 3], 1.)
compiled_fn(arr)
self.assertEqual(arr[1, 2, 3], 0.)
def _do_test(self, function, arg_types, *args, **kws):
_numba_compile = numba_compile(arg_types=arg_types) if arg_types is not None else numba_compile()
compiled_fn = _numba_compile(function)
self.assertEqual(compiled_fn(*args, **kws), function(*args, **kws))
def test_set_index_fn_0 (self):
arr = numpy.ones((4,4,4))
compiled_fn = numba_compile(arg_types = [['d']])(set_index_fn_0)
self.assertEqual(arr[1,2,3], 1.)
compiled_fn(arr)
self.assertEqual(arr[1,2,3], 0.)
def test_compiled_for_loop_fn_0(self):
test_data = numpy.array([1, 2, 3], dtype='l')
compiled_for_loop_fn = numba_compile(arg_types=[['l']])(for_loop_fn_0)
result = compiled_for_loop_fn(test_data)
self.assertEqual(result, 6)
self.assertEqual(result, for_loop_fn_0(testdata))
# to detect and desugar for loops over range/xrange/arange first.
i = 0
z_real = 0.
z_imag = 0.
while i < max_iters:
z_real_n = z_real * z_real - z_imag * z_imag + real_coord
z_imag = 2. * z_real * z_imag + imag_coord
z_real = z_real_n
if (z_real * z_real + z_imag * z_imag) >= 4:
return i
i += 1
return -1
try:
mandel_1c = numba_compile(arg_types=['d', 'd', 'i'],
ret_type='i')(mandel_1)
except:
if __debug__:
import traceback as tb
tb.print_exc()
mandel_1c = None
#@numba_compile(arg_types = ['d', 'd', 'd', 'i', [['b']], [[['b']]]])
def mandel_driver_1(min_x, max_x, min_y, nb_iterations, colors, image):
nb_colors = len(colors)
width = image.shape[0]
height = image.shape[1]
pixel_size = (max_x - min_x) / width
for x in range(width):
real = min_x + x * pixel_size
def test_while_loop_fn_3(self):
compiled_fn = numba_compile(arg_types=["l"])(while_loop_fn_3)
compiled_result = compiled_fn(3)
self.assertEqual(compiled_result, while_loop_fn_3(3))
self.assertEqual(compiled_result, 8.0)
def test_while_loop_fn_4(self):
compiled_fn = numba_compile(arg_types=["l", "l", "l"], ret_type="l")(while_loop_fn_4)
compiled_result = compiled_fn(1, 4, 1)
self.assertEqual(compiled_result, while_loop_fn_4(1, 4, 1))
self.assertEqual(compiled_result, 6)
def test_call_len(self):
testarr = numpy.arange(10.)
testfn = numba_compile(arg_types=[['d']], ret_type='l')(call_len)
self.assertEqual(testfn(testarr), 10)
def test_call_zeros_like(self):
testarr = numpy.array([1., 2, 3, 4, 5])
testfn = numba_compile(arg_types=[['d']],
ret_type=['d'])(call_zeros_like)
self.assertTrue((testfn(testarr) == numpy.zeros_like(testarr)).all())
def test_getattr_ndim_2(self):
test_data2 = numpy.array([[1., 2., 3.], [4., 5., 6.]])
compiled_fn2 = numba_compile(ret_type = 'i',
arg_types = [[['d']]])(get_ndarray_ndim)
self.assertEqual(compiled_fn2(test_data2), 2)
def _do_test(self, function, arg_types, *args, **kws):
_numba_compile = (numba_compile(arg_types = arg_types)
if arg_types is not None else numba_compile())
compiled_fn = _numba_compile(function)
self.assertEqual(compiled_fn(*args, **kws), function(*args, **kws))
def test_getattr_shape_1(self):
test_data = numpy.array([1., 2., 3.])
compiled_fn = numba_compile(ret_type = 'i%d*' % (_plat_bits // 8),
arg_types = [['d']])(get_ndarray_shape)
result = compiled_fn(test_data)
self.assertEqual(result[0], 3)
def test_while_loop_fn_4(self):
compiled_fn = numba_compile(arg_types = ['l', 'l', 'l'],
ret_type = 'l')(while_loop_fn_4)
compiled_result = compiled_fn(1, 4, 1)
self.assertEqual(compiled_result, while_loop_fn_4(1, 4, 1))
self.assertEqual(compiled_result, 6)
def test_get_complex_constant_fn (self):
compiled_get_complex_constant_fn = numba_compile(
arg_types = [], ret_type = 'D')(get_complex_constant_fn)
self.assertEqual(get_complex_constant_fn(),
compiled_get_complex_constant_fn())
def test_call_len(self):
testarr = numpy.arange(10.)
testfn = numba_compile(arg_types = [['d']], ret_type = 'l')(
call_len)
self.assertEqual(testfn(testarr), 10)
def test_if_fn_1(self):
if_fn_1c = numba_compile()(if_fn_1)
self.assertEqual(if_fn_1c(-1.), 42.)
self.assertEqual(if_fn_1c(1.), 22.)
def test_while_loop_fn_5(self):
compiled_fn = numba_compile(arg_types=["d", "d"])(while_loop_fn_5)
compiled_result = compiled_fn(3, 4)
self.assertEqual(compiled_result, while_loop_fn_5(3, 4))
self.assertEqual(compiled_result, 18.0)
def test_compiled_for_loop_fn_2(self):
compiled_for_loop_fn = numba_compile(arg_types=['i'],
ret_type='i')(for_loop_fn_2)
result = compiled_for_loop_fn(4)
self.assertEqual(result, 36)
self.assertEqual(result, for_loop_fn_2(4))
def test_compiled_for_loop_fn_1(self):
compiled_for_loop_fn = numba_compile(arg_types=['i', 'i', 'i'],
ret_type='i')(for_loop_fn_1)
result = compiled_for_loop_fn(1, 4, 1)
self.assertEqual(result, 6)
self.assertEqual(result, for_loop_fn_1(1, 4, 1))
def test_get_complex_constant_fn(self):
compiled_get_complex_constant_fn = numba_compile(
arg_types=[], ret_type='D')(get_complex_constant_fn)
self.assertEqual(get_complex_constant_fn(),
compiled_get_complex_constant_fn())
def test_compiled_for_loop_fn_1(self):
compiled_for_loop_fn = numba_compile(arg_types = ['i','i','i'],
ret_type = 'i')(for_loop_fn_1)
result = compiled_for_loop_fn(1, 4, 1)
self.assertEqual(result, 6)
self.assertEqual(result, for_loop_fn_1(1, 4, 1))
def test_while_loop_fn_5(self):
compiled_fn = numba_compile(arg_types = ['d', 'd'])(while_loop_fn_5)
compiled_result = compiled_fn(3, 4)
self.assertEqual(compiled_result, while_loop_fn_5(3, 4))
self.assertEqual(compiled_result, 18.)
def test_long(self):
func = numba_compile(ret_type=numba.long_,
arg_types=[numba.long_])(test_long)
self.assertEqual(func(-1), 42)
self.assertEqual(func(1), 22)
def test_while_loop_fn_3(self):
compiled_fn = numba_compile(arg_types = ['l'])(while_loop_fn_3)
compiled_result = compiled_fn(3)
self.assertEqual(compiled_result, while_loop_fn_3(3))
self.assertEqual(compiled_result, 8.)
def test_compiled_for_loop_fn_3(self):
compiled_for_loop_fn = numba_compile(arg_types = ['i'],
ret_type = 'i')(for_loop_fn_3)
result = compiled_for_loop_fn(3)
self.assertEqual(result, for_loop_fn_3(3))
self.assertEqual(result, 81)
def test_call_zeros_like(self):
testarr = numpy.array([1., 2, 3, 4, 5])
testfn = numba_compile(arg_types = [['d']], ret_type = ['d'])(
call_zeros_like)
self.assertTrue((testfn(testarr) == numpy.zeros_like(testarr)).all())
def test_int(self):
func = numba_compile(ret_type=numba.int_,
arg_types=[numba.int_])(test_int)
self.assertEqual(func(-1), 42)
self.assertEqual(func(1), 22)
def test_get_index_fn_0(self):
arr = numpy.ones((4, 4, 4))
arr[1, 2, 3] = 0.
compiled_fn = numba_compile(arg_types=[['d']])(get_index_fn_0)
self.assertEqual(compiled_fn(arr), 0.)
def test_long(self):
func = numba_compile(ret_type=numba.long_,
arg_types=[numba.long_])(test_long)
self.assertEqual(func(-1), 42)
self.assertEqual(func(1), 22)
def test_compiled_for_loop_fn_2(self):
compiled_for_loop_fn = numba_compile(arg_types = ['i'],
ret_type = 'i')(for_loop_fn_2)
result = compiled_for_loop_fn(4)
self.assertEqual(result, 36)
self.assertEqual(result, for_loop_fn_2(4))
def test_compiled_for_loop_fn_3(self):
compiled_for_loop_fn = numba_compile(arg_types=['i'],
ret_type='i')(for_loop_fn_3)
result = compiled_for_loop_fn(3)
self.assertEqual(result, for_loop_fn_3(3))
self.assertEqual(result, 81)
