def test_tensordot_reshape(dtype, num_charges):
np.random.seed(10)
R1 = 4
R2 = 4
q = np.random.randint(-5, 5, (num_charges, 10), dtype=np.int16)
charges1 = [
BaseCharge(q, charge_types=[U1Charge] * num_charges) for n in range(R1)
]
charges2 = [
BaseCharge(q, charge_types=[U1Charge] * num_charges) for n in range(R2)
]
flowsA = np.asarray([False] * R1)
flowsB = np.asarray([True] * R2)
A = BlockSparseTensor.random(
indices=[Index(charges1[n], flowsA[n]) for n in range(R1)],
dtype=dtype)
B = BlockSparseTensor.random(
indices=[Index(charges2[n], flowsB[n]) for n in range(R2)],
dtype=dtype)
Adense = A.todense().reshape((10, 10 * 10, 10))
Bdense = B.todense().reshape((10 * 10, 10, 10))
A = A.reshape((10, 10 * 10, 10))
B = B.reshape((10 * 10, 10, 10))
res = tensordot(A, B, ([0, 1], [2, 0]))
dense = np.tensordot(Adense, Bdense, ([0, 1], [2, 0]))
np.testing.assert_allclose(dense, res.todense())
def trace(tensor: BlockSparseTensor,
axes: Optional[Tuple[int, ...]] = None) -> BlockSparseTensor:
"""
Compute the trace of a matrix or tensor.
Args:
tensor: A `BlockSparseTensor`.
axes: The axes over which the trace should be computed.
Defaults to the last two indices of the tensor.
Returns:
BlockSparseTensor: The result of taking the trace.
"""
if tensor.ndim > 1:
if axes is None:
axes = (tensor.ndim - 2, tensor.ndim - 1)
if len(axes) != 2:
raise ValueError(f"`len(axes)` has to be 2, found `axes = {axes}`")
if not np.array_equal(tensor.flows[axes[0]],
np.logical_not(tensor.flows[axes[1]])):
raise ValueError(
f"trace indices for axes {axes} have non-matching flows.")
sparse_shape = tensor.sparse_shape
if sparse_shape[axes[0]].copy().flip_flow() != sparse_shape[axes[1]]:
raise ValueError(f"trace indices for axes {axes} are not matching")
#flatten the shape of `tensor`
out = tensor.reshape(
flatten([[tensor._charges[n].dim for n in o] for o in tensor._order]))
_, _, labels0 = np.intersect1d(
tensor._order[axes[0]], flatten(out._order), return_indices=True)
_, _, labels1 = np.intersect1d(
tensor._order[axes[1]], flatten(out._order), return_indices=True)
a0 = list(labels0[np.argsort(tensor._order[axes[0]])])
a1 = list(labels1[np.argsort(tensor._order[axes[1]])])
while len(a0) > 0:
i = a0.pop(0)
j = a1.pop(0)
identity = eye(
Index([out._charges[out._order[i][0]]],
[not out._flows[out._order[i][0]]]))
out = tensordot(out, identity, ([i, j], [0, 1])) # pytype: disable=wrong-arg-types
a0ar = np.asarray(a0)
mask_min = a0ar > np.min([i, j])
mask_max = a0ar > np.max([i, j])
a0ar[np.logical_and(mask_min, mask_max)] -= 2
a0ar[np.logical_xor(mask_min, mask_max)] -= 1
a1ar = np.asarray(a1)
mask_min = a1ar > np.min([i, j])
mask_max = a1ar > np.max([i, j])
a1ar[np.logical_and(mask_min, mask_max)] -= 2
a1ar[np.logical_xor(mask_min, mask_max)] -= 1
a0 = list(a0ar)
a1 = list(a1ar)
return out # pytype: disable=bad-return-type
raise ValueError("trace can only be taken for tensors with ndim > 1")
def test_tensordot_inner(R1, R2, dtype, num_charges):
np.random.seed(10)
DsA = np.random.randint(3, 5, R1)
Dscomm = np.random.randint(3, 5, 0)
DsB = np.random.randint(3, 5, R2)
A, B, indsA, indsB = get_contractable_tensors(R1, R2, 0, dtype,
num_charges, DsA, Dscomm,
DsB)
res = tensordot(A, B, (indsA, indsB))
dense_res = np.tensordot(A.todense(), B.todense(), (indsA, indsB))
np.testing.assert_allclose(dense_res, res.todense())
def test_tensordot_single_arg():
R = 3
dtype = np.float64
np.random.seed(10)
Ds = [10, 10, 10]
inds = [
Index(U1Charge.random(dimension=Ds[n], minval=-5, maxval=5), False)
for n in range(R)
]
A = BlockSparseTensor.random(inds, dtype=dtype)
res = tensordot(A, A.conj(), ([0]))
dense_res = np.tensordot(A.todense(), A.conj().todense(), ([0], [0]))
np.testing.assert_allclose(dense_res, res.todense())
def test_tensordot_outer(R1, R2, dtype, num_charges):
np.random.seed(10)
DsA = np.random.randint(3, 5, R1)
Dscomm = np.random.randint(3, 5, 0)
DsB = np.random.randint(3, 5, R2)
A, B, _, _ = get_contractable_tensors(R1, R2, 0, dtype, num_charges, DsA,
Dscomm, DsB)
res = tensordot(A, B, axes=0)
dense_res = np.tensordot(A.todense(), B.todense(), axes=0)
np.testing.assert_allclose(dense_res, res.todense())
for n in range(R1):
assert charge_equal(res.charges[n][0], A.charges[n][0])
for n in range(R1, R1 + R2):
assert charge_equal(res.charges[n][0], B.charges[n - R1][0])
def test_tensordot_empty_tensors(dtype, num_charges):
A, B, iA, iB = get_contractable_tensors(R1=4,
R2=4,
cont=2,
dtype=dtype,
num_charges=num_charges,
DsA=[10, 0],
Dscomm=[0, 4],
DsB=[8, 0])
free_inds_A = np.sort(list(set(np.arange(len(A.shape))) - set(iA)))
free_inds_B = np.sort(list(set(np.arange(len(B.shape))) - set(iB)))
res = tensordot(A, B, (iA, iB))
assert len(res.data) == 0
for n in range(2):
assert charge_equal(res.charges[n][0], A.charges[free_inds_A[n]][0])
for n in range(2, 4):
assert charge_equal(res.charges[n][0],
B.charges[free_inds_B[n - 2]][0])
def test_tensordot(R1, R2, cont, dtype, num_charges):
np.random.seed(10)
DsA = np.random.randint(5, 10, R1 - cont)
Dscomm = np.random.randint(5, 10, cont)
DsB = np.random.randint(5, 10, R2 - cont)
A, B, indsA, indsB = get_contractable_tensors(R1, R2, cont, dtype,
num_charges, DsA, Dscomm,
DsB)
res = tensordot(A, B, (indsA, indsB))
dense_res = np.tensordot(A.todense(), B.todense(), (indsA, indsB))
np.testing.assert_allclose(dense_res, res.todense())
free_inds_A = np.sort(list(set(np.arange(len(A.shape))) - set(indsA)))
free_inds_B = np.sort(list(set(np.arange(len(B.shape))) - set(indsB)))
for n, fiA in enumerate(free_inds_A):
assert charge_equal(res.charges[n][0], A.charges[fiA][0])
for n in range(len(free_inds_A), len(free_inds_A) + len(free_inds_B)):
assert charge_equal(res.charges[n][0],
B.charges[free_inds_B[n - len(free_inds_A)]][0])
def test_tensordot_inner_transpose(R1, R2, dtype, num_charges):
np.random.seed(10)
DsA = np.random.randint(3, 5, R1)
Dscomm = np.random.randint(3, 5, 0)
DsB = np.random.randint(3, 5, R2)
A, B, indsA, indsB = get_contractable_tensors(R1, R2, 0, dtype,
num_charges, DsA, Dscomm,
DsB)
orderA = np.arange(R1)
orderB = np.arange(R2)
np.random.shuffle(orderA)
np.random.shuffle(orderB)
A_ = A.transpose(orderA)
B_ = B.transpose(orderB)
_, indposA = np.unique(orderA, return_index=True)
_, indposB = np.unique(orderB, return_index=True)
indsA_ = indposA[indsA]
indsB_ = indposB[indsB]
res = tensordot(A_, B_, (indsA_, indsB_))
dense_res = np.tensordot(A_.todense(), B_.todense(), (indsA_, indsB_))
np.testing.assert_allclose(dense_res, res.todense())
def test_repr():
np.random.seed(10)
dtype = np.float64
D = 10
rank = 3
charges = [U1Charge.random(D, -2, 2) for _ in range(rank)]
flows = np.random.choice([True, False], size=rank, replace=True)
inds = [Index(c, f) for c, f in zip(charges, flows)]
T = ChargeArray.random(inds, dtype=dtype)
actual = T.__repr__()
expected = "ChargeArray\n shape: (10, 10, 10)\n " +\
" charge types: ['U1Charge']\n dtype: " +\
repr(T.dtype.name) + "\n flat flows: " + \
repr(list(flows)) + "\n order: " + repr(T._order)
assert actual == expected
res = tensordot(T, T.conj(), ([0, 1, 2], [0, 1, 2]))
actual = res.__repr__()
expected = "BlockSparseTensor\n shape: ()\n " +\
" charge types: no charge types (scalar)\n dtype: " +\
repr(res.dtype.name) + "\n flat flows: " + \
repr(list(res.flat_flows)) + "\n order: " + repr(res._order)
assert actual == expected
def test_tensordot_raises():
R1 = 3
R2 = 3
dtype = np.float64
np.random.seed(10)
Ds1 = np.arange(2, 2 + R1)
Ds2 = np.arange(2 + R1, 2 + R1 + R2)
is1 = [
Index(U1Charge.random(dimension=Ds1[n], minval=-5, maxval=5), False)
for n in range(R1)
]
is2 = [
Index(U1Charge.random(dimension=Ds2[n], minval=-5, maxval=5), False)
for n in range(R2)
]
A = BlockSparseTensor.random(is1, dtype=dtype)
B = BlockSparseTensor.random(is2, dtype=dtype)
with pytest.raises(ValueError):
tensordot(A, B, ([0, 1, 2, 3], [1, 2]))
with pytest.raises(ValueError):
tensordot(A, B, ([0, 1], [0, 1, 2, 3]))
with pytest.raises(ValueError):
tensordot(A, B, ([0], [1, 2]))
with pytest.raises(ValueError):
tensordot(A, B, [0, [1, 2]])
with pytest.raises(ValueError):
tensordot(A, B, ([0, 0], [1, 2]))
with pytest.raises(ValueError):
tensordot(A, B, ([0, 1], [1, 1]))
with pytest.raises(ValueError):
tensordot(A, B, ([0, 4], [1, 2]))
with pytest.raises(ValueError):
tensordot(A, B, ([0, 4], [1, 4]))
with pytest.raises(ValueError):
tensordot(A, B, ([0, 1], [0, 1]))
with pytest.raises(ValueError):
tensordot(A, A, ([0, 1], [0, 1]))
with pytest.raises(ValueError):
tensordot(A, A.conj(), ([0, 1], [1, 0]))
def test_tensordot_raises():
R1 = 3
R2 = 3
R3 = 3
dtype = np.float64
np.random.seed(10)
Ds1 = np.arange(2, 2 + R1)
Ds2 = np.arange(2 + R1, 2 + R1 + R2)
Ds3 = np.arange(2 + R1, 2 + R1 + R3)
is1 = [
Index(U1Charge.random(dimension=Ds1[n], minval=-5, maxval=5), False)
for n in range(R1)
]
is2 = [
Index(U1Charge.random(dimension=Ds2[n], minval=-5, maxval=5), False)
for n in range(R2)
]
is3 = [
Index(U1Charge.random(dimension=Ds3[n], minval=-5, maxval=5), False)
for n in range(R3)
]
A = BlockSparseTensor.random(is1, dtype=dtype)
B = BlockSparseTensor.random(is2, dtype=dtype)
C = BlockSparseTensor.random(is3, dtype=dtype)
with pytest.raises(ValueError, match="same length"):
tensordot(A, B, ([0, 1, 2, 3], [1, 2]))
with pytest.raises(ValueError, match="same length"):
tensordot(A, B, ([0, 1], [0, 1, 2, 3]))
with pytest.raises(ValueError, match="same length"):
tensordot(A, B, ([0], [1, 2]))
with pytest.raises(ValueError, match='invalid input'):
tensordot(A, B, [0, [1, 2]])
with pytest.raises(ValueError, match="incompatible elementary flows"):
tensordot(A, B, ([0, 0], [1, 2]))
with pytest.raises(ValueError):
tensordot(A, B, ([0, 1], [1, 1]))
with pytest.raises(ValueError, match="rank of `tensor1` is smaller than "):
tensordot(A, B, ([0, 4], [1, 2]))
with pytest.raises(ValueError, match="rank of `tensor2` is smaller than "):
tensordot(A, B, ([0, 1], [0, 4]))
with pytest.raises(ValueError):
tensordot(A, B, ([0, 4], [1, 4]))
with pytest.raises(ValueError):
tensordot(A, B, ([0, 1], [0, 1]))
with pytest.raises(ValueError):
tensordot(A, A, ([0, 1], [0, 1]))
with pytest.raises(ValueError):
tensordot(A, A.conj(), ([0, 1], [1, 0]))
with pytest.raises(ValueError,
match="is incompatible with `tensor1.shape"):
tensordot(A, A.conj(), ([0, 1, 2, 3], [0, 1, 2, 3]))
with pytest.raises(ValueError,
match="is incompatible with `tensor1.shape"):
tensordot(A, C, ([0, 1, 2, 3], [0, 1, 2, 3]))
def test_tensordot_raises():
R1 = 3
R2 = 3
R3 = 3
dtype = np.float64
np.random.seed(10)
Ds1 = np.arange(2, 2 + R1)
Ds2 = np.arange(2 + R1, 2 + R1 + R2)
Ds3 = np.arange(2 + R1, 2 + R1 + R3)
is1 = [
Index(U1Charge.random(dimension=Ds1[n], minval=-5, maxval=5), False)
for n in range(R1)
]
is2 = [
Index(U1Charge.random(dimension=Ds2[n], minval=-5, maxval=5), False)
for n in range(R2)
]
is3 = [
Index(U1Charge.random(dimension=Ds3[n], minval=-5, maxval=5), False)
for n in range(R3)
]
A = BlockSparseTensor.random(is1, dtype=dtype)
B = BlockSparseTensor.random(is2, dtype=dtype)
C = BlockSparseTensor.random(is3, dtype=dtype)
with pytest.raises(ValueError, match="same length"):
tensordot(A, B, ([0, 1, 2, 3], [1, 2]))
with pytest.raises(ValueError, match="same length"):
tensordot(A, B, ([0, 1], [0, 1, 2, 3]))
with pytest.raises(ValueError, match="same length"):
tensordot(A, B, ([0], [1, 2]))
with pytest.raises(ValueError, match='invalid input'):
tensordot(A, B, [0, [1, 2]])
with pytest.raises(ValueError, match="incompatible elementary flows"):
tensordot(A, B, ([0, 0], [1, 2]))
with pytest.raises(ValueError):
tensordot(A, B, ([0, 1], [1, 1]))
with pytest.raises(ValueError, match="rank of `tensor1` is smaller than "):
tensordot(A, B, ([0, 4], [1, 2]))
with pytest.raises(ValueError, match="rank of `tensor2` is smaller than "):
tensordot(A, B, ([0, 1], [0, 4]))
with pytest.raises(ValueError):
tensordot(A, B, ([0, 4], [1, 4]))
with pytest.raises(ValueError):
tensordot(A, B, ([0, 1], [0, 1]))
with pytest.raises(ValueError):
tensordot(A, A, ([0, 1], [0, 1]))
with pytest.raises(ValueError):
tensordot(A, A.conj(), ([0, 1], [1, 0]))
with pytest.raises(ValueError,
match="is incompatible with `tensor1.shape"):
tensordot(A, A.conj(), ([0, 1, 2, 3], [0, 1, 2, 3]))
with pytest.raises(ValueError,
match="is incompatible with `tensor1.shape"):
tensordot(A, C, ([0, 1, 2, 3], [0, 1, 2, 3]))
Ds1 = np.array([8, 9, 10, 11])
Ds2 = np.array([8, 9])
flows1 = [False] * len(Ds1)
flows2 = [False] * len(Ds2)
indices1 = [
Index(U1Charge.random(D, -2, 2), f) for D, f in zip(Ds1, flows1)
]
indices2 = [
Index(U1Charge.random(D, -2, 2), f) for D, f in zip(Ds2, flows2)
]
arr1 = BlockSparseTensor.random(indices1)
arr2 = BlockSparseTensor.random(indices2)
with pytest.raises(ValueError, match="axes2 = "):
tensordot(arr1, arr2, ([0, 1, 2], [0, 1, 2]))
Ds2 = np.array([8, 9, 2, 5, 11])
flows2 = [False] * len(Ds2)
indices1 = [
Index(U1Charge.random(D, -2, 2), f) for D, f in zip(Ds1, flows1)
]
indices2 = [
Index(U1Charge.random(D, -2, 2), f) for D, f in zip(Ds2, flows2)
]
arr1 = BlockSparseTensor.random(indices1)
arr2 = BlockSparseTensor.random(indices2).reshape(Ds1)
with pytest.raises(ValueError, match="incompatible elementary shapes "):
tensordot(arr1, arr2.conj(), ([2, 3], [2, 3]))
