def test_reduce_charges_non_trivial(num_charges): np.random.seed(10) left_charges = np.random.randint(-5, 6, (num_charges, 200), dtype=np.int16) right_charges = np.random.randint(-5, 6, (num_charges, 200), dtype=np.int16) target_charge = np.random.randint(-2, 3, (num_charges, 3), dtype=np.int16) charge_types = [U1Charge] * num_charges fused_charges = fuse_ndarray_charges(left_charges, right_charges, charge_types) dense_positions = reduce_charges([ BaseCharge(left_charges, charge_types=charge_types), BaseCharge(right_charges, charge_types=charge_types) ], [False, False], target_charge, return_locations=True) assert np.all( np.isin(np.squeeze(dense_positions[0].charges), np.squeeze(target_charge))) tmp = [] #pylint: disable=unsubscriptable-object for n in range(target_charge.shape[1]): #pylint: disable=no-member tmp.append( np.logical_and.reduce(fused_charges.T == target_charge[:, n][None, :], axis=1)) #pylint: disable=no-member mask = np.logical_or.reduce(tmp) np.testing.assert_allclose(dense_positions[1], np.nonzero(mask)[0])
def test_reduce_charges_2(): left_charges = np.asarray([[-2, 0, 1, 0, 0], [-3, 0, 2, 1, 0]]).astype(np.int16).T right_charges = np.asarray([[-1, 0, 2, 1], [-2, 2, 7, 0]]).astype(np.int16).T target_charge = np.zeros((1, 2), dtype=np.int16) fused_charges = fuse_ndarray_charges(left_charges, right_charges, [U1Charge, U1Charge]) dense_positions = reduce_charges([ BaseCharge(left_charges, charge_types=[U1Charge, U1Charge]), BaseCharge(right_charges, charge_types=[U1Charge, U1Charge]) ], [False, False], target_charge, return_locations=True) np.testing.assert_allclose(dense_positions[0].charges, 0) #pylint: disable=no-member np.testing.assert_allclose( dense_positions[1], np.nonzero(np.logical_and.reduce(fused_charges == target_charge, axis=1))[0])
def reduce_charges(charges: List[BaseCharge], flows: Union[np.ndarray, List[bool]], target_charges: np.ndarray, return_locations: Optional[bool] = False, strides: Optional[np.ndarray] = None) -> Any: """ Add quantum numbers arising from combining two or more charges into a single index, keeping only the quantum numbers that appear in `target_charges`. Equilvalent to using "combine_charges" followed by "reduce", but is generally much more efficient. Args: charges: List of `BaseCharge`, one for each leg of a tensor. flows: A list of bool, one for each leg of a tensor. with values `False` or `True` denoting inflowing and outflowing charge direction, respectively. target_charges: n-by-D array of charges which should be kept, with `n` the number of symmetries. return_locations: If `True` return the location of the kept values of the fused charges strides: Index strides with which to compute the retured locations of the kept elements. Defaults to trivial strides (based on row major order). Returns: BaseCharge: the fused index after reduction. np.ndarray: Locations of the fused BaseCharge charges that were kept. """ tensor_dims = [len(c) for c in charges] if len(charges) == 1: # reduce single index if strides is None: strides = np.array([1], dtype=SIZE_T) return charges[0].dual(flows[0]).reduce( target_charges, return_locations=return_locations, strides=strides[0]) # find size-balanced partition of charges partition = _find_best_partition(tensor_dims) # compute quantum numbers for each partition left_ind = fuse_charges(charges[:partition], flows[:partition]) right_ind = fuse_charges(charges[partition:], flows[partition:]) # compute combined qnums comb_qnums = fuse_ndarray_charges(left_ind.unique_charges, right_ind.unique_charges, charges[0].charge_types) #special case of empty charges #pylint: disable=unsubscriptable-object if (comb_qnums.shape[0] == 0) or (len(left_ind.charge_labels) == 0) or (len( right_ind.charge_labels) == 0): obj = charges[0].__new__(type(charges[0])) obj.__init__( np.empty((0, charges[0].num_symmetries), dtype=charges[0].dtype), np.empty(0, dtype=charges[0].label_dtype), charges[0].charge_types) if return_locations: return obj, np.empty(0, dtype=SIZE_T) return obj unique_comb_qnums, comb_labels = np.unique( comb_qnums, return_inverse=True, axis=0) num_unique = unique_comb_qnums.shape[0] # intersect combined qnums and target_charges reduced_qnums, label_to_unique, _ = intersect( unique_comb_qnums, target_charges, axis=0, return_indices=True) map_to_kept = -np.ones(num_unique, dtype=charges[0].label_dtype) map_to_kept[label_to_unique] = np.arange(len(label_to_unique)) # new_comb_labels is a matrix of shape # (left_ind.num_unique, right_ind.num_unique) # each row new_comb_labels[n,:] contains integers values. # Positions where values > 0 # denote labels of right-charges that are kept. new_comb_labels = map_to_kept[comb_labels].reshape( [left_ind.num_unique, right_ind.num_unique]) reduced_rows = [0] * left_ind.num_unique for n in range(left_ind.num_unique): temp_label = new_comb_labels[n, right_ind.charge_labels] reduced_rows[n] = temp_label[temp_label >= 0] reduced_labels = np.concatenate( [reduced_rows[n] for n in left_ind.charge_labels]) obj = charges[0].__new__(type(charges[0])) obj.__init__(reduced_qnums, reduced_labels, charges[0].charge_types) if return_locations: row_locs = [0] * left_ind.num_unique if strides is not None: # computed locations based on non-trivial strides row_pos = fuse_stride_arrays(tensor_dims[:partition], strides[:partition]) col_pos = fuse_stride_arrays(tensor_dims[partition:], strides[partition:]) for n in range(left_ind.num_unique): temp_label = new_comb_labels[n, right_ind.charge_labels] temp_keep = temp_label >= 0 if strides is not None: row_locs[n] = col_pos[temp_keep] else: row_locs[n] = np.where(temp_keep)[0] if strides is not None: reduced_locs = np.concatenate([ row_pos[n] + row_locs[left_ind.charge_labels[n]] for n in range(left_ind.dim) ]) else: reduced_locs = np.concatenate([ n * right_ind.dim + row_locs[left_ind.charge_labels[n]] for n in range(left_ind.dim) ]) return obj, reduced_locs return obj
def fuse_many_ndarray_charges(charges, charge_types): res = fuse_ndarray_charges(charges[0], charges[1], charge_types) for n in range(2, len(charges)): res = fuse_ndarray_charges(res, charges[n], charge_types) return res