def test_compute_sparse_lookup_non_ordered(flow):
np_flow = np.int(-(np.int(flow) - 0.5) * 2)
charge_labels = np.array([0, 0, 1, 5, 5, 0, 2, 3, 2, 3, 4, 0, 3, 3, 1, 5])
unique_charges = np.array([-1, 0, 1, -5, 7, 2])
np_targets = np.array([-1, 0, 2])
charges = [U1Charge(unique_charges, charge_labels=charge_labels)]
inds = np.nonzero(
np.isin((np_flow * unique_charges)[charge_labels], np_targets))[0]
targets = U1Charge(np_targets)
lookup, unique, labels = compute_sparse_lookup(charges, [flow], targets)
np.testing.assert_allclose(labels, np.sort(labels))
np.testing.assert_allclose(np.squeeze(unique.charges[:, lookup]),
(np_flow * unique_charges)[charge_labels][inds])
def test_compute_sparse_lookup():
q1 = np.array([-2, 0, -5, 7])
q2 = np.array([-3, 1, -2, 6, 2, -2])
expected_unique = np.array(
[-11, -8, -7, -6, -4, -3, -2, -1, 0, 1, 2, 3, 5, 6, 9, 10])
expected_labels_to_unique = np.array([7, 8, 9])
expected_lookup = np.array([9, 8, 8, 7, 9])
charges = [U1Charge(q1), U1Charge(q2)]
targets = U1Charge(np.array([-1, 0, 1]))
flows = [False, True]
lookup, unique, labels = compute_sparse_lookup(charges, flows, targets)
np.testing.assert_allclose(lookup, expected_lookup)
np.testing.assert_allclose(expected_unique, np.squeeze(unique.charges))
np.testing.assert_allclose(labels, expected_labels_to_unique)
def diag(tensor: ChargeArray) -> Any:
"""
Return a diagonal `BlockSparseTensor` from a `ChargeArray`, or
return the diagonal of a `BlockSparseTensor` as a `ChargeArray`.
For input of type `BlockSparseTensor`:
The full diagonal is obtained from finding the diagonal blocks of the
`BlockSparseTensor`, taking the diagonal elements of those and packing
the result into a ChargeArray. Note that the computed diagonal elements
are usually different from the diagonal elements obtained from
converting the `BlockSparseTensor` to dense storage and taking the diagonal.
Note that the flow of the resulting 1d `ChargeArray` object is `False`.
Args:
tensor: A `ChargeArray`.
Returns:
ChargeArray: A 1d `CharggeArray` containing the diagonal of `tensor`,
or a diagonal matrix of type `BlockSparseTensor` containing `tensor`
on its diagonal.
"""
if tensor.ndim > 2:
raise ValueError("`diag` currently only implemented for matrices, "
"found `ndim={}".format(tensor.ndim))
if not isinstance(tensor, BlockSparseTensor):
if tensor.ndim > 1:
raise ValueError(
"`diag` currently only implemented for `ChargeArray` with ndim=1, "
"found `ndim={}`".format(tensor.ndim))
flat_charges = tensor._charges + tensor._charges
flat_flows = list(tensor._flows) + list(np.logical_not(tensor._flows))
flat_order = list(tensor.flat_order) + list(
np.asarray(tensor.flat_order) + len(tensor._charges))
tr_partition = len(tensor._order[0])
blocks, charges, shapes = _find_transposed_diagonal_sparse_blocks(
flat_charges, flat_flows, tr_partition, flat_order)
data = np.zeros(
np.int64(np.sum(np.prod(shapes, axis=0))), dtype=tensor.dtype)
lookup, unique, labels = compute_sparse_lookup(tensor._charges,
tensor._flows, charges)
for n, block in enumerate(blocks):
label = labels[np.nonzero(unique == charges[n])[0][0]]
data[block] = np.ravel(
np.diag(tensor.data[np.nonzero(lookup == label)[0]]))
order = [
tensor._order[0],
list(np.asarray(tensor._order[0]) + len(tensor._charges))
]
new_charges = [tensor._charges[0].copy(), tensor._charges[0].copy()]
return BlockSparseTensor(
data,
charges=new_charges,
flows=list(tensor._flows) + list(np.logical_not(tensor._flows)),
order=order,
check_consistency=False)
flat_charges = tensor._charges
flat_flows = tensor._flows
flat_order = tensor.flat_order
tr_partition = len(tensor._order[0])
sparse_blocks, charges, block_shapes = _find_transposed_diagonal_sparse_blocks(
flat_charges, flat_flows, tr_partition, flat_order)
shapes = np.min(block_shapes, axis=0)
if len(sparse_blocks) > 0:
data = np.concatenate([
np.diag(np.reshape(tensor.data[sparse_blocks[n]], block_shapes[:, n]))
for n in range(len(sparse_blocks))
])
charge_labels = np.concatenate([
np.full(shapes[n], fill_value=n, dtype=np.int16)
for n in range(len(sparse_blocks))
])
else:
data = np.empty(0, dtype=tensor.dtype)
charge_labels = np.empty(0, dtype=np.int16)
newcharges = [charges[charge_labels]]
flows = [False]
return ChargeArray(data, newcharges, flows)
