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)))
