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}, {})