def test_reduce_charges_non_trivial(num_charges):
np.random.seed(10)
left_charges = np.random.randint(-5, 6, (num_charges, 200), dtype=np.int16)
right_charges = np.random.randint(-5,
6, (num_charges, 200),
dtype=np.int16)
target_charge = np.random.randint(-2, 3, (num_charges, 3), dtype=np.int16)
charge_types = [U1Charge] * num_charges
fused_charges = fuse_ndarray_charges(left_charges, right_charges,
charge_types)
dense_positions = reduce_charges([
BaseCharge(left_charges, charge_types=charge_types),
BaseCharge(right_charges, charge_types=charge_types)
], [False, False],
target_charge,
return_locations=True)
assert np.all(
np.isin(np.squeeze(dense_positions[0].charges),
np.squeeze(target_charge)))
tmp = []
#pylint: disable=unsubscriptable-object
for n in range(target_charge.shape[1]):
#pylint: disable=no-member
tmp.append(
np.logical_and.reduce(fused_charges.T == target_charge[:,
n][None, :],
axis=1))
#pylint: disable=no-member
mask = np.logical_or.reduce(tmp)
np.testing.assert_allclose(dense_positions[1], np.nonzero(mask)[0])
def test_reduce_charges_2():
left_charges = np.asarray([[-2, 0, 1, 0, 0], [-3, 0, 2, 1,
0]]).astype(np.int16).T
right_charges = np.asarray([[-1, 0, 2, 1], [-2, 2, 7, 0]]).astype(np.int16).T
target_charge = np.zeros((1, 2), dtype=np.int16)
fused_charges = fuse_ndarray_charges(left_charges, right_charges,
[U1Charge, U1Charge])
dense_positions = reduce_charges([
BaseCharge(left_charges, charge_types=[U1Charge, U1Charge]),
BaseCharge(right_charges, charge_types=[U1Charge, U1Charge])
], [False, False],
target_charge,
return_locations=True)
np.testing.assert_allclose(dense_positions[0].charges, 0)
#pylint: disable=no-member
np.testing.assert_allclose(
dense_positions[1],
np.nonzero(np.logical_and.reduce(fused_charges == target_charge,
axis=1))[0])
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 fuse_many_ndarray_charges(charges, charge_types):
res = fuse_ndarray_charges(charges[0], charges[1], charge_types)
for n in range(2, len(charges)):
res = fuse_ndarray_charges(res, charges[n], charge_types)
return res
