diff --git a/tensornetwork/block_tensor/block_tensor.py b/tensornetwork/block_tensor/block_tensor.py index 1ad942293..c39fa38e7 100644 --- a/tensornetwork/block_tensor/block_tensor.py +++ b/tensornetwork/block_tensor/block_tensor.py @@ -212,6 +212,116 @@ def retrieve_non_zero_diagonal_blocks( column_charges, return_counts=True) common_charges = np.intersect1d( unique_row_charges, -unique_column_charges, assume_unique=True) + + row_degeneracies = dict(zip(unique_row_charges, row_dims)) + column_degeneracies = dict(zip(unique_column_charges, column_dims)) + + mask = np.isin(column_charges, -common_charges) + masked_charges = column_charges[mask] + degeneracy_vector = np.empty(len(masked_charges), dtype=np.int64) + masks = {} + for c in common_charges: + mask = masked_charges == -c + masks[c] = mask + degeneracy_vector[mask] = row_degeneracies[c] + summed_degeneracies = np.cumsum(degeneracy_vector) + blocks = {} + + for c in common_charges: + a = np.expand_dims(summed_degeneracies[masks[c]] - row_degeneracies[c], 0) + b = np.expand_dims(np.arange(row_degeneracies[c]), 1) + if not return_data: + blocks[c] = [a + b, (row_degeneracies[c], column_degeneracies[-c])] + else: + blocks[c] = np.reshape(data[a + b], + (row_degeneracies[c], column_degeneracies[-c])) + return blocks + + +def retrieve_non_zero_diagonal_blocks_deprecated( + data: np.ndarray, + charges: List[np.ndarray], + flows: List[Union[bool, int]], + return_data: Optional[bool] = True) -> Dict: + """ + Given the meta data and underlying data of a symmetric matrix, compute + all diagonal blocks and return them in a dict. + Args: + data: An np.ndarray of the data. The number of elements in `data` + has to match the number of non-zero elements defined by `charges` + and `flows` + charges: List of np.ndarray, one for each leg. + Each np.ndarray `charges[leg]` is of shape `(D[leg],)`. + The bond dimension `D[leg]` can vary on each leg. + flows: A list of integers, one for each leg, + with values `1` or `-1`, denoting the flow direction + of the charges on each leg. `1` is inflowing, `-1` is outflowing + charge. + return_data: If `True`, the return dictionary maps quantum numbers `q` to + actual `np.ndarray` with the data. This involves a copy of data. + If `False`, the returned dict maps quantum numbers of a list + [locations, shape], where `locations` is an np.ndarray of type np.int64 + containing the locations of the tensor elements within A.data, i.e. + `A.data[locations]` contains the elements belonging to the tensor with + quantum numbers `(q,q). `shape` is the shape of the corresponding array. + + Returns: + dict: Dictionary mapping quantum numbers (integers) to either an np.ndarray + or a python list of locations and shapes, depending on the value of `return_data`. + """ + #TODO: this is currently way too slow!!!! + #Run the following benchmark for testing (typical MPS use case) + #retrieving the blocks is ~ 10 times as slow as multiplying them + + # D=4000 + # B=10 + # q1 = np.random.randint(0,B,D) + # q2 = np.asarray([0,1]) + # q3 = np.random.randint(0,B,D) + # i1 = Index(charges=q1,flow=1) + # i2 = Index(charges=q2,flow=1) + # i3 = Index(charges=q3,flow=-1) + # indices=[i1,i2,i3] + # A = BlockSparseTensor.random(indices=indices, dtype=np.complex128) + # A.reshape((D*2, D)) + # def multiply_blocks(blocks): + # for b in blocks.values(): + # np.dot(b.T, b) + # t1s=[] + # t2s=[] + # for n in range(10): + # print(n) + # t1 = time.time() + # b = A.get_diagonal_blocks() + # t1s.append(time.time() - t1) + # t1 = time.time() + # multiply_blocks(b) + # t2s.append(time.time() - t1) + # print('average retrieval time', np.average(t1s)) + # print('average multiplication time',np.average(t2s)) + + if len(charges) != 2: + raise ValueError("input has to be a two-dimensional symmetric matrix") + check_flows(flows) + if len(flows) != len(charges): + raise ValueError("`len(flows)` is different from `len(charges)`") + + row_charges = flows[0] * charges[0] # a list of charges on each row + column_charges = flows[1] * charges[1] # a list of charges on each column + + #get the unique charges + t1 = time.time() + unique_row_charges, row_dims = np.unique(row_charges, return_counts=True) + # print('finding unique row charges', time.time() - t1) + # t1 = time.time() + unique_column_charges, column_dims = np.unique( + column_charges, return_counts=True) + # print('finding unique column charges', time.time() - t1) + # t1 = time.time() + common_charges = np.intersect1d( + unique_row_charges, -unique_column_charges, assume_unique=True) + # print('finding unique intersections', time.time() - t1) + # t1 = time.time() #common_charges = np.intersect1d(row_charges, -column_charges) # for each matrix column find the number of non-zero elements in it @@ -223,21 +333,61 @@ def retrieve_non_zero_diagonal_blocks( column_degeneracies = dict(zip(unique_column_charges, column_dims)) number_of_seen_elements = 0 - idxs = {c: [] for c in common_charges} + #idxs = {c: [] for c in common_charges} + idxs = { + c: np.empty( + row_degeneracies[c] * column_degeneracies[-c], dtype=np.int64) + for c in common_charges + } + idxs_stops = {c: 0 for c in common_charges} + t1 = time.time() mask = np.isin(column_charges, -common_charges) - for charge in column_charges[mask]: - idxs[-charge].append( - np.arange(number_of_seen_elements, - row_degeneracies[-charge] + number_of_seen_elements)) + masked_charges = column_charges[mask] + print('finding mask', time.time() - t1) + # print(len(column_charges), len(masked_charges)) + t1 = time.time() + elements = {c: np.arange(row_degeneracies[c]) for c in common_charges} + for charge in masked_charges: + # idxs[-charge].append((number_of_seen_elements, + # row_degeneracies[-charge] + number_of_seen_elements)) + + idxs[-charge][ + idxs_stops[-charge]:idxs_stops[-charge] + + row_degeneracies[-charge]] = number_of_seen_elements + elements[-charge] + + # np.arange( + # number_of_seen_elements, + # row_degeneracies[-charge] + number_of_seen_elements) + number_of_seen_elements += row_degeneracies[-charge] + idxs_stops[-charge] += row_degeneracies[-charge] + print('getting start and stop', time.time() - t1) + # t1 = time.time() + # for charge in masked_charges: + # tmp = np.arange(number_of_seen_elements, + # row_degeneracies[-charge] + number_of_seen_elements) + # number_of_seen_elements += row_degeneracies[-charge] + # print('running the partial loop', time.time() - t1) + + ####################################################################################### + #looks like this takes pretty long for rectangular matrices where shape[1] >> shape[0] + #it's mostly np.arange that causes the overhead. + # t1 = time.time() + # for charge in masked_charges: + # idxs[-charge].append( + # np.arange(number_of_seen_elements, + # row_degeneracies[-charge] + number_of_seen_elements)) + # number_of_seen_elements += row_degeneracies[-charge] + # print('running the full loop', time.time() - t1) + ####################################################################################### blocks = {} if not return_data: for c, idx in idxs.items(): - num_elements = np.sum([len(t) for t in idx]) - indexes = np.empty(num_elements, dtype=np.int64) - np.concatenate(idx, out=indexes) - blocks[c] = [indexes, (row_degeneracies[c], column_degeneracies[-c])] + #num_elements = np.sum([len(t) for t in idx]) + #indexes = np.empty(num_elements, dtype=np.int64) + #np.concatenate(idx, out=indexes) + blocks[c] = [idx, (row_degeneracies[c], column_degeneracies[-c])] return blocks for c, idx in idxs.items(): @@ -246,7 +396,6 @@ def retrieve_non_zero_diagonal_blocks( np.concatenate(idx, out=indexes) blocks[c] = np.reshape(data[indexes], (row_degeneracies[c], column_degeneracies[-c])) - return blocks