def test_reduce_charges_non_trivial(num_charges):
np.random.seed(10)
left_charges = np.random.randint(-5, 6, (200, num_charges), dtype=np.int16)
right_charges = np.random.randint(-5,
6, (200, num_charges),
dtype=np.int16)
target_charge = np.random.randint(-2, 3, (3, num_charges), 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[0]):
#pylint: disable=no-member
tmp.append(
np.logical_and.reduce(
fused_charges == 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 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 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(dim=1).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)
