def test_fuse_stride_arrays():
dims = np.asarray([2, 3, 4, 5])
strides = np.asarray([120, 60, 20, 5, 1])
actual = fuse_stride_arrays(dims, strides)
expected = fuse_ndarrays([
np.arange(0, strides[n] * dims[n], strides[n], dtype=np.uint32)
for n in range(len(dims))
])
np.testing.assert_allclose(actual, expected)
def reduce_charges(charges: List[BaseCharge],
flows: Union[np.ndarray, List[bool]],
target_charges: np.ndarray,
return_locations: Optional[bool] = False,
strides: Optional[np.ndarray] = None) -> Any:
"""
Add quantum numbers arising from combining two or more charges into a
single index, keeping only the quantum numbers that appear in
`target_charges`. Equilvalent to using "combine_charges" followed
by "reduce", but is generally much more efficient.
Args:
charges: List of `BaseCharge`, one for each leg of a
tensor.
flows: A list of bool, one for each leg of a tensor.
with values `False` or `True` denoting inflowing and
outflowing charge direction, respectively.
target_charges: n-by-D array of charges which should be kept,
with `n` the number of symmetries.
return_locations: If `True` return the location of the kept
values of the fused charges
strides: Index strides with which to compute the
retured locations of the kept elements. Defaults to trivial strides
(based on row major order).
Returns:
BaseCharge: the fused index after reduction.
np.ndarray: Locations of the fused BaseCharge charges that were kept.
"""
tensor_dims = [len(c) for c in charges]
if len(charges) == 1:
# reduce single index
if strides is None:
strides = np.array([1], dtype=SIZE_T)
return charges[0].dual(flows[0]).reduce(
target_charges, return_locations=return_locations, strides=strides[0])
# find size-balanced partition of charges
partition = _find_best_partition(tensor_dims)
# compute quantum numbers for each partition
left_ind = fuse_charges(charges[:partition], flows[:partition])
right_ind = fuse_charges(charges[partition:], flows[partition:])
# compute combined qnums
comb_qnums = fuse_ndarray_charges(left_ind.unique_charges,
right_ind.unique_charges,
charges[0].charge_types)
#special case of empty charges
#pylint: disable=unsubscriptable-object
if (comb_qnums.shape[0] == 0) or (len(left_ind.charge_labels) == 0) or (len(
right_ind.charge_labels) == 0):
obj = charges[0].__new__(type(charges[0]))
obj.__init__(
np.empty((0, charges[0].num_symmetries), dtype=charges[0].dtype),
np.empty(0, dtype=charges[0].label_dtype), charges[0].charge_types)
if return_locations:
return obj, np.empty(0, dtype=SIZE_T)
return obj
unique_comb_qnums, comb_labels = np.unique(
comb_qnums, return_inverse=True, axis=0)
num_unique = unique_comb_qnums.shape[0]
# intersect combined qnums and target_charges
reduced_qnums, label_to_unique, _ = intersect(
unique_comb_qnums, target_charges, axis=0, return_indices=True)
map_to_kept = -np.ones(num_unique, dtype=charges[0].label_dtype)
map_to_kept[label_to_unique] = np.arange(len(label_to_unique))
# new_comb_labels is a matrix of shape
# (left_ind.num_unique, right_ind.num_unique)
# each row new_comb_labels[n,:] contains integers values.
# Positions where values > 0
# denote labels of right-charges that are kept.
new_comb_labels = map_to_kept[comb_labels].reshape(
[left_ind.num_unique, right_ind.num_unique])
reduced_rows = [0] * left_ind.num_unique
for n in range(left_ind.num_unique):
temp_label = new_comb_labels[n, right_ind.charge_labels]
reduced_rows[n] = temp_label[temp_label >= 0]
reduced_labels = np.concatenate(
[reduced_rows[n] for n in left_ind.charge_labels])
obj = charges[0].__new__(type(charges[0]))
obj.__init__(reduced_qnums, reduced_labels, charges[0].charge_types)
if return_locations:
row_locs = [0] * left_ind.num_unique
if strides is not None:
# computed locations based on non-trivial strides
row_pos = fuse_stride_arrays(tensor_dims[:partition], strides[:partition])
col_pos = fuse_stride_arrays(tensor_dims[partition:], strides[partition:])
for n in range(left_ind.num_unique):
temp_label = new_comb_labels[n, right_ind.charge_labels]
temp_keep = temp_label >= 0
if strides is not None:
row_locs[n] = col_pos[temp_keep]
else:
row_locs[n] = np.where(temp_keep)[0]
if strides is not None:
reduced_locs = np.concatenate([
row_pos[n] + row_locs[left_ind.charge_labels[n]]
for n in range(left_ind.dim)
])
else:
reduced_locs = np.concatenate([
n * right_ind.dim + row_locs[left_ind.charge_labels[n]]
for n in range(left_ind.dim)
])
return obj, reduced_locs
return obj
def test_find_transposed_diagonal_sparse_blocks(num_charges, order, D):
order = list(order)
num_legs = len(order)
np.random.seed(10)
np_charges = [
np.random.randint(-5, 5, (num_charges, D), dtype=np.int16)
for _ in range(num_legs)
]
tr_charge_list = []
charge_list = []
for c in range(num_charges):
tr_charge_list.append(
fuse_ndarrays(
[np_charges[order[n]][c, :] for n in range(num_legs)]))
charge_list.append(
fuse_ndarrays([np_charges[n][c, :] for n in range(num_legs)]))
tr_fused = np.stack(tr_charge_list, axis=0)
fused = np.stack(charge_list, axis=0)
dims = [c.shape[1] for c in np_charges]
strides = _get_strides(dims)
transposed_linear_positions = fuse_stride_arrays(
dims, [strides[o] for o in order])
left_charges = np.stack([
fuse_ndarrays(
[np_charges[order[n]][c, :] for n in range(num_legs // 2)])
for c in range(num_charges)
],
axis=0)
right_charges = np.stack([
fuse_ndarrays([
np_charges[order[n]][c, :] for n in range(num_legs // 2, num_legs)
]) for c in range(num_charges)
],
axis=0)
#pylint: disable=no-member
mask = np.logical_and.reduce(fused.T == np.zeros((1, num_charges)), axis=1)
nz = np.nonzero(mask)[0]
dense_to_sparse = np.empty(len(mask), dtype=np.int64)
dense_to_sparse[mask] = np.arange(len(nz))
#pylint: disable=no-member
tr_mask = np.logical_and.reduce(tr_fused.T == np.zeros((1, num_charges)),
axis=1)
tr_nz = np.nonzero(tr_mask)[0]
tr_linear_locs = transposed_linear_positions[tr_nz]
# pylint: disable=no-member
left_inds, _ = np.divmod(tr_nz, right_charges.shape[1])
left = left_charges[:, left_inds]
unique_left = np.unique(left, axis=1)
blocks = []
for n in range(unique_left.shape[1]):
ul = unique_left[:, n][None, :]
#pylint: disable=no-member
blocks.append(dense_to_sparse[tr_linear_locs[np.nonzero(
np.logical_and.reduce(left.T == ul, axis=1))[0]]])
charges = [
BaseCharge(c, charge_types=[U1Charge] * num_charges)
for c in np_charges
]
flows = [False] * num_legs
bs, cs, ss = _find_transposed_diagonal_sparse_blocks(
charges, flows, tr_partition=num_legs // 2, order=order)
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)
