def test_split_node_qr_unitarity(dtype, num_charges):
np.random.seed(10)
a = tn.Node(get_square_matrix(50, num_charges, dtype=dtype),
backend='symmetric')
q, r = tn.split_node_qr(a, [a[0]], [a[1]])
r[0] | q[1]
qbar = tn.linalg.node_linalg.conj(q)
q[1] ^ qbar[1]
u1 = q @ qbar
qbar[0] ^ q[0]
u2 = qbar @ q
blocks, _, shapes = _find_diagonal_sparse_blocks(u1.tensor.flat_charges,
u1.tensor.flat_flows,
len(u1.tensor._order[0]))
for n, block in enumerate(blocks):
np.testing.assert_almost_equal(
np.reshape(u1.tensor.data[block], shapes[:, n]),
np.eye(N=shapes[0, n], M=shapes[1, n]))
blocks, _, shapes = _find_diagonal_sparse_blocks(u2.tensor.flat_charges,
u2.tensor.flat_flows,
len(u2.tensor._order[0]))
for n, block in enumerate(blocks):
np.testing.assert_almost_equal(
np.reshape(u2.tensor.data[block], shapes[:, n]),
np.eye(N=shapes[0, n], M=shapes[1, n]))
def test_pinv(dtype, num_charges):
np.random.seed(10)
R = 2
D = 10
charge = BaseCharge(np.random.randint(-5,
6, (D, num_charges),
dtype=np.int16),
charge_types=[U1Charge] * num_charges)
flows = [True, False]
A = BlockSparseTensor.random([Index(charge, flows[n]) for n in range(R)],
(-0.5, 0.5),
dtype=dtype)
invA = pinv(A)
left_eye = invA @ A
blocks, _, shapes = _find_diagonal_sparse_blocks(left_eye.flat_charges,
left_eye.flat_flows, 1)
for n, block in enumerate(blocks):
t = np.reshape(left_eye.data[block], shapes[:, n])
assert np.linalg.norm(t - np.eye(t.shape[0], t.shape[1])) < 1E-12
right_eye = A @ invA
blocks, _, shapes = _find_diagonal_sparse_blocks(right_eye.flat_charges,
right_eye.flat_flows, 1)
for n, block in enumerate(blocks):
t = np.reshape(right_eye.data[block], shapes[:, n])
assert np.linalg.norm(t - np.eye(t.shape[0], t.shape[1])) < 1E-12
def test_get_diag(dtype, num_charges, Ds, flow):
np.random.seed(10)
np_flow = -np.int((np.int(flow) - 0.5) * 2)
indices = [
Index(
BaseCharge(np.random.randint(-2, 3, (Ds[n], num_charges)),
charge_types=[U1Charge] * num_charges), flow)
for n in range(2)
]
arr = BlockSparseTensor.random(indices, dtype=dtype)
fused = fuse_charges(arr.flat_charges, arr.flat_flows)
inds = np.nonzero(fused == np.zeros((1, num_charges), dtype=np.int16))[0]
# pylint: disable=no-member
left, _ = np.divmod(inds, Ds[1])
unique = np.unique(np_flow * (indices[0]._charges[0].charges[left, :]),
axis=0)
diagonal = diag(arr)
sparse_blocks, _, block_shapes = _find_diagonal_sparse_blocks(
arr.flat_charges, arr.flat_flows, 1)
data = np.concatenate([
np.diag(np.reshape(arr.data[sparse_blocks[n]], block_shapes[:, n]))
for n in range(len(sparse_blocks))
])
np.testing.assert_allclose(data, diagonal.data)
np.testing.assert_allclose(unique, diagonal.flat_charges[0].unique_charges)
def contiguous(self,
permutation: Optional[Union[Tuple, List,
np.ndarray]] = None,
inplace: Optional[bool] = False) -> Any:
"""
Transpose the tensor data in place such that the linear order
of the elements in `BlockSparseTensor.data` corresponds to the
current order of tensor indices.
Consider a tensor with current order given by `_order=[[1,2],[3],[0]]`,
i.e. `data` was initialized according to order [0,1,2,3], and the tensor
has since been reshaped and transposed. The linear oder of `data` does not
match the desired order [1,2,3,0] of the tensor. `contiguous` fixes this
by permuting `data` into this order, transposing `_charges` and `_flows`,
and changing `_order` to `[[0,1],[2],[3]]`.
Args:
permutation: An optional alternative order to be used to transposed the
tensor. If `None` defaults to `BlockSparseTensor.permutation`.
"""
flat_charges = self._charges
flat_flows = self._flows
if permutation is None:
permutation = self.flat_order
if np.array_equal(permutation, np.arange(len(permutation))):
return self
tr_partition = _find_best_partition(
[flat_charges[n].dim for n in permutation])
tr_sparse_blocks, tr_charges, _ = _find_transposed_diagonal_sparse_blocks(
flat_charges, flat_flows, tr_partition, permutation)
sparse_blocks, charges, _ = _find_diagonal_sparse_blocks(
[flat_charges[n] for n in permutation],
[flat_flows[n] for n in permutation], tr_partition)
data = np.empty(len(self.data), dtype=self.dtype)
for n, sparse_block in enumerate(sparse_blocks):
ind = np.nonzero(tr_charges == charges[n])[0][0]
perm = tr_sparse_blocks[ind]
data[sparse_block] = self.data[perm]
_, inds = np.unique(permutation, return_index=True)
new_flat_order = inds[self.flat_order]
tmp = np.append(0, np.cumsum([len(o) for o in self._order]))
order = [
list(new_flat_order[tmp[n]:tmp[n + 1]])
for n in range(len(tmp) - 1)
]
charges = [self._charges[o] for o in permutation]
flows = [self._flows[o] for o in permutation]
if not inplace:
return BlockSparseTensor(data,
charges=charges,
flows=flows,
order=order,
check_consistency=False)
self.data = data
self._order = order
self._charges = charges
self._flows = flows
return self
def test_find_diagonal_sparse_blocks(num_legs, num_charges):
np.random.seed(10)
np_charges = [
np.random.randint(-5, 5, (60, num_charges), dtype=np.int16)
for _ in range(num_legs)
]
fused = np.stack([
fuse_ndarrays([np_charges[n][:, c]
for n in range(num_legs)])
for c in range(num_charges)
],
axis=1)
left_charges = np.stack([
fuse_ndarrays([np_charges[n][:, c]
for n in range(num_legs // 2)])
for c in range(num_charges)
],
axis=1)
right_charges = np.stack([
fuse_ndarrays(
[np_charges[n][:, c]
for n in range(num_legs // 2, num_legs)])
for c in range(num_charges)
],
axis=1)
#pylint: disable=no-member
nz = np.nonzero(
np.logical_and.reduce(fused == np.zeros((1, num_charges)), axis=1))[0]
linear_locs = np.arange(len(nz))
# pylint: disable=no-member
left_inds, _ = np.divmod(nz, right_charges.shape[0])
left = left_charges[left_inds, :]
unique_left = np.unique(left, axis=0)
blocks = []
for n in range(unique_left.shape[0]):
ul = unique_left[n, :][None, :]
#pylint: disable=no-member
blocks.append(linear_locs[np.nonzero(
np.logical_and.reduce(left == ul, axis=1))[0]])
charges = [
BaseCharge(left_charges, charge_types=[U1Charge] * num_charges),
BaseCharge(right_charges, charge_types=[U1Charge] * num_charges)
]
bs, cs, ss = _find_diagonal_sparse_blocks(charges, [False, False], 1)
np.testing.assert_allclose(cs.charges, unique_left)
for b1, b2 in zip(blocks, bs):
assert np.all(b1 == b2)
assert np.sum(np.prod(ss, axis=0)) == np.sum([len(b) for b in bs])
np.testing.assert_allclose(unique_left, cs.charges)
def test_eye(dtype, num_charges, D):
charge = BaseCharge(np.random.randint(-5,
6, (D, num_charges),
dtype=np.int16),
charge_types=[U1Charge] * num_charges)
flow = False
index = Index(charge, flow)
A = eye(index, dtype=dtype)
blocks, _, shapes = _find_diagonal_sparse_blocks(A.flat_charges,
A.flat_flows, 1)
for n, block in enumerate(blocks):
t = np.reshape(A.data[block], shapes[:, n])
np.testing.assert_almost_equal(t, np.eye(t.shape[0], t.shape[1]))
def test_create_diag(dtype, num_charges, flow):
np.random.seed(10)
D = 200
index = Index(
BaseCharge(np.random.randint(-2, 3, (D, num_charges)),
charge_types=[U1Charge] * num_charges), flow)
arr = ChargeArray.random([index], dtype=dtype)
diagarr = diag(arr)
dense = np.ravel(diagarr.todense())
np.testing.assert_allclose(np.sort(dense[dense != 0.0]),
np.sort(diagarr.data[diagarr.data != 0.0]))
sparse_blocks, charges, block_shapes = _find_diagonal_sparse_blocks(
diagarr.flat_charges, diagarr.flat_flows, 1)
for n, block in enumerate(sparse_blocks):
shape = block_shapes[:, n]
block_diag = np.diag(np.reshape(diagarr.data[block], shape))
np.testing.assert_allclose(
arr.data[np.squeeze((index._charges[0] * flow) == charges[n])],
block_diag)
def outerproduct(tensor1: BlockSparseTensor,
tensor2: BlockSparseTensor) -> BlockSparseTensor:
"""
Compute the outer product of two `BlockSparseTensor`.
The first `tensor1.ndim` indices of the resulting tensor are the
indices of `tensor1`, the last `tensor2.ndim` indices are those
of `tensor2`.
Args:
tensor1: A tensor.
tensor2: A tensor.
Returns:
BlockSparseTensor: The result of taking the outer product.
"""
final_charges = tensor1._charges + tensor2._charges
final_flows = list(tensor1._flows) + list(tensor2._flows)
order2 = [
list(np.asarray(s) + len(tensor1._charges)) for s in tensor2._order
]
data = np.zeros(compute_num_nonzero(final_charges, final_flows),
dtype=tensor1.dtype)
if ((len(tensor1.data) > 0) and
(len(tensor2.data) > 0)) and (len(data) > 0):
# find the location of the zero block in the output
final_block_maps, final_block_charges, _ = _find_diagonal_sparse_blocks(
final_charges, final_flows, len(tensor1._charges))
index = np.nonzero(
final_block_charges == final_block_charges.identity_charges(
dim=1))[0][0]
data[final_block_maps[index].ravel()] = np.outer(
tensor1.data, tensor2.data).ravel()
return BlockSparseTensor(data,
charges=final_charges,
flows=final_flows,
order=tensor1._order + order2,
check_consistency=False)
def eye(column_index: Index,
row_index: Optional[Index] = None,
dtype: Optional[Type[np.number]] = None) -> BlockSparseTensor:
"""
Return an identity matrix.
Args:
column_index: The column index of the matrix.
row_index: The row index of the matrix.
dtype: The dtype of the matrix.
Returns:
BlockSparseTensor
"""
if row_index is None:
row_index = column_index.copy().flip_flow()
if dtype is None:
dtype = np.float64
blocks, _, shapes = _find_diagonal_sparse_blocks(
column_index.flat_charges + row_index.flat_charges,
column_index.flat_flows + row_index.flat_flows,
len(column_index.flat_charges))
data = np.empty(np.int64(np.sum(np.prod(shapes, axis=0))), dtype=dtype)
for n, block in enumerate(blocks):
data[block] = np.ravel(np.eye(shapes[0, n], shapes[1, n], dtype=dtype))
order = [list(np.arange(0, len(column_index.flat_charges)))] + [
list(
np.arange(
len(column_index.flat_charges),
len(column_index.flat_charges) + len(row_index.flat_charges)))
]
return BlockSparseTensor(
data=data,
charges=column_index.flat_charges + row_index.flat_charges,
flows=column_index.flat_flows + row_index.flat_flows,
order=order,
check_consistency=False)
