def test_reduce_1d():
q = np.array([0, 1, 2, 0, 6, 1, -9, 0, -7])
Q = BaseCharge(charges=q)
target_charge = np.array([0, 1])
expected = np.array([[0, 1, 0, 1, 0]]).T
res, locs = Q.reduce(target_charge, return_locations=True)
np.testing.assert_allclose(res.charges, expected)
np.testing.assert_allclose(locs, [0, 1, 3, 5, 7])
def test_reduce_integer():
q = np.array([0, 1, 2, 0, 6, 1, -9, 0, -7])
Q = BaseCharge(charges=q)
target_charge = 0
expected = np.zeros((3, 1))
res, locs = Q.reduce(target_charge, return_locations=True)
np.testing.assert_allclose(res.charges, expected)
np.testing.assert_allclose(locs, [0, 3, 7])
def test_reduce():
q = np.array([[0, 1, 2, 0, 6, 1, -9, 0, -7], [2, 3, 4, -1, 4, 3, 1, 2, 0]])
Q = BaseCharge(charges=q)
target_charge = np.array([[0, 1, 6, -12], [2, 3, 4, 16]])
expected = np.array([[0, 1, 6, 1, 0], [2, 3, 4, 3, 2]])
res, locs = Q.reduce(target_charge, return_locations=True)
np.testing.assert_allclose(res.charges, expected)
np.testing.assert_allclose(locs, [0, 1, 4, 5, 7])
