entropy normalization arg (#3369)

polars/polars-core/src/series/ops/log.rs

@@ -26,14 +26,19 @@ impl Series {
     /// Compute the entropy as `-sum(pk * log(pk)`.
     /// where `pk` are discrete probabilities.
     #[cfg_attr(docsrs, doc(cfg(feature = "log")))]
-    pub fn entropy(&self, base: f64) -> Option<f64> {
+    pub fn entropy(&self, base: f64, normalize: bool) -> Option<f64> {
         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
+                let pk = if normalize {
+                    let sum = pk.sum_as_series();
+
+                    if sum.get(0).extract::<f64>()? != 1.0 {
+                        pk / &sum
+                    } else {
+                        pk.clone()
+                    }
                 } else {
                     pk.clone()
                 };
@@ -44,7 +49,7 @@ impl Series {
             _ => self
                 .cast(&DataType::Float64)
                 .ok()
-                .and_then(|s| s.entropy(base)),
+                .and_then(|s| s.entropy(base, normalize)),
         }
     }
 }

polars/polars-lazy/src/dsl/mod.rs

@@ -1810,9 +1810,9 @@ impl Expr {
     #[cfg_attr(docsrs, doc(cfg(feature = "log")))]
     /// Compute the entropy as `-sum(pk * log(pk)`.
     /// where `pk` are discrete probabilities.
-    pub fn entropy(self, base: f64) -> Self {
+    pub fn entropy(self, base: f64, normalize: bool) -> Self {
         self.apply(
-            move |s| Ok(Series::new(s.name(), [s.entropy(base)])),
+            move |s| Ok(Series::new(s.name(), [s.entropy(base, normalize)])),
             GetOutput::map_dtype(|dt| {
                 if matches!(dt, DataType::Float32) {
                     DataType::Float32

py-polars/polars/internals/expr.py

@@ -2884,19 +2884,20 @@ def log(self, base: float = math.e) -> "Expr":
         """
         return wrap_expr(self._pyexpr.log(base))
 
-    def entropy(self, base: float = math.e) -> "Expr":
+    def entropy(self, base: float = math.e, normalize: bool = False) -> "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
             Given base, defaults to `e`
+        normalize
+            Normalize pk if it doesn't sum to 1.
+
         """
-        return wrap_expr(self._pyexpr.entropy(base))
+        return wrap_expr(self._pyexpr.entropy(base, normalize))
 
     # Below are the namespaces defined. Keep these at the end of the definition of Expr, as to not confuse mypy with
     # the type annotation `str` with the namespace "str"

py-polars/polars/internals/series.py

@@ -947,25 +947,32 @@ def unique_counts(self) -> "Series":
         """
         return pli.select(pli.lit(self).unique_counts()).to_series()
 
-    def entropy(self, base: float = math.e) -> Optional[float]:
+    def entropy(self, base: float = math.e, normalize: bool = False) -> 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.
 
+        Parameters
+        ----------
+        base
+            Given base, defaults to `e`
+        normalize
+            Normalize pk if it doesn't sum to 1.
+
         Examples
         --------
 
         >>> a = pl.Series([0.99, 0.005, 0.005])
-        >>> a.entropy()
+        >>> a.entropy(normalize=True)
         0.06293300616044681
         >>> b = pl.Series([0.65, 0.10, 0.25])
-        >>> b.entropy()
+        >>> b.entropy(normalize=True)
         0.8568409950394724
 
         """
-        return pli.select(pli.lit(self).entropy(base)).to_series()[0]
+        return pli.select(pli.lit(self).entropy(base, normalize)).to_series()[0]
 
     @property
     def name(self) -> str:

py-polars/src/lazy/dsl.rs

@@ -1358,8 +1358,8 @@ impl PyExpr {
         self.inner.clone().log(base).into()
     }
 
-    pub fn entropy(&self, base: f64) -> Self {
-        self.inner.clone().entropy(base).into()
+    pub fn entropy(&self, base: f64, normalize: bool) -> Self {
+        self.inner.clone().entropy(base, normalize).into()
     }
     pub fn hash(&self, seed: usize) -> Self {
         self.inner.clone().hash(seed).into()

py-polars/tests/test_exprs.py

@@ -173,7 +173,9 @@ def test_entropy() -> None:
     )
 
     assert (
-        df.groupby("group", maintain_order=True).agg(pl.col("id").entropy())
+        df.groupby("group", maintain_order=True).agg(
+            pl.col("id").entropy(normalize=True)
+        )
     ).frame_equal(
         pl.DataFrame(
             {"group": ["A", "B"], "id": [1.0397207708399179, 1.371381017771811]}