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