From 41e420841b6fcc91695a9a2817b2ac0e0ae95b29 Mon Sep 17 00:00:00 2001 From: Albert Julius Liu Date: Mon, 30 Oct 2023 21:35:17 -0700 Subject: [PATCH] add `population.entropy()` --- src/icepool/population/base.py | 10 ++++++++++ tests/statistics_test.py | 5 +++++ 2 files changed, 15 insertions(+) diff --git a/src/icepool/population/base.py b/src/icepool/population/base.py index f55546e5..c4039043 100644 --- a/src/icepool/population/base.py +++ b/src/icepool/population/base.py @@ -683,6 +683,16 @@ def excess_kurtosis( self: 'Population[numbers.Rational] | Population[float]') -> float: return self.standardized_moment(4) - 3.0 + def entropy(self, base: float = 2.0) -> float: + """The entropy of a random sample from this population. + + Args: + base: The logarithm base to use. Default is 2.0, which gives the + entropy in bits. + """ + return -sum( + p * math.log(p, base) for p in self.probabilities() if p > 0.0) + # Joint statistics. class _Marginals(Sequence): diff --git a/tests/statistics_test.py b/tests/statistics_test.py index 71fc2432..18fdedd1 100644 --- a/tests/statistics_test.py +++ b/tests/statistics_test.py @@ -1,6 +1,7 @@ import icepool import pytest +from icepool import Die from fractions import Fraction @@ -41,3 +42,7 @@ def test_min_quantile(): def test_max_quantile(): assert icepool.d6.quantile_low(100) == 6 assert icepool.d6.quantile_high(100) == 6 + + +def test_entropy_with_zeros(): + assert Die({1: 1, 2: 0, 3: 1}).entropy() == 1.0