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_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_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_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])
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_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_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_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_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_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_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_multiple_terms():
n, m, o = symbols('n m o', 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 in (s0 + s1 + s2 + s3[:-1], s0 + s1 + s3 + s2[:-1],
s0 + s2 + s1 + s3[:-1], s0 + s2 + s3 + s1[:-1],
s0 + s3 + s1 + s2[:-1], s0 + s3 + s2 + s1[:-1])
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_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_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_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_Indexed_properties():
i, j = symbols('i j', integer=True)
A = Indexed('A', i, j)
assert A.rank == 2
assert A.indices == (i, j)
assert A.base == IndexedBase('A')
assert A.ranges == [None, None]
pytest.raises(IndexException, lambda: A.shape)
n, m = symbols('n m', integer=True)
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(m, n)
pytest.raises(IndexException, lambda: Indexed("A", Idx(i, m), Idx(j)).shape)
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])}
}
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 in (expected % {'rhs': 'y(i) + A(i, j)*x(j)'},
expected % {'rhs': 'y(i) + x(j)*A(i, j)'},
expected % {'rhs': 'x(j)*A(i, j) + y(i)'},
expected % {'rhs': 'A(i, j)*x(j) + y(i)'})
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)
k = Idx('k', m)
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
pytest.raises(NotImplementedError,
lambda: ccode(A[k, k]*A[i, j]*x[j], assign_to=y[i]))
def test_get_indices_mul():
x = IndexedBase('x')
y = IndexedBase('y')
i, j = Idx('i'), Idx('j')
assert get_indices(x[j] * y[i]) == ({i, j}, {})
assert get_indices(x[i] * y[j]) == ({i, j}, {})
def test_get_indices_exceptions():
x = IndexedBase('x')
y = IndexedBase('y')
i, j = Idx('i'), Idx('j')
pytest.raises(IndexConformanceException, lambda: get_indices(x[i] + y[j]))
def test_Indexed_subs():
i, j, k = symbols('i j k', integer=True)
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_scalar_broadcast():
x = IndexedBase('x')
y = IndexedBase('y')
i = Idx('i')
assert get_indices(x[i] + y[i, i]) == ({i}, {})
def test_IndexedBase_construction():
pytest.raises(TypeError, lambda: IndexedBase(1))
def test_get_indices_Indexed():
x = IndexedBase('x')
i, j = Idx('i'), Idx('j')
assert get_indices(x[i, j]) == ({i, j}, {})
assert get_indices(x[j, i]) == ({j, i}, {})
