fix entropy (#3070)

pola-rs · Apr 6, 2022 · 0363e97 · 0363e97
1 parent 800a450
commit 0363e97
Show file tree

Hide file tree

Showing 5 changed files with 40 additions and 19 deletions.
diff --git a/polars/polars-core/src/series/ops/log.rs b/polars/polars-core/src/series/ops/log.rs
@@ -27,8 +27,24 @@ impl Series {
     /// where `pk` are discrete probabilities.
     #[cfg_attr(docsrs, doc(cfg(feature = "log")))]
     pub fn entropy(&self, base: f64) -> Option<f64> {
-        let pk = self;
-        let log_pk = pk.log(base);
-        (pk * &log_pk).sum::<f64>().map(|v| -v)
+        match self.dtype() {
+            DataType::Float32 | DataType::Float64 => {
+                let pk = self;
+                let sum = pk.sum_as_series();
+
+                let pk = if sum.get(0).extract::<f64>()? != 1.0 {
+                    pk / &sum
+                } else {
+                    pk.clone()
+                };
+
+                let log_pk = pk.log(base);
+                (&pk * &log_pk).sum::<f64>().map(|v| -v)
+            }
+            _ => self
+                .cast(&DataType::Float64)
+                .ok()
+                .and_then(|s| s.entropy(base)),
+        }
     }
 }
diff --git a/polars/polars-lazy/src/dsl/mod.rs b/polars/polars-lazy/src/dsl/mod.rs
@@ -2141,7 +2141,11 @@ impl Expr {
                 }
             }),
         )
-        .with_fmt("entropy")
+        .with_function_options(|mut options| {
+            options.fmt_str = "entropy";
+            options.auto_explode = true;
+            options
+        })
     }
 
     #[cfg(feature = "strings")]

diff --git a/py-polars/polars/internals/expr.py b/py-polars/polars/internals/expr.py
@@ -2879,6 +2879,8 @@ def entropy(self, base: float = math.e) -> "Expr":
         Compute the entropy as `-sum(pk * log(pk)`.
         where `pk` are discrete probabilities.
 
+        This routine will normalize pk if they don’t sum to 1.
+
         Parameters
         ----------
         base

diff --git a/py-polars/polars/internals/series.py b/py-polars/polars/internals/series.py
@@ -933,6 +933,8 @@ def entropy(self, base: float = math.e) -> Optional[float]:
         """
         Compute the entropy as `-sum(pk * log(pk)`.
         where `pk` are discrete probabilities.
+
+        This routine will normalize pk if they don’t sum to 1.
         """
         return pli.select(pli.lit(self).entropy(base)).to_series()[0]
 

diff --git a/py-polars/tests/test_exprs.py b/py-polars/tests/test_exprs.py
@@ -165,23 +165,20 @@ def test_unique_counts() -> None:
 
 
 def test_entropy() -> None:
-    df = pl.DataFrame({"id": [1, 1, 2, 2, 3]})
+    df = pl.DataFrame(
+        {
+            "group": ["A", "A", "A", "B", "B", "B", "B"],
+            "id": [1, 2, 1, 4, 5, 4, 6],
+        }
+    )
+
     assert (
-        df.select(
-            [
-                (
-                    -(
-                        pl.col("id").unique_counts()
-                        / pl.count()
-                        * (pl.col("id").unique_counts() / pl.count()).log()
-                    ).sum()
-                ).alias("e0"),
-                ((pl.col("id").unique_counts() / pl.count()).entropy()).alias("e1"),
-            ]
-        ).rows()
-        == [(1.0549201679861442, 1.0549201679861442)]
+        df.groupby("group", maintain_order=True).agg(pl.col("id").entropy())
+    ).frame_equal(
+        pl.DataFrame(
+            {"group": ["A", "B"], "id": [1.0397207708399179, 1.371381017771811]}
+        )
     )
-    assert df["id"].entropy() == -6.068425588244111
 
 
 def test_dot_in_groupby() -> None: