Ensure KDE is normalized

src/arviz_stats/base/density.py

@@ -1,6 +1,7 @@
 # pylint: disable=invalid-name,too-many-lines
 """Density estimation functions for ArviZ."""
 
+import math
 import warnings
 
 import numpy as np
@@ -238,7 +239,7 @@ def _a1inv(self, x):  # pylint: disable=no-self-use
         return 1 / (x**3 - 4 * x**2 + 3 * x)
 
     def _kappa_mle(self, x):
-        mean = self._circular_mean(x)
+        mean = self.circular_mean(x)
         kappa = self._a1inv(np.mean(np.cos(x - mean)))
         return kappa
 
@@ -474,6 +475,10 @@ def kde_linear(
 
         if cumulative:
             pdf = pdf.cumsum() / pdf.sum()
+        else:
+            # explicitly normalize to 1
+            bin_width = grid_edges[1] - grid_edges[0]
+            pdf /= pdf.sum() * bin_width
 
         return grid, pdf, bw
 
@@ -495,9 +500,7 @@ def kde_convolution(self, x, bw, grid_edges, grid_counts, grid_len, bound_correc
 
         grid = (grid_edges[1:] + grid_edges[:-1]) / 2
 
-        kernel_n = int(bw * 2 * np.pi)
-        if kernel_n == 0:
-            kernel_n = 1
+        kernel_n = 2 * math.ceil(4 * bw) + 1
 
         kernel = gaussian(kernel_n, bw)
 
@@ -633,7 +636,7 @@ def kde_circular(
             raise ValueError(f"Numeric `bw` must be positive.\nInput: {bw:.4f}.")
         if isinstance(bw, str):
             if bw == "taylor":
-                bw = self._bw_taylor(x)
+                bw = self.bw_taylor(x)
             else:
                 raise ValueError(f"`bw` must be a positive numeric or `taylor`, not {bw}")
         bw *= bw_fct

tests/base/test_stats.py

@@ -30,6 +30,31 @@ def multivariable_log_likelihood(centered_eight):
     return centered_eight
 
 
+@pytest.mark.parametrize(
+    "kde_kwargs",
+    [
+        {},
+        {
+            "adaptive": True,
+        },
+        {"circular": True},
+    ],
+    ids=["default", "adaptive", "circular"],
+)
+@pytest.mark.parametrize("bound_correction", [True, False])
+def test_kde_is_normalized(bound_correction, kde_kwargs):
+    rng = np.random.default_rng(43)
+    if kde_kwargs.get("circular", False):
+        data = rng.vonmises(np.pi, 1, (1_000, 100))
+    else:
+        data = rng.normal(size=(1_000, 100))
+    sample = ndarray_to_dataarray(data, "x", sample_dims=["sample"])
+    kde = sample.azstats.kde(dims="sample", bound_correction=bound_correction, **kde_kwargs)
+    dx = kde.sel(plot_axis="x").diff(dim="kde_dim")
+    density_norm = kde.sel(plot_axis="y").sum(dim="kde_dim") * dx
+    assert_array_almost_equal(density_norm, 1, 6)
+
+
 def test_hdi_idata(centered_eight):
     accessor = centered_eight.posterior.ds.azstats
     result = accessor.hdi()