class TestSize:
@pytest.mark.parametrize("dim", (-1, 3))
def test_wrong_dim(self, dim):
a = zn.Array([[1, 2], [2, 3], [3, 4]])
with pytest.raises(
ValueError,
match=f"Array has 0\\.\\.2 dimensions, but {dim} were specified"
):
a.size(dim)
@pytest.mark.parametrize("array", [
zn.Array([[1, 2], [2, 3], [3, 4]]),
zn.Array(zn.var(int), shape=(3, 2))
])
@pytest.mark.parametrize("dim, expected", ((0, 3), (1, 2)))
def test_dim(self, array, dim, expected):
class MyModel(zn.Model):
def __init__(self, array):
self.a = array
self.s = self.a.size(dim)
self.s2 = self.a.size(dim) * 2 # test arithmetic operations
self.s_bigger_2 = self.a.size(
dim) > 2 # test compare operations
assert self.s.type is int
assert self.s2.type is int
assert self.s_bigger_2.type is int
result = MyModel(array).solve_satisfy()
assert result["s"] == expected
assert result["s_bigger_2"] == (expected > 2)
assert result["s2"] == (expected * 2)
def test_slice_wrong_start(self, start):
array = zn.Array(zn.var(int), shape=(3, 2, 4))
with pytest.raises(
ValueError,
match=re.escape(
f"start({start}) should be smaller then stop(10)")):
_ = array[1, start:10]
def test_array_and_op():
array = zn.Array(zn.var(int), shape=(3, 4))
v = zn.var(int)
iter_var, op = get_iter_var_and_op(array, v + 1)
assert op.op == _Op_code.add
assert isinstance(op.params[1], int) and op.params[1] == 1
assert isinstance(op.params[0], zn.var) and op.params[0] is v
def test_wrong_dim(self, dim):
a = zn.Array([[1, 2], [2, 3], [3, 4]])
with pytest.raises(
ValueError,
match=f"Array has 0\\.\\.2 dimensions, but {dim} were specified"
):
a.size(dim)
def test_slice_and_func():
def fn(par):
return 2 - par
array = zn.Array(zn.var(int), shape=(3, 4))
iter_var, op = get_iter_var_and_op(array[2:], fn)
assert iter_var.name == "par"
assert iter_var.type is int
assert op.op == _Op_code.sub
assert isinstance(op.params[0], int) and op.params[0] == 2
assert op.params[1].name == "par"
def create_array(name, ndims):
p = zn.Array(zn.var(range(5)), shape=(5,) * ndims)
p._name = name
return p
def __init__(self):
self.a = zn.Array([[1, 2, 3], [4, 5, 6]])
self.start = zn.par(1)
self.b = zn.sum(self.a[:, self.start + 1:self.start * 3])
def __init__(self, array):
self.a = zn.Array(array)
def test_var_array_without_shape():
with pytest.raises(ValueError, match="shape wasn't specified"):
zn.Array(zn.var(int))
def __init__(self, array):
self.a = zn.Array(array)
self.s1 = zn.sum(self.a[0, 1:2])
self.s2 = zn.sum(self.a[2:, 1:])
self.s3 = zn.sum(self.a[2:, 2:])
def __init__(self, array):
self.a = zn.Array(array)
with pytest.raises(
ValueError,
match="All elements of the array should be the same type"):
MyModel(array)
def test_var_array_without_shape():
with pytest.raises(ValueError, match="shape wasn't specified"):
zn.Array(zn.var(int))
@pytest.mark.parametrize(
"array", (zn.Array([[1], [2]]), zn.Array(zn.var(int), shape=(1, 2))))
@pytest.mark.parametrize("indexes", ((2, 0), (0, 2)))
def test_array_index_error(array, indexes):
with pytest.raises(IndexError):
array[indexes]
@pytest.mark.parametrize("array", (((1, 0, 1), (1, 3)), [[1, 3], [1, 1, 0]],
(((1, 1), 1), )))
def test_different_length(array):
class MyModel(zn.Model):
def __init__(self, array):
self.a = zn.Array(array)
with pytest.raises(
ValueError,
def __init__(self, low, high):
self.a = zn.Array(zn.var(range(1, 10)), shape=40)
self.n = zn.var(range(low, high))
self.constraints = [zn.ndistinct(self.a) == self.n]
def test_neg_slice_2d(self, pos):
a = zn.Array([[1], [2], [3]])
with pytest.raises(ValueError,
match="Negative indexes are not supported for now"):
_ = a[pos]
def test_step(self):
a = zn.Array([1, 2, 3])
with pytest.raises(
ValueError,
match="step other then 1 isn't supported, but it is 2"):
_ = a[1:3:2]
def test_neg_int(self):
a = zn.Array([1, 2, 3])
with pytest.raises(ValueError,
match="Negative indexes are not supported for now"):
_ = a[-1]
def test_tuple_1d(self):
a = zn.Array([1, 2, 3])
with pytest.raises(
ValueError,
match="Array has 1 dimensions but 2 were specified"):
_ = a[1, 0]
def test_int_1d(self):
a = zn.Array([1, 2, 3])
a = a[1]
assert a.pos == (1, )
class TestPos:
def test_int_1d(self):
a = zn.Array([1, 2, 3])
a = a[1]
assert a.pos == (1, )
def test_tuple_1d(self):
a = zn.Array([1, 2, 3])
with pytest.raises(
ValueError,
match="Array has 1 dimensions but 2 were specified"):
_ = a[1, 0]
def test_neg_int(self):
a = zn.Array([1, 2, 3])
with pytest.raises(ValueError,
match="Negative indexes are not supported for now"):
_ = a[-1]
def test_step(self):
a = zn.Array([1, 2, 3])
with pytest.raises(
ValueError,
match="step other then 1 isn't supported, but it is 2"):
_ = a[1:3:2]
@pytest.mark.parametrize(
"pos",
[slice(-1, 2), slice(1, -2), slice(-1, -2)])
def test_neg_slice_2d(self, pos):
a = zn.Array([[1], [2], [3]])
with pytest.raises(ValueError,
match="Negative indexes are not supported for now"):
_ = a[pos]
@pytest.mark.parametrize("start", (10, 15))
def test_slice_wrong_start(self, start):
array = zn.Array(zn.var(int), shape=(3, 2, 4))
with pytest.raises(
ValueError,
match=re.escape(
f"start({start}) should be smaller then stop(10)")):
_ = array[1, start:10]
@pytest.mark.parametrize("pos, expected",
[((0, 1), (0, 1)),
((0, slice(1, 3)), (0, slice(1, 3, 1)))])
def test_2d(self, pos, expected):
a = zn.Array([[1, 2], [2, 3], [3, 4]])
a = a[pos]
assert a.pos == expected
@pytest.mark.parametrize("pos", (1, slice(1, 2)))
@pytest.mark.parametrize("array", (zn.Array(
[[1, 2], [2, 3], [3, 4]]), zn.Array(zn.var(int), shape=(3, 2, 4))))
def test_added_last_dim(self, pos, array):
ndim = array.ndims()
array = array[pos]
assert len(array.pos) == ndim
assert array.pos[1].start == 0
assert array.pos[1].stop.op == _Op_code.size
assert array.pos[1].step == 1
def test_2d(self, pos, expected):
a = zn.Array([[1, 2], [2, 3], [3, 4]])
a = a[pos]
assert a.pos == expected
def __init__(self, array):
self.a = zn.Array(array)
self.s = zn.sum(self.a)
self.s1 = zn.sum(self.a[:, 1:])