def test_find_best_partition(): with pytest.raises(ValueError): _find_best_partition([5])
def _find_transposed_diagonal_sparse_blocks( charges: List[BaseCharge], flows: Union[np.ndarray, List[bool]], tr_partition: int, order: Optional[Union[List, np.ndarray]] = None ) -> Tuple[List, BaseCharge, np.ndarray]: """ Find the diagonal blocks of a transposed tensor with meta-data `charges` and `flows`. `charges` and `flows` are the charges and flows of the untransposed tensor, `order` is the final transposition, and `tr_partition` is the partition of the transposed tensor according to which the diagonal blocks should be found. Args: charges: List of `BaseCharge`, one for each leg of a tensor. flows: A list of bool, one for each leg of a tensor. with values `False` or `True` denoting inflowing and outflowing charge direction, respectively. tr_partition: Location of the transposed tensor partition (i.e. such that the tensor is viewed as a matrix between `charges[order[:partition]]` and `charges[order[partition:]]`). order: Order with which to permute the tensor axes. Returns: block_maps (List[np.ndarray]): list of integer arrays, which each containing the location of a symmetry block in the data vector. block_qnums (BaseCharge): The charges of the corresponding blocks. block_dims (np.ndarray): 2-by-m array of matrix dimensions of each block. """ flows = np.asarray(flows) cacher = get_cacher() if cacher.do_caching: hash_val = _to_string(charges, flows, tr_partition, order) if hash_val in cacher.cache: return cacher.cache[hash_val] if np.array_equal(order, None) or (np.array_equal( np.array(order), np.arange(len(charges)))): # no transpose order return _find_diagonal_sparse_blocks(charges, flows, tr_partition) # general case: non-trivial transposition is required num_inds = len(charges) tensor_dims = np.array([charges[n].dim for n in range(num_inds)], dtype=int) strides = np.append(np.flip(np.cumprod(np.flip(tensor_dims[1:]))), 1) # compute qnums of row/cols in original tensor orig_partition = _find_best_partition(tensor_dims) orig_width = np.prod(tensor_dims[orig_partition:]) orig_unique_row_qnums = compute_unique_fused_charges( charges[:orig_partition], flows[:orig_partition]) orig_unique_col_qnums, orig_col_degen = compute_fused_charge_degeneracies( charges[orig_partition:], np.logical_not(flows[orig_partition:])) orig_block_qnums, row_map, col_map = intersect( orig_unique_row_qnums.unique_charges, orig_unique_col_qnums.unique_charges, axis=0, return_indices=True) orig_num_blocks = orig_block_qnums.shape[0] if orig_num_blocks == 0: # special case: trivial number of non-zero elements obj = charges[0].__new__(type(charges[0])) obj.__init__( np.empty((0, charges[0].num_symmetries), dtype=charges[0].dtype), np.arange(0, dtype=charges[0].label_dtype), charges[0].charge_types) return [], obj, np.empty((2, 0), dtype=SIZE_T) orig_row_ind = fuse_charges(charges[:orig_partition], flows[:orig_partition]) orig_col_ind = fuse_charges(charges[orig_partition:], np.logical_not(flows[orig_partition:])) inv_row_map = -np.ones(orig_unique_row_qnums.unique_charges.shape[0], dtype=charges[0].label_dtype) inv_row_map[row_map] = np.arange(len(row_map), dtype=charges[0].label_dtype) all_degens = np.append(orig_col_degen[col_map], 0)[inv_row_map[orig_row_ind.charge_labels]] all_cumul_degens = np.cumsum(np.insert(all_degens[:-1], 0, 0)).astype(SIZE_T) dense_to_sparse = np.empty(orig_width, dtype=SIZE_T) for n in range(orig_num_blocks): dense_to_sparse[orig_col_ind.charge_labels == col_map[n]] = np.arange( orig_col_degen[col_map[n]], dtype=SIZE_T) # define properties of new tensor resulting from transposition new_strides = strides[order] new_row_charges = [charges[n] for n in order[:tr_partition]] new_col_charges = [charges[n] for n in order[tr_partition:]] new_row_flows = flows[order[:tr_partition]] new_col_flows = flows[order[tr_partition:]] if tr_partition == 0: # special case: reshape into row vector # compute qnums of row/cols in transposed tensor unique_col_qnums, new_col_degen = compute_fused_charge_degeneracies( new_col_charges, np.logical_not(new_col_flows)) identity_charges = charges[0].identity_charges(dim=1) block_qnums, new_row_map, new_col_map = intersect( identity_charges.unique_charges, unique_col_qnums.unique_charges, axis=0, return_indices=True) block_dims = np.array([[1], new_col_degen[new_col_map]], dtype=SIZE_T) num_blocks = 1 col_ind, col_locs = reduce_charges(new_col_charges, np.logical_not(new_col_flows), block_qnums, return_locations=True, strides=new_strides[tr_partition:]) # find location of blocks in transposed tensor (w.r.t positions in original) #pylint: disable=no-member orig_row_posR, orig_col_posR = np.divmod( col_locs[col_ind.charge_labels == 0], orig_width) block_maps = [(all_cumul_degens[orig_row_posR] + dense_to_sparse[orig_col_posR]).ravel()] obj = charges[0].__new__(type(charges[0])) obj.__init__( block_qnums, np.arange(block_qnums.shape[0], dtype=charges[0].label_dtype), charges[0].charge_types) elif tr_partition == len(charges): # special case: reshape into col vector # compute qnums of row/cols in transposed tensor unique_row_qnums, new_row_degen = compute_fused_charge_degeneracies( new_row_charges, new_row_flows) identity_charges = charges[0].identity_charges(dim=1) block_qnums, new_row_map, new_col_map = intersect( unique_row_qnums.unique_charges, identity_charges.unique_charges, axis=0, return_indices=True) block_dims = np.array([new_row_degen[new_row_map], [1]], dtype=SIZE_T) num_blocks = 1 row_ind, row_locs = reduce_charges(new_row_charges, new_row_flows, block_qnums, return_locations=True, strides=new_strides[:tr_partition]) # find location of blocks in transposed tensor (w.r.t positions in original) #pylint: disable=no-member orig_row_posL, orig_col_posL = np.divmod( row_locs[row_ind.charge_labels == 0], orig_width) block_maps = [(all_cumul_degens[orig_row_posL] + dense_to_sparse[orig_col_posL]).ravel()] obj = charges[0].__new__(type(charges[0])) obj.__init__( block_qnums, np.arange(block_qnums.shape[0], dtype=charges[0].label_dtype), charges[0].charge_types) else: unique_row_qnums, new_row_degen = compute_fused_charge_degeneracies( new_row_charges, new_row_flows) unique_col_qnums, new_col_degen = compute_fused_charge_degeneracies( new_col_charges, np.logical_not(new_col_flows)) block_qnums, new_row_map, new_col_map = intersect( unique_row_qnums.unique_charges, unique_col_qnums.unique_charges, axis=0, return_indices=True) block_dims = np.array( [new_row_degen[new_row_map], new_col_degen[new_col_map]], dtype=SIZE_T) num_blocks = len(new_row_map) row_ind, row_locs = reduce_charges(new_row_charges, new_row_flows, block_qnums, return_locations=True, strides=new_strides[:tr_partition]) col_ind, col_locs = reduce_charges(new_col_charges, np.logical_not(new_col_flows), block_qnums, return_locations=True, strides=new_strides[tr_partition:]) block_maps = [0] * num_blocks for n in range(num_blocks): #pylint: disable=no-member orig_row_posL, orig_col_posL = np.divmod( row_locs[row_ind.charge_labels == n], orig_width) #pylint: disable=no-member orig_row_posR, orig_col_posR = np.divmod( col_locs[col_ind.charge_labels == n], orig_width) block_maps[n] = ( all_cumul_degens[np.add.outer(orig_row_posL, orig_row_posR)] + dense_to_sparse[np.add.outer(orig_col_posL, orig_col_posR)]).ravel() obj = charges[0].__new__(type(charges[0])) obj.__init__( block_qnums, np.arange(block_qnums.shape[0], dtype=charges[0].label_dtype), charges[0].charge_types) if cacher.do_caching: cacher.cache[hash_val] = (block_maps, obj, block_dims) return cacher.cache[hash_val] return block_maps, obj, block_dims
def test_find_best_partition_raises(): d = [5, 4, 5, 2, 6, 8] p = _find_best_partition(d) assert p == 3
def reduce_charges(charges: List[BaseCharge], flows: Union[np.ndarray, List[bool]], target_charges: np.ndarray, return_locations: Optional[bool] = False, strides: Optional[np.ndarray] = None) -> Any: """ Add quantum numbers arising from combining two or more charges into a single index, keeping only the quantum numbers that appear in `target_charges`. Equilvalent to using "combine_charges" followed by "reduce", but is generally much more efficient. Args: charges: List of `BaseCharge`, one for each leg of a tensor. flows: A list of bool, one for each leg of a tensor. with values `False` or `True` denoting inflowing and outflowing charge direction, respectively. target_charges: n-by-D array of charges which should be kept, with `n` the number of symmetries. return_locations: If `True` return the location of the kept values of the fused charges strides: Index strides with which to compute the retured locations of the kept elements. Defaults to trivial strides (based on row major order). Returns: BaseCharge: the fused index after reduction. np.ndarray: Locations of the fused BaseCharge charges that were kept. """ tensor_dims = [len(c) for c in charges] if len(charges) == 1: # reduce single index if strides is None: strides = np.array([1], dtype=SIZE_T) return charges[0].dual(flows[0]).reduce( target_charges, return_locations=return_locations, strides=strides[0]) # find size-balanced partition of charges partition = _find_best_partition(tensor_dims) # compute quantum numbers for each partition left_ind = fuse_charges(charges[:partition], flows[:partition]) right_ind = fuse_charges(charges[partition:], flows[partition:]) # compute combined qnums comb_qnums = fuse_ndarray_charges(left_ind.unique_charges, right_ind.unique_charges, charges[0].charge_types) #special case of empty charges #pylint: disable=unsubscriptable-object if (comb_qnums.shape[0] == 0) or (len(left_ind.charge_labels) == 0) or (len(right_ind.charge_labels) == 0): obj = charges[0].__new__(type(charges[0])) obj.__init__( np.empty((0, charges[0].num_symmetries), dtype=charges[0].dtype), np.empty(0, dtype=charges[0].label_dtype), charges[0].charge_types) if return_locations: return obj, np.empty(0, dtype=SIZE_T) return obj unique_comb_qnums, comb_labels = unique(comb_qnums, return_inverse=True) num_unique = unique_comb_qnums.shape[0] # intersect combined qnums and target_charges reduced_qnums, label_to_unique, _ = intersect(unique_comb_qnums, target_charges, axis=0, return_indices=True) map_to_kept = -np.ones(num_unique, dtype=charges[0].label_dtype) map_to_kept[label_to_unique] = np.arange(len(label_to_unique)) # new_comb_labels is a matrix of shape # (left_ind.num_unique, right_ind.num_unique) # each row new_comb_labels[n,:] contains integers values. # Positions where values > 0 # denote labels of right-charges that are kept. new_comb_labels = map_to_kept[comb_labels].reshape( [left_ind.num_unique, right_ind.num_unique]) reduced_rows = [0] * left_ind.num_unique for n in range(left_ind.num_unique): temp_label = new_comb_labels[n, right_ind.charge_labels] reduced_rows[n] = temp_label[temp_label >= 0] reduced_labels = np.concatenate( [reduced_rows[n] for n in left_ind.charge_labels]) obj = charges[0].__new__(type(charges[0])) obj.__init__(reduced_qnums, reduced_labels, charges[0].charge_types) if return_locations: row_locs = [0] * left_ind.num_unique if strides is not None: # computed locations based on non-trivial strides row_pos = fuse_stride_arrays(tensor_dims[:partition], strides[:partition]) col_pos = fuse_stride_arrays(tensor_dims[partition:], strides[partition:]) for n in range(left_ind.num_unique): temp_label = new_comb_labels[n, right_ind.charge_labels] temp_keep = temp_label >= 0 if strides is not None: row_locs[n] = col_pos[temp_keep] else: row_locs[n] = np.where(temp_keep)[0] if strides is not None: reduced_locs = np.concatenate([ row_pos[n] + row_locs[left_ind.charge_labels[n]] for n in range(left_ind.dim) ]) else: reduced_locs = np.concatenate([ n * right_ind.dim + row_locs[left_ind.charge_labels[n]] for n in range(left_ind.dim) ]) return obj, reduced_locs return obj