def withNewIndex(factor, indexDict, indexList):
'Takes an expression, an indexDict as the one created by serchIndexInFactor, and list of (new?) indecis. Builds a new expression, same as factor, but with the indecis in indexList'
if tensor.isTensor(factor):
factor = factor.withNewIndex(
*[indexList[i] for i in indexDict['index'] ])
if 'args' in indexDict:
argsList = list(factor.args)
argsDict = indexDict['args']
for (j,arg) in enumerate(argsList):
if j in argsDict:
argsList[j] = withNewIndex(arg, argsDict[j], indexList)
factor = type(factor)(*argsList)
return factor
def serch(exp,indexList,indexPos):
indexDict = {}
if tensor.isTensor(exp):
expIndex = []
for ind in exp.index:
indexList.append(ind)
expIndex.append(indexPos)
indexPos += 1
indexDict['index'] = expIndex
if getattr(exp,'args',False):
argsDict = {}
for (j,arg) in enumerate(exp.args):
argDict,indexList,indexPos = serch(arg,indexList,indexPos)
if argDict: argsDict[j]=argDict
if argsDict: indexDict['args']=argsDict
return indexDict, indexList, indexPos
def subsIndex(exp,oldIndex,newIndex):
'''
Substitute oldIndex for newIndex in exp.
If oldIndex is given as a string, all indecis with that name will be replaced.
If oldIndex is not a sring, every index equal to oldIndex will be replaced.
'''
if not tensor.isTensor(exp):
if not getattr(exp, 'args', False):
return exp
return type(exp)(*[subsIndex(arg,oldIndex,newIndex)
for arg in exp.args])
newIndexList = []
for ind in exp.index:
if ind==oldIndex or getattr(ind,'name',None)==oldIndex:
newIndexList.append(newIndex)
else:
newIndexList.append(ind)
if not getattr(exp,'args',False):
return exp.withNewIndex(*newIndexList)
return type(type(exp))(*newIndexList)(
*[subsIndex(arg,oldIndex,newIndex) for arg in exp.args])