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)
