def test_IndexedBase_subs():
i, j, k = symbols('i j k', integer=True)
A = IndexedBase(a)
B = IndexedBase(b)
C = IndexedBase(c)
assert A[i] == B[i].subs({b: a})
assert isinstance(C[1].subs({C: Dict({1: 2})}), type(A[1]))
def test_ccode_loops_matrix_vector():
n, m = symbols('n m', integer=True)
A = IndexedBase('A')
x = IndexedBase('x')
y = IndexedBase('y')
i = Idx('i', m)
j = Idx('j', n)
s = ('for (int i=0; i<m; i++){\n'
' y[i] = 0;\n'
'}\n'
'for (int i=0; i<m; i++){\n'
' for (int j=0; j<n; j++){\n'
' y[i] = x[j]*A[%s] + y[i];\n' % (i * n + j) + ' }\n'
'}')
c = ccode(A[i, j] * x[j], assign_to=y[i])
assert c == s
pytest.raises(ValueError, lambda: ccode(A[i, j] * x[j], assign_to=x[j]))
s2 = ('for (int i=0; i<m; i++){\n'
' y[i] = 0;\n'
'}\n'
'for (int i=0; i<m; i++){\n'
' for (int i=0; i<m; i++){\n'
' y[i] = y[i] + A[m*i + i];\n'
' }\n}')
c = ccode(A[i, i], assign_to=y[i])
assert c == s2
def test_Indexed_subs():
i, j, k = symbols('i j k', integer=True)
a, b = symbols('a b')
A = IndexedBase(a)
B = IndexedBase(b)
assert A[i, j] == B[i, j].subs(b, a)
assert A[i, j] == A[i, k].subs(k, j)
def test_m_tensor_loops_multiple_contractions():
# see comments in previous test about vectorizing
n, m, o, p = symbols('n m o p', integer=True)
A = IndexedBase('A')
B = IndexedBase('B')
y = IndexedBase('y')
i = Idx('i', m)
j = Idx('j', n)
k = Idx('k', o)
l = Idx('l', p)
result, = codegen(('tensorthing', Eq(y[i], B[j, k, l] * A[i, j, k, l])),
'Octave',
header=False,
empty=False)
source = result[1]
expected = ('function y = tensorthing(A, B, m, n, o, p)\n'
' for i = 1:m\n'
' y(i) = 0;\n'
' end\n'
' for i = 1:m\n'
' for j = 1:n\n'
' for k = 1:o\n'
' for l = 1:p\n'
' y(i) = y(i) + B(j, k, l).*A(i, j, k, l);\n'
' end\n'
' end\n'
' end\n'
' end\n'
'end\n')
assert source == expected
def test_ccode_loops_multiple_contractions():
n, m, o, p = symbols('n m o p', integer=True)
a = IndexedBase('a')
b = IndexedBase('b')
y = IndexedBase('y')
i = Idx('i', m)
j = Idx('j', n)
k = Idx('k', o)
l = Idx('l', p)
s = ('for (int i=0; i<m; i++){\n'
' y[i] = 0;\n'
'}\n'
'for (int i=0; i<m; i++){\n'
' for (int j=0; j<n; j++){\n'
' for (int k=0; k<o; k++){\n'
' for (int l=0; l<p; l++){\n'
' y[i] = y[i] + b[%s]*a[%s];\n' %
(j * o * p + k * p + l, i * n * o * p + j * o * p + k * p + l) +
' }\n'
' }\n'
' }\n'
'}')
c = ccode(b[j, k, l] * a[i, j, k, l], assign_to=y[i])
assert c == s
def test_ccode_loops_addfactor():
n, m, o, p = symbols('n m o p', integer=True)
a = IndexedBase('a')
b = IndexedBase('b')
c = IndexedBase('c')
y = IndexedBase('y')
i = Idx('i', m)
j = Idx('j', n)
k = Idx('k', o)
l = Idx('l', p)
s = ('for (int i=0; i<m; i++){\n'
' y[i] = 0;\n'
'}\n'
'for (int i=0; i<m; i++){\n'
' for (int j=0; j<n; j++){\n'
' for (int k=0; k<o; k++){\n'
' for (int l=0; l<p; l++){\n'
' y[i] = (a[%s] + b[%s])*c[%s] + y[i];\n' %
(i * n * o * p + j * o * p + k * p + l, i * n * o * p + j * o * p +
k * p + l, j * o * p + k * p + l) + ' }\n'
' }\n'
' }\n'
'}')
c = ccode((a[i, j, k, l] + b[i, j, k, l]) * c[j, k, l], assign_to=y[i])
assert c == s
def test_contraction_structure_Mul_and_Pow():
x = IndexedBase('x')
y = IndexedBase('y')
i, j, k = Idx('i'), Idx('j'), Idx('k')
i_ji = x[i]**(y[j] * x[i])
assert get_contraction_structure(i_ji) == {None: {i_ji}}
ij_i = (x[i] * y[j])**(y[i])
assert get_contraction_structure(ij_i) == {None: {ij_i}}
j_ij_i = x[j] * (x[i] * y[j])**(y[i])
assert get_contraction_structure(j_ij_i) == {(j, ): {j_ij_i}}
j_i_ji = x[j] * x[i]**(y[j] * x[i])
assert get_contraction_structure(j_i_ji) == {(j, ): {j_i_ji}}
ij_exp_kki = x[i] * y[j] * exp(y[i] * y[k, k])
result = get_contraction_structure(ij_exp_kki)
expected = {
(i, ): {ij_exp_kki},
ij_exp_kki: [{
None: {exp(y[i] * y[k, k])},
exp(y[i] * y[k, k]): [{
None: {y[i] * y[k, k]},
y[i] * y[k, k]: [{
(k, ): {y[k, k]}
}]
}]
}]
}
assert result == expected
def test_loops():
n, m = symbols('n,m', integer=True)
A = IndexedBase('A')
x = IndexedBase('x')
y = IndexedBase('y')
i = Idx('i', m)
j = Idx('j', n)
expected = ('do i = 1, m\n'
' y(i) = 0\n'
'end do\n'
'do i = 1, m\n'
' do j = 1, n\n'
' y(i) = %(rhs)s\n'
' end do\n'
'end do')
code = fcode(A[i, j] * x[j], assign_to=y[i], source_format='free')
assert (code == expected % {
'rhs': 'y(i) + A(i, j)*x(j)'
} or code == expected % {
'rhs': 'y(i) + x(j)*A(i, j)'
} or code == expected % {
'rhs': 'x(j)*A(i, j) + y(i)'
} or code == expected % {
'rhs': 'A(i, j)*x(j) + y(i)'
})
def test_ccode_loops_multiple_terms():
n, m, o, p = symbols('n m o p', integer=True)
a = IndexedBase('a')
b = IndexedBase('b')
c = IndexedBase('c')
y = IndexedBase('y')
i = Idx('i', m)
j = Idx('j', n)
k = Idx('k', o)
s0 = ('for (int i=0; i<m; i++){\n' ' y[i] = 0;\n' '}\n')
s1 = ('for (int i=0; i<m; i++){\n'
' for (int j=0; j<n; j++){\n'
' for (int k=0; k<o; k++){\n'
' y[i] = b[j]*b[k]*c[%s] + y[i];\n' %
(i * n * o + j * o + k) + ' }\n'
' }\n'
'}\n')
s2 = ('for (int i=0; i<m; i++){\n'
' for (int k=0; k<o; k++){\n'
' y[i] = b[k]*a[%s] + y[i];\n' % (i * o + k) + ' }\n'
'}\n')
s3 = ('for (int i=0; i<m; i++){\n'
' for (int j=0; j<n; j++){\n'
' y[i] = b[j]*a[%s] + y[i];\n' % (i * n + j) + ' }\n'
'}\n')
c = ccode(b[j] * a[i, j] + b[k] * a[i, k] + b[j] * b[k] * c[i, j, k],
assign_to=y[i])
assert (c == s0 + s1 + s2 + s3[:-1] or c == s0 + s1 + s3 + s2[:-1]
or c == s0 + s2 + s1 + s3[:-1] or c == s0 + s2 + s3 + s1[:-1]
or c == s0 + s3 + s1 + s2[:-1] or c == s0 + s3 + s2 + s1[:-1])
def test_IndexedBase_subs():
i, j, k = symbols('i j k', integer=True)
a, b, c = symbols('a b c')
A = IndexedBase(a)
B = IndexedBase(b)
C = IndexedBase(c)
assert A[i] == B[i].subs(b, a)
assert isinstance(C[1].subs(C, Dict({1: 2})), type(A[1]))
def test_fcode_Indexed_without_looking_for_contraction():
len_y = 5
y = IndexedBase('y', shape=(len_y, ))
x = IndexedBase('x', shape=(len_y, ))
Dy = IndexedBase('Dy', shape=(len_y - 1, ))
i = Idx('i', len_y - 1)
e = Eq(Dy[i], (y[i + 1] - y[i]) / (x[i + 1] - x[i]))
code0 = fcode(e.rhs, assign_to=e.lhs, contract=False)
assert code0.endswith('Dy(i) = (y(i + 1) - y(i))/(x(i + 1) - x(i))')
def test_ccode_Indexed_without_looking_for_contraction():
len_y = 5
y = IndexedBase('y', shape=(len_y, ))
x = IndexedBase('x', shape=(len_y, ))
Dy = IndexedBase('Dy', shape=(len_y - 1, ))
i = Idx('i', len_y - 1)
e = Eq(Dy[i], (y[i + 1] - y[i]) / (x[i + 1] - x[i]))
code0 = ccode(e.rhs, assign_to=e.lhs, contract=False)
assert code0 == 'Dy[i] = (y[%s] - y[i])/(x[%s] - x[i]);' % (i + 1, i + 1)
def test_cse_Indexed():
len_y = 5
y = IndexedBase('y', shape=(len_y, ))
x = IndexedBase('x', shape=(len_y, ))
i = Idx('i', len_y - 1)
expr1 = (y[i + 1] - y[i]) / (x[i + 1] - x[i])
expr2 = 1 / (x[i + 1] - x[i])
replacements, reduced_exprs = cse([expr1, expr2])
assert len(replacements) > 0
def test_IndexedBase_shape():
i, j, m, n = symbols('i j m n', integer=True)
a = IndexedBase('a', shape=(m, m))
b = IndexedBase('a', shape=(m, n))
assert b.shape == Tuple(m, n)
assert a[i, j] != b[i, j]
assert a[i, j] == b[i, j].subs(n, m)
assert b.func(*b.args) == b
assert b[i, j].func(*b[i, j].args) == b[i, j]
pytest.raises(IndexException, lambda: b[i])
pytest.raises(IndexException, lambda: b[i, i, j])
def test_get_contraction_structure_complex():
x = IndexedBase('x')
y = IndexedBase('y')
A = IndexedBase('A')
i, j, k = Idx('i'), Idx('j'), Idx('k')
expr1 = y[i] + A[i, j] * x[j]
d1 = {None: {y[i]}, (j, ): {A[i, j] * x[j]}}
assert get_contraction_structure(expr1) == d1
expr2 = expr1 * A[k, i] + x[k]
d2 = {None: {x[k]}, (i, ): {expr1 * A[k, i]}, expr1 * A[k, i]: [d1]}
assert get_contraction_structure(expr2) == d2
def test_m_loops():
# Note: an Octave programmer would probably vectorize this across one or
# more dimensions. Also, size(A) would be used rather than passing in m
# and n. Perhaps users would expect us to vectorize automatically here?
# Or is it possible to represent such things using IndexedBase?
n, m = symbols('n m', integer=True)
A = IndexedBase('A')
x = IndexedBase('x')
y = IndexedBase('y')
i = Idx('i', m)
j = Idx('j', n)
result, = codegen(('mat_vec_mult', Eq(y[i], A[i, j] * x[j])),
'Octave',
header=False,
empty=False)
source = result[1]
expected = ('function y = mat_vec_mult(A, m, n, x)\n'
' for i = 1:m\n'
' y(i) = 0;\n'
' end\n'
' for i = 1:m\n'
' for j = 1:n\n'
' y(i) = %(rhs)s + y(i);\n'
' end\n'
' end\n'
'end\n')
assert (source == expected % {
'rhs': 'A(%s, %s).*x(j)' % (i, j)
} or source == expected % {
'rhs': 'x(j).*A(%s, %s)' % (i, j)
})
result, = codegen(('mat_vec_mult', Eq(y[i], A[i, j] * x[j])),
'Octave',
header=False,
empty=False,
argument_sequence=[x, A, m, n])
source = result[1]
expected = ('function y = mat_vec_mult(x, A, m, n)\n'
' for i = 1:m\n'
' y(i) = 0;\n'
' end\n'
' for i = 1:m\n'
' for j = 1:n\n'
' y(i) = %(rhs)s + y(i);\n'
' end\n'
' end\n'
'end\n')
assert (source == expected % {
'rhs': 'A(%s, %s).*x(j)' % (i, j)
} or source == expected % {
'rhs': 'x(j).*A(%s, %s)' % (i, j)
})
def test_dummy_loops():
i, m = symbols('i m', integer=True, cls=Dummy)
x = IndexedBase('x')
y = IndexedBase('y')
i = Idx(i, m)
expected = (
'do i_%(icount)i = 1, m_%(mcount)i\n'
' y(i_%(icount)i) = x(i_%(icount)i)\n'
'end do'
) % {'icount': i.label.dummy_index, 'mcount': m.dummy_index}
code = fcode(x[i], assign_to=y[i], source_format='free')
assert code == expected
def test_dummy_loops():
i, m = symbols('i m', integer=True, cls=Dummy)
x = IndexedBase('x')
y = IndexedBase('y')
i = Idx(i, m)
expected = (
'for (int i_%(icount)i=0; i_%(icount)i<m_%(mcount)i; i_%(icount)i++){\n'
' y[i_%(icount)i] = x[i_%(icount)i];\n'
'}'
) % {'icount': i.label.dummy_index, 'mcount': m.dummy_index}
code = ccode(x[i], assign_to=y[i])
assert code == expected
def test_contraction_structure_simple_Pow():
x = IndexedBase('x')
y = IndexedBase('y')
i, j = Idx('i'), Idx('j')
ii_jj = x[i, i]**y[j, j]
assert get_contraction_structure(ii_jj) == {
None: {ii_jj},
ii_jj: [{
(i, ): {x[i, i]}
}, {
(j, ): {y[j, j]}
}]
}
def test_ccode_Indexed():
n, m, o = symbols('n m o', integer=True)
i, j, k = Idx('i', n), Idx('j', m), Idx('k', o)
p = CCodePrinter()
p._not_c = set()
x = IndexedBase('x')[j]
assert p._print_Indexed(x) == 'x[j]'
A = IndexedBase('A')[i, j]
assert p._print_Indexed(A) == 'A[%s]' % (m * i + j)
B = IndexedBase('B')[i, j, k]
assert p._print_Indexed(B) == 'B[%s]' % (i * o * m + j * o + k)
assert p._not_c == set()
def test_get_indices_Pow():
x = IndexedBase('x')
y = IndexedBase('y')
A = IndexedBase('A')
i, j, k = Idx('i'), Idx('j'), Idx('k')
assert get_indices(Pow(x[i], y[j])) == ({i, j}, {})
assert get_indices(Pow(x[i, k], y[j, k])) == ({i, j, k}, {})
assert get_indices(Pow(A[i, k], y[k] + A[k, j] * x[j])) == ({i, k}, {})
assert get_indices(Pow(2, x[i])) == get_indices(exp(x[i]))
# test of a design decision, this may change:
assert get_indices(Pow(x[i], 2)) == ({
i,
}, {})
def test_IndexedBase_shape():
i, j, m, n = symbols('i j m n', integer=True)
a = IndexedBase('a', shape=(m, m))
b = IndexedBase('a', shape=(m, n))
assert b.shape == Tuple(m, n)
assert a[i, j] != b[i, j]
assert a[i, j] == b[i, j].subs({n: m})
assert b.func(b.label, shape=b.shape) == b
assert b[i, j].func(*b[i, j].args) == b[i, j]
pytest.raises(IndexException, lambda: b[i])
pytest.raises(IndexException, lambda: b[i, i, j])
F = IndexedBase('F', shape=m)
assert F.shape == Tuple(m)
assert F[i].subs({i: j}) == F[j]
pytest.raises(IndexException, lambda: F[i, j])
def test_Indexed_constructor():
i, j = symbols('i j', integer=True)
A = Indexed('A', i, j)
assert A == Indexed(Symbol('A'), i, j)
assert A == Indexed(IndexedBase('A'), i, j)
pytest.raises(TypeError, lambda: Indexed(A, i, j))
pytest.raises(IndexException, lambda: Indexed("A"))
def test_Assignment():
x, y = symbols('x, y')
A = MatrixSymbol('A', 3, 1)
mat = Matrix([1, 2, 3])
B = IndexedBase('B')
n = symbols('n', integer=True)
i = Idx('i', n)
# Here we just do things to show they don't error
Assignment(x, y)
Assignment(x, 0)
Assignment(A, mat)
Assignment(A[1, 0], 0)
Assignment(A[1, 0], x)
Assignment(B[i], x)
Assignment(B[i], 0)
# Here we test things to show that they error
# Matrix to scalar
pytest.raises(ValueError, lambda: Assignment(B[i], A))
pytest.raises(ValueError, lambda: Assignment(B[i], mat))
pytest.raises(ValueError, lambda: Assignment(x, mat))
pytest.raises(ValueError, lambda: Assignment(x, A))
pytest.raises(ValueError, lambda: Assignment(A[1, 0], mat))
# Scalar to matrix
pytest.raises(ValueError, lambda: Assignment(A, x))
pytest.raises(ValueError, lambda: Assignment(A, 0))
# Non-atomic lhs
pytest.raises(TypeError, lambda: Assignment(mat, A))
pytest.raises(TypeError, lambda: Assignment(0, x))
pytest.raises(TypeError, lambda: Assignment(x * x, 1))
pytest.raises(TypeError, lambda: Assignment(A + A, mat))
pytest.raises(TypeError, lambda: Assignment(B, 0))
def test_symbolic_indexing():
M = ImmutableDenseNDimArray([[x, y], [z, w]])
Mij = M[i, j]
assert isinstance(Mij, Indexed)
Ms = ImmutableSparseNDimArray([[2, 3 * x], [4, 5]])
msij = Ms[i, j]
assert isinstance(msij, Indexed)
for oi, oj in [(0, 0), (0, 1), (1, 0), (1, 1)]:
assert Mij.subs({i: oi, j: oj}) == M[oi, oj]
assert msij.subs({i: oi, j: oj}) == Ms[oi, oj]
A = IndexedBase('A', (0, 2))
assert A[0, 0].subs({A: M}) == x
assert A[i, j].subs({A: M}) == M[i, j]
assert M[i, j].subs({M: A}) == A[i, j]
assert isinstance(M[3 * i - 2, j], Indexed)
assert M[3 * i - 2, j].subs({i: 1, j: 0}) == M[1, 0]
assert isinstance(M[i, 0], Indexed)
assert M[i, 0].subs({i: 0}) == M[0, 0]
assert M[0, i].subs({i: 1}) == M[0, 1]
Mo = ImmutableDenseNDimArray([1, 2, 3])
assert Mo[i].subs({i: 1}) == 2
Mos = ImmutableSparseNDimArray([1, 2, 3])
assert Mos[i].subs({i: 1}) == 2
pytest.raises(ValueError, lambda: M[i, 2])
pytest.raises(ValueError, lambda: M[i, -1])
pytest.raises(ValueError, lambda: M[2, i])
pytest.raises(ValueError, lambda: M[-1, i])
pytest.raises(ValueError, lambda: Ms[i, 2])
pytest.raises(ValueError, lambda: Ms[i, -1])
pytest.raises(ValueError, lambda: Ms[2, i])
pytest.raises(ValueError, lambda: Ms[-1, i])
def test_contraction_structure_Add_in_Pow():
x = IndexedBase('x')
y = IndexedBase('y')
i, j = Idx('i'), Idx('j')
s_ii_jj_s = (1 + x[i, i])**(1 + y[j, j])
expected = {
None: {s_ii_jj_s},
s_ii_jj_s: [{
None: {1},
(i, ): {x[i, i]}
}, {
None: {1},
(j, ): {y[j, j]}
}]
}
result = get_contraction_structure(s_ii_jj_s)
assert result == expected
def test_Indexed_coeff():
N = Symbol('N', integer=True)
len_y = N
i = Idx('i', len_y - 1)
y = IndexedBase('y', shape=(len_y, ))
a = (1 / y[i + 1] * y[i]).coeff(y[i])
b = (y[i] / y[i + 1]).coeff(y[i])
assert a == b
def test_IndexedBase_sugar():
i, j = symbols('i j', integer=True)
A1 = Indexed(a, i, j)
A2 = IndexedBase(a)
assert A1 == A2[i, j]
assert A1 == A2[(i, j)]
assert A1 == A2[[i, j]]
assert A1 == A2[Tuple(i, j)]
assert all(a.is_Integer for a in A2[1, 0].args[1:])
def test_indexed_is_constant():
A = IndexedBase("A")
i, j, k = symbols("i,j,k")
assert not A[i].is_constant()
assert A[i].is_constant(j)
assert not A[1 + 2 * i, k].is_constant()
assert not A[1 + 2 * i, k].is_constant(i)
assert A[1 + 2 * i, k].is_constant(j)
assert not A[1 + 2 * i, k].is_constant(k)
def test_ccode_loops_add():
n, m = symbols('n m', integer=True)
A = IndexedBase('A')
x = IndexedBase('x')
y = IndexedBase('y')
z = IndexedBase('z')
i = Idx('i', m)
j = Idx('j', n)
s = ('for (int i=0; i<m; i++){\n'
' y[i] = x[i] + z[i];\n'
'}\n'
'for (int i=0; i<m; i++){\n'
' for (int j=0; j<n; j++){\n'
' y[i] = x[j]*A[%s] + y[i];\n' % (i * n + j) + ' }\n'
'}')
c = ccode(A[i, j] * x[j] + x[i] + z[i], assign_to=y[i])
assert c == s