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