def test_toTensor(self): self.showTestCaseBegin("diagonal tensor toTensor") a = DiagonalTensor(shape = (3, 3, 3), labels = ['a', 'b', 'c'], data = np.array([1, 2, 3])) aTensor = a.toTensor() aRealTensor = np.zeros((3, 3, 3)) for i in range(3): aRealTensor[(i, i, i)] = i + 1 self.assertTrue(funcs.floatArrayEqual(aTensor, aRealTensor)) self.showTestCaseEnd("diagonal tensor toTensor")
def test_toTensorLike(self): a = DiagonalTensor(shape = (3, 3), tensorLikeFlag = False) self.assertIsNotNone(a.a) aLike = a.toTensorLike() self.assertIsNone(aLike.a) self.assertTrue(aLike.tensorLikeFlag) self.assertTupleEqual(aLike.shape, (3, 3)) self.assertListEqual(aLike.labels, ['a', 'b'])
def test_diagonalTensorLike(self): diagonalTensor = DiagonalTensor(shape = (3, 3), tensorLikeFlag = True) self.assertEqual(diagonalTensor.dim, 2) self.assertTupleEqual(diagonalTensor.shape, (3, 3)) self.assertListEqual(diagonalTensor.labels, ['a', 'b']) self.assertIsNone(diagonalTensor.a) diagonalTensor = DiagonalTensor(data = np.zeros((3, 3)), tensorLikeFlag = True) # only data is given self.assertTrue(diagonalTensor.diagonalFlag) self.assertEqual(diagonalTensor.dim, 2) self.assertTupleEqual(diagonalTensor.shape, (3, 3)) self.assertListEqual(diagonalTensor.labels, ['a', 'b']) self.assertIsNone(diagonalTensor.a) diagonalTensor = DiagonalTensor(data = np.ones((3, 3)), labels = ['up', 'down'], tensorLikeFlag = True) # only data is given self.assertTrue(diagonalTensor.diagonalFlag) self.assertEqual(diagonalTensor.dim, 2) self.assertTupleEqual(diagonalTensor.shape, (3, 3)) self.assertListEqual(diagonalTensor.labels, ['up', 'down']) self.assertIsNone(diagonalTensor.a) self.assertRaises(AssertionError, diagonalTensor.norm) self.assertRaises(AssertionError, diagonalTensor.trace) self.assertRaises(AssertionError, diagonalTensor.toTensor) self.assertRaises(AssertionError, lambda: diagonalTensor.toMatrix(rows = None, cols = None)) self.assertRaises(AssertionError, diagonalTensor.toVector)
def test_DiagonalTensorCopy(self): self.showTestCaseBegin("diagonal tensor copy") aData = np.ones((3, 3)) a = DiagonalTensor(shape = (3, 3), labels = ['a', 'b'], data = aData) aData[(0, 0)] = 2.0 self.assertEqual(a.a[0], 1.0) a = DiagonalTensor(shape = (3, 3), labels = ['a', 'b']) b = a.copy() b.a[0] = 2.0 self.assertEqual(a.a[0], 1.0) self.assertEqual(b.a[0], 2.0) self.assertEqual(a.diagonalFlag, b.diagonalFlag) self.assertEqual(a.tensorLikeFlag, b.tensorLikeFlag) a.renameLabel('a', 'c') self.assertListEqual(b.labels, ['a', 'b']) self.showTestCaseEnd('diagonal tensor copy')
def test_diagonalTensor(self): self.showTestCaseBegin("diagonal tensor") diagonalTensor = DiagonalTensor(shape = (3, 3)) self.assertEqual(diagonalTensor.dim, 2) self.assertTupleEqual(diagonalTensor.shape, (3, 3)) self.assertListEqual(diagonalTensor.labels, ['a', 'b']) diagonalTensor = DiagonalTensor(data = np.zeros((3, 3))) # only data is given self.assertTrue(diagonalTensor.diagonalFlag) self.assertEqual(diagonalTensor.dim, 2) self.assertTupleEqual(diagonalTensor.shape, (3, 3)) self.assertListEqual(diagonalTensor.labels, ['a', 'b']) diagonalTensor = DiagonalTensor(data = np.ones((3, 3)), labels = ['up', 'down']) # only data is given self.assertTrue(diagonalTensor.diagonalFlag) self.assertEqual(diagonalTensor.dim, 2) self.assertTupleEqual(diagonalTensor.shape, (3, 3)) self.assertListEqual(diagonalTensor.labels, ['up', 'down']) self.assertEqual(diagonalTensor.norm(), np.sqrt(3.0)) self.assertEqual(diagonalTensor.trace(), 3.0) self.showTestCaseEnd("diagonal tensor")
def nothingErrorFunc(): _ = DiagonalTensor()
def dataShapeErrorFunc3(): _ = DiagonalTensor(data = np.zeros((2, 2, 3)))
def dataNoneErrorFunc(): _ = DiagonalTensor(labels = ['a', 'b', 'c'], data = None)
def dataShapeErrorFunc2(): _ = DiagonalTensor(labels = ['a', 'b', 'c'], data = np.zeros((2, 2, 3)))
def labelsSizeNotEqualFunc(): _ = DiagonalTensor(legs = [legA, legB], labels = ['a'])
def shapeOrderNotEqualFunc(): _ = DiagonalTensor(legs = [legA, legB], shape = (6, 5))
def labelsShortFunc(): _ = DiagonalTensor(shape = (2, 2), labels = ['a1'])
def labelsOrderNotEqualFunc(): _ = DiagonalTensor(legs = [legA, legB], labels = ['b', 'a'])
def dimensionless1DDataErrorFunc2(): _ = DiagonalTensor(labels = [], data = np.zeros(3))
def shapeNotEqualFunc(): _ = DiagonalTensor(shape = (2, 3))
def labelsShapeNotCompatibleFunc(): _ = DiagonalTensor(legs = [legA, legB], labels = ['a'], data = np.zeros((5, 5)))
def dataShapeNotEqualFunc(): _ = DiagonalTensor(legs = [legA, legB], data = np.zeros((5, 6)))
def dataDimNotEqualFunc(): _ = DiagonalTensor(legs = [legA, legB], data = np.zeros((5, 6, 7)))
def test_extremeContraction(self): self.showTestCaseBegin("diagonal tensor extreme contraction") aData = np.random.random_sample(2) aTensorData = np.array([[[aData[0], 0], [0, 0]], [[0, 0], [0, aData[1]]]]) bData = np.random.random_sample((2, 2)) a = DiagonalTensor(shape = (2, 2, 2), labels = ['a1', 'a2', 'a3'], data = aData) b = Tensor(shape = (2, 2), labels = ['b1', 'b2'], data = bData) makeLink('a1', 'b1', a, b) makeLink('a3', 'b2', a, b) res1 = a @ b a = Tensor(shape = (2, 2, 2), labels = ['a1', 'a2', 'a3'], data = aTensorData) b = Tensor(shape = (2, 2), labels = ['b1', 'b2'], data = bData) makeLink('a1', 'b1', a, b) makeLink('a3', 'b2', a, b) res2 = b @ a self.assertTrue(funcs.compareLists(list(res1.labels), ['a2'])) self.assertTrue(funcs.compareLists(list(res2.labels), ['a2'])) # self.assertListEqual(list(res1.a), list(res2.a)) self.assertTrue(funcs.floatArrayEqual(res1.a, res2.a)) aData = np.random.random_sample(2) aTensorData = np.array([[aData[0], 0], [0, aData[1]]]) bData = np.random.random_sample((2, 2)) a = DiagonalTensor(shape = (2, 2), labels = ['a1', 'a2'], data = aData) b = Tensor(shape = (2, 2), labels = ['b1', 'b2'], data = bData) makeLink('a1', 'b1', a, b) makeLink('a2', 'b2', a, b) res1 = a @ b a = Tensor(shape = (2, 2), labels = ['a1', 'a2'], data = aTensorData) b = Tensor(shape = (2, 2), labels = ['b1', 'b2'], data = bData) makeLink('a1', 'b1', a, b) makeLink('a2', 'b2', a, b) res2 = b @ a self.assertTrue(funcs.compareLists(res1.labels, [])) self.assertTrue(funcs.compareLists(res2.labels, [])) # self.assertListEqual(list(res1.single(), list(res2.a)) self.assertTrue(funcs.floatEqual(res1.single(), res2.single())) aData = np.random.random_sample(2) aTensorData = np.array([[[aData[0], 0], [0, 0]], [[0, 0], [0, aData[1]]]]) bData = np.random.random_sample((2, 2, 2)) a = DiagonalTensor(shape = (2, 2, 2), labels = ['a1', 'a2', 'a3'], data = aData) b = Tensor(shape = (2, 2, 2), labels = ['b1', 'b2', 'b3'], data = bData) makeLink('a1', 'b1', a, b) makeLink('a3', 'b2', a, b) res1 = a @ b a = Tensor(shape = (2, 2, 2), labels = ['a1', 'a2', 'a3'], data = aTensorData) b = Tensor(shape = (2, 2, 2), labels = ['b1', 'b2', 'b3'], data = bData) makeLink('a1', 'b1', a, b) makeLink('a3', 'b2', a, b) res2 = b @ a self.assertTrue(funcs.compareLists(list(res1.labels), ['a2', 'b3'])) self.assertTrue(funcs.compareLists(list(res2.labels), ['a2', 'b3'])) # print(res1.labels, res2.labels) # print(res1.a, res2.a) res2.reArrange(res1.labels) self.assertTrue(funcs.floatArrayEqual(res1.a, res2.a)) # self.assertListEqual(list(np.ravel(res1.a)), list(np.ravel(res2.a))) aData = np.random.random_sample(2) aTensorData = np.array([[[aData[0], 0], [0, 0]], [[0, 0], [0, aData[1]]]]) bData = np.random.random_sample((2, 2, 2)) a = DiagonalTensor(shape = (2, 2, 2), labels = ['a1', 'a2', 'a3'], data = aData) b = Tensor(shape = (2, 2, 2), labels = ['b1', 'b2', 'b3'], data = bData) makeLink('a1', 'b1', a, b) # makeLink('a3', 'b2', a, b) res1 = a @ b a = Tensor(shape = (2, 2, 2), labels = ['a1', 'a2', 'a3'], data = aTensorData) b = Tensor(shape = (2, 2, 2), labels = ['b1', 'b2', 'b3'], data = bData) makeLink('a1', 'b1', a, b) # makeLink('a3', 'b2', a, b) res2 = b @ a # print('ndEye(2, 2) = {}'.format(funcs.ndEye(2, 2))) # print('ndEye(1, 2) = {}'.format(funcs.ndEye(1, 2))) # print('ndEye(3, 2) = {}'.format(funcs.ndEye(3, 2))) self.assertTrue(funcs.compareLists(list(res1.labels), ['a2', 'a3', 'b2', 'b3'])) self.assertTrue(funcs.compareLists(list(res2.labels), ['a2', 'a3', 'b2', 'b3'])) # print(res1.labels, res2.labels) res2.reArrange(res1.labels) # print('res1 = {}, res2 = {}'.format(res1.a, res2.a)) self.assertTrue(funcs.floatArrayEqual(res1.a, res2.a)) # self.assertListEqual(list(np.ravel(res1.a)), list(np.ravel(res2.a))) aData = np.random.random_sample(2) aTensorData = np.array([[[aData[0], 0], [0, 0]], [[0, 0], [0, aData[1]]]]) bData = np.random.random_sample((2, 4, 7)) a = DiagonalTensor(shape = (2, 2, 2), labels = ['a1', 'a2', 'a3'], data = aData) b = Tensor(shape = (2, 4, 7), labels = ['b1', 'b2', 'b3'], data = bData) makeLink('a1', 'b1', a, b) # makeLink('a3', 'b2', a, b) res1 = a @ b a = Tensor(shape = (2, 2, 2), labels = ['a1', 'a2', 'a3'], data = aTensorData) b = Tensor(shape = (2, 4, 7), labels = ['b1', 'b2', 'b3'], data = bData) makeLink('a1', 'b1', a, b) # makeLink('a3', 'b2', a, b) res2 = b @ a # print('ndEye(2, 2) = {}'.format(funcs.ndEye(2, 2))) # print('ndEye(1, 2) = {}'.format(funcs.ndEye(1, 2))) # print('ndEye(3, 2) = {}'.format(funcs.ndEye(3, 2))) self.assertTrue(funcs.compareLists(list(res1.labels), ['a2', 'a3', 'b2', 'b3'])) self.assertTrue(funcs.compareLists(list(res2.labels), ['a2', 'a3', 'b2', 'b3'])) # print(res1.labels, res2.labels) res2.reArrange(res1.labels) # print('res1 = {}, res2 = {}'.format(res1.a, res2.a)) self.assertTrue(funcs.floatArrayEqual(res1.a, res2.a)) self.showTestCaseEnd("diagonal tensor extreme contraction")
def test_diagonalRename(self): a = DiagonalTensor(shape = (3, 3, 3), labels = ['abc', 'def', 'abc']) a.renameLabels(['abc', 'abc'], ['abc1', 'abc2']) self.assertTrue(funcs.compareLists(a.labels, ['abc1', 'abc2', 'def']))
def test_contraction(self): self.showTestCaseBegin("diagonal tensor contraction") # print('Begin test diagonalTensor contraction.') a = DiagonalTensor(shape = (2, 2), labels = ['a', 'b']) b = Tensor(shape = (2, ), labels = ['x']) c = Tensor(shape = (2, ), labels = ['y']) makeLink('a', 'x', a, b) makeLink('b', 'y', a, c) seq = generateOptimalSequence([a, b, c], typicalDim = 10) # print('optimal sequence = {}'.format(seq)) prod, cost = contractAndCostWithSequence([a, b, c], seq = seq) # print('cost = {}'.format(cost)) # prod = contractTensorList([a, b, c], outProductWarning = False) self.assertTrue(funcs.compareLists(prod.labels, [])) self.assertListEqual(seq, [(0, 2), (1, 0)]) self.assertEqual(cost, 4.0) # if we use Tensor instead of DiagonalTensor for a # then the cost should be 12.0, and the order should be (1, 2), (0, 1) # the optimal cost of diagonal tensors can be achieved if we use diagonal nature for contraction a = DiagonalTensor(shape = (2, 2, 2), labels = ['a', 'b', 'c']) b = DiagonalTensor(shape = (2, 2), labels = ['x', 'y']) makeLink('a', 'x', a, b) prod, cost = contractAndCostWithSequence([a, b]) self.assertEqual(cost, 2) self.assertTrue(funcs.compareLists(prod.labels, ['b', 'c', 'y'])) aData = np.array([[[1, 0], [0, 0]], [[0, 0], [0, 3]]]) bData = np.random.random_sample(2) cData = np.random.random_sample(2) a = DiagonalTensor(data = aData, labels = ['a', 'b', 'c']) b = Tensor(data = bData, labels = ['x']) c = Tensor(data = cData, labels = ['y']) makeLink('a', 'x', a, b) makeLink('b', 'y', a, c) res1, _ = contractAndCostWithSequence([a, b, c]) # print('seq = {}'.format(generateOptimalSequence([a, b, c]))) a = Tensor(data = aData, labels = ['a', 'b', 'c']) b = Tensor(data = bData, labels = ['x']) c = Tensor(data = cData, labels = ['y']) makeLink('a', 'x', a, b) makeLink('b', 'y', a, c) res2, _ = contractAndCostWithSequence([a, b, c]) # self.assertListEqual(list(res1.a), list(res2.a)) self.assertTrue(funcs.floatArrayEqual(res1.a, res2.a)) # print(cost1, cost2) # print(res1.a, res2.a) # test diagonal tensor contraction a = DiagonalTensor(shape = (2, 2), labels = ['a1', 'a2']) b = DiagonalTensor(shape = (2, 2, 2), labels = ['b1', 'b2', 'b3']) makeLink('a1', 'b1', a, b) res = a @ b self.assertTupleEqual(res.shape, (2, 2, 2)) self.assertEqual(res.dim, 3) self.assertTrue(res.diagonalFlag) self.assertTrue((res.a == np.ones(2)).all()) # test for diagonal * diagonal contraction cost(just O(length)) a = DiagonalTensor(shape = (2, 2), labels = ['a1', 'a2']) b = DiagonalTensor(shape = 2, labels = ['b1', 'b2']) # deduce dim makeLink('a1', 'b2', a, b) cost, _ = contractCost(a, b) self.assertEqual(cost, 2.0) res, cost = contractAndCostWithSequence([a, b]) self.assertEqual(res.dim, 2) self.assertEqual(res._length, 2) self.assertTupleEqual(res.shape, (2, 2)) self.assertEqual(cost, 2.0) self.assertTrue(res.diagonalFlag) a = DiagonalTensor(shape = (2, 2), labels = ['a1', 'a2']) b = Tensor(shape = (2, 3, 3), labels = ['b1', 'b2', 'b3']) # deduce dim makeLink('a1', 'b1', a, b) cost, _ = contractCost(a, b) self.assertEqual(cost, 18.0) res, cost = contractAndCostWithSequence([a, b]) # print(res) self.assertEqual(res.dim, 3) self.assertTrue(funcs.compareLists(list(res.shape), [2, 3, 3])) self.assertEqual(cost, 18.0) self.showTestCaseEnd("diagonal tensor contraction")
def test_outerProduct(self): a = Tensor(shape=(2, ), labels=['a']) b = Tensor(shape=(2, ), labels=['b']) op = contractTwoTensors(a, b, outProductWarning=False) self.assertTrue(funcs.compareLists(op.labels, ['a', 'b'])) a = Tensor(shape=(2, 2, 2), labels=['a', 'b', 'c']) b = Tensor(shape=(2, ), labels=['x']) c = Tensor(shape=(2, ), labels=['y']) makeLink('a', 'x', a, b) makeLink('b', 'y', a, c) prod = contractTensorList([a, b, c], outProductWarning=False) self.assertTrue(funcs.compareLists(prod.labels, ['c'])) dataA = np.random.random_sample((2, 2)) dataB = np.random.random_sample((3, 3)) a = Tensor(shape=(2, 2), labels=['a1', 'a2'], data=dataA) b = Tensor(shape=(3, 3), labels=['b1', 'b2'], data=dataB) prod = contractTensorList([a, b], outProductWarning=False) prod.reArrange(['a1', 'a2', 'b1', 'b2']) res = np.zeros((2, 2, 3, 3)) for i in range(2): for j in range(2): for k in range(3): for l in range(3): res[(i, j, k, l)] = dataA[(i, j)] * dataB[(k, l)] # print(res, prod.a) self.assertTrue(funcs.floatArrayEqual(res, prod.a)) a = Tensor(shape=(2, 2), labels=['a1', 'a2'], data=dataA) b = DiagonalTensor(shape=(3, 3), labels=['b1', 'b2'], data=np.diag(dataB)) prod = contractTensorList([a, b], outProductWarning=False) prodData = prod.toTensor(['a1', 'a2', 'b1', 'b2']) # prod.reArrange(['a1', 'a2', 'b1', 'b2']) res = np.zeros((2, 2, 3, 3)) for i in range(2): for j in range(2): for k in range(3): # for l in range(3): res[(i, j, k, k)] = dataA[(i, j)] * dataB[(k, k)] # print(res, prod.a) self.assertTrue(funcs.floatArrayEqual(res, prodData)) dataA = np.random.random_sample((2, 2)) dataB = np.random.random_sample(3) a = Tensor(shape=(2, 2), labels=['a1', 'a2'], data=dataA) b = DiagonalTensor(shape=(3, 3, 3), labels=['b1', 'b2', 'b3'], data=dataB) prod = contractTensorList([a, b], outProductWarning=False) prodData = prod.toTensor(['a1', 'a2', 'b1', 'b2', 'b3']) # prod.reArrange(['a1', 'a2', 'b1', 'b2']) res = np.zeros((2, 2, 3, 3, 3)) for i in range(2): for j in range(2): for k in range(3): # for l in range(3): res[(i, j, k, k, k)] = dataA[(i, j)] * dataB[k] # print(res, prod.a) self.assertTrue(funcs.floatArrayEqual(res, prodData))
def labelsLongFunc(): _ = DiagonalTensor(shape = (2, 2), labels = ['a', 'b', 'c'])
def shapeSizeNotEqualFunc(): _ = DiagonalTensor(legs = [legA, legB], shape = (5, 6, 7))
def dataDimErrorFunc(): _ = DiagonalTensor(shape = (2, 2), data = np.zeros((2, 2, 2)))
def dataShapeErrorFunc(): _ = DiagonalTensor(shape = (2, 2), data = np.zeros((2, 3))) # no error in 9be9325, newly added
def test_diagonalSumOut(self): a = DiagonalTensor(shape = (3, 3, 3), labels = ['abc', 'def', 'abc']) a.sumOutLegByLabel(['abc', 'abc']) self.assertTrue(funcs.compareLists(a.labels, ['def']))
def dataDimErrorFunc2(): _ = DiagonalTensor(labels = ['a', 'b', 'c'], data = np.zeros((2, 2)))
def contractTwoTensors(ta, tb, bonds=None, outProductWarning=True): """ Calculate the result of contraction of two tensors. Parameters ---------- ta, tb : Tensor bonds : list of Bond, optional If given, then contract only over the given bonds but not all bonds shared. outProductWarning : bool, default True Whether to make a warning message of outer product, used for debug. Returns ------- Tensor The contraction result of ta and tb. """ # contract between bonds(if not, then find the shared legs) # this requires that: tensor contraction happens in-place # after we prepare some tensor, we must create tensors to make links if (not ta.diagonalFlag) and (tb.diagonalFlag): return contractTwoTensors(tb, ta, bonds=bonds, outProductWarning=outProductWarning) if (bonds is None): bonds = shareBonds(ta, tb) if (ta.tensorLikeFlag != tb.tensorLikeFlag): raise TypeError( funcs.errorMessage( "ta and tb must be the same type(tensor/tensorlike): {} and {} gotten." .format(ta.typeName, tb.typeName), location='CTL.tensor.contract.contractTwoTensors')) tensorLikeContract = ta.tensorLikeFlag if (len(bonds) == 0): if (outProductWarning): warnings.warn( '{} and {} do not share same label, do out product'.format( ta, tb), RuntimeWarning) # aMatrix = ta.toMatrix(rows = ta.legs, cols = []) # bMatrix = tb.toMatrix(rows = [], cols = tb.legs) # data = xplib.xp.matmul(aMatrix, bMatrix) labels = ta.labels + tb.labels shape = ta.shape + tb.shape legs = ta.legs + tb.legs if (ta.diagonalFlag and tb.diagonalFlag): if (tensorLikeContract): return DiagonalTensor(labels=labels, data=None, shape=shape, legs=legs, tensorLikeFlag=True) else: return DiagonalTensor(labels=labels, data=ta.a * tb.a, shape=shape, legs=legs) elif (ta.diagonalFlag): if (tensorLikeContract): return Tensor(labels=labels, data=None, shape=shape, legs=legs, tensorLikeFlag=True) else: data = xplib.xp.zeros(shape, dtype=ta.a.dtype) einsumStr = ('j' * ta.dim) + '...->j...' outerData = xplib.xp.multiply.outer(ta.a, tb.a) xplib.xp.einsum(einsumStr, data)[...] = outerData return Tensor(labels=labels, data=data, shape=shape, legs=legs) else: if (tensorLikeContract): return Tensor(labels=labels, data=None, shape=shape, legs=legs, tensorLikeFlag=True) else: return Tensor(labels=labels, data=xplib.xp.multiply.outer(ta.a, tb.a), shape=shape, legs=legs) # aVector = ta.toVector() # bVector = tb.toVector() # data = xplib.xp.outer(aVector, bVector) # data = xplib.xp.reshape(data, shape) # return Tensor(labels = labels, data = data, legs = legs) contractALegs = [bond.legs[0] for bond in bonds] contractBLegs = [bond.legs[1] for bond in bonds] taRemainLegs = ta.complementLegs(contractALegs) tbRemainLegs = tb.complementLegs(contractBLegs) newLegs = taRemainLegs + tbRemainLegs newShape = tuple([leg.dim for leg in newLegs]) if (ta.diagonalFlag) and (tb.diagonalFlag): # return a diagonal tensor if (tensorLikeContract): return DiagonalTensor(shape=newShape, data=None, legs=newLegs, tensorLikeFlag=True) if (len(newLegs) != 0): return DiagonalTensor(shape=newShape, data=ta.a * tb.a, legs=newLegs) else: return DiagonalTensor(data=xplib.xp.array(xplib.xp.sum(ta.a * tb.a))) if (ta.diagonalFlag): if (tensorLikeContract): return Tensor(shape=newShape, data=None, legs=newLegs, tensorLikeFlag=True) # then tb is not diagonal tensor # 1. calculate the core with broadcast # 2. calculate the real tensor with xplib.xp.outer # how to broadcast? # we need to broadcast from the end(instead of the first dimension) # so we need to make the contract legs to the end # then we need to take diagonal from these dimensions tb.moveLegsToFront(tbRemainLegs) dim = len(contractBLegs) l = contractBLegs[0].dim # print('contract A legs: {}'.format(len(contractALegs))) # print('contract B legs: {}'.format(len(contractBLegs))) remADim = len(taRemainLegs) einsumStr = '...' + ('j' * dim) + '->...j' # print('ta.a = {}, tb.a = {}'.format(ta.a, tb.a)) data = xplib.xp.einsum(einsumStr, tb.a) * ta.a # print('einsum str = {}'.format(einsumStr)) # print('einsum b = {}'.format(xplib.xp.einsum(einsumStr, tb.a))) # print('data = {}'.format(data)) if (remADim == 0): newData = xplib.xp.sum(data, axis=-1) else: # print('remADim = {}, data.shape = {}'.format(remADim, data.shape)) newData = xplib.xp.zeros(data.shape + (data.shape[-1], ) * (remADim - 1), dtype=data.dtype) # print('newData.shape = {}'.format(newData.shape)) einsumDiagStr = '...' + ('j' * (remADim)) + '->...j' xplib.xp.einsum(einsumDiagStr, newData)[...] = data # newData = xplib.xp.einsum(eimsumDiagStr, data) # print(funcs.ndEye(remADim - 1, l)) # newData = xplib.xp.multiply.outer(data, funcs.ndEye(remADim - 1, l)) # print('newData = {}'.format(newData)) # print('newData = {}'.format(newData)) # print('shape = {}, data = {}, legs = {}'.format(newShape, newData, newLegs)) newLegs = tbRemainLegs + taRemainLegs newShape = tuple([leg.dim for leg in newLegs]) return Tensor(shape=newShape, data=newData, legs=newLegs) if (tensorLikeContract): return Tensor(shape=newShape, data=None, legs=newLegs, tensorLikeFlag=True) dataA = ta.toMatrix(rows=None, cols=contractALegs) dataB = tb.toMatrix(rows=contractBLegs, cols=None) newData = xplib.xp.matmul(dataA, dataB) newData = xplib.xp.reshape(newData, newShape) return Tensor(shape=newShape, data=newData, legs=newLegs)
def legDimNotEqualFunc(): _ = DiagonalTensor(legs = [legA, legBError])