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(): 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) fused_charges = fuse_ndarrays([left_charges, right_charges]) dense_positions = reduce_charges( [U1Charge(left_charges), U1Charge(right_charges)], [False, False], target_charge, return_locations=True) np.testing.assert_allclose(dense_positions[0].charges, 0) np.testing.assert_allclose( dense_positions[1], np.nonzero(fused_charges == target_charge[0, 0])[0])
def fromdense(cls, indices: List[Index], array: np.ndarray) -> "BlockSparseTensor": """ Initialize a BlockSparseTensor from a dense array. Args: indices: A list of `Index` objects. array: A numpy array. Returns: BlockSparseTensors: A Tensor initialized from the elements of `array` at the positions where `indices` fuse to the identity charge. """ shape = [i.dim for i in indices] if not np.array_equal(shape, array.shape): raise ValueError( f"Cannot initialize an BlockSparseTensor of shape {shape}" f" from an array of shape {array.shape}") tmp = np.append(0, np.cumsum([len(i.flat_charges) for i in indices])) order = [ list(np.arange(tmp[n], tmp[n + 1])) for n in range(len(tmp) - 1) ] charges = [] flows = [] for i in indices: charges.extend(i.flat_charges) flows.extend(i.flat_flows) _, locs = reduce_charges( charges=charges, flows=flows, target_charges=charges[0].identity_charges.unique_charges, return_locations=True) ar = np.ravel(array) data = ar[locs] return cls(data=data, charges=charges, flows=flows, order=order, check_consistency=False)
def test_reduce_charges_2(): left_charges = np.asarray([[-2, 0, 1, 0, 0], [-3, 0, 2, 1, 0]]).astype(np.int16) right_charges = np.asarray([[-1, 0, 2, 1], [-2, 2, 7, 0]]).astype(np.int16) target_charge = np.zeros((2, 1), 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.T == target_charge.T, axis=1))[0])