diff --git a/tensornetwork/block_tensor/block_tensor.py b/tensornetwork/block_tensor/block_tensor.py index 5c7d2ed93..8ab2f3898 100644 --- a/tensornetwork/block_tensor/block_tensor.py +++ b/tensornetwork/block_tensor/block_tensor.py @@ -27,6 +27,25 @@ Tensor = Any +def get_flat_order(indices, order): + flat_charges, _ = get_flat_meta_data(indices) + flat_labels = np.arange(len(flat_charges)) + cum_num_legs = np.append(0, np.cumsum([len(i.flat_charges) for i in indices])) + flat_order = np.concatenate( + [flat_labels[cum_num_legs[n]:cum_num_legs[n + 1]] for n in order]) + + return flat_order + + +def get_flat_meta_data(indices): + charges = [] + flows = [] + for i in indices: + flows.extend(i.flat_flows) + charges.extend(i.flat_charges) + return charges, flows + + def fuse_stride_arrays(dims: np.ndarray, strides: np.ndarray) -> np.ndarray: return fuse_ndarrays([ np.arange(0, strides[n] * dims[n], strides[n], dtype=np.uint32) @@ -193,11 +212,10 @@ def compute_num_nonzero(charges: List[BaseCharge], flows: List[bool]) -> int: charges, flows) res = accumulated_charges == accumulated_charges.identity_charges nz_inds = np.nonzero(res)[0] - if len(nz_inds) == 0: - raise ValueError( - "given leg-charges `charges` and flows `flows` are incompatible " - "with a symmetric tensor") - return np.squeeze(accumulated_degeneracies[nz_inds][0]) + + if len(nz_inds) > 0: + return np.squeeze(accumulated_degeneracies[nz_inds][0]) + return 0 def reduce_charges( @@ -326,7 +344,7 @@ def _find_diagonal_sparse_blocks( partition: int) -> (np.ndarray, np.ndarray, np.ndarray): """ Find the location of all non-trivial symmetry blocks from the data vector of - of SymTensor (when viewed as a matrix across some prescribed index + of BlockSparseTensor (when viewed as a matrix across some prescribed index bi-partition). Args: charges (List[SymIndex]): list of SymIndex. @@ -355,7 +373,11 @@ def _find_diagonal_sparse_blocks( if partition == len(flows): block_dims = np.flipud(block_dims) - return block_maps, block_qnums, block_dims + obj = charges[0].__new__(type(charges[0])) + obj.__init__(block_qnums, np.arange(block_qnums.shape[1], dtype=np.int16), + charges[0].charge_types) + + return block_maps, obj, block_dims else: unique_row_qnums, row_degen = compute_fused_charge_degeneracies( @@ -419,12 +441,9 @@ def _find_transposed_diagonal_sparse_blocks( tr_partition: int, order: np.ndarray = None) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: """ - Find the location of all non-trivial symmetry blocks from the data vector of - of SymTensor after transposition (and then viewed as a matrix across some - prescribed index bi-tr_partition). Produces and equivalent result to - retrieve_blocks acting on a transposed SymTensor, but is much faster. + Args: - charges (List[SymIndex]): list of SymIndex. + charges (List[BaseCharge]): List of charges. flows (np.ndarray): vector of bools describing index orientations. tr_partition (int): location of tensor partition (i.e. such that the tensor is viewed as a matrix between first partition charges and @@ -444,18 +463,115 @@ def _find_transposed_diagonal_sparse_blocks( # 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 + return [], np.array([], dtype=np.uint32), np.array([], dtype=np.uint32) + + 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=np.int16) + for n in range(len(row_map)): + inv_row_map[row_map[n]] = n + + 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(np.uint32) + dense_to_sparse = np.zeros(orig_width, dtype=np.uint32) + 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=np.uint32) + + # 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=np.uint32) + 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) + 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=np.int16), + 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=np.uint32) + 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) + 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=np.int16), + charges[0].charge_types) else: - # 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) - - # 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:]] unique_row_qnums, new_row_degen = compute_fused_charge_degeneracies( new_row_charges, new_row_flows) @@ -471,6 +587,7 @@ def _find_transposed_diagonal_sparse_blocks( [new_row_degen[new_row_map], new_col_degen[new_col_map]], dtype=np.uint32) num_blocks = len(new_row_map) + row_ind, row_locs = reduce_charges( new_row_charges, new_row_flows, @@ -484,38 +601,6 @@ def _find_transposed_diagonal_sparse_blocks( block_qnums, return_locations=True, strides=new_strides[tr_partition:]) - 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] - - 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=np.int16) - for n in range(len(row_map)): - inv_row_map[row_map[n]] = n - - 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(np.uint32) - dense_to_sparse = np.zeros(orig_width, dtype=np.uint32) - 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=np.uint32) block_maps = [0] * num_blocks for n in range(num_blocks): @@ -530,7 +615,7 @@ def _find_transposed_diagonal_sparse_blocks( obj.__init__(block_qnums, np.arange(block_qnums.shape[1], dtype=np.int16), charges[0].charge_types) - return block_maps, obj, block_dims + return block_maps, obj, block_dims class BlockSparseTensor: @@ -540,25 +625,9 @@ class BlockSparseTensor: The class currently only supports a single U(1) symmetry and only numpy.ndarray. - Attributes: - * self.data: A 1d np.ndarray storing the underlying - data of the tensor - * self.charges: A list of `np.ndarray` of shape - (D,), where D is the bond dimension. Once we go beyond - a single U(1) symmetry, this has to be updated. - - * self.flows: A list of integers of length `k`. - `self.flows` determines the flows direction of charges - on each leg of the tensor. A value of `-1` denotes - outflowing charge, a value of `1` denotes inflowing - charge. - The tensor data is stored in self.data, a 1d np.ndarray. """ - def copy(self): - return BlockSparseTensor(self.data.copy(), [i.copy() for i in self.indices]) - def __init__(self, data: np.ndarray, indices: List[Index]) -> None: """ Args: @@ -568,7 +637,8 @@ def __init__(self, data: np.ndarray, indices: List[Index]) -> None: indices: List of `Index` objecst, one for each leg. """ self.indices = indices - num_non_zero_elements = compute_num_nonzero(self.charges, self.flows) + num_non_zero_elements = compute_num_nonzero(self.flat_charges, + self.flat_flows) if num_non_zero_elements != len(data.flat): raise ValueError("number of tensor elements {} defined " @@ -578,18 +648,25 @@ def __init__(self, data: np.ndarray, indices: List[Index]) -> None: self.data = np.asarray(data.flat) #do not copy data + def copy(self): + return BlockSparseTensor(self.data.copy(), [i.copy() for i in self.indices]) + def todense(self) -> np.ndarray: """ Map the sparse tensor to dense storage. """ out = np.asarray(np.zeros(self.dense_shape, dtype=self.dtype).flat) - charges = self.charges + charges = self.flat_charges out[np.nonzero( - fuse_charges(charges, self.flows) == charges[0].identity_charges) + fuse_charges(charges, self.flat_flows) == charges[0].identity_charges) [0]] = self.data return np.reshape(out, self.dense_shape) + @property + def ndim(self): + return len(self.indices) + @classmethod def randn(cls, indices: List[Index], dtype: Optional[Type[np.number]] = None) -> "BlockSparseTensor": @@ -601,8 +678,7 @@ def randn(cls, indices: List[Index], Returns: BlockSparseTensor """ - charges = [i.charges for i in indices] - flows = [i.flow for i in indices] + charges, flows = get_flat_meta_data(indices) num_non_zero_elements = compute_num_nonzero(charges, flows) backend = backend_factory.get_backend('numpy') data = backend.randn((num_non_zero_elements,), dtype=dtype) @@ -619,8 +695,7 @@ def ones(cls, indices: List[Index], Returns: BlockSparseTensor """ - charges = [i.charges for i in indices] - flows = [i.flow for i in indices] + charges, flows = get_flat_meta_data(indices) num_non_zero_elements = compute_num_nonzero(charges, flows) backend = backend_factory.get_backend('numpy') data = backend.ones((num_non_zero_elements,), dtype=dtype) @@ -637,8 +712,7 @@ def zeros(cls, indices: List[Index], Returns: BlockSparseTensor """ - charges = [i.charges for i in indices] - flows = [i.flow for i in indices] + charges, flows = get_flat_meta_data(indices) num_non_zero_elements = compute_num_nonzero(charges, flows) backend = backend_factory.get_backend('numpy') data = backend.zeros((num_non_zero_elements,), dtype=dtype) @@ -655,9 +729,7 @@ def random(cls, indices: List[Index], Returns: BlockSparseTensor """ - charges = [i.charges for i in indices] - flows = [i.flow for i in indices] - + charges, flows = get_flat_meta_data(indices) num_non_zero_elements = compute_num_nonzero(charges, flows) dtype = dtype if dtype is not None else np.float64 @@ -683,7 +755,7 @@ def dense_shape(self) -> Tuple: Returns: Tuple: A tuple of `int`. """ - return tuple([i.dimension for i in self.indices]) + return tuple([i.dim for i in self.indices]) @property def shape(self) -> Tuple: @@ -706,6 +778,20 @@ def flows(self): def charges(self): return [i.charges for i in self.indices] + @property + def flat_charges(self): + flat = [] + for i in self.indices: + flat.extend(i.flat_charges) + return flat + + @property + def flat_flows(self): + flat = [] + for i in self.indices: + flat.extend(i.flat_flows) + return flat + def transpose( self, order: Union[List[int], np.ndarray], @@ -724,9 +810,10 @@ def transpose( #check for trivial permutation if np.all(order == np.arange(len(order))): - return self - flat_indices, flat_charges, flat_flows, _, flat_order = flatten_meta_data( - self.indices, order) + return BlockSparseTensor(self.data, self.indices) + flat_charges, flat_flows = get_flat_meta_data(self.indices) + flat_order = get_flat_order(self.indices, order) + print(flat_order) tr_partition = _find_best_partition( [len(flat_charges[n]) for n in flat_order]) @@ -742,25 +829,8 @@ def transpose( ind = np.nonzero(tr_charges == charges[n])[0][0] permutation = tr_sparse_blocks[ind] data[sparse_block] = self.data[permutation] - self.indices = [self.indices[o] for o in order] - self.data = data - return self - def reset_shape(self) -> None: - """ - Bring the tensor back into its elementary shape. - """ - self.indices = self.get_elementary_indices() - - def get_elementary_indices(self) -> List: - """ - Compute the elementary indices of the array. - """ - elementary_indices = [] - for i in self.indices: - elementary_indices.extend(i.get_elementary_indices()) - - return elementary_indices + return BlockSparseTensor(data, [self.indices[o] for o in order]) def reshape(self, shape: Union[Iterable[Index], Iterable[int]]) -> None: """ @@ -801,52 +871,38 @@ def reshape(self, shape: Union[Iterable[Index], Iterable[int]]) -> None: new_shape = [] for s in shape: if isinstance(s, Index): - new_shape.append(s.dimension) + new_shape.append(s.dim) else: new_shape.append(s) + # a few simple checks if np.prod(new_shape) != np.prod(self.dense_shape): raise ValueError("A tensor with {} elements cannot be " "reshaped into a tensor with {} elements".format( - np.prod(self.shape), np.prod(self.dense_shape))) - - #keep a copy of the old indices for the case where reshaping fails - #FIXME: this is pretty hacky! - indices = [i.copy() for i in self.indices] - flat_indices = [] - for i in indices: - flat_indices.extend(i.get_elementary_indices()) - - def raise_error(): - #if this error is raised then `shape` is incompatible - #with the elementary indices. We then reset the shape - #to what is was before the call to `reshape`. - # self.indices = index_copy - raise ValueError("The shape {} is incompatible with the " - "elementary shape {} of the tensor.".format( - new_shape, - tuple([e.dimension for e in flat_indices]))) + np.prod(self.shape), np.prod(new_shape))) + + flat_charges, flat_flows = get_flat_meta_data(self.indices) + flat_dims = [f.dim for f in flat_charges] + partitions = [0] for n in range(len(new_shape)): - if new_shape[n] > flat_indices[n].dimension: - while new_shape[n] > flat_indices[n].dimension: - #fuse flat_indices - i1, i2 = flat_indices.pop(n), flat_indices.pop(n) - #note: the resulting flow is set to one since the flow - #is multiplied into the charges. As a result the tensor - #will then be invariant in any case. - flat_indices.insert(n, fuse_index_pair(i1, i2)) - if flat_indices[n].dimension > new_shape[n]: - raise_error() - elif new_shape[n] < flat_indices[n].dimension: - raise_error() - #at this point the first len(new_shape) flat_indices of the tensor - #match the `new_shape`. - while len(new_shape) < len(flat_indices): - i2, i1 = flat_indices.pop(), flat_indices.pop() - flat_indices.append(fuse_index_pair(i1, i2)) - - result = BlockSparseTensor(data=self.data, indices=flat_indices) + tmp = np.nonzero(np.cumprod(flat_dims) == new_shape[n])[0] + if len(tmp) == 0: + raise ValueError("The shape {} is incompatible with the " + "elementary shape {} of the tensor.".format( + new_shape, tuple([e.dim for e in flat_charges]))) + + partitions.append(tmp[0] + 1) + flat_dims = flat_dims[partitions[-1]:] + partitions = np.cumsum(partitions) + new_flat_charges = [] + new_flat_flows = [] + for n in range(1, len(partitions)): + new_flat_charges.append(flat_charges[partitions[n - 1]:partitions[n]]) + new_flat_flows.append(flat_flows[partitions[n - 1]:partitions[n]]) + + indices = [Index(c, f) for c, f in zip(new_flat_charges, new_flat_flows)] + result = BlockSparseTensor(data=self.data, indices=indices) return result @@ -901,15 +957,40 @@ def transpose(tensor: BlockSparseTensor, Returns: BlockSparseTensor: The transposed tensor. """ - result = tensor.copy() - result.transpose(order) - return result + return tensor.transpose() + + +def outerproduct(tensor1: BlockSparseTensor, + tensor2: BlockSparseTensor) -> BlockSparseTensor: + """ + Compute the outer product of two BlockSparseTensor. + Args: + tensor1: A tensor. + tensor2: A tensor. + Returns: + BlockSparseTensor: The result of taking the outer product. + """ + + final_charges = tensor1.flat_charges + tensor2.flat_charges + final_flows = tensor1.flat_flows + tensor2.flat_flows + data = np.zeros( + compute_num_nonzero(final_charges, final_flows), dtype=tensor1.dtype) + if ((len(tensor1.data) > 0) and (len(tensor2.data) > 0)) and (len(data) > 0): + # find the location of the zero block in the output + final_block_maps, final_block_charges, final_block_dims = _find_diagonal_sparse_blocks( + final_charges, final_flows, len(tensor1.flat_charges)) + index = np.nonzero( + final_block_charges == final_block_charges.identity_charges)[0][0] + data[final_block_maps[index].ravel()] = np.outer(tensor1.data, + tensor2.data).ravel() + + return BlockSparseTensor(data, tensor1.indices + tensor2.indices) def tensordot( tensor1: BlockSparseTensor, tensor2: BlockSparseTensor, - axes: Sequence[Sequence[int]], + axes: Optional[Union[Sequence[Sequence[int]], int]] = 2, final_order: Optional[Union[List, np.ndarray]] = None) -> BlockSparseTensor: """ Contract two `BlockSparseTensor`s along `axes`. @@ -922,6 +1003,14 @@ def tensordot( BlockSparseTensor: The result of the tensor contraction. """ + + 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)): + axes = [np.array(axes, dtype=np.int16), np.array(axes, dtype=np.int16)] axes1 = axes[0] axes2 = axes[1] if not np.all(np.unique(axes1) == np.sort(axes1)): @@ -931,6 +1020,22 @@ def tensordot( raise ValueError( "Some values in axes[1] = {} appear more than once!".format(axes2n)) + if len(axes1) == 0: + res = outerproduct(tensor1, tensor2) + if final_order is not None: + return res.transpose(final_order) + return res + + if (len(axes1) == tensor1.ndim) and (len(axes2) == tensor2.ndim): + isort = np.argsort(axes1) + data = np.dot(tensor1.data, + tensor2.transpose(np.asarray(axes2)[isort]).data) + if len(tensor1.indices[0].flat_charges) > 0: + identity_charges = tensor1.indices[0].flat_charges[0].identity_charges + + return BlockSparseTensor( + data=data, indices=[Index(identity_charges, flow=False)]) + if max(axes1) >= len(tensor1.shape): raise ValueError( "rank of `tensor1` is smaller than `max(axes1) = {}.`".format( @@ -940,22 +1045,27 @@ def tensordot( raise ValueError( "rank of `tensor2` is smaller than `max(axes2) = {}`".format( max(axes1))) - elementary_1, elementary_2 = [], [] + + contr_flows_1 = [] + contr_flows_2 = [] + contr_charges_1 = [] + contr_charges_2 = [] for a in axes1: - elementary_1.extend(tensor1.indices[a].get_elementary_indices()) + contr_flows_1.extend(tensor1.indices[a].flat_flows) + contr_charges_1.extend(tensor1.indices[a].flat_charges) for a in axes2: - elementary_2.extend(tensor2.indices[a].get_elementary_indices()) + contr_flows_2.extend(tensor2.indices[a].flat_flows) + contr_charges_2.extend(tensor2.indices[a].flat_charges) - if len(elementary_2) != len(elementary_1): - raise ValueError("axes1 and axes2 have incompatible elementary" - " shapes {} and {}".format(elementary_1, elementary_2)) + if len(contr_charges_2) != len(contr_charges_1): + raise ValueError( + "axes1 and axes2 have incompatible elementary" + " shapes {} and {}".format([e.dim for e in contr_charges_1], + [e.dim for e in contr_charges_2])) if not np.all( - np.array([i.flow for i in elementary_1]) == np.array( - [not i.flow for i in elementary_2])): + np.asarray(contr_flows_1) == np.logical_not(np.asarray(contr_flows_2))): raise ValueError("axes1 and axes2 have incompatible elementary" - " flows {} and {}".format( - np.array([i.flow for i in elementary_1]), - np.array([i.flow for i in elementary_2]))) + " flows {} and {}".format(contr_flows_1, contr_flows_2)) free_axes1 = sorted(set(np.arange(len(tensor1.shape))) - set(axes1)) free_axes2 = sorted(set(np.arange(len(tensor2.shape))) - set(axes2)) @@ -975,43 +1085,31 @@ def tensordot( new_order1 = free_axes1 + list(axes1) new_order2 = list(axes2) + free_axes2 - contr_flat_indices_1 = [] - for n in axes1: - contr_flat_indices_1.extend(tensor1.indices[n].get_elementary_indices()) + flat_charges_1, flat_flows_1 = get_flat_meta_data(tensor1.indices) + flat_charges_2, flat_flows_2 = get_flat_meta_data(tensor2.indices) - contr_flat_indices_2 = [] - for n in axes2: - contr_flat_indices_2.extend(tensor2.indices[n].get_elementary_indices()) - #get the flattened indices for the output tensor - left_indices = [] - right_indices = [] + flat_order_1 = get_flat_order(tensor1.indices, new_order1) + flat_order_2 = get_flat_order(tensor2.indices, new_order2) + + left_charges = [] + right_charges = [] + left_flows = [] + right_flows = [] + free_indices = [] for n in free_axes1: - left_indices.extend(tensor1.indices[n].get_elementary_indices()) + free_indices.append(tensor1.indices[n]) + left_charges.extend(tensor1.indices[n].flat_charges) + left_flows.extend(tensor1.indices[n].flat_flows) for n in free_axes2: - right_indices.extend(tensor2.indices[n].get_elementary_indices()) - - indices = left_indices + right_indices - - flat_charges1 = [i.charges for i in left_indices - ] + [i.charges for i in contr_flat_indices_1] - flat_flows1 = [i.flow for i in left_indices - ] + [i.flow for i in contr_flat_indices_1] + free_indices.append(tensor2.indices[n]) + right_charges.extend(tensor2.indices[n].flat_charges) + right_flows.extend(tensor2.indices[n].flat_flows) - flat_charges2 = [i.charges for i in contr_flat_indices_2 - ] + [i.charges for i in right_indices] - - flat_flows2 = [i.flow for i in contr_flat_indices_2 - ] + [i.flow for i in right_indices] - - flat_order1 = new_flat_order(tensor1.indices, new_order1) - flat_order2 = new_flat_order(tensor2.indices, new_order2) - - tr_partition1 = len(left_indices) - tr_partition2 = len(contr_flat_indices_2) tr_sparse_blocks_1, charges1, shapes_1 = _find_transposed_diagonal_sparse_blocks( - flat_charges1, flat_flows1, tr_partition1, flat_order1) + 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_charges2, flat_flows2, tr_partition2, flat_order2) + 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, @@ -1023,14 +1121,14 @@ def tensordot( if final_order is not None: #in this case we view the result of the diagonal multiplication #as a transposition of the final tensor - final_indices = [indices[n] for n in final_order] + final_indices = [free_indices[n] for n in final_order] _, reverse_order = np.unique(final_order, return_index=True) - - flat_final_indices, flat_final_charges, flat_final_flows, flat_final_strides, flat_final_order, tr_partition = flatten_meta_data( - final_indices, reverse_order, len(free_axes1)) + flat_reversed_order = get_flat_order(final_indices, reverse_order) + flat_final_charges, flat_final_flows = get_flat_meta_data(final_indices) sparse_blocks_final, charges_final, shapes_final = _find_transposed_diagonal_sparse_blocks( - flat_final_charges, flat_final_flows, tr_partition, flat_final_order) + flat_final_charges, flat_final_flows, len(left_charges), + flat_reversed_order) num_nonzero_elements = np.sum([len(v) for v in sparse_blocks_final]) data = np.zeros( @@ -1055,10 +1153,10 @@ def tensordot( return BlockSparseTensor(data=data, indices=final_indices) else: #Note: `cs` may contain charges that are not present in `common_charges` - charges = [i.charges for i in indices] - flows = [i.flow for i in indices] + charges = left_charges + right_charges + flows = left_flows + right_flows sparse_blocks, cs, shapes = _find_diagonal_sparse_blocks( - charges, flows, len(left_indices)) + charges, flows, len(left_charges)) num_nonzero_elements = np.sum([len(v) for v in sparse_blocks]) #Note that empty is not a viable choice here. data = np.zeros( @@ -1077,40 +1175,4 @@ def tensordot( tensor1.data[tr_sparse_blocks_1[n1].reshape( shapes_1[:, n1])] @ tensor2.data[tr_sparse_blocks_2[n2].reshape( shapes_2[:, n2])]).ravel() - return BlockSparseTensor(data=data, indices=indices) - - -def new_flat_order(indices, order): - elementary_indices = {} - flat_elementary_indices = [] - for n in range(len(indices)): - elementary_indices[n] = indices[n].get_elementary_indices() - flat_elementary_indices.extend(elementary_indices[n]) - flat_index_list = np.arange(len(flat_elementary_indices)) - cum_num_legs = np.append( - 0, np.cumsum([len(elementary_indices[n]) for n in range(len(indices))])) - - flat_order = np.concatenate( - [flat_index_list[cum_num_legs[n]:cum_num_legs[n + 1]] for n in order]) - - return flat_order - - -def flatten_meta_data(indices, order): - elementary_indices = {} - flat_elementary_indices = [] - for n in range(len(indices)): - elementary_indices[n] = indices[n].get_elementary_indices() - flat_elementary_indices.extend(elementary_indices[n]) - flat_index_list = np.arange(len(flat_elementary_indices)) - cum_num_legs = np.append( - 0, np.cumsum([len(elementary_indices[n]) for n in range(len(indices))])) - - flat_charges = [i.charges for i in flat_elementary_indices] - flat_flows = [i.flow for i in flat_elementary_indices] - flat_dims = [len(c) for c in flat_charges] - flat_strides = _get_strides(flat_dims) - flat_order = np.concatenate( - [flat_index_list[cum_num_legs[n]:cum_num_legs[n + 1]] for n in order]) - - return flat_elementary_indices, flat_charges, flat_flows, flat_strides, flat_order + return BlockSparseTensor(data=data, indices=free_indices) diff --git a/tensornetwork/block_tensor/block_tensor_test.py b/tensornetwork/block_tensor/block_tensor_test.py index 2c8ada9c0..adc77fc6a 100644 --- a/tensornetwork/block_tensor/block_tensor_test.py +++ b/tensornetwork/block_tensor/block_tensor_test.py @@ -8,6 +8,52 @@ np_dtypes = [np.float32, np.float16, np.float64, np.complex64, np.complex128] +def get_contractable_tensors(R1, R2, cont): + DsA = np.random.randint(5, 10, R1) + DsB = np.random.randint(5, 10, R2) + assert R1 >= cont + assert R2 >= cont + chargesA = [ + U1Charge(np.random.randint(-5, 5, DsA[n])) for n in range(R1 - cont) + ] + commoncharges = [ + U1Charge(np.random.randint(-5, 5, DsA[n + R1 - cont])) + for n in range(cont) + ] + chargesB = [ + U1Charge(np.random.randint(-5, 5, DsB[n])) for n in range(R2 - cont) + ] + #contracted indices + indsA = np.random.choice(np.arange(R1), cont, replace=False) + indsB = np.random.choice(np.arange(R2), cont, replace=False) + + flowsA = np.full(R1, False, dtype=np.bool) + flowsB = np.full(R2, False, dtype=np.bool) + flowsB[indsB] = True + + indicesA = [None for _ in range(R1)] + indicesB = [None for _ in range(R2)] + for n in range(len(indsA)): + indicesA[indsA[n]] = Index(commoncharges[n], flowsA[indsA[n]]) + indicesB[indsB[n]] = Index(commoncharges[n], flowsB[indsB[n]]) + compA = list(set(np.arange(R1)) - set(indsA)) + compB = list(set(np.arange(R2)) - set(indsB)) + + for n in range(len(compA)): + indicesA[compA[n]] = Index(chargesA[n], flowsA[compA[n]]) + for n in range(len(compB)): + indicesB[compB[n]] = Index(chargesB[n], flowsB[compB[n]]) + indices_final = [] + for n in sorted(compA): + indices_final.append(indicesA[n]) + for n in sorted(compB): + indices_final.append(indicesB[n]) + shapes = tuple([i.dim for i in indices_final]) + A = BlockSparseTensor.random(indices=indicesA) + B = BlockSparseTensor.random(indices=indicesB) + return A, B, indsA, indsB + + @pytest.mark.parametrize("dtype", np_dtypes) def test_block_sparse_init(dtype): D = 10 #bond dimension @@ -23,17 +69,16 @@ def test_block_sparse_init(dtype): Index(charges=charges[n], flow=flows[n], name='index{}'.format(n)) for n in range(rank) ] - num_elements = compute_num_nonzero([i.charges for i in indices], - [i.flow for i in indices]) + num_elements = compute_num_nonzero(charges, flows) A = BlockSparseTensor.random(indices=indices, dtype=dtype) assert A.dtype == dtype for r in range(rank): - assert A.indices[r].name == 'index{}'.format(r) + assert A.indices[r].name[0] == 'index{}'.format(r) assert A.dense_shape == tuple([D] * rank) assert len(A.data) == num_elements -def test_find_dense_positions(): +def test_reduce_charges(): left_charges = np.asarray([-2, 0, 1, 0, 0]).astype(np.int16) right_charges = np.asarray([-1, 0, 2, 1]).astype(np.int16) target_charge = np.zeros((1, 1), dtype=np.int16) @@ -49,7 +94,7 @@ def test_find_dense_positions(): def test_transpose(): R = 4 - Ds = [20, 3, 4, 5] + Ds = np.random.randint(10, 20, R) final_order = np.arange(R) np.random.shuffle(final_order) charges = [U1Charge(np.random.randint(-5, 5, Ds[n])) for n in range(R)] @@ -58,253 +103,94 @@ def test_transpose(): A = BlockSparseTensor.random(indices=indices) Adense = A.todense() dense_res = np.transpose(Adense, final_order) - A.transpose(final_order) - np.testing.assert_allclose(dense_res, A.todense()) + B = A.transpose(final_order) + np.testing.assert_allclose(dense_res, B.todense()) -def test_tensordot(): +def test_reshape(): R = 4 - DsA = [10, 12, 14, 16] - DsB = [14, 16, 18, 20] - chargesA = [U1Charge(np.random.randint(-5, 5, DsA[n])) for n in range(R // 2)] - commoncharges = [ - U1Charge(np.random.randint(-5, 5, DsA[n + R // 2])) for n in range(R // 2) - ] - chargesB = [ - U1Charge(np.random.randint(-5, 5, DsB[n + R // 2])) for n in range(R // 2) - ] - indsA = np.random.choice(np.arange(R), R // 2, replace=False) - indsB = np.random.choice(np.arange(R), R // 2, replace=False) - flowsA = np.full(R, False, dtype=np.bool) - flowsB = np.full(R, False, dtype=np.bool) + Ds = [3, 4, 5, 6] + charges = [U1Charge(np.random.randint(-5, 5, Ds[n])) for n in range(R)] + flows = np.full(R, fill_value=False, dtype=np.bool) + indices = [Index(charges[n], flows[n]) for n in range(R)] + A = BlockSparseTensor.random(indices=indices) + B = A.reshape([Ds[0] * Ds[1], Ds[2], Ds[3]]) + Adense = A.todense() + Bdense = Adense.reshape([Ds[0] * Ds[1], Ds[2], Ds[3]]) + np.testing.assert_allclose(Bdense, B.todense()) - flowsB[indsB] = True - indicesA = [None for _ in range(R)] - indicesB = [None for _ in range(R)] - for n in range(len(indsA)): - indicesA[indsA[n]] = Index(commoncharges[n], flowsA[indsA[n]]) - indicesB[indsB[n]] = Index(commoncharges[n], flowsB[indsB[n]]) - compA = list(set(np.arange(R)) - set(indsA)) - compB = list(set(np.arange(R)) - set(indsB)) - for n in range(len(compA)): - indicesA[compA[n]] = Index(chargesA[n], flowsA[compA[n]]) - indicesB[compB[n]] = Index(chargesB[n], flowsB[compB[n]]) - indices_final = [] - for n in sorted(compA): - indices_final.append(indicesA[n]) - for n in sorted(compB): - indices_final.append(indicesB[n]) - shapes = tuple([i.dimension for i in indices_final]) - A = BlockSparseTensor.random(indices=indicesA) - B = BlockSparseTensor.random(indices=indicesB) +def test_reshape_transpose(): + R = 4 + Ds = [3, 4, 5, 6] + charges = [U1Charge(np.random.randint(-5, 5, Ds[n])) for n in range(R)] + flows = np.full(R, fill_value=False, dtype=np.bool) + indices = [Index(charges[n], flows[n]) for n in range(R)] + A = BlockSparseTensor.random(indices=indices) + B = A.reshape([Ds[0] * Ds[1], Ds[2], Ds[3]]).transpose([2, 0, 1]) + dense = A.todense().reshape([Ds[0] * Ds[1], Ds[2], + Ds[3]]).transpose([2, 0, 1]) + np.testing.assert_allclose(dense, B.todense()) - final_order = np.arange(R) + +@pytest.mark.parametrize("R1, R2, cont", [(4, 4, 2), (4, 3, 3), (3, 4, 3)]) +def test_tensordot(R1, R2, cont): + A, B, indsA, indsB = get_contractable_tensors(R1, R2, cont) + res = tensordot(A, B, (indsA, indsB)) + dense_res = np.tensordot(A.todense(), B.todense(), (indsA, indsB)) + np.testing.assert_allclose(dense_res, res.todense()) + + +def test_tensordot_reshape(): + R1 = 4 + R2 = 4 + + q = np.random.randint(-5, 5, 10, dtype=np.int16) + charges1 = [U1Charge(q) for n in range(R1)] + charges2 = [U1Charge(q) for n in range(R2)] + flowsA = np.asarray([False] * R1) + flowsB = np.asarray([True] * R2) + A = BlockSparseTensor.random(indices=[ + Index(charges1[n], flowsA[n], name='a{}'.format(n)) for n in range(R1) + ]) + B = BlockSparseTensor.random(indices=[ + Index(charges2[n], flowsB[n], name='b{}'.format(n)) for n in range(R2) + ]) + + Adense = A.todense().reshape((10, 10 * 10, 10)) + Bdense = B.todense().reshape((10 * 10, 10, 10)) + + A = A.reshape((10, 10 * 10, 10)) + B = B.reshape((10 * 10, 10, 10)) + + res = tensordot(A, B, ([0, 1], [2, 0])) + dense = np.tensordot(Adense, Bdense, ([0, 1], [2, 0])) + np.testing.assert_allclose(dense, res.todense()) + + +@pytest.mark.parametrize("R1, R2, cont", [(4, 4, 2), (4, 3, 3), (3, 4, 3)]) +def test_tensordot_final_order(R1, R2, cont): + A, B, indsA, indsB = get_contractable_tensors(R1, R2, cont) + final_order = np.arange(R1 + R2 - 2 * cont) np.random.shuffle(final_order) - Adense = A.todense() - Bdense = B.todense() + res = tensordot(A, B, (indsA, indsB), final_order=final_order) dense_res = np.transpose( - np.tensordot(Adense, Bdense, (indsA, indsB)), final_order) + np.tensordot(A.todense(), B.todense(), (indsA, indsB)), final_order) + np.testing.assert_allclose(dense_res, res.todense()) - res = tensordot(A, B, (indsA, indsB), final_order=final_order) + +@pytest.mark.parametrize("R1, R2", [(2, 2), (3, 3), (4, 4), (1, 1)]) +def test_tensordot_inner(R1, R2): + + A, B, indsA, indsB = get_contractable_tensors(R1, R2, 0) + res = tensordot(A, B, (indsA, indsB)) + dense_res = np.tensordot(A.todense(), B.todense(), (indsA, indsB)) np.testing.assert_allclose(dense_res, res.todense()) -# def test_find_dense_positions_2(): -# D = 40 #bond dimension -# B = 4 #number of blocks -# dtype = np.int16 #the dtype of the quantum numbers -# rank = 4 -# flows = np.asarray([1 for _ in range(rank)]) -# flows[-2::] = -1 -# charges = [ -# np.random.randint(-B // 2, B // 2 + 1, D).astype(dtype) -# for _ in range(rank) -# ] -# indices = [ -# Index( -# charges=U1Charge(charges[n]), flow=flows[n], name='index{}'.format(n)) -# for n in range(rank) -# ] -# n1 = compute_num_nonzero([i.charges for i in indices], -# [i.flow for i in indices]) - -# i01 = indices[0] * indices[1] -# i23 = indices[2] * indices[3] -# positions = find_dense_positions([i01.charges, i23.charges], [1, 1], -# U1Charge(np.asarray([0]))) -# assert len(positions[0]) == n1 - -# def test_find_sparse_positions(): -# D = 40 #bond dimension -# B = 4 #number of blocks -# dtype = np.int16 #the dtype of the quantum numbers -# rank = 4 -# flows = np.asarray([1 for _ in range(rank)]) -# flows[-2::] = -1 -# charges = [ -# np.random.randint(-B // 2, B // 2 + 1, D).astype(dtype) -# for _ in range(rank) -# ] -# indices = [ -# Index( -# charges=U1Charge(charges[n]), flow=flows[n], name='index{}'.format(n)) -# for n in range(rank) -# ] -# n1 = compute_num_nonzero([i.charges for i in indices], -# [i.flow for i in indices]) -# i01 = indices[0] * indices[1] -# i23 = indices[2] * indices[3] -# unique_row_charges = np.unique(i01.charges.charges) -# unique_column_charges = np.unique(i23.charges.charges) -# common_charges = np.intersect1d( -# unique_row_charges, -unique_column_charges, assume_unique=True) -# blocks = find_sparse_positions([i01.charges, i23.charges], [1, 1], -# target_charges=U1Charge(np.asarray([0]))) -# assert sum([len(v) for v in blocks.values()]) == n1 -# np.testing.assert_allclose(np.sort(blocks[0]), np.arange(n1)) - -# def test_find_sparse_positions_2(): -# D = 1000 #bond dimension -# B = 4 #number of blocks -# dtype = np.int16 #the dtype of the quantum numbers -# charges = np.random.randint(-B // 2, B // 2 + 1, D).astype(dtype) -# index = Index(charges=U1Charge(charges), flow=1, name='index0') -# targets = np.asarray([-1, 0, 1]) -# blocks = find_sparse_positions([index.charges], [index.flow], -# target_charges=U1Charge(targets)) - -# inds = np.isin(charges, targets) -# relevant_charges = charges[inds] -# blocks_ = {t: np.nonzero(relevant_charges == t)[0] for t in targets} -# assert np.all( -# np.asarray(list(blocks.keys())) == np.asarray(list(blocks_.keys()))) -# for k in blocks.keys(): -# assert np.all(blocks[k] == blocks_[k]) - -# def test_find_sparse_positions_3(): -# D = 40 #bond dimension -# B = 4 #number of blocks -# dtype = np.int16 #the dtype of the quantum numbers -# flows = [1, -1] - -# rank = len(flows) -# charges = [ -# np.random.randint(-B // 2, B // 2 + 1, D).astype(dtype) -# for _ in range(rank) -# ] -# indices = [ -# Index( -# charges=U1Charge(charges[n]), flow=flows[n], name='index{}'.format(n)) -# for n in range(rank) -# ] -# i1, i2 = indices -# common_charges = np.intersect1d(i1.charges.charges, i2.charges.charges) -# row_locations = find_sparse_positions( -# charges=[i1.charges, i2.charges], -# flows=flows, -# target_charges=U1Charge(common_charges)) -# fused = (i1 * i2).charges -# relevant = fused.charges[np.isin(fused.charges, common_charges)] -# for k, v in row_locations.items(): -# np.testing.assert_allclose(np.nonzero(relevant == k)[0], np.sort(v)) - -# # def test_dense_transpose(): -# # Ds = [10, 11, 12] #bond dimension -# # rank = len(Ds) -# # flows = np.asarray([1 for _ in range(rank)]) -# # flows[-2::] = -1 -# # charges = [U1Charge(np.zeros(Ds[n], dtype=np.int16)) for n in range(rank)] -# # indices = [ -# # Index(charges=charges[n], flow=flows[n], name='index{}'.format(n)) -# # for n in range(rank) -# # ] -# # A = BlockSparseTensor.random(indices=indices, dtype=np.float64) -# # B = np.transpose(np.reshape(A.data.copy(), Ds), (1, 0, 2)) -# # A.transpose((1, 0, 2)) -# # np.testing.assert_allclose(A.data, B.flat) - -# # B = np.transpose(np.reshape(A.data.copy(), [11, 10, 12]), (1, 0, 2)) -# # A.transpose((1, 0, 2)) - -# # np.testing.assert_allclose(A.data, B.flat) - -# @pytest.mark.parametrize("R", [1, 2]) -# def test_find_diagonal_dense_blocks(R): -# rs = [U1Charge(np.random.randint(-4, 4, 50)) for _ in range(R)] -# cs = [U1Charge(np.random.randint(-4, 4, 50)) for _ in range(R)] -# charges = rs + cs - -# left_fused = fuse_charges(charges[0:R], [1] * R) -# right_fused = fuse_charges(charges[R:], [1] * R) -# left_unique = left_fused.unique() -# right_unique = right_fused.unique() -# zero = left_unique.zero_charge -# blocks = {} -# rdim = len(right_fused) -# for lu in left_unique: -# linds = np.nonzero(left_fused == lu)[0] -# rinds = np.nonzero(right_fused == lu * (-1))[0] -# if (len(linds) > 0) and (len(rinds) > 0): -# blocks[lu] = fuse_ndarrays([linds * rdim, rinds]) -# comm, blocks_ = _find_diagonal_dense_blocks(rs, cs, [1] * R, [1] * R) -# for n in range(len(comm)): -# assert np.all(blocks[comm.charges[n]] == blocks_[n][0]) - -# # #@pytest.mark.parametrize("dtype", np_dtypes) -# # def test_find_diagonal_dense_blocks_2(): -# # R = 1 -# # rs = [U1Charge(np.random.randint(-4, 4, 50)) for _ in range(R)] -# # cs = [U1Charge(np.random.randint(-4, 4, 50)) for _ in range(R)] -# # charges = rs + cs - -# # left_fused = fuse_charges(charges[0:R], [1] * R) -# # right_fused = fuse_charges(charges[R:], [1] * R) -# # left_unique = left_fused.unique() -# # right_unique = right_fused.unique() -# # zero = left_unique.zero_charge -# # blocks = {} -# # rdim = len(right_fused) -# # for lu in left_unique: -# # linds = np.nonzero(left_fused == lu)[0] -# # rinds = np.nonzero(right_fused == lu * (-1))[0] -# # if (len(linds) > 0) and (len(rinds) > 0): -# # blocks[lu] = fuse_ndarrays([linds * rdim, rinds]) -# # comm, blocks_ = _find_diagonal_dense_blocks(rs, cs, [1] * R, [1] * R) -# # for n in range(len(comm)): -# # assert np.all(blocks[comm.charges[n]] == blocks_[n][0]) - -# @pytest.mark.parametrize("R", [1, 2]) -# def test_find_diagonal_dense_blocks_transposed(R): -# order = np.arange(2 * R) -# np.random.shuffle(order) -# rs = [U1Charge(np.random.randint(-4, 4, 50)) for _ in range(R)] -# cs = [U1Charge(np.random.randint(-4, 4, 40)) for _ in range(R)] -# charges = rs + cs -# dims = np.asarray([len(c) for c in charges]) -# strides = np.flip(np.append(1, np.cumprod(np.flip(dims[1::])))) -# stride_arrays = [np.arange(dims[n]) * strides[n] for n in range(2 * R)] - -# left_fused = fuse_charges([charges[n] for n in order[0:R]], [1] * R) -# right_fused = fuse_charges([charges[n] for n in order[R:]], [1] * R) -# lstrides = fuse_ndarrays([stride_arrays[n] for n in order[0:R]]) -# rstrides = fuse_ndarrays([stride_arrays[n] for n in order[R:]]) - -# left_unique = left_fused.unique() -# right_unique = right_fused.unique() -# blocks = {} -# rdim = len(right_fused) -# for lu in left_unique: -# linds = np.nonzero(left_fused == lu)[0] -# rinds = np.nonzero(right_fused == lu * (-1))[0] -# if (len(linds) > 0) and (len(rinds) > 0): -# tmp = fuse_ndarrays([linds * rdim, rinds]) -# blocks[lu] = _find_values_in_fused(tmp, lstrides, rstrides) - -# comm, blocks_ = _find_diagonal_dense_blocks([charges[n] for n in order[0:R]], -# [charges[n] for n in order[R:]], -# [1] * R, [1] * R, -# row_strides=strides[order[0:R]], -# column_strides=strides[order[R:]]) -# for n in range(len(comm)): -# assert np.all(blocks[comm.charges[n]] == blocks_[n][0]) +@pytest.mark.parametrize("R1, R2", [(2, 2), (2, 1), (1, 2), (1, 1)]) +def test_tensordot_outer(R1, R2): + A, B, indsA, indsB = get_contractable_tensors(R1, R2, 0) + res = tensordot(A, B, axes=0) + dense_res = np.tensordot(A.todense(), B.todense(), axes=0) + np.testing.assert_allclose(dense_res, res.todense()) diff --git a/tensornetwork/block_tensor/charge.py b/tensornetwork/block_tensor/charge.py index e5cf1d094..3aa223492 100644 --- a/tensornetwork/block_tensor/charge.py +++ b/tensornetwork/block_tensor/charge.py @@ -121,6 +121,13 @@ def dual(self, take_dual: Optional[bool] = False) -> np.ndarray: return obj return self + def copy(self): + obj = self.__new__(type(self)) + obj.__init__( + charges=self.unique_charges.copy(), + charge_labels=self.charge_labels.copy(), + charge_types=self.charge_types) + @property def charges(self): return self.unique_charges[:, self.charge_labels] diff --git a/tensornetwork/block_tensor/index.py b/tensornetwork/block_tensor/index.py index fc2d033d9..111924b4d 100644 --- a/tensornetwork/block_tensor/index.py +++ b/tensornetwork/block_tensor/index.py @@ -16,7 +16,7 @@ from __future__ import division from __future__ import print_function import numpy as np -from tensornetwork.block_tensor.charge import BaseCharge +from tensornetwork.block_tensor.charge import BaseCharge, fuse_charges import copy from typing import List, Union, Any, Optional, Tuple, Text @@ -29,47 +29,27 @@ class Index: """ def __init__(self, - charges: BaseCharge, - flow: int, - name: Optional[Text] = None, - left_child: Optional["Index"] = None, - right_child: Optional["Index"] = None): - self._charges = charges #ChargeCollection([charges]) + charges: Union[List[BaseCharge], BaseCharge], + flow: Union[List[int], int], + name: Optional[Union[List[Text], Text]] = None) -> None: + if isinstance(charges, BaseCharge): + charges = [charges] + self._charges = charges + if isinstance(flow, (np.bool_, bool, np.bool)): + flow = [flow] + if not all([isinstance(f, (np.bool_, np.bool, bool)) for f in flow]): + raise TypeError("flows have to be boolean") self.flow = flow - self.left_child = left_child - self.right_child = right_child + if isinstance(name, str): + name = [name] self.name = name def __repr__(self): - return str(self.dimension) + return str(self.dim) @property - def is_leave(self): - return (self.left_child is None) and (self.right_child is None) - - @property - def dimension(self): - return np.prod([len(i.charges) for i in self.get_elementary_indices()]) - - def _copy_helper(self, index: "Index", copied_index: "Index") -> None: - """ - Helper function for copy - """ - if index.left_child != None: - left_copy = Index( - charges=copy.deepcopy(index.left_child.charges), - flow=copy.deepcopy(index.left_child.flow), - name=copy.deepcopy(index.left_child.name)) - - copied_index.left_child = left_copy - self._copy_helper(index.left_child, left_copy) - if index.right_child != None: - right_copy = Index( - charges=copy.deepcopy(index.right_child.charges), - flow=copy.deepcopy(index.right_child.flow), - name=copy.deepcopy(index.right_child.name)) - copied_index.right_child = right_copy - self._copy_helper(index.right_child, right_copy) + def dim(self): + return np.prod([i.dim for i in self._charges]) def copy(self): """ @@ -78,30 +58,28 @@ def copy(self): `Index` are copied as well. """ index_copy = Index( - charges=copy.deepcopy(self._charges), + charges=[c.copy() for c in self._charges], flow=copy.deepcopy(self.flow), - name=self.name) - - self._copy_helper(self, index_copy) + name=copy.deepcopy(self.name)) return index_copy - def _leave_helper(self, index: "Index", leave_list: List) -> None: - if index.left_child: - self._leave_helper(index.left_child, leave_list) - if index.right_child: - self._leave_helper(index.right_child, leave_list) - if (index.left_child is None) and (index.right_child is None): - leave_list.append(index) + @property + def flat_charges(self) -> List: + """ + Returns: + List: A list containing the elementary indices (the leaves) + of `Index`. + """ + return self._charges - def get_elementary_indices(self) -> List: + @property + def flat_flows(self) -> List: """ Returns: List: A list containing the elementary indices (the leaves) of `Index`. """ - leave_list = [] - self._leave_helper(self, leave_list) - return leave_list + return self.flow def __mul__(self, index: "Index") -> "Index": """ @@ -114,9 +92,117 @@ def __mul__(self, index: "Index") -> "Index": @property def charges(self): - if self.is_leave: - return self._charges - return self.left_child.charges * self.left_child.flow + self.right_child.charges * self.right_child.flow + return fuse_charges(self.flat_charges, self.flat_flows) + + """ + An index class to store indices of a symmetric tensor. + An index keeps track of all its childs by storing references + to them (i.e. it is a binary tree). + """ + + +# class Index: +# """ +# An index class to store indices of a symmetric tensor. +# An index keeps track of all its childs by storing references +# to them (i.e. it is a binary tree). +# """ + +# def __init__(self, +# charges: BaseCharge, +# flow: int, +# name: Optional[Text] = None, +# left_child: Optional["Index"] = None, +# right_child: Optional["Index"] = None): +# self._charges = charges #ChargeCollection([charges]) +# self.flow = flow +# self.left_child = left_child +# self.right_child = right_child +# self.name = name + +# def __repr__(self): +# return str(self.dimension) + +# @property +# def is_leave(self): +# return (self.left_child is None) and (self.right_child is None) + +# @property +# def dimension(self): +# return np.prod([len(i.charges) for i in self.get_elementary_indices()]) + +# def _copy_helper(self, index: "Index", copied_index: "Index") -> None: +# """ +# Helper function for copy +# """ +# if index.left_child != None: +# left_copy = Index( +# charges=copy.deepcopy(index.left_child.charges), +# flow=copy.deepcopy(index.left_child.flow), +# name=copy.deepcopy(index.left_child.name)) + +# copied_index.left_child = left_copy +# self._copy_helper(index.left_child, left_copy) +# if index.right_child != None: +# right_copy = Index( +# charges=copy.deepcopy(index.right_child.charges), +# flow=copy.deepcopy(index.right_child.flow), +# name=copy.deepcopy(index.right_child.name)) +# copied_index.right_child = right_copy +# self._copy_helper(index.right_child, right_copy) + +# def copy(self): +# """ +# Returns: +# Index: A deep copy of `Index`. Note that all children of +# `Index` are copied as well. +# """ +# index_copy = Index( +# charges=copy.deepcopy(self._charges), +# flow=copy.deepcopy(self.flow), +# name=self.name) + +# self._copy_helper(self, index_copy) +# return index_copy + +# def _leave_helper(self, index: "Index", leave_list: List) -> None: +# if index.left_child: +# self._leave_helper(index.left_child, leave_list) +# if index.right_child: +# self._leave_helper(index.right_child, leave_list) +# if (index.left_child is None) and (index.right_child is None): +# leave_list.append(index) + +# def get_elementary_indices(self) -> List: +# """ +# Returns: +# List: A list containing the elementary indices (the leaves) +# of `Index`. +# """ +# leave_list = [] +# self._leave_helper(self, leave_list) +# return leave_list + +# def __mul__(self, index: "Index") -> "Index": +# """ +# Merge `index` and self into a single larger index. +# The flow of the resulting index is set to 1. +# Flows of `self` and `index` are multiplied into +# the charges upon fusing.n +# """ +# return fuse_index_pair(self, index) + +# @property +# def charges(self): +# if self.is_leave: +# return self._charges +# return self.left_child.charges * self.left_child.flow + self.right_child.charges * self.right_child.flow + +# """ +# An index class to store indices of a symmetric tensor. +# An index keeps track of all its childs by storing references +# to them (i.e. it is a binary tree). +# """ def fuse_index_pair(left_index: Index, @@ -137,7 +223,8 @@ def fuse_index_pair(left_index: Index, "index1 and index2 are the same object. Can only fuse distinct objects") return Index( - charges=None, flow=flow, left_child=left_index, right_child=right_index) + charges=left_index.flat_charges + right_index.flat_charges, + flow=left_index.flat_flows + right_index.flat_flows) def fuse_indices(indices: List[Index], flow: Optional[int] = False) -> Index: