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_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_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_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_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])) == ({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 tt():
bc = VectorFieldBase()
br = OneFormFieldBase(bc)
i = Idx('i', 2)
x = VIB("x", [bc], {(0, ): 2, (1, ): 2})
y = VIB("y", [br], {(0, ): 2, (1, ): 2})
res = x[i] * y[i]
print(res)
assert (res == 8)
j = Idx('j', 2)
A = VIB("A", [bc, bc], {(0, 0): 3})
x
res = get_contraction_structure(A[i, j] * y[i])
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_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 __getitem__(self, indices, **kw_args):
self.checkShapeCompatibility(indices)
if not (is_sequence(indices)):
# only one index
index = indices
if self.shape and len(self.shape) != 1:
raise IncompatibleShapeException()
if type(index) is int:
return (self.data[(index, )])
else:
return VI(self, index, **kw_args)
else:
if self.shape and len(self.shape) != len(indices):
raise IndexException("Rank mismatch.")
if all(type(k) is int for k in indices):
# all indices are integers
return (self.data[indices])
elif all(type(i) is Idx or type(i) is int for i in indices):
# some indices are symbolic some are integers
# first handle the integer indices
fixedIndexPositions = [
i for i in range(len(indices))
if not (type(indices[i]) is Idx)
]
symbolicIndexPositions = [
i for i in range(len(indices)) if type(indices[i]) is Idx
]
freeIndices = [i for i in indices if type(i) is Idx]
freeKeys = [
k for k in self.data.keys()
if all(k[p] == indices[p] for p in fixedIndexPositions)
]
if self.bases:
remainingBases = [
b for p, b in enumerate(self.bases)
if p in symbolicIndexPositions
]
else:
remainingBases = None
newTensorIndexSet = TensorIndexSet("intermediate",
remainingBases)
newTensorIndexSet.data = {}
for k in freeKeys:
newTensorIndexSet.data[extractIndices(
k, symbolicIndexPositions)] = self.data[k]
# now check the remaining part for contraction
expr = VI(newTensorIndexSet, *freeIndices, **kw_args)
cs = get_contraction_structure(expr)
csk = list(cs.keys())
if None in csk:
#no contraction
return (expr)
else:
# compute the contracted indexed object or number.
# Since we are in __getitem__ we assume that we deal here
# only with one VI instance not with a many term expression
dummySuspects = csk[0]
#freeIndices=indices
for d in dummySuspects:
cs = get_contraction_structure(expr)
csk = cs.keys()
# there could be multiple contractions like x[_i,^i,^i,_i]
# which is the same as x[_i,^i,^j,_j] (or x[_i,^j,^i,_j]
# but we do not implement this because it is not very clear
dummyPositions = [
p for p, ind in enumerate(freeIndices) if ind == d
]
# find pairs
n = len(dummyPositions)
if n != 2:
# exclude things like [i,i,i,j] where one index occures more than 2 times
# rendering the contraction ambigous
raise (ContractionException(d, n))
p0 = dummyPositions[0]
p1 = dummyPositions[1]
b0 = newTensorIndexSet.bases[p0].dual
b1 = newTensorIndexSet.bases[p1].dual
if b1 != b0.dual:
# implement only natural pairing (up and down indices)
raise (ContractionIncompatibleBaseException(
d, p0, p1, b0, b1))
newBases = [
b for i, b in enumerate(newTensorIndexSet.bases)
if not (i in dummyPositions)
]
nonDummyPositions = [
p for p, ind in enumerate(freeIndices) if ind != d
]
if len(newBases) == 0:
# the result will be a scalar
res = sum([
newTensorIndexSet.data[k]
for k in newTensorIndexSet.data.keys()
if k[p0] == k[p1]
])
return (res)
else:
newData = {}
newKeys = set([
deleteIndices(k, dummyPositions)
for k in newTensorIndexSet.data.keys()
])
for nk in newKeys:
newData[nk] = simplify(
sum([
newTensorIndexSet.data[k]
for k in newTensorIndexSet.data.keys()
if k[p0] == k[p1] and deleteIndices(
k, dummyPositions) == nk
]))
newTensorIndexSet.data = newData
newTensorIndexSet.bases = newBases
freeIndices = [
i for p, i in enumerate(freeIndices)
if p in nonDummyPositions
]
expr = VI(newTensorIndexSet, *freeIndices,
**kw_args)
return (expr)
from sympy.tensor.index_methods import get_contraction_structure from sympy import symbols, default_sort_key from sympy.tensor import IndexedBase, Idx x, y = map(IndexedBase, ['x', 'y']) i, j = map(Idx, ['i', 'j']) print(type(x[i] * y[j])) get_contraction_structure(x[i] * y[j])