def test_trimmed_mean_corner_cases(self):
td = TDigest()
mean = td.trimmed_mean(0, 100)
assert mean == 0
td.update(1)
mean = td.trimmed_mean(0, 100)
assert mean == 1
td.update(1000)
mean = td.trimmed_mean(0, 100)
assert mean == 500.5
def test_trimmed_mean_negative(self):
td = TDigest()
for i in range(100):
td.update(random.random())
for i in range(10):
td.update(i * 100)
mean = td.trimmed_mean(1, 99)
assert mean >= 0
def test_trimmed_mean(self, percentile_range, data_size):
p1 = percentile_range[0]
p2 = percentile_range[1]
t = TDigest()
x = random.random(size=data_size)
t.batch_update(x)
tm_actual = t.trimmed_mean(p1, p2)
tm_expected = x[
bitwise_and(x >= percentile(x, p1), x <= percentile(x, p2))
].mean()
testing.assert_allclose(tm_actual, tm_expected, rtol=0.01, atol=0.01)