def test_reduce_charges_non_trivial_2(): np.random.seed(10) left_charges1 = np.random.randint(-5, 5, 200, dtype=np.int16) left_charges2 = np.random.randint(-5, 5, 200, dtype=np.int16) right_charges1 = np.random.randint(-5, 5, 200, dtype=np.int16) right_charges2 = np.random.randint(-5, 5, 200, dtype=np.int16) target_charge = np.array([[-2, 0, 3], [-1, 1, 0]]).astype(np.int16) fused_charges1 = fuse_ndarrays([left_charges1, right_charges1]) fused_charges2 = fuse_ndarrays([left_charges2, right_charges2]) dense_positions = reduce_charges([ U1Charge(left_charges1) @ U1Charge(left_charges2), U1Charge(right_charges1) @ U1Charge(right_charges2) ], [False, False], target_charge, return_locations=True) masks = [] assert np.all(dense_positions[0].isin(target_charge)) #pylint: disable=unsubscriptable-object for n in range(target_charge.shape[1]): mask1 = np.isin(fused_charges1, np.squeeze(target_charge[0, n])) mask2 = np.isin(fused_charges2, np.squeeze(target_charge[1, n])) masks.append(np.logical_and(mask1, mask2)) #pylint: disable=no-member np.testing.assert_allclose( np.nonzero(np.logical_or.reduce(masks))[0], dense_positions[1])
def test_find_diagonal_sparse_blocks(num_legs, num_charges): np.random.seed(10) np_charges = [ np.random.randint(-5, 5, (60, num_charges), dtype=np.int16) for _ in range(num_legs) ] fused = np.stack([ fuse_ndarrays([np_charges[n][:, c] for n in range(num_legs)]) for c in range(num_charges) ], axis=1) left_charges = np.stack([ fuse_ndarrays([np_charges[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[n][:, c] for n in range(num_legs // 2, num_legs)]) for c in range(num_charges) ], axis=1) #pylint: disable=no-member nz = np.nonzero( np.logical_and.reduce(fused == np.zeros((1, num_charges)), axis=1))[0] linear_locs = np.arange(len(nz)) # pylint: disable=no-member left_inds, _ = np.divmod(nz, right_charges.shape[0]) left = left_charges[left_inds, :] unique_left = np.unique(left, axis=0) blocks = [] for n in range(unique_left.shape[0]): ul = unique_left[n, :][None, :] #pylint: disable=no-member blocks.append(linear_locs[np.nonzero( np.logical_and.reduce(left == ul, axis=1))[0]]) charges = [ BaseCharge(left_charges, charge_types=[U1Charge] * num_charges), BaseCharge(right_charges, charge_types=[U1Charge] * num_charges) ] print(left_charges) print(right_charges) bs, cs, ss = _find_diagonal_sparse_blocks(charges, [False, False], 1) 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 test_fuse_stride_arrays(): dims = np.asarray([2, 3, 4, 5]) strides = np.asarray([120, 60, 20, 5, 1]) actual = fuse_stride_arrays(dims, strides) expected = fuse_ndarrays([ np.arange(0, strides[n] * dims[n], strides[n], dtype=np.uint32) for n in range(len(dims)) ]) np.testing.assert_allclose(actual, expected)
def test_compute_unique_fused_charges(): np.random.seed(10) qs = [np.random.randint(-3, 3, 100) for _ in range(3)] charges = [U1Charge(q) for q in qs] flows = [False, True, False] np_flows = [1, -1, 1] unique = compute_unique_fused_charges(charges, flows) fused = fuse_ndarrays([qs[n] * np_flows[n] for n in range(3)]) exp_unique = np.unique(fused) np.testing.assert_allclose(np.squeeze(unique.charges), exp_unique)
def test_compute_num_nonzero(): np.random.seed(10) qs = [np.random.randint(-3, 3, 100) for _ in range(3)] charges = [U1Charge(q) for q in qs] flows = [False, True, False] np_flows = [1, -1, 1] fused = fuse_ndarrays([qs[n] * np_flows[n] for n in range(3)]) nz1 = compute_num_nonzero(charges, flows) nz2 = len(np.nonzero(fused == 0)[0]) assert nz1 == nz2
def test_fuse_charges(): num_charges = 5 B = 6 D = 10 np_charges = [ np.random.randint(-B // 2, B // 2 + 1, D).astype(np.int16) for _ in range(num_charges) ] charges = [U1Charge(c) for c in np_charges] flows = [True, False, True, False, True] np_flows = np.ones(5, dtype=np.int16) np_flows[flows] = -1 fused = fuse_charges(charges, flows) np_fused = fuse_ndarrays([c * f for c, f in zip(np_charges, np_flows)]) np.testing.assert_allclose(np.squeeze(fused.charges), np_fused)
def fuse_stride_arrays(dims: Union[List[int], np.ndarray], strides: Union[List[int], np.ndarray]) -> np.ndarray: """ Compute linear positions of tensor elements of a tensor with dimensions `dims` according to `strides`. Args: dims: An np.ndarray of (original) tensor dimensions. strides: An np.ndarray of (possibly permuted) strides. Returns: np.ndarray: Linear positions of tensor elements according to `strides`. """ return fuse_ndarrays([ np.arange(0, strides[n] * dims[n], strides[n], dtype=SIZE_T) for n in range(len(dims)) ])
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 fuse_many_charges(num_charges, num_charge_types, seed, D, B, use_flows=False): np.random.seed(seed) if use_flows: flows = np.random.choice([True, False], num_charges, replace=True) else: flows = np.asarray([False] * num_charges) np_flows = np.ones(num_charges, dtype=np.int16) np_flows[flows] = -1 charges = [ get_charges(-B // 2, B // 2, D, num_charge_types) for _ in range(num_charges) ] fused = [ fuse_ndarrays( [charges[n][m] * np_flows[n] for n in range(num_charges)]) for m in range(num_charge_types) ] final_charges = [U1Charge(charges[n][0]) for n in range(num_charges)] for n in range(num_charges): for m in range(1, num_charge_types): final_charges[n] = final_charges[n] @ U1Charge(charges[n][m]) np_target_charges = np.random.randint(-B, B, num_charge_types, dtype=np.int16) target_charges = [ U1Charge(np.array([np_target_charges[n]])) for n in range(num_charge_types) ] target = target_charges[0] for m in range(1, num_charge_types): target = target @ target_charges[m] final = final_charges[0] * flows[0] for n in range(1, num_charges): final = final + final_charges[n] * flows[n] nz_1 = np.nonzero(final == target)[0] masks = [fused[m] == target.charges[m, 0] for m in range(num_charge_types)] #pylint: disable=no-member nz_2 = np.nonzero(np.logical_and.reduce(masks))[0] return nz_1, nz_2
def test_reduce_charges_non_trivial(): np.random.seed(10) left_charges = np.random.randint(-5, 5, 200, dtype=np.int16) right_charges = np.random.randint(-5, 5, 200, dtype=np.int16) target_charge = np.array([[-2, 0, 3]]).astype(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) assert np.all( np.isin(np.squeeze(dense_positions[0].charges), np.squeeze(target_charge))) mask = np.isin(fused_charges, np.squeeze(target_charge)) np.testing.assert_allclose(dense_positions[1], np.nonzero(mask)[0])
def test_fuse_ndarrays(): d1 = np.asarray([0, 1]) d2 = np.asarray([2, 3, 4]) fused = fuse_ndarrays([d1, d2]) np.testing.assert_allclose(fused, [2, 3, 4, 3, 4, 5])
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, (num_charges, D), 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=0) fused = np.stack(charge_list, axis=0) dims = [c.shape[1] 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=0) 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=0) #pylint: disable=no-member mask = np.logical_and.reduce(fused.T == 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.T == 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[1]) left = left_charges[:, left_inds] unique_left = np.unique(left, axis=1) blocks = [] for n in range(unique_left.shape[1]): ul = unique_left[:, n][None, :] #pylint: disable=no-member blocks.append(dense_to_sparse[tr_linear_locs[np.nonzero( np.logical_and.reduce(left.T == 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)