def compute_sparse_lookup(
charges: List[BaseCharge], flows: Union[np.ndarray, List[bool]],
target_charges: BaseCharge
) -> Tuple[np.ndarray, BaseCharge, np.ndarray]:
"""
Compute lookup table for how dense index positions map
to sparse index positions, treating only those elements as non-zero
whose charges fuse to `target_charges`.
Args:
charges: List of `BaseCharge` objects.
flows: A list of `bool`; the flow directions.
target_charges: A `BaseCharge`; the target charges for which
the fusion of `charges` is non-zero.
Returns:
lookup: An np.ndarray of positive numbers between `0` and
`len(unique_charges)`. The position of values `n` in `lookup`
are positions with charge values `unique_charges[n]`.
unique_charges: The unique charges of fusion of `charges`
label_to_unique: The integer labels of the unique charges.
"""
fused_charges = fuse_charges(charges, flows)
unique_charges, inverse = unique(fused_charges.charges,
return_inverse=True)
_, label_to_unique, _ = intersect(unique_charges,
target_charges.charges,
return_indices=True)
# _, label_to_unique, _ = unique_charges.intersect(
# target_charges, return_indices=True)
tmp = np.full(unique_charges.shape[0],
fill_value=-1,
dtype=charges[0].label_dtype)
obj = charges[0].__new__(type(charges[0]))
obj.__init__(charges=unique_charges,
charge_labels=None,
charge_types=charges[0].charge_types)
tmp[label_to_unique] = label_to_unique
lookup = tmp[inverse]
lookup = lookup[lookup >= 0]
return lookup, obj, np.sort(label_to_unique)
def __getitem__(self, n: Union[List[int], np.ndarray, int]) -> "BaseCharge":
"""
Return the charge-element at position `n`, wrapped into a `BaseCharge`
object.
Args:
n: An integer or `np.ndarray`.
Returns:
BaseCharge: The charges at `n`.
"""
if isinstance(n, (np.integer, int)):
n = np.asarray([n])
n = np.asarray(n)
obj = self.__new__(type(self))
if self._unique_charges is not None:
labels = self.charge_labels[n]
unique_labels, new_labels = unique(labels, return_inverse=True)
unique_charges = self.unique_charges[unique_labels, :]
obj.__init__(unique_charges, new_labels, self.charge_types)
return obj
obj.__init__(
self._charges[n, :], charge_labels=None, charge_types=self.charge_types)
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, (D, num_charges), 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=1)
fused = np.stack(charge_list, axis=1)
dims = [c.shape[0] 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=1)
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=1)
#pylint: disable=no-member
mask = np.logical_and.reduce(fused == 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 == 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[0])
left = left_charges[left_inds, :]
unique_left = unique(left)
blocks = []
for n in range(unique_left.shape[0]):
ul = unique_left[n, :][None, :]
#pylint: disable=no-member
blocks.append(dense_to_sparse[tr_linear_locs[np.nonzero(
np.logical_and.reduce(left == 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)
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 = unique(comb_qnums, return_inverse=True)
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 charge_labels(self):
if self._charge_labels is None:
self._unique_charges, self._charge_labels = unique(
self.charges, return_inverse=True)
self._charges = None
return self._charge_labels
def unique(
self, #pylint: disable=inconsistent-return-statements
return_index: bool = False,
return_inverse: bool = False,
return_counts: bool = False) -> Any:
"""
Compute the unique charges in `BaseCharge`.
See unique for a more detailed explanation. This function
does the same but instead of a np.ndarray, it returns the unique
elements (not neccessarily sorted in standard order) in a `BaseCharge`
object.
Args:
return_index: If `True`, also return the indices of `self.charges`
(along the specified axis,
if provided, or in the flattened array) that result in the unique array.
return_inverse: If `True`, also return the indices of the unique array
(for the specified
axis, if provided) that can be used to reconstruct `self.charges`.
return_counts: If `True`, also return the number of times each unique
item appears in `self.charges`.
Returns:
BaseCharge: The sorted unique values.
np.ndarray: The indices of the first occurrences of the unique values
in the original array. Only provided if `return_index` is True.
np.ndarray: The indices to reconstruct the original array from the
unique array. Only provided if `return_inverse` is True.
np.ndarray: The number of times each of the unique values comes up in the
original array. Only provided if `return_counts` is True.
"""
obj = self.__new__(type(self))
if self._charges is not None:
tmp = unique(self._charges,
return_index=return_index,
return_inverse=return_inverse,
return_counts=return_counts)
if any([return_index, return_inverse, return_counts]):
unique_charges = tmp[0]
obj.__init__(charges=unique_charges,
charge_labels=np.arange(unique_charges.shape[0],
dtype=self.label_dtype),
charge_types=self.charge_types)
tmp[0] = obj
else:
obj.__init__(charges=tmp,
charge_labels=np.arange(tmp.shape[0],
dtype=self.label_dtype),
charge_types=self.charge_types)
tmp = obj
return tmp
if self._unique_charges is not None:
if not return_index:
obj.__init__(charges=self._unique_charges,
charge_labels=np.arange(
self._unique_charges.shape[0],
dtype=self.label_dtype),
charge_types=self.charge_types)
out = [obj]
if return_inverse:
out.append(self._charge_labels)
if return_counts:
_, cnts = unique(self._charge_labels, return_counts=True)
out.append(cnts)
if len(out) > 1:
return out
return out[0]
tmp = unique(self._charge_labels,
return_index=return_index,
return_inverse=return_inverse,
return_counts=return_counts)
unique_charges = self._unique_charges[tmp[0], :]
obj.__init__(charges=unique_charges,
charge_labels=np.arange(unique_charges.shape[0],
dtype=self.label_dtype),
charge_types=self.charge_types)
tmp[0] = obj
return tmp
