def intersect(self,
other,
assume_unique=False,
return_indices=False) -> Any:
"""
Compute the intersection of `self` with `other`. See also np.intersect1d.
Args:
other: A BaseCharge object.
assume_unique: If `True` assume that elements are unique.
return_indices: If `True`, return index-labels.
Returns:
If `return_indices=True`:
BaseCharge
np.ndarray: The indices of the first occurrences of the
common values in `self`.
np.ndarray: The indices of the first occurrences of the
common values in `other`.
If `return_indices=False`:
BaseCharge
"""
if isinstance(other, type(self)):
out = intersect(self.charges,
other.charges,
assume_unique=assume_unique,
axis=0,
return_indices=return_indices)
else:
if other.ndim == 1:
other = other[:, None]
out = intersect(self.charges,
np.asarray(other),
axis=0,
assume_unique=assume_unique,
return_indices=return_indices)
obj = self.__new__(type(self))
if return_indices:
obj.__init__(
charges=out[0],
charge_labels=np.arange(out[0].shape[0],
dtype=self.label_dtype),
charge_types=self.charge_types,
)
return obj, out[1], out[2]
obj.__init__(
charges=out,
charge_labels=np.arange(out.shape[0], dtype=self.label_dtype),
charge_types=self.charge_types,
)
return obj
def test_intersect_6(dtype):
a = np.array([[0, 2], [1, 3], [2, 4]], dtype=dtype)
b = np.array([[0, 2], [-2, 3], [6, 4], [2, 4]], dtype=dtype)
out, la, lb = intersect(a, b, axis=0, 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_4(dtype):
a = np.array([0, 1, 2, 3, 4], dtype=dtype)
b = np.array([0, -1, 4], dtype=dtype)
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_ndarray_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_1d(dtype):
a = np.random.randint(-5, 5, 10, dtype=dtype)
b = np.random.randint(-2, 2, 8, dtype=dtype)
out, la, lb = intersect(a, b, axis=0, return_indices=True)
out_, la_, lb_ = np.intersect1d(a, b, return_indices=True)
np.testing.assert_allclose(out, out_)
np.testing.assert_allclose(la, la_)
np.testing.assert_allclose(lb, lb_)
def test_intersect_ndarray_raises():
np.random.seed(10)
a = np.random.randint(0, 10, (4, 5, 1))
b = np.random.randint(0, 10, (4, 6))
with pytest.raises(ValueError, match="array ndims"):
intersect(a, b, axis=0)
a = np.random.randint(0, 10, (4, 5))
b = np.random.randint(0, 10, (4, 6))
with pytest.raises(ValueError, match="array widths"):
intersect(a, b, axis=0)
with pytest.raises(NotImplementedError, match="intersection can only"):
intersect(a, b, axis=2)
d = np.random.randint(0, 10, (3, 7, 3))
e = np.random.randint(0, 10, (3, 7, 3))
with pytest.raises(NotImplementedError, match="_intersect_ndarray is only"):
intersect(d, e, axis=1)
def reduce(self,
target_charges: Union[int, np.ndarray],
return_locations: bool = False,
strides: Optional[int] = 1) -> Any:
"""
Reduce the dimension of a
charge to keep only the charge values that intersect target_charges
Args:
target_charges: array of unique charges to keep.
return_locations: If `True`, also return the locations of
target values within `BaseCharge`.
strides: An optional stride value.
Returns:
BaseCharge: charge of reduced dimension.
np.ndarray: If `return_locations = True`; the index locations
of target values.
"""
if isinstance(target_charges, (np.integer, int)):
target_charges = np.asarray([target_charges], dtype=self.dtype)
if target_charges.ndim == 1:
target_charges = target_charges[:, None]
target_charges = np.asarray(target_charges, dtype=self.dtype)
# find intersection of index charges and target charges
reduced_charges, label_to_unique, _ = intersect(self.unique_charges,
target_charges,
axis=0,
return_indices=True)
num_unique = len(label_to_unique)
# construct the map to the reduced charges
map_to_reduced = np.full(self.dim,
fill_value=-1,
dtype=self.label_dtype)
map_to_reduced[label_to_unique] = np.arange(num_unique,
dtype=self.label_dtype)
# construct the map to the reduced charges
reduced_ind_labels = map_to_reduced[self.charge_labels]
reduced_locs = reduced_ind_labels >= 0
new_ind_labels = reduced_ind_labels[reduced_locs].astype(
self.label_dtype)
obj = self.__new__(type(self))
obj.__init__(reduced_charges, new_ind_labels, self.charge_types)
if return_locations:
return obj, strides * np.flatnonzero(reduced_locs).astype(
np.uint32)
return obj
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 _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)
cacher = get_cacher()
if cacher.do_caching:
hash_val = _to_string(charges, flows, tr_partition, order)
if hash_val in cacher.cache:
return cacher.cache[hash_val]
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=0,
return_indices=True)
orig_num_blocks = orig_block_qnums.shape[0]
if orig_num_blocks == 0:
# special case: trivial number of non-zero elements
obj = charges[0].__new__(type(charges[0]))
obj.__init__(
np.empty((0, charges[0].num_symmetries), 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[0],
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(dim=1)
block_qnums, new_row_map, new_col_map = intersect(
identity_charges.unique_charges,
unique_col_qnums.unique_charges,
axis=0,
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[0], 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(dim=1)
block_qnums, new_row_map, new_col_map = intersect(
unique_row_qnums.unique_charges,
identity_charges.unique_charges,
axis=0,
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[0], 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=0,
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[0], dtype=charges[0].label_dtype),
charges[0].charge_types)
if cacher.do_caching:
cacher.cache[hash_val] = (block_maps, obj, block_dims)
return cacher.cache[hash_val]
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.
"""
cacher = get_cacher()
if cacher.do_caching:
hash_val = _to_string(charges, flows, partition, list(range(len(charges))))
if hash_val in cacher.cache:
return cacher.cache[hash_val]
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(dim=1).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(1, 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=0,
return_indices=True)
num_blocks = block_qnums.shape[0]
if num_blocks == 0:
obj = charges[0].__new__(type(charges[0]))
obj.__init__(
np.zeros((0, charges[0].num_symmetries), 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[0], dtype=charges[0].label_dtype),
charges[0].charge_types)
if cacher.do_caching:
cacher.cache[hash_val] = (block_maps, obj, block_dims)
return cacher.cache[hash_val]
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[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_intersect_5(dtype):
a = np.array([[0, 2], [1, 3], [2, 4]], dtype=dtype)
b = np.array([[0, 2], [-2, 3], [6, 6]], dtype=dtype)
out = intersect(a, b, axis=0)
np.testing.assert_allclose(np.array([[0, 2]]), out)
def test_intersect_3(dtype):
a = np.array([0, 1, 2, 3, 4], dtype=dtype)
b = np.array([0, -1, 4], dtype=dtype)
out = intersect(a, b)
np.testing.assert_allclose([0, 4], out)
def test_intersect_1(dtype):
a = np.array([[0, 1, 2], [2, 3, 4]], dtype=dtype)
b = np.array([[0, -2, 6], [2, 3, 4]], dtype=dtype)
out = intersect(a, b, axis=1)
np.testing.assert_allclose(np.array([[0], [2]]), out)
def tensordot(
tensor1: BlockSparseTensor,
tensor2: BlockSparseTensor,
axes: Optional[Union[Sequence[Sequence[int]], 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))
#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 (returns an ndim=0 tensor)
if (len(axes1) == tensor1.ndim) and (len(axes2) == tensor2.ndim):
t1 = tensor1.transpose(axes1).contiguous()
t2 = tensor2.transpose(axes2).contiguous()
return BlockSparseTensor(data=np.dot(t1.data, t2.data),
charges=[],
flows=[],
order=[],
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._flows
flat_charges_2, flat_flows_2 = tensor2._charges, tensor2._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( #pylint: disable=line-too-long
flat_charges_1, flat_flows_1, len(left_charges), flat_order_1)
tr_sparse_blocks_2, charges2, shapes_2 = _find_transposed_diagonal_sparse_blocks( #pylint: disable=line-too-long
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=0,
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_transposed_diagonal_sparse_blocks(
charges, flows, len(left_charges), list(range(len(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=0,
return_indices=True)[1]
for n in range(common_charges.shape[0]):
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_ndarray_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)
