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])) == ({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_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_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_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_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_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: {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: {S.One},
(i, ): {x[i] * y[i]}
}
assert get_contraction_structure(x[i]**y[i]) == {None: {x[i]**y[i]}}
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_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_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]}
}]
}
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_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
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_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_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]}
}]
})
