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