class ArrayTests(ArrayTestsBase, unittest.TestCase):
A = Array(Long, (2, 3, 2))
B = Array(Long, (12, ))
def test_repr(self):
self.assertEqual(repr(self.A), "Array(Long, shape=(2, 3, 2))")
self.assertEqual(repr(self.B), "Array(Long, shape=(12,))")
def test_str(self):
self.assertEqual(str(self.A), "<Array [2 x [3 x [2 x Long]]]>")
self.assertEqual(str(self.B), "<Array [12 x Long]>")
def test_and_eval(self):
"""Don't evaluate `and`ed expressions if not necessary."""
from nitrous.types.array import Array
Bool1 = Array(Bool, (1,))
@function(Bool, b=Bool1)
def side_effect(b):
b[0] = True
return True
@function(Bool, a=Bool, b=Bool1)
def and_(a, b):
return a and side_effect(b)
m = module([and_])
# First value is true; b should evaluate as well
x = (Bool.c_type * 1)(False)
m.and_(True, x)
self.assertTrue(x[0])
# First value is false; should skip next term.
x = (Bool.c_type * 1)(False)
m.and_(False, x)
self.assertFalse(x[0])
def test_or_eval(self):
"""Don't evaluate `or`ed expressions if not necessary."""
from nitrous.types.array import Array
Bool1 = Array(Bool, (1,))
@function(Bool, b=Bool1)
def side_effect(b):
b[0] = True
return True
@function(Bool, a=Bool, b=Bool1)
def or_(a, b):
return a or side_effect(b)
m = module([or_])
# First value is true; should skip b
x = (Bool.c_type * 1)(False)
m.or_(True, x)
self.assertFalse(x[0])
# First value is false; should evaluate b
x = (Bool.c_type * 1)(False)
m.or_(False, x)
self.assertTrue(x[0])
def test_for_range(self):
"""More advanced loop ranges."""
import ctypes
Long8 = Array(Long, (8, ))
@function(Long, data=Long8, start=Long, end=Long)
def loop_1(data, start, end):
for i in range(start, end):
data[i] = i
return 0
@function(Long, data=Long8, start=Long, end=Long, step=Long)
def loop_2(data, start, end, step):
for i in range(start, end, step):
data[i] = i
return 0
m = module([loop_1, loop_2])
data = (ctypes.c_long * 8)()
m.loop_1(data, 2, 7)
self.assertEqual(list(data), [0, 0, 2, 3, 4, 5, 6, 0])
data = (ctypes.c_long * 8)()
m.loop_2(data, 2, 7, 2)
self.assertEqual(list(data), [0, 0, 2, 0, 4, 0, 6, 0])
def test_assign_wrong_type(self):
@function(x=Array(Long, shape=(1, )), y=Double)
def f(x, y):
x[0] = y
message = ">>> x\[0\] = y"
with self.assertRaisesRegexp(TypeError, message):
module([f])
def test_aug(self):
"""Augmented assignment."""
@function(Long, a=Long, b=Array(Long, (1, )))
def f(a, b):
a += 5
b[0] += 7
return a
m = module([f])
b = (Long.c_type * 1)(5)
self.assertEqual(m.f(6, b), 11)
self.assertEqual(b[0], 12)
def test_unpack_subscript(self):
"""Unpacking tuples into array elements"""
A, B, C = range(3)
@function(Long, d=Array(Long, (3, )))
def foo(d):
d[C], d[A], d[B] = d[A], d[B], d[C]
return d[B] * 10 + d[A] * 100 + d[C] * 1000
m = module([foo])
d = (Long.c_type * 3)(5, 6, 2)
self.assertEqual(m.foo(d), 5000 + 600 + 20)
def test_load_element(self):
import ctypes
@function(Long, data=Array(Long, (5, )), i=Long)
def get_i(data, i):
e = data[i]
return e
m = module([get_i])
dtype = (ctypes.c_long * 5)
data = dtype(0, 10, 20, 30, 40)
for i in range(5):
self.assertEqual(m.get_i(data, i), data[i])
def test_nd_index(self):
import ctypes
@function(Long, y=Array(Long, shape=(2, 3, 2)), i=Long, j=Long, k=Long)
def x(y, i, j, k):
return y[i, j, k]
m = module([x])
dtype = (((ctypes.c_long * 2) * 3) * 2)
data = dtype(((1, 2), (3, 4), (5, 6)), ((7, 8), (9, 10), (11, 12)))
for i in range(2):
for j in range(3):
for k in range(2):
self.assertEqual(m.x(data, i, j, k), data[i][j][k])
def test_init_2d(self):
"""Multi-dimensional array initialization."""
from nitrous.types import Double
Double2x2 = Array(Double, (2, 2))
@function(Double2x2, x=Double, y=Double, z=Double, w=Double)
def make_2x2(x, y, z, w):
return Double2x2(((x, y), (z, w)))
m = module([make_2x2])
c = m.make_2x2(1.0, 2.0, 3.0, 4.0)
self.assertEqual(c[0][0], 1.0)
self.assertEqual(c[0][1], 2.0)
self.assertEqual(c[1][0], 3.0)
self.assertEqual(c[1][1], 4.0)
def test_for(self):
"""Simple for loop given range stop value."""
import ctypes
@function(Long, data=Array(Long, (6, )), n=Long)
def loop_1(data, n):
for i in range(n):
data[i] = (i if i < 3 else 99)
return 0
m = module([loop_1])
# 5 + 1 elements to check stop correctness
data = (ctypes.c_long * 6)()
m.loop_1(data, 5)
self.assertEqual(list(data), [0, 1, 2, 99, 99, 0])
def test_while(self):
import ctypes
@function(Long, data=Array(Long, (6, )), n=Long)
def loop_1(data, n):
i = 0
while i < n:
data[i] = (i if i < 3 else 99)
i += 1
return 0
m = module([loop_1])
# 5 + 1 elements to check stop correctness
data = (ctypes.c_long * 6)()
m.loop_1(data, 5)
self.assertEqual(list(data), [0, 1, 2, 99, 99, 0])
def test_alloc_return(self):
"""Allocate array and pass back through return value."""
from nitrous.types import Double
Coord = Array(Double, (3, ))
@function(Coord, x=Double, y=Double, z=Double)
def make_coord(x, y, z):
return Coord((x, y, z))
@function(Coord, x=Double, y=Double, z=Double)
def make_coord_2(x, y, z):
return make_coord(x, y, z)
m = module([make_coord, make_coord_2])
c = m.make_coord_2(1.0, 2.0, 3.0)
self.assertEqual(tuple(c), (1.0, 2.0, 3.0))
def test_double_while(self):
import ctypes
@function(data=Array(Long, (9, )), n=Long)
def loop_1(data, n):
i = 0
while i < n:
j = 0
while j < n:
data[i * n + j] = i + j
j += 1
i += 1
m = module([loop_1])
data = (ctypes.c_long * 9)()
expected = [0, 1, 2, 1, 2, 3, 2, 3, 4]
m.loop_1(data, 3)
self.assertEqual(list(data), expected)
def test_double_for(self):
import ctypes
@function(Long, data=Array(Long, (9, )), n=Long)
def loop_1(data, n):
j = 0
for i in range(n):
for j in range(n):
data[i * n + j] = i + j
# Check availability and correctness of
# loop variables outside the loop.
return i * j
m = module([loop_1])
data = (ctypes.c_long * 9)()
expected = [0, 1, 2, 1, 2, 3, 2, 3, 4]
self.assertEqual(m.loop_1(data, 3), 4)
self.assertEqual(list(data), expected)
def test_store_element(self):
import ctypes
@function(Long, data=Array(Long, (5, )), i=Long, e=Long)
def set_i(data, i, e):
data[i] = e
return 0
m = module([set_i])
dtype = (ctypes.c_long * 5)
data = dtype(0, 0, 0, 0, 0)
m.set_i(data, 1, 2)
m.set_i(data, 2, 5)
m.set_i(data, 4, 10)
expected = (0, 2, 5, 0, 10)
for i in range(5):
self.assertEqual(data[i], expected[i])
import unittest
import numpy as np
from nitrous.function import function
from nitrous.types import Double, Index
from nitrous.types.array import Array, FastSlice, Slice, Any
DoubleNx3 = Slice(Double, shape=(Any, 3))
DoubleN = Slice(Double)
DoubleNArray = FastSlice(Double)
Double3 = Array(Double, (3, ))
X, Y, Z = range(3)
@function(d=DoubleNx3, out=DoubleNArray, n=Index, m=Index)
def sum_1(d, out, n, m):
for k in range(m):
for i in range(n):
for j in range(i + 1, n):
out[X] += d[i, X] * d[j, X]
out[Y] += d[i, Y] * d[j, Y]
out[Z] += d[i, Z] * d[j, Z]
@function(d=DoubleNx3, out=DoubleN, n=Index, m=Index)
def sum_2(d, out, n, m):
for k in range(m):