def test_intersect_4():
a = np.array([0, 1, 2, 3, 4])
b = np.array([0, -1, 4])
out, la, lb = intersect(a, b, return_indices=True)
np.testing.assert_allclose([0, 4], out)
np.testing.assert_allclose(la, [0, 4])
np.testing.assert_allclose(lb, [0, 2])
def test_intersect_2():
a = np.array([[0, 1, 2], [2, 3, 4]])
b = np.array([[0, -2, 6, 2], [2, 3, 4, 4]])
out, la, lb = intersect(a, b, axis=1, return_indices=True)
np.testing.assert_allclose(np.array([[0, 2], [2, 4]]), out)
np.testing.assert_allclose(la, [0, 2])
np.testing.assert_allclose(lb, [0, 3])
def test_intersect_raises():
np.random.seed(10)
a = np.random.randint(0, 10, (4, 5))
b = np.random.randint(0, 10, (4, 6))
with pytest.raises(ValueError):
intersect(a, b, axis=0)
c = np.random.randint(0, 10, (3, 7))
with pytest.raises(ValueError):
intersect(a, c, axis=1)
with pytest.raises(NotImplementedError):
intersect(a, c, axis=2)
d = np.random.randint(0, 10, (3, 7, 3))
e = np.random.randint(0, 10, (3, 7, 3))
with pytest.raises(NotImplementedError):
intersect(d, e, axis=1)
def tensordot(
tensor1: BlockSparseTensor,
tensor2: BlockSparseTensor,
axes: Optional[Union[Sequence[Sequence[int]],
int]] = 2) -> BlockSparseTensor:
"""
Contract two `BlockSparseTensor`s along `axes`.
Args:
tensor1: First tensor.
tensor2: Second tensor.
axes: The axes to contract.
Returns:
BlockSparseTensor: The result of the tensor contraction.
"""
#process scalar input for `axes`
if isinstance(axes, (np.integer, int)):
axes = [
np.arange(tensor1.ndim - axes, tensor1.ndim, dtype=np.int16),
np.arange(0, axes, dtype=np.int16)
]
elif isinstance(axes[0], (np.integer, int)):
if len(axes) > 1:
raise ValueError(
"invalid input `axes = {}` to tensordot".format(axes))
axes = [np.array(axes, dtype=np.int16), np.array(axes, dtype=np.int16)]
axes1 = axes[0]
axes2 = axes[1]
if len(axes1) != len(axes2):
raise ValueError(
"`axes1 = {}` and `axes2 = {}` have to be of same length. ".format(
axes1, axes2))
if len(axes1) > len(tensor1.shape):
raise ValueError(
"`axes1 = {}` is incompatible with `tensor1.shape = {}. ".format(
axes1, tensor1.shape))
if len(axes2) > len(tensor2.shape):
raise ValueError(
"`axes2 = {}` is incompatible with `tensor2.shape = {}. ".format(
axes2, tensor2.shape))
if not np.all(np.unique(axes1) == np.sort(axes1)):
raise ValueError(
"Some values in axes[0] = {} appear more than once!".format(axes1))
if not np.all(np.unique(axes2) == np.sort(axes2)):
raise ValueError(
"Some values in axes[1] = {} appear more than once!".format(axes2))
#special case outer product
if len(axes1) == 0:
return outerproduct(tensor1, tensor2)
#more checks
if max(axes1) >= len(tensor1.shape):
raise ValueError(
"rank of `tensor1` is smaller than `max(axes1) = {}.`".format(
max(axes1)))
if max(axes2) >= len(tensor2.shape):
raise ValueError(
"rank of `tensor2` is smaller than `max(axes2) = {}`".format(
max(axes1)))
contr_flows_1 = []
contr_flows_2 = []
contr_charges_1 = []
contr_charges_2 = []
for a in axes1:
contr_flows_1.extend(tensor1._flows[tensor1._order[a]])
contr_charges_1.extend(
[tensor1._charges[n] for n in tensor1._order[a]])
for a in axes2:
contr_flows_2.extend(tensor2._flows[tensor2._order[a]])
contr_charges_2.extend(
[tensor2._charges[n] for n in tensor2._order[a]])
if len(contr_charges_2) != len(contr_charges_1):
raise ValueError(
"`axes1 = {}` and `axes2 = {}` have incompatible elementary"
" shapes {} and {}".format(axes1, axes2,
[e.dim for e in contr_charges_1],
[e.dim for e in contr_charges_2]))
if not np.all(
np.asarray(contr_flows_1) == np.logical_not(
np.asarray(contr_flows_2))):
raise ValueError(
"`axes1 = {}` and `axes2 = {}` have incompatible elementary"
" flows {} and {}".format(axes1, axes2, contr_flows_1,
contr_flows_2))
charge_check = [
charge_equal(c1, c2)
for c1, c2 in zip(contr_charges_1, contr_charges_2)
]
if not np.all(charge_check):
inds = np.nonzero(np.logical_not(charge_check))[0]
raise ValueError(
"`axes = {}` of tensor1 and `axes = {}` of tensor2 have incompatible charges"
" {} and {}".format(
np.array(axes1)[inds],
np.array(axes2)[inds], [contr_charges_1[i] for i in inds],
[contr_charges_2[i] for i in inds]))
#checks finished
#special case inner product
if (len(axes1) == tensor1.ndim) and (len(axes2) == tensor2.ndim):
t1 = tensor1.transpose(axes1).transpose_data()
t2 = tensor2.transpose(axes2).transpose_data()
data = np.dot(t1.data, t2.data)
charge = tensor1._charges[0]
final_charge = charge.__new__(type(charge))
final_charge.__init__(np.empty((charge.num_symmetries, 0),
dtype=np.int16),
charge_labels=np.empty(0, dtype=np.int16),
charge_types=charge.charge_types)
return BlockSparseTensor(data=data,
charges=[final_charge],
flows=[False],
order=[[0]],
check_consistency=False)
#in all other cases we perform a regular tensordot
free_axes1 = sorted(set(np.arange(tensor1.ndim)) - set(axes1))
free_axes2 = sorted(set(np.arange(tensor2.ndim)) - set(axes2))
new_order1 = [tensor1._order[n]
for n in free_axes1] + [tensor1._order[n] for n in axes1]
new_order2 = [tensor2._order[n]
for n in axes2] + [tensor2._order[n] for n in free_axes2]
flat_order_1 = flatten(new_order1)
flat_order_2 = flatten(new_order2)
flat_charges_1, flat_flows_1 = tensor1._charges, tensor1.flat_flows
flat_charges_2, flat_flows_2 = tensor2._charges, tensor2.flat_flows
left_charges = []
right_charges = []
left_flows = []
right_flows = []
left_order = []
right_order = []
s = 0
for n in free_axes1:
left_charges.extend([tensor1._charges[o] for o in tensor1._order[n]])
left_order.append(list(np.arange(s, s + len(tensor1._order[n]))))
s += len(tensor1._order[n])
left_flows.extend([tensor1._flows[o] for o in tensor1._order[n]])
s = 0
for n in free_axes2:
right_charges.extend([tensor2._charges[o] for o in tensor2._order[n]])
right_order.append(
list(len(left_charges) + np.arange(s, s + len(tensor2._order[n]))))
s += len(tensor2._order[n])
right_flows.extend([tensor2._flows[o] for o in tensor2._order[n]])
tr_sparse_blocks_1, charges1, shapes_1 = _find_transposed_diagonal_sparse_blocks(
flat_charges_1, flat_flows_1, len(left_charges), flat_order_1)
tr_sparse_blocks_2, charges2, shapes_2 = _find_transposed_diagonal_sparse_blocks(
flat_charges_2, flat_flows_2, len(contr_charges_2), flat_order_2)
common_charges, label_to_common_1, label_to_common_2 = intersect(
charges1.unique_charges,
charges2.unique_charges,
axis=1,
return_indices=True)
#Note: `cs` may contain charges that are not present in `common_charges`
charges = left_charges + right_charges
flows = left_flows + right_flows
sparse_blocks, cs, _ = _find_diagonal_sparse_blocks(
charges, flows, len(left_charges))
num_nonzero_elements = np.int64(np.sum([len(v) for v in sparse_blocks]))
#Note that empty is not a viable choice here.
data = np.zeros(num_nonzero_elements,
dtype=np.result_type(tensor1.dtype, tensor2.dtype))
label_to_common_final = intersect(cs.unique_charges,
common_charges,
axis=1,
return_indices=True)[1]
for n in range(common_charges.shape[1]):
n1 = label_to_common_1[n]
n2 = label_to_common_2[n]
nf = label_to_common_final[n]
data[sparse_blocks[nf].ravel()] = np.ravel(
np.matmul(
tensor1.data[tr_sparse_blocks_1[n1].reshape(shapes_1[:, n1])],
tensor2.data[tr_sparse_blocks_2[n2].reshape(shapes_2[:, n2])]))
res = BlockSparseTensor(data=data,
charges=charges,
flows=flows,
order=left_order + right_order,
check_consistency=False)
return res
def test_intersect_3():
a = np.array([0, 1, 2, 3, 4])
b = np.array([0, -1, 4])
out = intersect(a, b)
np.testing.assert_allclose([0, 4], out)
def test_intersect_1():
a = np.array([[0, 1, 2], [2, 3, 4]])
b = np.array([[0, -2, 6], [2, 3, 4]])
out = intersect(a, b, axis=1)
np.testing.assert_allclose(np.array([[0], [2]]), out)
def _find_transposed_diagonal_sparse_blocks(
charges: List[BaseCharge],
flows: Union[np.ndarray, List[bool]],
tr_partition: int,
order: Optional[Union[List, np.ndarray]] = None
) -> Tuple[List, BaseCharge, np.ndarray]:
"""
Find the diagonal blocks of a transposed tensor with
meta-data `charges` and `flows`. `charges` and `flows`
are the charges and flows of the untransposed tensor,
`order` is the final transposition, and `tr_partition`
is the partition of the transposed tensor according to
which the diagonal blocks should be found.
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.
tr_partition: Location of the transposed tensor partition (i.e. such that the
tensor is viewed as a matrix between `charges[order[:partition]]` and
`charges[order[partition:]]`).
order: Order with which to permute the tensor axes.
Returns:
block_maps (List[np.ndarray]): list of integer arrays, which each
containing the location of a symmetry block in the data vector.
block_qnums (BaseCharge): The charges of the corresponding blocks.
block_dims (np.ndarray): 2-by-m array of matrix dimensions of each block.
"""
flows = np.asarray(flows)
if np.array_equal(order, None) or (np.array_equal(
np.array(order), np.arange(len(charges)))):
# no transpose order
return _find_diagonal_sparse_blocks(charges, flows, tr_partition)
# general case: non-trivial transposition is required
num_inds = len(charges)
tensor_dims = np.array([charges[n].dim for n in range(num_inds)],
dtype=int)
strides = np.append(np.flip(np.cumprod(np.flip(tensor_dims[1:]))), 1)
# compute qnums of row/cols in original tensor
orig_partition = _find_best_partition(tensor_dims)
orig_width = np.prod(tensor_dims[orig_partition:])
orig_unique_row_qnums = compute_unique_fused_charges(
charges[:orig_partition], flows[:orig_partition])
orig_unique_col_qnums, orig_col_degen = compute_fused_charge_degeneracies(
charges[orig_partition:], np.logical_not(flows[orig_partition:]))
orig_block_qnums, row_map, col_map = intersect(
orig_unique_row_qnums.unique_charges,
orig_unique_col_qnums.unique_charges,
axis=1,
return_indices=True)
orig_num_blocks = orig_block_qnums.shape[1]
if orig_num_blocks == 0:
# special case: trivial number of non-zero elements
obj = charges[0].__new__(type(charges[0]))
obj.__init__(
np.empty((charges[0].num_symmetries, 0), dtype=charges[0].dtype),
np.arange(0, dtype=charges[0].label_dtype),
charges[0].charge_types)
return [], obj, np.empty((2, 0), dtype=SIZE_T)
orig_row_ind = fuse_charges(charges[:orig_partition],
flows[:orig_partition])
orig_col_ind = fuse_charges(charges[orig_partition:],
np.logical_not(flows[orig_partition:]))
inv_row_map = -np.ones(orig_unique_row_qnums.unique_charges.shape[1],
dtype=charges[0].label_dtype)
inv_row_map[row_map] = np.arange(len(row_map),
dtype=charges[0].label_dtype)
all_degens = np.append(orig_col_degen[col_map],
0)[inv_row_map[orig_row_ind.charge_labels]]
all_cumul_degens = np.cumsum(np.insert(all_degens[:-1], 0,
0)).astype(SIZE_T)
dense_to_sparse = np.empty(orig_width, dtype=SIZE_T)
for n in range(orig_num_blocks):
dense_to_sparse[orig_col_ind.charge_labels == col_map[n]] = np.arange(
orig_col_degen[col_map[n]], dtype=SIZE_T)
# define properties of new tensor resulting from transposition
new_strides = strides[order]
new_row_charges = [charges[n] for n in order[:tr_partition]]
new_col_charges = [charges[n] for n in order[tr_partition:]]
new_row_flows = flows[order[:tr_partition]]
new_col_flows = flows[order[tr_partition:]]
if tr_partition == 0:
# special case: reshape into row vector
# compute qnums of row/cols in transposed tensor
unique_col_qnums, new_col_degen = compute_fused_charge_degeneracies(
new_col_charges, np.logical_not(new_col_flows))
identity_charges = charges[0].identity_charges
block_qnums, new_row_map, new_col_map = intersect(
identity_charges.unique_charges,
unique_col_qnums.unique_charges,
axis=1,
return_indices=True)
block_dims = np.array([[1], new_col_degen[new_col_map]], dtype=SIZE_T)
num_blocks = 1
col_ind, col_locs = reduce_charges(new_col_charges,
np.logical_not(new_col_flows),
block_qnums,
return_locations=True,
strides=new_strides[tr_partition:])
# find location of blocks in transposed tensor (w.r.t positions in original)
#pylint: disable=no-member
orig_row_posR, orig_col_posR = np.divmod(
col_locs[col_ind.charge_labels == 0], orig_width)
block_maps = [(all_cumul_degens[orig_row_posR] +
dense_to_sparse[orig_col_posR]).ravel()]
obj = charges[0].__new__(type(charges[0]))
obj.__init__(
block_qnums,
np.arange(block_qnums.shape[1], dtype=charges[0].label_dtype),
charges[0].charge_types)
elif tr_partition == len(charges):
# special case: reshape into col vector
# compute qnums of row/cols in transposed tensor
unique_row_qnums, new_row_degen = compute_fused_charge_degeneracies(
new_row_charges, new_row_flows)
identity_charges = charges[0].identity_charges
block_qnums, new_row_map, new_col_map = intersect(
unique_row_qnums.unique_charges,
identity_charges.unique_charges,
axis=1,
return_indices=True)
block_dims = np.array([new_row_degen[new_row_map], [1]], dtype=SIZE_T)
num_blocks = 1
row_ind, row_locs = reduce_charges(new_row_charges,
new_row_flows,
block_qnums,
return_locations=True,
strides=new_strides[:tr_partition])
# find location of blocks in transposed tensor (w.r.t positions in original)
#pylint: disable=no-member
orig_row_posL, orig_col_posL = np.divmod(
row_locs[row_ind.charge_labels == 0], orig_width)
block_maps = [(all_cumul_degens[orig_row_posL] +
dense_to_sparse[orig_col_posL]).ravel()]
obj = charges[0].__new__(type(charges[0]))
obj.__init__(
block_qnums,
np.arange(block_qnums.shape[1], dtype=charges[0].label_dtype),
charges[0].charge_types)
else:
unique_row_qnums, new_row_degen = compute_fused_charge_degeneracies(
new_row_charges, new_row_flows)
unique_col_qnums, new_col_degen = compute_fused_charge_degeneracies(
new_col_charges, np.logical_not(new_col_flows))
block_qnums, new_row_map, new_col_map = intersect(
unique_row_qnums.unique_charges,
unique_col_qnums.unique_charges,
axis=1,
return_indices=True)
block_dims = np.array(
[new_row_degen[new_row_map], new_col_degen[new_col_map]],
dtype=SIZE_T)
num_blocks = len(new_row_map)
row_ind, row_locs = reduce_charges(new_row_charges,
new_row_flows,
block_qnums,
return_locations=True,
strides=new_strides[:tr_partition])
col_ind, col_locs = reduce_charges(new_col_charges,
np.logical_not(new_col_flows),
block_qnums,
return_locations=True,
strides=new_strides[tr_partition:])
block_maps = [0] * num_blocks
for n in range(num_blocks):
#pylint: disable=no-member
orig_row_posL, orig_col_posL = np.divmod(
row_locs[row_ind.charge_labels == n], orig_width)
#pylint: disable=no-member
orig_row_posR, orig_col_posR = np.divmod(
col_locs[col_ind.charge_labels == n], orig_width)
block_maps[n] = (
all_cumul_degens[np.add.outer(orig_row_posL, orig_row_posR)] +
dense_to_sparse[np.add.outer(orig_col_posL,
orig_col_posR)]).ravel()
obj = charges[0].__new__(type(charges[0]))
obj.__init__(
block_qnums,
np.arange(block_qnums.shape[1], dtype=charges[0].label_dtype),
charges[0].charge_types)
return block_maps, obj, block_dims
def _find_diagonal_sparse_blocks(
charges: List[BaseCharge], flows: Union[np.ndarray, List[bool]],
partition: int) -> Tuple[List, BaseCharge, np.ndarray]:
"""
Find the location of all non-trivial symmetry blocks from the data vector of
of BlockSparseTensor (when viewed as a matrix across some prescribed index
bi-partition).
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.
partition: location of tensor partition (i.e. such that the
tensor is viewed as a matrix between `charges[:partition]` and
the remaining charges).
Returns:
block_maps (List[np.ndarray]): list of integer arrays, which each
containing the location of a symmetry block in the data vector.
block_qnums (BaseCharge): The charges of the corresponding blocks.n
block, with 'n' the number of symmetries and 'm' the number of blocks.
block_dims (np.ndarray): 2-by-m array of matrix dimensions of each block.
"""
num_inds = len(charges)
if partition in (0, num_inds):
# special cases (matrix of trivial height or width)
num_nonzero = compute_num_nonzero(charges, flows)
block_maps = [np.arange(0, num_nonzero, dtype=SIZE_T).ravel()]
block_qnums = charges[0].identity_charges.charges
block_dims = np.array([[1], [num_nonzero]])
if partition == len(flows):
block_dims = np.flipud(block_dims)
obj = charges[0].__new__(type(charges[0]))
obj.__init__(block_qnums, np.arange(0, dtype=charges[0].label_dtype),
charges[0].charge_types)
return block_maps, obj, block_dims
unique_row_qnums, row_degen = compute_fused_charge_degeneracies(
charges[:partition], flows[:partition])
unique_col_qnums, col_degen = compute_fused_charge_degeneracies(
charges[partition:], np.logical_not(flows[partition:]))
block_qnums, row_to_block, col_to_block = intersect(
unique_row_qnums.unique_charges,
unique_col_qnums.unique_charges,
axis=1,
return_indices=True)
num_blocks = block_qnums.shape[1]
if num_blocks == 0:
obj = charges[0].__new__(type(charges[0]))
obj.__init__(
np.zeros((charges[0].num_symmetries, 0), dtype=charges[0].dtype),
np.arange(0, dtype=charges[0].label_dtype),
charges[0].charge_types)
return [], obj, np.empty((2, 0), dtype=SIZE_T)
# calculate number of non-zero elements in each row of the matrix
row_ind = reduce_charges(charges[:partition], flows[:partition],
block_qnums)
row_num_nz = col_degen[col_to_block[row_ind.charge_labels]]
cumulate_num_nz = np.insert(np.cumsum(row_num_nz[0:-1]), 0,
0).astype(SIZE_T)
# calculate mappings for the position in datavector of each block
if num_blocks < 15:
# faster method for small number of blocks
row_locs = np.concatenate([
(row_ind.charge_labels == n) for n in range(num_blocks)
]).reshape(num_blocks, row_ind.dim)
else:
# faster method for large number of blocks
row_locs = np.zeros([num_blocks, row_ind.dim], dtype=bool)
row_locs[row_ind.charge_labels,
np.arange(row_ind.dim)] = np.ones(row_ind.dim, dtype=bool)
block_dims = np.array(
[[row_degen[row_to_block[n]], col_degen[col_to_block[n]]]
for n in range(num_blocks)],
dtype=SIZE_T).T
#pylint: disable=unsubscriptable-object
block_maps = [
np.ravel(cumulate_num_nz[row_locs[n, :]][:, None] +
np.arange(block_dims[1, n])[None, :])
for n in range(num_blocks)
]
obj = charges[0].__new__(type(charges[0]))
obj.__init__(block_qnums,
np.arange(block_qnums.shape[1], dtype=charges[0].label_dtype),
charges[0].charge_types)
return block_maps, obj, block_dims
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[1] == 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((charges[0].num_symmetries, 0), 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=1)
num_unique = unique_comb_qnums.shape[1]
# intersect combined qnums and target_charges
reduced_qnums, label_to_unique, _ = intersect(unique_comb_qnums,
target_charges,
axis=1,
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