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_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_Idx_subs():
i, a, b = symbols('i a b', integer=True)
assert Idx(i, a).subs(a, b) == Idx(i, b)
assert Idx(i, a).subs(i, b) == Idx(b, a)
assert Idx(i).subs(i, 2) == Idx(2)
assert Idx(i, a).subs(a, 2) == Idx(i, 2)
assert Idx(i, (a, b)).subs(i, 2) == Idx(2, (a, b))
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_Idx_func_args():
i, a, b = symbols('i a b', integer=True)
ii = Idx(i)
assert ii.func(*ii.args) == ii
ii = Idx(i, a)
assert ii.func(*ii.args) == ii
ii = Idx(i, (a, b))
assert ii.func(*ii.args) == ii
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_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_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_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_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_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_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
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_Indexed_shape_precedence():
i, j = symbols('i j', integer=True)
o, p = symbols('o p', integer=True)
n, m = symbols('n m', integer=True)
a = IndexedBase('a', shape=(o, p))
assert a.shape == Tuple(o, p)
assert Indexed(
a, Idx(i, m), Idx(j, n)).ranges == [Tuple(0, m - 1), Tuple(0, n - 1)]
assert Indexed(a, Idx(i, m), Idx(j, n)).shape == Tuple(o, p)
assert Indexed(
a, Idx(i, m), Idx(j)).ranges == [Tuple(0, m - 1), Tuple(None, None)]
assert Indexed(a, Idx(i, m), Idx(j)).shape == Tuple(o, p)
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_get_contraction_structure_basic():
x = IndexedBase('x')
y = IndexedBase('y')
i, j = Idx('i'), Idx('j')
f = Function('f')
assert get_contraction_structure(x[i] * y[j]) == {None: {x[i] * y[j]}}
assert get_contraction_structure(x[i] + y[j]) == {None: {x[i], y[j]}}
assert get_contraction_structure(x[i] * y[i]) == {(i, ): {x[i] * y[i]}}
assert get_contraction_structure(1 + x[i] * y[i]) == {
None: {1},
(i, ): {x[i] * y[i]}
}
assert get_contraction_structure(x[i]**y[i]) == {None: {x[i]**y[i]}}
assert (get_contraction_structure(f(x[i, i])) == {
None: {f(x[i, i])},
f(x[i, i]): [{
(i, ): {x[i, i]}
}]
})
def test_contraction_structure_Pow_in_Pow():
x = IndexedBase('x')
y = IndexedBase('y')
z = IndexedBase('z')
i, j, k = Idx('i'), Idx('j'), Idx('k')
ii_jj_kk = x[i, i]**y[j, j]**z[k, k]
expected = {
None: {ii_jj_kk},
ii_jj_kk: [{
(i, ): {x[i, i]}
}, {
None: {y[j, j]**z[k, k]},
y[j, j]**z[k, k]: [{
(j, ): {y[j, j]}
}, {
(k, ): {z[k, k]}
}]
}]
}
assert get_contraction_structure(ii_jj_kk) == expected
def test_get_indices_add():
x = IndexedBase('x')
y = IndexedBase('y')
A = IndexedBase('A')
i, j, k = Idx('i'), Idx('j'), Idx('k')
assert get_indices(x[i] + 2 * y[i]) == ({
i,
}, {})
assert get_indices(y[i] + 2 * A[i, j] * x[j]) == ({
i,
}, {})
assert get_indices(y[i] + 2 * (x[i] + A[i, j] * x[j])) == ({
i,
}, {})
assert get_indices(y[i] + x[i] * (A[j, j] + 1)) == ({
i,
}, {})
assert get_indices(y[i] + x[i] * x[j] * (y[j] + A[j, k] * x[k])) == ({
i,
}, {})
def test_ccode_inline_function():
g = implemented_function('g', Lambda(x, 2 * x))
assert ccode(g(x)) == '2*x'
g = implemented_function('g', Lambda(x, 2 * x / Catalan))
assert ccode(
g(x)) == 'double const Catalan = %s;\n2*x/Catalan' % Catalan.evalf()
A = IndexedBase('A')
i = Idx('i', symbols('n', integer=True))
g = implemented_function('g', Lambda(x, x * (1 + x) * (2 + x)))
assert ccode(g(A[i]),
assign_to=A[i]) == ('for (int i=0; i<n; i++){\n'
' A[i] = (A[i] + 1)*(A[i] + 2)*A[i];\n'
'}')
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_Idx_func_args():
i, a, b = symbols('i a b', integer=True)
ii = Idx(i)
assert ii.func(*ii.args) == ii
ii = Idx(i, a)
assert ii.func(*ii.args) == ii
ii = Idx(i, (a, b))
assert ii.func(*ii.args) == ii
def test_inline_function():
g = implemented_function('g', Lambda(x, 2 * x))
assert fcode(g(x)) == ' 2*x'
g = implemented_function('g', Lambda(x, 2 * pi / x))
assert fcode(g(x)) == (' parameter (pi = 3.14159265358979d0)\n'
' 2*pi/x')
A = IndexedBase('A')
i = Idx('i', symbols('n', integer=True))
g = implemented_function('g', Lambda(x, x * (1 + x) * (2 + x)))
assert fcode(
g(A[i]),
assign_to=A[i]) == (' do i = 1, n\n'
' A(i) = (A(i) + 1)*(A(i) + 2)*A(i)\n'
' end do')
def test_ufunc_support():
f = Function('f')
g = Function('g')
x = IndexedBase('x')
y = IndexedBase('y')
i, j = Idx('i'), Idx('j')
assert get_indices(f(x[i])) == ({i}, {})
assert get_indices(f(x[i], y[j])) == ({i, j}, {})
assert get_indices(f(y[i]) * g(x[i])) == (set(), {})
assert get_indices(f(a, x[i])) == ({i}, {})
assert get_indices(f(a, y[i], x[j]) * g(x[i])) == ({j}, {})
assert get_indices(g(f(x[i]))) == ({i}, {})
assert get_contraction_structure(f(x[i])) == {None: {f(x[i])}}
assert get_contraction_structure(f(y[i]) * g(x[i])) == {
(i, ): {f(y[i]) * g(x[i])}
}
assert get_contraction_structure(f(y[i]) * g(f(x[i]))) == {
(i, ): {f(y[i]) * g(f(x[i]))}
}
assert get_contraction_structure(f(x[j], y[i]) * g(x[i])) == {
(i, ): {f(x[j], y[i]) * g(x[i])}
}