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