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 todense(self) -> np.ndarray:
"""
Map the sparse tensor to dense storage.
"""
if len(self.shape) == 0:
return self.data
out = np.asarray(np.zeros(self.shape, dtype=self.dtype).flat)
out[np.nonzero(
fuse_charges(self._charges, self._flows) ==
self._charges[0].identity_charges)[0]] = self.data
result = np.reshape(out, [c.dim for c in self._charges])
flat_order = flatten(self._order)
return result.transpose(flat_order).reshape(self.shape)
def tensordot(
tensor1: BlockSparseTensor,
tensor2: BlockSparseTensor,
axes: Optional[Union[Sequence[Sequence[int]],
int]] = 2) -> BlockSparseTensor:
"""
Contract two `BlockSparseTensor`s along `axes`.
Args:
tensor1: First tensor.
tensor2: Second tensor.
axes: The axes to contract.
Returns:
BlockSparseTensor: The result of the tensor contraction.
"""
#process scalar input for `axes`
if isinstance(axes, (np.integer, int)):
axes = [
np.arange(tensor1.ndim - axes, tensor1.ndim, dtype=np.int16),
np.arange(0, axes, dtype=np.int16)
]
elif isinstance(axes[0], (np.integer, int)):
if len(axes) > 1:
raise ValueError(
"invalid input `axes = {}` to tensordot".format(axes))
axes = [np.array(axes, dtype=np.int16), np.array(axes, dtype=np.int16)]
axes1 = axes[0]
axes2 = axes[1]
if len(axes1) != len(axes2):
raise ValueError(
"`axes1 = {}` and `axes2 = {}` have to be of same length. ".format(
axes1, axes2))
if len(axes1) > len(tensor1.shape):
raise ValueError(
"`axes1 = {}` is incompatible with `tensor1.shape = {}. ".format(
axes1, tensor1.shape))
if len(axes2) > len(tensor2.shape):
raise ValueError(
"`axes2 = {}` is incompatible with `tensor2.shape = {}. ".format(
axes2, tensor2.shape))
if not np.all(np.unique(axes1) == np.sort(axes1)):
raise ValueError(
"Some values in axes[0] = {} appear more than once!".format(axes1))
if not np.all(np.unique(axes2) == np.sort(axes2)):
raise ValueError(
"Some values in axes[1] = {} appear more than once!".format(axes2))
#special case outer product
if len(axes1) == 0:
return outerproduct(tensor1, tensor2)
#more checks
if max(axes1) >= len(tensor1.shape):
raise ValueError(
"rank of `tensor1` is smaller than `max(axes1) = {}.`".format(
max(axes1)))
if max(axes2) >= len(tensor2.shape):
raise ValueError(
"rank of `tensor2` is smaller than `max(axes2) = {}`".format(
max(axes1)))
contr_flows_1 = []
contr_flows_2 = []
contr_charges_1 = []
contr_charges_2 = []
for a in axes1:
contr_flows_1.extend(tensor1._flows[tensor1._order[a]])
contr_charges_1.extend(
[tensor1._charges[n] for n in tensor1._order[a]])
for a in axes2:
contr_flows_2.extend(tensor2._flows[tensor2._order[a]])
contr_charges_2.extend(
[tensor2._charges[n] for n in tensor2._order[a]])
if len(contr_charges_2) != len(contr_charges_1):
raise ValueError(
"`axes1 = {}` and `axes2 = {}` have incompatible elementary"
" shapes {} and {}".format(axes1, axes2,
[e.dim for e in contr_charges_1],
[e.dim for e in contr_charges_2]))
if not np.all(
np.asarray(contr_flows_1) == np.logical_not(
np.asarray(contr_flows_2))):
raise ValueError(
"`axes1 = {}` and `axes2 = {}` have incompatible elementary"
" flows {} and {}".format(axes1, axes2, contr_flows_1,
contr_flows_2))
charge_check = [
charge_equal(c1, c2)
for c1, c2 in zip(contr_charges_1, contr_charges_2)
]
if not np.all(charge_check):
inds = np.nonzero(np.logical_not(charge_check))[0]
raise ValueError(
"`axes = {}` of tensor1 and `axes = {}` of tensor2 have incompatible charges"
" {} and {}".format(
np.array(axes1)[inds],
np.array(axes2)[inds], [contr_charges_1[i] for i in inds],
[contr_charges_2[i] for i in inds]))
#checks finished
#special case inner product
if (len(axes1) == tensor1.ndim) and (len(axes2) == tensor2.ndim):
t1 = tensor1.transpose(axes1).transpose_data()
t2 = tensor2.transpose(axes2).transpose_data()
data = np.dot(t1.data, t2.data)
charge = tensor1._charges[0]
final_charge = charge.__new__(type(charge))
final_charge.__init__(np.empty((charge.num_symmetries, 0),
dtype=np.int16),
charge_labels=np.empty(0, dtype=np.int16),
charge_types=charge.charge_types)
return BlockSparseTensor(data=data,
charges=[final_charge],
flows=[False],
order=[[0]],
check_consistency=False)
#in all other cases we perform a regular tensordot
free_axes1 = sorted(set(np.arange(tensor1.ndim)) - set(axes1))
free_axes2 = sorted(set(np.arange(tensor2.ndim)) - set(axes2))
new_order1 = [tensor1._order[n]
for n in free_axes1] + [tensor1._order[n] for n in axes1]
new_order2 = [tensor2._order[n]
for n in axes2] + [tensor2._order[n] for n in free_axes2]
flat_order_1 = flatten(new_order1)
flat_order_2 = flatten(new_order2)
flat_charges_1, flat_flows_1 = tensor1._charges, tensor1.flat_flows
flat_charges_2, flat_flows_2 = tensor2._charges, tensor2.flat_flows
left_charges = []
right_charges = []
left_flows = []
right_flows = []
left_order = []
right_order = []
s = 0
for n in free_axes1:
left_charges.extend([tensor1._charges[o] for o in tensor1._order[n]])
left_order.append(list(np.arange(s, s + len(tensor1._order[n]))))
s += len(tensor1._order[n])
left_flows.extend([tensor1._flows[o] for o in tensor1._order[n]])
s = 0
for n in free_axes2:
right_charges.extend([tensor2._charges[o] for o in tensor2._order[n]])
right_order.append(
list(len(left_charges) + np.arange(s, s + len(tensor2._order[n]))))
s += len(tensor2._order[n])
right_flows.extend([tensor2._flows[o] for o in tensor2._order[n]])
tr_sparse_blocks_1, charges1, shapes_1 = _find_transposed_diagonal_sparse_blocks(
flat_charges_1, flat_flows_1, len(left_charges), flat_order_1)
tr_sparse_blocks_2, charges2, shapes_2 = _find_transposed_diagonal_sparse_blocks(
flat_charges_2, flat_flows_2, len(contr_charges_2), flat_order_2)
common_charges, label_to_common_1, label_to_common_2 = intersect(
charges1.unique_charges,
charges2.unique_charges,
axis=1,
return_indices=True)
#Note: `cs` may contain charges that are not present in `common_charges`
charges = left_charges + right_charges
flows = left_flows + right_flows
sparse_blocks, cs, _ = _find_diagonal_sparse_blocks(
charges, flows, len(left_charges))
num_nonzero_elements = np.int64(np.sum([len(v) for v in sparse_blocks]))
#Note that empty is not a viable choice here.
data = np.zeros(num_nonzero_elements,
dtype=np.result_type(tensor1.dtype, tensor2.dtype))
label_to_common_final = intersect(cs.unique_charges,
common_charges,
axis=1,
return_indices=True)[1]
for n in range(common_charges.shape[1]):
n1 = label_to_common_1[n]
n2 = label_to_common_2[n]
nf = label_to_common_final[n]
data[sparse_blocks[nf].ravel()] = np.ravel(
np.matmul(
tensor1.data[tr_sparse_blocks_1[n1].reshape(shapes_1[:, n1])],
tensor2.data[tr_sparse_blocks_2[n2].reshape(shapes_2[:, n2])]))
res = BlockSparseTensor(data=data,
charges=charges,
flows=flows,
order=left_order + right_order,
check_consistency=False)
return res
def flat_order(self) -> List:
"""
The flattened `ChargeArray._oder`.
"""
return flatten(self._order)
def test_flatten():
listoflist = [[1, 2], [3, 4], [5]]
flat = flatten(listoflist)
np.testing.assert_allclose(flat, [1, 2, 3, 4, 5])
