def test_ccode_loops_multiple_contractions():
from diofant.tensor import IndexedBase, Idx
from diofant import symbols
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_m_tensor_loops_multiple_contractions():
# see comments in previous test about vectorizing
from diofant.tensor import IndexedBase, Idx
from diofant import symbols
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_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_jscode_loops_addfactor():
from diofant.tensor import IndexedBase, Idx
from diofant import symbols
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 (var i=0; i<m; i++){\n'
' y[i] = 0;\n'
'}\n'
'for (var i=0; i<m; i++){\n'
' for (var j=0; j<n; j++){\n'
' for (var k=0; k<o; k++){\n'
' for (var 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 = jscode((a[i, j, k, l] + b[i, j, k, l])*c[j, k, l], assign_to=y[i])
assert c == s
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?
from diofant.tensor import IndexedBase, Idx
from diofant import symbols
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)
})
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():
from diofant.tensor import IndexedBase, Idx
from diofant import symbols
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_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_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_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 (var 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 = jscode(x[i], assign_to=y[i])
assert code == expected
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_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_jscode_Indexed():
from diofant.tensor import IndexedBase, Idx
from diofant import symbols
n, m, o = symbols('n m o', integer=True)
i, j, k = Idx('i', n), Idx('j', m), Idx('k', o)
p = JavascriptCodePrinter()
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_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_jscode_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 (var i=0; i<m; i++){\n'
' y[i] = 0;\n'
'}\n'
'for (var i=0; i<m; i++){\n'
' for (var j=0; j<n; j++){\n'
' y[i] = x[j]*A[n*i + j] + y[i];\n'
' }\n'
'}'
)
c = jscode(A[i, j]*x[j], assign_to=y[i])
assert c == s
def test_jscode_inline_function():
x = symbols('x')
g = implemented_function('g', Lambda(x, 2*x))
assert jscode(g(x)) == "2*x"
g = implemented_function('g', Lambda(x, 2*x/Catalan))
assert jscode(g(x)) == "var Catalan = %s;\n2*x/Catalan" % Catalan.n()
A = IndexedBase('A')
i = Idx('i', symbols('n', integer=True))
g = implemented_function('g', Lambda(x, x*(1 + x)*(2 + x)))
assert jscode(g(A[i]), assign_to=A[i]) == (
"for (var i=0; i<n; i++){\n"
" A[i] = (A[i] + 1)*(A[i] + 2)*A[i];\n"
"}"
)
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_jscode_loops_add():
from diofant.tensor import IndexedBase, Idx
from diofant import symbols
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 (var i=0; i<m; i++){\n'
' y[i] = x[i] + z[i];\n'
'}\n'
'for (var i=0; i<m; i++){\n'
' for (var j=0; j<n; j++){\n'
' y[i] = x[j]*A[n*i + j] + y[i];\n'
' }\n'
'}'
)
c = jscode(A[i, j]*x[j] + x[i] + z[i], assign_to=y[i])
assert c == s