def test_iterative_weighted_nd():
    data = np.array([[-10], [3], [5], [7], [8], [50], [100], [5000]])
    weights = np.array([1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0])
    y_init = [np.sum(data * np.expand_dims(weights, -1)) / np.sum(weights)]
    y, _ = weiszfeld_nd(data, y_init, weights=weights)
    print(y_init, y)
    assert 6.99 < y[-1] < 7.01
def test_iterative_nd():
    data = np.array([[-10], [3], [5], [7], [8], [50], [5000]])
    y_init = np.mean(data, axis=0)
    y, _ = weiszfeld_nd(data, y_init)
    print(y_init, y)
    assert 6.99 < y < 7.01