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])