def test_cache(): D = 10 mpsinds = [ Index(U1Charge(np.random.randint(-5, 5, D, dtype=np.int16)), False), Index(U1Charge(np.random.randint(-5, 5, D, dtype=np.int16)), False), Index(U1Charge(np.random.randint(-5, 5, D, dtype=np.int16)), False), Index(U1Charge(np.random.randint(-5, 5, D, dtype=np.int16)), True) ] A = BlockSparseTensor.random(mpsinds) B = A.conj() res_charges = [ A.flat_charges[2], A.flat_charges[3], B.flat_charges[2], B.flat_charges[3] ] res_flows = [ A.flat_flows[2], A.flat_flows[3], B.flat_flows[2], B.flat_flows[3] ] enable_caching() ncon([A, B], [[1, 2, -1, -2], [1, 2, -3, -4]], backend='symmetric') cacher = get_cacher() sA = _to_string(A.flat_charges, A.flat_flows, 2, [2, 3, 0, 1]) sB = _to_string(B.flat_charges, B.flat_flows, 2, [0, 1, 2, 3]) sC = _to_string(res_charges, res_flows, 2, [0, 1, 2, 3]) blocksA, chargesA, dimsA = _find_transposed_diagonal_sparse_blocks( A.flat_charges, A.flat_flows, 2, [2, 3, 0, 1]) blocksB, chargesB, dimsB = _find_transposed_diagonal_sparse_blocks( B.flat_charges, B.flat_flows, 2, [0, 1, 2, 3]) blocksC, chargesC, dimsC = _find_transposed_diagonal_sparse_blocks( res_charges, res_flows, 2, [0, 1, 2, 3]) assert sA in cacher.cache assert sB in cacher.cache assert sC in cacher.cache for b1, b2 in zip(cacher.cache[sA][0], blocksA): np.testing.assert_allclose(b1, b2) for b1, b2 in zip(cacher.cache[sB][0], blocksB): np.testing.assert_allclose(b1, b2) for b1, b2 in zip(cacher.cache[sC][0], blocksC): np.testing.assert_allclose(b1, b2) assert charge_equal(cacher.cache[sA][1], chargesA) assert charge_equal(cacher.cache[sB][1], chargesB) assert charge_equal(cacher.cache[sC][1], chargesC) np.testing.assert_allclose(cacher.cache[sA][2], dimsA) np.testing.assert_allclose(cacher.cache[sB][2], dimsB) np.testing.assert_allclose(cacher.cache[sC][2], dimsC) disable_caching() clear_cache()
def pinv(matrix: BlockSparseTensor, rcond: Optional[float] = 1E-15, hermitian: Optional[bool] = False) -> BlockSparseTensor: """ Compute the Moore-Penrose pseudo inverse of `matrix`. Args: rcond: Pseudo inverse cutoff. Returns: BlockSparseTensor: The pseudo inverse of `matrix`. """ if matrix.ndim != 2: raise ValueError("`pinv` can only be taken for matrices, " "found tensor.ndim={}".format(matrix.ndim)) flat_charges = matrix._charges flat_flows = matrix._flows flat_order = matrix.flat_order tr_partition = len(matrix._order[0]) blocks, _, shapes = _find_transposed_diagonal_sparse_blocks( flat_charges, flat_flows, tr_partition, flat_order) data = np.empty(np.sum(np.prod(shapes, axis=0)), dtype=matrix.dtype) for n, block in enumerate(blocks): data[block] = np.ravel( np.linalg.pinv( np.reshape(matrix.data[block], shapes[:, n]), rcond=rcond, hermitian=hermitian).T) return BlockSparseTensor( data=data, charges=matrix._charges, flows=np.logical_not(matrix._flows), order=matrix._order, check_consistency=False).transpose((1, 0))#pytype: disable=bad-return-type
def inv(matrix: BlockSparseTensor) -> BlockSparseTensor: """ Compute the matrix inverse of `matrix`. Returns: BlockSparseTensor: The inverse of `matrix`. """ if matrix.ndim != 2: raise ValueError("`inv` can only be taken for matrices, " "found tensor.ndim={}".format(matrix.ndim)) flat_charges = matrix._charges flat_flows = matrix._flows flat_order = matrix.flat_order tr_partition = len(matrix._order[0]) blocks, _, shapes = _find_transposed_diagonal_sparse_blocks( flat_charges, flat_flows, tr_partition, flat_order) data = np.empty(np.sum(np.prod(shapes, axis=0)), dtype=matrix.dtype) for n, block in enumerate(blocks): data[block] = np.ravel( np.linalg.inv(np.reshape(matrix.data[block], shapes[:, n])).T) return BlockSparseTensor( data=data, charges=matrix._charges, flows=np.logical_not(matrix._flows), order=matrix._order, check_consistency=False).transpose((1, 0))#pytype: disable=bad-return-type
def eig(matrix: BlockSparseTensor) -> Tuple[ChargeArray, BlockSparseTensor]: """ Compute the eigen decomposition of an `M` by `M` matrix `matrix`. Args: matrix: A matrix (i.e. a rank-2 tensor) of type `BlockSparseTensor` Returns: (ChargeArray,BlockSparseTensor): The eigenvalues and eigenvectors """ if matrix.ndim != 2: raise NotImplementedError( "eig currently supports only rank-2 tensors.") flat_charges = matrix._charges flat_flows = matrix._flows flat_order = matrix.flat_order tr_partition = len(matrix._order[0]) blocks, charges, shapes = _find_transposed_diagonal_sparse_blocks( flat_charges, flat_flows, tr_partition, flat_order) eigvals = [] v_blocks = [] for n, block in enumerate(blocks): e, v = np.linalg.eig(np.reshape(matrix.data[block], shapes[:, n])) eigvals.append(e) v_blocks.append(v) tmp_labels = [ np.full(len(eigvals[n]), fill_value=n, dtype=np.int16) for n in range(len(eigvals)) ] if len(tmp_labels) > 0: eigvalscharge_labels = np.concatenate(tmp_labels) else: eigvalscharge_labels = np.empty(0, dtype=np.int16) eigvalscharge = charges[eigvalscharge_labels] if len(eigvals) > 0: all_eigvals = np.concatenate(eigvals) else: all_eigvals = np.empty(0, dtype=get_real_dtype(matrix.dtype)) E = ChargeArray(all_eigvals, [eigvalscharge], [False]) charges_v = [eigvalscharge ] + [matrix._charges[o] for o in matrix._order[0]] order_v = [[0]] + [list(np.arange(1, len(matrix._order[0]) + 1))] flows_v = [True] + [matrix._flows[o] for o in matrix._order[0]] if len(v_blocks) > 0: all_v_blocks = np.concatenate([np.ravel(v.T) for v in v_blocks]) else: all_v_blocks = np.empty(0, dtype=matrix.dtype) V = BlockSparseTensor(all_v_blocks, charges=charges_v, flows=flows_v, order=order_v, check_consistency=False).transpose() return E, V #pytype: disable=bad-return-type
def contiguous(self, permutation: Optional[Union[Tuple, List, np.ndarray]] = None, inplace: Optional[bool] = False) -> Any: """ Transpose the tensor data in place such that the linear order of the elements in `BlockSparseTensor.data` corresponds to the current order of tensor indices. Consider a tensor with current order given by `_order=[[1,2],[3],[0]]`, i.e. `data` was initialized according to order [0,1,2,3], and the tensor has since been reshaped and transposed. The linear oder of `data` does not match the desired order [1,2,3,0] of the tensor. `contiguous` fixes this by permuting `data` into this order, transposing `_charges` and `_flows`, and changing `_order` to `[[0,1],[2],[3]]`. Args: permutation: An optional alternative order to be used to transposed the tensor. If `None` defaults to `BlockSparseTensor.permutation`. """ flat_charges = self._charges flat_flows = self._flows if permutation is None: permutation = self.flat_order if np.array_equal(permutation, np.arange(len(permutation))): return self tr_partition = _find_best_partition( [flat_charges[n].dim for n in permutation]) tr_sparse_blocks, tr_charges, _ = _find_transposed_diagonal_sparse_blocks( flat_charges, flat_flows, tr_partition, permutation) sparse_blocks, charges, _ = _find_diagonal_sparse_blocks( [flat_charges[n] for n in permutation], [flat_flows[n] for n in permutation], tr_partition) data = np.empty(len(self.data), dtype=self.dtype) for n, sparse_block in enumerate(sparse_blocks): ind = np.nonzero(tr_charges == charges[n])[0][0] perm = tr_sparse_blocks[ind] data[sparse_block] = self.data[perm] _, inds = np.unique(permutation, return_index=True) new_flat_order = inds[self.flat_order] tmp = np.append(0, np.cumsum([len(o) for o in self._order])) order = [ list(new_flat_order[tmp[n]:tmp[n + 1]]) for n in range(len(tmp) - 1) ] charges = [self._charges[o] for o in permutation] flows = [self._flows[o] for o in permutation] if not inplace: return BlockSparseTensor(data, charges=charges, flows=flows, order=order, check_consistency=False) self.data = data self._order = order self._charges = charges self._flows = flows return self
def eye_sym(charges, flows, pivot=1): ind_temp = [Index(charges[n], flows[n]) for n in range(len(flows))] arr = BT.zeros(ind_temp) sparse_blocks, cs, dims = _find_transposed_diagonal_sparse_blocks( charges, flows, pivot, list(range(len(charges)))) for k in range(len(cs)): arr.data[sparse_blocks[k]] = np.eye(N=dims[0, k], M=dims[1, k]).flatten() return arr
def diag(tensor: ChargeArray) -> Any: """ Return a diagonal `BlockSparseTensor` from a `ChargeArray`, or return the diagonal of a `BlockSparseTensor` as a `ChargeArray`. For input of type `BlockSparseTensor`: The full diagonal is obtained from finding the diagonal blocks of the `BlockSparseTensor`, taking the diagonal elements of those and packing the result into a ChargeArray. Note that the computed diagonal elements are usually different from the diagonal elements obtained from converting the `BlockSparseTensor` to dense storage and taking the diagonal. Note that the flow of the resulting 1d `ChargeArray` object is `False`. Args: tensor: A `ChargeArray`. Returns: ChargeArray: A 1d `CharggeArray` containing the diagonal of `tensor`, or a diagonal matrix of type `BlockSparseTensor` containing `tensor` on its diagonal. """ if tensor.ndim > 2: raise ValueError("`diag` currently only implemented for matrices, " "found `ndim={}".format(tensor.ndim)) if not isinstance(tensor, BlockSparseTensor): if tensor.ndim > 1: raise ValueError( "`diag` currently only implemented for `ChargeArray` with ndim=1, " "found `ndim={}`".format(tensor.ndim)) flat_charges = tensor._charges + tensor._charges flat_flows = list(tensor._flows) + list(np.logical_not(tensor._flows)) flat_order = list(tensor.flat_order) + list( np.asarray(tensor.flat_order) + len(tensor._charges)) tr_partition = len(tensor._order[0]) blocks, charges, shapes = _find_transposed_diagonal_sparse_blocks( flat_charges, flat_flows, tr_partition, flat_order) data = np.zeros( np.int64(np.sum(np.prod(shapes, axis=0))), dtype=tensor.dtype) lookup, unique, labels = compute_sparse_lookup(tensor._charges, tensor._flows, charges) for n, block in enumerate(blocks): label = labels[np.nonzero(unique == charges[n])[0][0]] data[block] = np.ravel( np.diag(tensor.data[np.nonzero(lookup == label)[0]])) order = [ tensor._order[0], list(np.asarray(tensor._order[0]) + len(tensor._charges)) ] new_charges = [tensor._charges[0].copy(), tensor._charges[0].copy()] return BlockSparseTensor( data, charges=new_charges, flows=list(tensor._flows) + list(np.logical_not(tensor._flows)), order=order, check_consistency=False) flat_charges = tensor._charges flat_flows = tensor._flows flat_order = tensor.flat_order tr_partition = len(tensor._order[0]) sparse_blocks, charges, block_shapes = _find_transposed_diagonal_sparse_blocks(#pylint: disable=line-too-long flat_charges, flat_flows, tr_partition, flat_order) shapes = np.min(block_shapes, axis=0) if len(sparse_blocks) > 0: data = np.concatenate([ np.diag(np.reshape(tensor.data[sparse_blocks[n]], block_shapes[:, n])) for n in range(len(sparse_blocks)) ]) charge_labels = np.concatenate([ np.full(shapes[n], fill_value=n, dtype=np.int16) for n in range(len(sparse_blocks)) ]) else: data = np.empty(0, dtype=tensor.dtype) charge_labels = np.empty(0, dtype=np.int16) newcharges = [charges[charge_labels]] flows = [False] return ChargeArray(data, newcharges, flows)
def qr(matrix: BlockSparseTensor, mode: Optional[Text] = 'reduced') -> Any: """ Compute the qr decomposition of an `M` by `N` matrix `matrix`. The matrix is factorized into `q*r`, with `q` an orthogonal matrix and `r` an upper triangular matrix. Args: matrix: A matrix (i.e. a rank-2 tensor) of type `BlockSparseTensor` mode : Can take values {'reduced', 'complete', 'r', 'raw'}. If K = min(M, N), then * 'reduced' : returns q, r with dimensions (M, K), (K, N) (default) * 'complete' : returns q, r with dimensions (M, M), (M, N) * 'r' : returns r only with dimensions (K, N) Returns: (BlockSparseTensor,BlockSparseTensor): If mode = `reduced` or `complete` BlockSparseTensor: If mode = `r`. """ if mode == 'raw': raise NotImplementedError('mode `raw` currenntly not supported') if matrix.ndim != 2: raise NotImplementedError("qr currently supports only rank-2 tensors.") flat_charges = matrix._charges flat_flows = matrix._flows flat_order = matrix.flat_order tr_partition = len(matrix._order[0]) blocks, charges, shapes = _find_transposed_diagonal_sparse_blocks( flat_charges, flat_flows, tr_partition, flat_order) q_blocks = [] r_blocks = [] for n, block in enumerate(blocks): out = np.linalg.qr(np.reshape(matrix.data[block], shapes[:, n]), mode) if mode in ('reduced', 'complete'): q_blocks.append(out[0]) r_blocks.append(out[1]) elif mode == 'r': r_blocks.append(out) else: raise ValueError('unknown value {} for input `mode`'.format(mode)) tmp_r_charge_labels = [ np.full(r_blocks[n].shape[0], fill_value=n, dtype=np.int16) for n in range(len(r_blocks)) ] if len(tmp_r_charge_labels) > 0: left_r_charge_labels = np.concatenate(tmp_r_charge_labels) else: left_r_charge_labels = np.empty(0, dtype=np.int16) left_r_charge = charges[left_r_charge_labels] charges_r = [left_r_charge] + [matrix._charges[o] for o in matrix._order[1]] flows_r = [False] + [matrix._flows[o] for o in matrix._order[1]] order_r = [[0]] + [list(np.arange(1, len(matrix._order[1]) + 1))] if len(r_blocks) > 0: all_r_blocks = np.concatenate([np.ravel(r) for r in r_blocks]) else: all_r_blocks = np.empty(0, dtype=matrix.dtype) R = BlockSparseTensor( all_r_blocks, charges=charges_r, flows=flows_r, order=order_r, check_consistency=False) if mode in ('reduced', 'complete'): tmp_right_q_charge_labels = [ np.full(q_blocks[n].shape[1], fill_value=n, dtype=np.int16) for n in range(len(q_blocks)) ] if len(tmp_right_q_charge_labels) > 0: right_q_charge_labels = np.concatenate(tmp_right_q_charge_labels) else: right_q_charge_labels = np.empty(0, dtype=np.int16) right_q_charge = charges[right_q_charge_labels] charges_q = [ right_q_charge, ] + [matrix._charges[o] for o in matrix._order[0]] order_q = [[0]] + [list(np.arange(1, len(matrix._order[0]) + 1))] flows_q = [True] + [matrix._flows[o] for o in matrix._order[0]] if len(q_blocks) > 0: all_q_blocks = np.concatenate([np.ravel(q.T) for q in q_blocks]) else: all_q_blocks = np.empty(0, dtype=matrix.dtype) return BlockSparseTensor( all_q_blocks, charges=charges_q, flows=flows_q, order=order_q, check_consistency=False).transpose((1, 0)), R return R
def svd(matrix: BlockSparseTensor, full_matrices: Optional[bool] = True, compute_uv: Optional[bool] = True, hermitian: Optional[bool] = False) -> Any: """ Compute the singular value decomposition of `matrix`. The matrix if factorized into `u * s * vh`, with `u` and `vh` the left and right singular vectors of `matrix`, and `s` its singular values. Args: matrix: A matrix (i.e. an order-2 tensor) of type `BlockSparseTensor` full_matrices: If `True`, expand `u` and `v` to square matrices If `False` return the "economic" svd, i.e. `u.shape[1]=s.shape[0]` and `v.shape[0]=s.shape[1]` compute_uv: If `True`, return `u` and `v`. hermitian: If `True`, assume hermiticity of `matrix`. Returns: If `compute_uv` is `True`: Three BlockSparseTensors `U,S,V`. If `compute_uv` is `False`: A BlockSparseTensors `S` containing the singular values. """ if matrix.ndim != 2: raise NotImplementedError("svd currently supports only tensors of order 2.") flat_charges = matrix._charges flat_flows = matrix._flows flat_order = matrix.flat_order tr_partition = len(matrix._order[0]) blocks, charges, shapes = _find_transposed_diagonal_sparse_blocks( flat_charges, flat_flows, tr_partition, flat_order) u_blocks = [] singvals = [] v_blocks = [] for n, block in enumerate(blocks): out = np.linalg.svd( np.reshape(matrix.data[block], shapes[:, n]), full_matrices, compute_uv, hermitian) if compute_uv: u_blocks.append(out[0]) singvals.append(out[1]) v_blocks.append(out[2]) else: singvals.append(out) tmp_labels = [ np.full(len(singvals[n]), fill_value=n, dtype=np.int16) for n in range(len(singvals)) ] if len(tmp_labels) > 0: left_singval_charge_labels = np.concatenate(tmp_labels) else: left_singval_charge_labels = np.empty(0, dtype=np.int16) left_singval_charge = charges[left_singval_charge_labels] if len(singvals) > 0: all_singvals = np.concatenate(singvals) else: all_singvals = np.empty(0, dtype=get_real_dtype(matrix.dtype)) S = ChargeArray(all_singvals, [left_singval_charge], [False]) if compute_uv: #define the new charges on the two central bonds tmp_left_labels = [ np.full(u_blocks[n].shape[1], fill_value=n, dtype=np.int16) for n in range(len(u_blocks)) ] if len(tmp_left_labels) > 0: left_charge_labels = np.concatenate(tmp_left_labels) else: left_charge_labels = np.empty(0, dtype=np.int16) tmp_right_labels = [ np.full(v_blocks[n].shape[0], fill_value=n, dtype=np.int16) for n in range(len(v_blocks)) ] if len(tmp_right_labels) > 0: right_charge_labels = np.concatenate(tmp_right_labels) else: right_charge_labels = np.empty(0, dtype=np.int16) new_left_charge = charges[left_charge_labels] new_right_charge = charges[right_charge_labels] charges_u = [new_left_charge ] + [matrix._charges[o] for o in matrix._order[0]] order_u = [[0]] + [list(np.arange(1, len(matrix._order[0]) + 1))] flows_u = [True] + [matrix._flows[o] for o in matrix._order[0]] charges_v = [new_right_charge ] + [matrix._charges[o] for o in matrix._order[1]] flows_v = [False] + [matrix._flows[o] for o in matrix._order[1]] order_v = [[0]] + [list(np.arange(1, len(matrix._order[1]) + 1))] # We fill in data into the transposed U # note that transposing is essentially free if len(u_blocks) > 0: all_u_blocks = np.concatenate([np.ravel(u.T) for u in u_blocks]) all_v_blocks = np.concatenate([np.ravel(v) for v in v_blocks]) else: all_u_blocks = np.empty(0, dtype=matrix.dtype) all_v_blocks = np.empty(0, dtype=matrix.dtype) return BlockSparseTensor( all_u_blocks, charges=charges_u, flows=flows_u, order=order_u, check_consistency=False).transpose((1, 0)), S, BlockSparseTensor( all_v_blocks, charges=charges_v, flows=flows_v, order=order_v, check_consistency=False) return S
def svd(bt, tensor: BlockSparseTensor, pivot_axis: int, max_singular_values: Optional[int] = None, max_truncation_error: Optional[float] = None, relative: Optional[bool] = False ) -> Tuple[Tensor, Tensor, Tensor, Tensor]: """ Computes the singular value decomposition (SVD) of a tensor. See tensornetwork.backends.tensorflow.decompositions for details. """ left_dims = tensor.shape[:pivot_axis] right_dims = tensor.shape[pivot_axis:] matrix = bt.reshape(tensor, [np.prod(left_dims), np.prod(right_dims)]) flat_charges = matrix._charges flat_flows = matrix._flows flat_order = matrix.flat_order tr_partition = len(matrix._order[0]) blocks, charges, shapes = _find_transposed_diagonal_sparse_blocks( flat_charges, flat_flows, tr_partition, flat_order) u_blocks = [] singvals = [] v_blocks = [] for n, b in enumerate(blocks): out = np.linalg.svd(np.reshape(matrix.data[b], shapes[:, n]), full_matrices=False, compute_uv=True) u_blocks.append(out[0]) singvals.append(out[1]) v_blocks.append(out[2]) orig_num_singvals = np.int64(np.sum([len(s) for s in singvals])) discarded_singvals = np.zeros(0, dtype=get_real_dtype(tensor.dtype)) if (max_singular_values is not None) and (max_singular_values >= orig_num_singvals): max_singular_values = None if (max_truncation_error is not None) or (max_singular_values is not None): max_D = np.max([len(s) for s in singvals]) if len(singvals) > 0 else 0 #extend singvals of all blocks into a matrix by padding each block with 0 if len(singvals) > 0: extended_singvals = np.stack([ np.append(s, np.zeros(max_D - len(s), dtype=s.dtype)) for s in singvals ], axis=1) else: extended_singvals = np.empty((0, 0), dtype=get_real_dtype(tensor.dtype)) extended_flat_singvals = np.ravel(extended_singvals) #sort singular values inds = np.argsort(extended_flat_singvals, kind='stable') discarded_inds = np.zeros(0, dtype=SIZE_T) if inds.shape[0] > 0: maxind = inds[-1] else: maxind = 0 if max_truncation_error is not None: if relative and (len(singvals) > 0): max_truncation_error = max_truncation_error * np.max( [s[0] for s in singvals]) kept_inds_mask = np.sqrt( np.cumsum(np.square( extended_flat_singvals[inds]))) > max_truncation_error trunc_inds_mask = np.logical_not(kept_inds_mask) discarded_inds = inds[trunc_inds_mask] inds = inds[kept_inds_mask] if max_singular_values is not None: #if the original number of non-zero singular values #is smaller than `max_singular_values` we need to reset #`max_singular_values` (we were filling in 0.0 into singular #value blocks to facilitate trunction steps, thus we could end up #with more singular values than originally there). if max_singular_values > orig_num_singvals: max_singular_values = orig_num_singvals if max_singular_values < len(inds): discarded_inds = np.append(discarded_inds, inds[:(-1) * max_singular_values]) inds = inds[(-1) * max_singular_values::] if len(inds) == 0: #special case of truncation to 0 dimension; warnings.warn("svd_decomposition truncated to 0 dimensions. " "Adjusting to `max_singular_values = 1`") inds = np.asarray([maxind]) if extended_singvals.shape[1] > 0: #pylint: disable=no-member keep = np.divmod(inds, extended_singvals.shape[1]) else: keep = (np.zeros(1, dtype=SIZE_T), np.zeros(1, dtype=SIZE_T)) newsingvals = [ extended_singvals[keep[0][keep[1] == n], keep[1][keep[1] == n]][::-1] for n in range(extended_singvals.shape[1]) ] discarded_singvals = extended_flat_singvals[discarded_inds] singvals = newsingvals if len(singvals) > 0: left_singval_charge_labels = np.concatenate([ np.full(singvals[n].shape[0], fill_value=n, dtype=np.int16) for n in range(len(singvals)) ]) all_singvals = np.concatenate(singvals) #define the new charges on the two central bonds left_charge_labels = np.concatenate([ np.full(len(singvals[n]), fill_value=n, dtype=np.int16) for n in range(len(u_blocks)) ]) right_charge_labels = np.concatenate([ np.full(len(singvals[n]), fill_value=n, dtype=np.int16) for n in range(len(v_blocks)) ]) all_ublocks = np.concatenate([ np.ravel(np.transpose(u_blocks[n][:, 0:len(singvals[n])])) for n in range(len(u_blocks)) ]) all_vblocks = np.concatenate([ np.ravel(v_blocks[n][0:len(singvals[n]), :]) for n in range(len(v_blocks)) ]) else: left_singval_charge_labels = np.empty(0, dtype=np.int16) all_singvals = np.empty(0, dtype=get_real_dtype(tensor.dtype)) left_charge_labels = np.empty(0, dtype=np.int16) right_charge_labels = np.empty(0, dtype=np.int16) all_ublocks = np.empty(0, dtype=get_real_dtype(tensor.dtype)) all_vblocks = np.empty(0, dtype=get_real_dtype(tensor.dtype)) left_singval_charge = charges[left_singval_charge_labels] S = ChargeArray(all_singvals, [left_singval_charge], [False]) new_left_charge = charges[left_charge_labels] new_right_charge = charges[right_charge_labels] #get the indices of the new tensors U,S and V charges_u = [new_left_charge ] + [matrix._charges[o] for o in matrix._order[0]] order_u = [[0]] + [list(np.arange(1, len(matrix._order[0]) + 1))] flows_u = [True] + [matrix._flows[o] for o in matrix._order[0]] charges_v = [new_right_charge ] + [matrix._charges[o] for o in matrix._order[1]] flows_v = [False] + [matrix._flows[o] for o in matrix._order[1]] order_v = [[0]] + [list(np.arange(1, len(matrix._order[1]) + 1))] #We fill in data into the transposed U U = BlockSparseTensor(all_ublocks, charges=charges_u, flows=flows_u, order=order_u, check_consistency=False).transpose((1, 0)) V = BlockSparseTensor(all_vblocks, charges=charges_v, flows=flows_v, order=order_v, check_consistency=False) left_shape = left_dims + (S.shape[0], ) right_shape = (S.shape[0], ) + right_dims return U.reshape(left_shape), S, V.reshape( right_shape), discarded_singvals[discarded_singvals > 0.0]
def _eigh_free( matrix: BlockSparseTensor, which: Optional[Text] = 'LM', full_sort: Optional[bool] = True, threshold: Optional[float] = None, max_kept: Optional[int] = None, UPLO: Optional[Text] = 'L') -> Tuple[ChargeArray, BlockSparseTensor]: """ Compute eigenvectors and eigenvalues of a hermitian matrix where the output charge order is free. """ # reshape into matrix if needed pivot = matrix.ndim // 2 m_shape = matrix.shape matrix = matrix.reshape( [np.prod(m_shape[:pivot]), np.prod(m_shape[pivot:])]) if max_kept is None: max_kept = matrix.shape[0] max_kept = min(max_kept, matrix.shape[0]) if threshold is None: if which == 'LM' or which == 'LA': threshold = -float('inf') elif which == 'SM' or which == 'SA': threshold = float('inf') # compute info about each block flat_charges = matrix._charges flat_flows = matrix._flows flat_order = matrix.flat_order tr_partition = len(matrix._order[0]) blocks, charges, shapes = _find_transposed_diagonal_sparse_blocks( flat_charges, flat_flows, tr_partition, flat_order) num_blocks = len(blocks) # diagonalize each block eigvals = [0] * num_blocks v_blocks = [0] * num_blocks for n, block in enumerate(blocks): etemp, vtemp = np.linalg.eigh( np.reshape(matrix.data[block], shapes[:, n]), UPLO) # sort within each block if which == 'LM': ord_temp = np.flip(np.argsort(abs(etemp))) elif which == 'LA': ord_temp = np.flip(np.argsort(etemp)) elif which == 'SM': ord_temp = np.argsort(abs(etemp)) elif which == 'SA': ord_temp = np.argsort(etemp) eigvals[n] = etemp[ord_temp] v_blocks[n] = (vtemp[:, ord_temp].T) # combine and sort eigenvalues from all symmetry blocks tmp_labels = [ np.full(len(eigvals[n]), fill_value=n, dtype=np.int16) for n in range(len(eigvals)) ] tmp_degens = [ np.arange(len(eigvals[n]), dtype=np.int16) for n in range(len(eigvals)) ] all_eigvals = np.concatenate(eigvals) all_labels = np.concatenate(tmp_labels) if which == 'LM': eig_ord = np.flip(np.argsort(np.abs(all_eigvals))) num_kept = min(sum(np.abs(all_eigvals) >= threshold), max_kept) elif which == 'LA': eig_ord = np.flip(np.argsort(all_eigvals)) num_kept = min(sum(all_eigvals >= threshold), max_kept) elif which == 'SM': eig_ord = np.argsort(np.abs(all_eigvals)) num_kept = min(sum(np.abs(all_eigvals) <= threshold), max_kept) elif which == 'SA': eig_ord = np.argsort(all_eigvals) num_kept = min(sum(all_eigvals <= threshold), max_kept) if full_sort: e_labels = np.concatenate(tmp_labels)[eig_ord[:num_kept]] e_degens = np.concatenate(tmp_degens)[eig_ord[:num_kept]] e_charge = charges[e_labels] E = ChargeArray(all_eigvals[eig_ord[:num_kept]], [e_charge], [False]) else: num_per_block = [ sum(all_labels[eig_ord[:num_kept]] == n) for n in range(num_blocks) ] e_labels = np.concatenate([ np.full(num_per_block[n], fill_value=n, dtype=np.int16) for n in range(num_blocks) ]) e_degens = np.concatenate([ np.arange(num_per_block[n], dtype=np.int16) for n in range(num_blocks) ]) e_charge = charges[e_labels] new_eigvals = np.concatenate( [eigvals[n][:num_per_block[n]] for n in range(num_blocks)]) E = ChargeArray(new_eigvals, [e_charge], [False]) charges_v = [e_charge] + [matrix._charges[o] for o in matrix._order[0]] order_v = [[0]] + [list(np.arange(1, len(matrix._order[0]) + 1))] flows_v = [True] + [matrix._flows[o] for o in matrix._order[0]] all_v_blocks = np.concatenate( [v_blocks[e_labels[n]][e_degens[n], :] for n in range(num_kept)]) fin_shape = [*m_shape[:pivot], num_kept] V = BlockSparseTensor( all_v_blocks, charges=charges_v, flows=flows_v, order=order_v, check_consistency=False).transpose().reshape(fin_shape) return E, V
def _eigh_fixed( matrix: BlockSparseTensor, link_charges: ChargeArray, which: Optional[Text] = 'LM', UPLO: Optional[Text] = 'L') -> Tuple[ChargeArray, BlockSparseTensor]: """ Compute eigenvectors and eigenvalues of a hermitian matrix where the output charge order is fixed to match that of `link_charges`. """ # reshape into matrix if needed pivot = matrix.ndim // 2 m_shape = matrix.shape matrix = matrix.reshape( [np.prod(m_shape[:pivot]), np.prod(m_shape[pivot:])]) # compute info about each block flat_charges = matrix._charges flat_flows = matrix._flows flat_order = matrix.flat_order tr_partition = len(matrix._order[0]) blocks, charges, shapes = _find_transposed_diagonal_sparse_blocks( flat_charges, flat_flows, tr_partition, flat_order) # intersect between link charges and block charges link_uni, link_pos, link_counts = link_charges.unique(return_inverse=True, return_counts=True) _, blk_common, link_common = charges.intersect(link_uni, return_indices=True) # diagonalize each block eigvals = [] v_blocks = [] for n, m in enumerate(blk_common): e, v = np.linalg.eigh(np.reshape(matrix.data[blocks[m]], shapes[:, m]), UPLO) # sort within each block if which == 'SA': blk_sort = np.argsort(e) elif which == 'LA': blk_sort = np.flip(np.argsort(e)) elif which == 'SM': blk_sort = np.argsort(np.abs(e)) elif which == 'LM': blk_sort = np.flip(np.argsort(np.abs(e))) eigvals.append(e[blk_sort]) v_blocks.append(v[:, blk_sort].T) link_degens = np.zeros(np.size(link_pos), dtype=np.int64) for n, count in enumerate(link_counts): link_degens[link_pos == n] = np.arange(count, dtype=np.int64) all_eigvals = np.zeros(link_pos.size, dtype=matrix.dtype) for n in range(len(link_pos)): all_eigvals[n] = eigvals[link_pos[n]][link_degens[n]] e_charge = charges[blk_common[link_pos]] E = ChargeArray(all_eigvals, [e_charge], [False]) charges_v = [e_charge] + [matrix._charges[o] for o in matrix._order[0]] order_v = [[0]] + [list(np.arange(1, len(matrix._order[0]) + 1))] flows_v = [True] + [matrix._flows[o] for o in matrix._order[0]] if len(v_blocks) > 0: all_v_blocks = np.concatenate([ v_blocks[link_pos[n]][link_degens[n], :] for n in range(len(all_eigvals)) ]) else: all_v_blocks = np.empty(0, dtype=matrix.dtype) fin_shape = [*m_shape[:pivot], len(all_eigvals)] V = BlockSparseTensor( all_v_blocks, charges=charges_v, flows=flows_v, order=order_v, check_consistency=False).transpose().reshape(fin_shape) return E, V
def test_find_transposed_diagonal_sparse_blocks(num_charges, order, D): order = list(order) num_legs = len(order) np.random.seed(10) np_charges = [ np.random.randint(-5, 5, (D, num_charges), dtype=np.int16) for _ in range(num_legs) ] tr_charge_list = [] charge_list = [] for c in range(num_charges): tr_charge_list.append( fuse_ndarrays( [np_charges[order[n]][:, c] for n in range(num_legs)])) charge_list.append( fuse_ndarrays([np_charges[n][:, c] for n in range(num_legs)])) tr_fused = np.stack(tr_charge_list, axis=1) fused = np.stack(charge_list, axis=1) dims = [c.shape[0] for c in np_charges] strides = _get_strides(dims) transposed_linear_positions = fuse_stride_arrays( dims, [strides[o] for o in order]) left_charges = np.stack([ fuse_ndarrays( [np_charges[order[n]][:, c] for n in range(num_legs // 2)]) for c in range(num_charges) ], axis=1) right_charges = np.stack([ fuse_ndarrays([ np_charges[order[n]][:, c] for n in range(num_legs // 2, num_legs) ]) for c in range(num_charges) ], axis=1) #pylint: disable=no-member mask = np.logical_and.reduce(fused == np.zeros((1, num_charges)), axis=1) nz = np.nonzero(mask)[0] dense_to_sparse = np.empty(len(mask), dtype=np.int64) dense_to_sparse[mask] = np.arange(len(nz)) #pylint: disable=no-member tr_mask = np.logical_and.reduce(tr_fused == np.zeros((1, num_charges)), axis=1) tr_nz = np.nonzero(tr_mask)[0] tr_linear_locs = transposed_linear_positions[tr_nz] # pylint: disable=no-member left_inds, _ = np.divmod(tr_nz, right_charges.shape[0]) left = left_charges[left_inds, :] unique_left = unique(left) blocks = [] for n in range(unique_left.shape[0]): ul = unique_left[n, :][None, :] #pylint: disable=no-member blocks.append(dense_to_sparse[tr_linear_locs[np.nonzero( np.logical_and.reduce(left == ul, axis=1))[0]]]) charges = [ BaseCharge(c, charge_types=[U1Charge] * num_charges) for c in np_charges ] flows = [False] * num_legs bs, cs, ss = _find_transposed_diagonal_sparse_blocks( charges, flows, tr_partition=num_legs // 2, order=order) np.testing.assert_allclose(cs.charges, unique_left) for b1, b2 in zip(blocks, bs): assert np.all(b1 == b2) assert np.sum(np.prod(ss, axis=0)) == np.sum([len(b) for b in bs]) np.testing.assert_allclose(unique_left, cs.charges)
def 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