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: set([i_ji])} ij_i = (x[i] * y[j])**(y[i]) assert get_contraction_structure(ij_i) == {None: set([ij_i])} j_ij_i = x[j] * (x[i] * y[j])**(y[i]) assert get_contraction_structure(j_ij_i) == {(j, ): set([j_ij_i])} j_i_ji = x[j] * x[i]**(y[j] * x[i]) assert get_contraction_structure(j_i_ji) == {(j, ): set([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, ): set([ij_exp_kki]), ij_exp_kki: [{ None: set([exp(y[i] * y[k, k])]), exp(y[i] * y[k, k]): [{ None: set([y[i] * y[k, k]]), y[i] * y[k, k]: [{ (k, ): set([y[k, k]]) }] }] }] } assert result == expected
def test_contraction_structure_Add_in_Pow(): x = IndexedBase('x') y = IndexedBase('y') i, j, k = Idx('i'), Idx('j'), Idx('k') s_ii_jj_s = (1 + x[i, i])**(1 + y[j, j]) expected = { None: set([s_ii_jj_s]), s_ii_jj_s: [{ None: set([S.One]), (i, ): set([x[i, i]]) }, { None: set([S.One]), (j, ): set([y[j, j]]) }] } result = get_contraction_structure(s_ii_jj_s) assert result == expected s_ii_jk_s = (1 + x[i, i])**(1 + y[j, k]) expected_2 = { None: set([(x[i, i] + 1)**(y[j, k] + 1)]), s_ii_jk_s: [{ None: set([S.One]), (i, ): set([x[i, i]]) }] } result_2 = get_contraction_structure(s_ii_jk_s) assert result_2 == expected_2
def test_ufunc_support(): f = Function('f') g = Function('g') x = IndexedBase('x') y = IndexedBase('y') i, j, k = Idx('i'), Idx('j'), Idx('k') a = symbols('a') assert get_indices(f(x[i])) == (set([i]), {}) assert get_indices(f(x[i], y[j])) == (set([i, j]), {}) assert get_indices(f(y[i]) * g(x[i])) == (set(), {}) assert get_indices(f(a, x[i])) == (set([i]), {}) assert get_indices(f(a, y[i], x[j]) * g(x[i])) == (set([j]), {}) assert get_indices(g(f(x[i]))) == (set([i]), {}) assert get_contraction_structure(f(x[i])) == {None: set([f(x[i])])} assert get_contraction_structure(f(y[i]) * g(x[i])) == { (i, ): set([f(y[i]) * g(x[i])]) } assert get_contraction_structure(f(y[i]) * g(f(x[i]))) == { (i, ): set([f(y[i]) * g(f(x[i]))]) } assert get_contraction_structure(f(x[j], y[i]) * g(x[i])) == { (i, ): set([f(x[j], y[i]) * g(x[i])]) }
def test_issue_18604(): m = symbols("m") assert Idx("i", m).name == 'i' assert Idx("i", m).lower == 0 assert Idx("i", m).upper == m - 1 m = symbols("m", real=False) raises(TypeError, lambda: Idx("i", m))
def test_ufunc_support(): f = Function("f") g = Function("g") x = IndexedBase("x") y = IndexedBase("y") i, j = Idx("i"), Idx("j") a = symbols("a") assert get_indices(f(x[i])) == (set([i]), {}) assert get_indices(f(x[i], y[j])) == (set([i, j]), {}) assert get_indices(f(y[i]) * g(x[i])) == (set(), {}) assert get_indices(f(a, x[i])) == (set([i]), {}) assert get_indices(f(a, y[i], x[j]) * g(x[i])) == (set([j]), {}) assert get_indices(g(f(x[i]))) == (set([i]), {}) assert get_contraction_structure(f(x[i])) == {None: set([f(x[i])])} assert get_contraction_structure(f(y[i]) * g(x[i])) == { (i, ): set([f(y[i]) * g(x[i])]) } assert get_contraction_structure(f(y[i]) * g(f(x[i]))) == { (i, ): set([f(y[i]) * g(f(x[i]))]) } assert get_contraction_structure(f(x[j], y[i]) * g(x[i])) == { (i, ): set([f(x[j], y[i]) * g(x[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: set([ii_jj_kk]), ii_jj_kk: [ { (i, ): set([x[i, i]]) }, { None: set([y[j, j]**z[k, k]]), y[j, j]**z[k, k]: [{ (j, ): set([y[j, j]]) }, { (k, ): set([z[k, k]]) }], }, ], } assert get_contraction_structure(ii_jj_kk) == expected
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_basic(): x = IndexedBase('x') y = IndexedBase('y') i, j = Idx('i'), Idx('j') assert get_contraction_structure(x[i]*y[j]) == {None: set([x[i]*y[j]])} assert get_contraction_structure(x[i] + y[j]) == {None: set([x[i], y[j]])} assert get_contraction_structure(x[i]*y[i]) == {(i,): set([x[i]*y[i]])} assert get_contraction_structure(1 + x[i]*y[i]) == {None: set([S.One]), (i,): set([x[i]*y[i]])} assert get_contraction_structure(x[i]**y[i]) == {None: set([x[i]**y[i]])}
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]) == (set([i,]), {}) assert get_indices(y[i] + 2*A[i, j]*x[j]) == (set([i,]), {}) assert get_indices(y[i] + 2*(x[i] + A[i, j]*x[j])) == (set([i,]), {}) assert get_indices(y[i] + x[i]*(A[j, j] + 1)) == (set([i,]), {}) assert get_indices(y[i] + x[i]*x[j]*(y[j] + A[j, k]*x[k])) == (set([i,]), {})
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: set([y[i]]), (j,): set([A[i, j]*x[j]])} assert get_contraction_structure(expr1) == d1 expr2 = expr1*A[k, i] + x[k] d2 = {None: set([x[k]]), (i,): set([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_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])) == (set([i,j]), {}) assert get_indices(Pow(x[i, k], y[j, k])) == (set([i, j, k]), {}) assert get_indices(Pow(A[i, k], y[k] + A[k, j]*x[j])) == (set([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)) == (set([i,]), {})
def test_contraction_structure_simple_Pow(): x = IndexedBase('x') y = IndexedBase('y') i, j, k = Idx('i'), Idx('j'), Idx('k') ii_jj = x[i, i]**y[j, j] assert get_contraction_structure(ii_jj) == { None: set([ii_jj]), ii_jj: [ {(i,): set([x[i, i]])}, {(j,): set([y[j, j]])} ] }
def test_get_contraction_structure_basic(): x = IndexedBase("x") y = IndexedBase("y") i, j = Idx("i"), Idx("j") assert get_contraction_structure(x[i] * y[j]) == {None: set([x[i] * y[j]])} assert get_contraction_structure(x[i] + y[j]) == {None: set([x[i], y[j]])} assert get_contraction_structure(x[i] * y[i]) == { (i, ): set([x[i] * y[i]]) } assert get_contraction_structure(1 + x[i] * y[i]) == { None: set([S.One]), (i, ): set([x[i] * y[i]]), } assert get_contraction_structure(x[i]**y[i]) == {None: set([x[i]**y[i]])}
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])) == (set([i, j]), {}) assert get_indices(Pow(x[i, k], y[j, k])) == (set([i, j, k]), {}) assert get_indices(Pow(A[i, k], y[k] + A[k, j] * x[j])) == (set([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)) == (set([ i, ]), {})
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: set([y[i]]), (j, ): set([A[i, j] * x[j]])} assert get_contraction_structure(expr1) == d1 expr2 = expr1 * A[k, i] + x[k] d2 = { None: set([x[k]]), (i, ): set([expr1 * A[k, i]]), expr1 * A[k, i]: [d1] } assert get_contraction_structure(expr2) == d2
def test_Idx_inequalities_current_fails(): i14 = Idx("i14", (1, 4)) assert S(5) >= i14 assert S(5) > i14 assert not (S(5) <= i14) assert not (S(5) < i14)
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_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_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]) == (set([ i, ]), {}) assert get_indices(y[i] + 2 * A[i, j] * x[j]) == (set([ i, ]), {}) assert get_indices(y[i] + 2 * (x[i] + A[i, j] * x[j])) == (set([ i, ]), {}) assert get_indices(y[i] + x[i] * (A[j, j] + 1)) == (set([ i, ]), {}) assert get_indices(y[i] + x[i] * x[j] * (y[j] + A[j, k] * x[k])) == (set([ i, ]), {})
def test_contraction_structure_simple_Pow(): x = IndexedBase('x') y = IndexedBase('y') i, j, k = Idx('i'), Idx('j'), Idx('k') 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]}} ] } ii_jk = x[i, i]**y[j, k] assert get_contraction_structure(ii_jk) == { None: {x[i, i]**y[j, k]}, x[i, i]**y[j, k]: [ {(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_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: set([ii_jj_kk]), ii_jj_kk: [ {(i,): set([x[i, i]])}, { None: set([y[j, j]**z[k, k]]), y[j, j]**z[k, k]: [ {(j,): set([y[j, j]])}, {(k,): set([z[k, k]])} ] } ] } assert get_contraction_structure(ii_jj_kk) == expected
def test_contraction_structure_simple_Pow(): x = IndexedBase("x") y = IndexedBase("y") i, j, k = Idx("i"), Idx("j"), Idx("k") ii_jj = x[i, i]**y[j, j] assert get_contraction_structure(ii_jj) == { None: set([ii_jj]), ii_jj: [{ (i, ): set([x[i, i]]) }, { (j, ): set([y[j, j]]) }], } ii_jk = x[i, i]**y[j, k] assert get_contraction_structure(ii_jk) == { None: set([x[i, i]**y[j, k]]), x[i, i]**y[j, k]: [{ (i, ): set([x[i, 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_contraction_structure_Add_in_Pow(): x = IndexedBase('x') y = IndexedBase('y') i, j, k = Idx('i'), Idx('j'), Idx('k') s_ii_jj_s = (1 + x[i, i])**(1 + y[j, j]) expected = { None: {s_ii_jj_s}, s_ii_jj_s: [ {None: {S.One}, (i,): {x[i, i]}}, {None: {S.One}, (j,): {y[j, j]}} ] } result = get_contraction_structure(s_ii_jj_s) assert result == expected s_ii_jk_s = (1 + x[i, i]) ** (1 + y[j, k]) expected_2 = { None: {(x[i, i] + 1)**(y[j, k] + 1)}, s_ii_jk_s: [ {None: {S.One}, (i,): {x[i, i]}} ] } result_2 = get_contraction_structure(s_ii_jk_s) assert result_2 == expected_2
def test_Indexed_properties(): i, j = symbols('i j', integer=True) A = Indexed('A', i, j) assert A.name == 'A[i, j]' assert A.rank == 2 assert A.indices == (i, j) assert A.base == IndexedBase('A') assert A.ranges == [None, None] 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) raises(IndexException, lambda: Indexed("A", Idx(i, m), Idx(j)).shape)
def ufuncify(args, expr, language=None, backend='numpy', tempdir=None, flags=None, verbose=False, helpers=None): """Generates a binary function that supports broadcasting on numpy arrays. Parameters ---------- args : iterable Either a Symbol or an iterable of symbols. Specifies the argument sequence for the function. expr A SymPy expression that defines the element wise operation. language : string, optional If supplied, (options: 'C' or 'F95'), specifies the language of the generated code. If ``None`` [default], the language is inferred based upon the specified backend. backend : string, optional Backend used to wrap the generated code. Either 'numpy' [default], 'cython', or 'f2py'. tempdir : string, optional Path to directory for temporary files. If this argument is supplied, the generated code and the wrapper input files are left intact in the specified path. flags : iterable, optional Additional option flags that will be passed to the backend verbose : bool, optional If True, autowrap will not mute the command line backends. This can be helpful for debugging. helpers : iterable, optional Used to define auxillary expressions needed for the main expr. If the main expression needs to call a specialized function it should be put in the ``helpers`` iterable. Autowrap will then make sure that the compiled main expression can link to the helper routine. Items should be tuples with (<funtion_name>, <sympy_expression>, <arguments>). It is mandatory to supply an argument sequence to helper routines. Note ---- The default backend ('numpy') will create actual instances of ``numpy.ufunc``. These support ndimensional broadcasting, and implicit type conversion. Use of the other backends will result in a "ufunc-like" function, which requires equal length 1-dimensional arrays for all arguments, and will not perform any type conversions. References ---------- [1] http://docs.scipy.org/doc/numpy/reference/ufuncs.html Examples ======== >>> from sympy.utilities.autowrap import ufuncify >>> from sympy.abc import x, y >>> import numpy as np >>> f = ufuncify((x, y), y + x**2) >>> type(f) <class 'numpy.ufunc'> >>> f([1, 2, 3], 2) array([ 3., 6., 11.]) >>> f(np.arange(5), 3) array([ 3., 4., 7., 12., 19.]) For the F2Py and Cython backends, inputs are required to be equal length 1-dimensional arrays. The F2Py backend will perform type conversion, but the Cython backend will error if the inputs are not of the expected type. >>> f_fortran = ufuncify((x, y), y + x**2, backend='F2Py') >>> f_fortran(1, 2) array([ 3.]) >>> f_fortran(np.array([1, 2, 3]), np.array([1.0, 2.0, 3.0])) array([ 2., 6., 12.]) >>> f_cython = ufuncify((x, y), y + x**2, backend='Cython') >>> f_cython(1, 2) # doctest: +ELLIPSIS Traceback (most recent call last): ... TypeError: Argument '_x' has incorrect type (expected numpy.ndarray, got int) >>> f_cython(np.array([1.0]), np.array([2.0])) array([ 3.]) """ if isinstance(args, Symbol): args = (args, ) else: args = tuple(args) if language: _validate_backend_language(backend, language) else: language = _infer_language(backend) helpers = helpers if helpers else () flags = flags if flags else () if backend.upper() == 'NUMPY': # maxargs is set by numpy compile-time constant NPY_MAXARGS # If a future version of numpy modifies or removes this restriction # this variable should be changed or removed maxargs = 32 helps = [] for name, expr, args in helpers: helps.append(make_routine(name, expr, args)) code_wrapper = UfuncifyCodeWrapper(C99CodeGen("ufuncify"), tempdir, flags, verbose) if not isinstance(expr, (list, tuple)): expr = [expr] if len(expr) == 0: raise ValueError('Expression iterable has zero length') if (len(expr) + len(args)) > maxargs: raise ValueError( 'Cannot create ufunc with more than {0} total arguments: got {1} in, {2} out' .format(maxargs, len(args), len(expr))) routines = [ make_routine('autofunc{}'.format(idx), exprx, args) for idx, exprx in enumerate(expr) ] return code_wrapper.wrap_code(routines, helpers=helps) else: # Dummies are used for all added expressions to prevent name clashes # within the original expression. y = IndexedBase(Dummy()) m = Dummy(integer=True) i = Idx(Dummy(integer=True), m) f = implemented_function(Dummy().name, Lambda(args, expr)) # For each of the args create an indexed version. indexed_args = [IndexedBase(Dummy(str(a))) for a in args] # Order the arguments (out, args, dim) args = [y] + indexed_args + [m] args_with_indices = [a[i] for a in indexed_args] return autowrap(Eq(y[i], f(*args_with_indices)), language, backend, tempdir, args, flags, verbose, helpers)
def test_sympy__tensor__indexed__Indexed(): from sympy.tensor.indexed import Indexed, Idx assert _test_args(Indexed('A', Idx('i'), Idx('j')))
def test_sympy__tensor__indexed__Idx(): from sympy.tensor.indexed import Idx assert _test_args(Idx('test')) assert _test_args(Idx(1, (0, 10)))