from __future__ import annotations from dataclasses import dataclass from typing import Any, Callable import numpy as np from numpy import ndarray import pandas as pd from pandas import DataFrame try: from scipy.stats import gaussian_kde _no_scipy = False except ImportError: from seaborn.external.kde import gaussian_kde _no_scipy = True from seaborn._core.groupby import GroupBy from seaborn._core.scales import Scale from seaborn._stats.base import Stat @dataclass class KDE(Stat): """ Compute a univariate kernel density estimate. Parameters ---------- bw_adjust : float Factor that multiplicatively scales the value chosen using `bw_method`. Increasing will make the curve smoother. See Notes. bw_method : string, scalar, or callable Method for determining the smoothing bandwidth to use. Passed directly to :class:`scipy.stats.gaussian_kde`; see there for options. common_norm : bool or list of variables If `True`, normalize so that the areas of all curves sums to 1. If `False`, normalize each curve independently. If a list, defines variable(s) to group by and normalize within. common_grid : bool or list of variables If `True`, all curves will share the same evaluation grid. If `False`, each evaluation grid is independent. If a list, defines variable(s) to group by and share a grid within. gridsize : int or None Number of points in the evaluation grid. If None, the density is evaluated at the original datapoints. cut : float Factor, multiplied by the kernel bandwidth, that determines how far the evaluation grid extends past the extreme datapoints. When set to 0, the curve is truncated at the data limits. cumulative : bool If True, estimate a cumulative distribution function. Requires scipy. Notes ----- The *bandwidth*, or standard deviation of the smoothing kernel, is an important parameter. Much like histogram bin width, using the wrong bandwidth can produce a distorted representation. Over-smoothing can erase true features, while under-smoothing can create false ones. The default uses a rule-of-thumb that works best for distributions that are roughly bell-shaped. It is a good idea to check the default by varying `bw_adjust`. Because the smoothing is performed with a Gaussian kernel, the estimated density curve can extend to values that may not make sense. For example, the curve may be drawn over negative values when data that are naturally positive. The `cut` parameter can be used to control the evaluation range, but datasets that have many observations close to a natural boundary may be better served by a different method. Similar distortions may arise when a dataset is naturally discrete or "spiky" (containing many repeated observations of the same value). KDEs will always produce a smooth curve, which could be misleading. The units on the density axis are a common source of confusion. While kernel density estimation produces a probability distribution, the height of the curve at each point gives a density, not a probability. A probability can be obtained only by integrating the density across a range. The curve is normalized so that the integral over all possible values is 1, meaning that the scale of the density axis depends on the data values. If scipy is installed, its cython-accelerated implementation will be used. Examples -------- .. include:: ../docstrings/objects.KDE.rst """ bw_adjust: float = 1 bw_method: str | float | Callable[[gaussian_kde], float] = "scott" common_norm: bool | list[str] = True common_grid: bool | list[str] = True gridsize: int | None = 200 cut: float = 3 cumulative: bool = False def __post_init__(self): if self.cumulative and _no_scipy: raise RuntimeError("Cumulative KDE evaluation requires scipy") def _check_var_list_or_boolean(self, param: str, grouping_vars: Any) -> None: """Do input checks on grouping parameters.""" value = getattr(self, param) if not ( isinstance(value, bool) or (isinstance(value, list) and all(isinstance(v, str) for v in value)) ): param_name = f"{self.__class__.__name__}.{param}" raise TypeError(f"{param_name} must be a boolean or list of strings.") self._check_grouping_vars(param, grouping_vars, stacklevel=3) def _fit(self, data: DataFrame, orient: str) -> gaussian_kde: """Fit and return a KDE object.""" # TODO need to handle singular data fit_kws: dict[str, Any] = {"bw_method": self.bw_method} if "weight" in data: fit_kws["weights"] = data["weight"] kde = gaussian_kde(data[orient], **fit_kws) kde.set_bandwidth(kde.factor * self.bw_adjust) return kde def _get_support(self, data: DataFrame, orient: str) -> ndarray: """Define the grid that the KDE will be evaluated on.""" if self.gridsize is None: return data[orient].to_numpy() kde = self._fit(data, orient) bw = np.sqrt(kde.covariance.squeeze()) gridmin = data[orient].min() - bw * self.cut gridmax = data[orient].max() + bw * self.cut return np.linspace(gridmin, gridmax, self.gridsize) def _fit_and_evaluate( self, data: DataFrame, orient: str, support: ndarray ) -> DataFrame: """Transform single group by fitting a KDE and evaluating on a support grid.""" empty = pd.DataFrame(columns=[orient, "weight", "density"], dtype=float) if len(data) < 2: return empty try: kde = self._fit(data, orient) except np.linalg.LinAlgError: return empty if self.cumulative: s_0 = support[0] density = np.array([kde.integrate_box_1d(s_0, s_i) for s_i in support]) else: density = kde(support) weight = data["weight"].sum() return pd.DataFrame({orient: support, "weight": weight, "density": density}) def _transform( self, data: DataFrame, orient: str, grouping_vars: list[str] ) -> DataFrame: """Transform multiple groups by fitting KDEs and evaluating.""" empty = pd.DataFrame(columns=[*data.columns, "density"], dtype=float) if len(data) < 2: return empty try: support = self._get_support(data, orient) except np.linalg.LinAlgError: return empty grouping_vars = [x for x in grouping_vars if data[x].nunique() > 1] if not grouping_vars: return self._fit_and_evaluate(data, orient, support) groupby = GroupBy(grouping_vars) return groupby.apply(data, self._fit_and_evaluate, orient, support) def __call__( self, data: DataFrame, groupby: GroupBy, orient: str, scales: dict[str, Scale], ) -> DataFrame: if "weight" not in data: data = data.assign(weight=1) data = data.dropna(subset=[orient, "weight"]) # Transform each group separately grouping_vars = [str(v) for v in data if v in groupby.order] if not grouping_vars or self.common_grid is True: res = self._transform(data, orient, grouping_vars) else: if self.common_grid is False: grid_vars = grouping_vars else: self._check_var_list_or_boolean("common_grid", grouping_vars) grid_vars = [v for v in self.common_grid if v in grouping_vars] res = ( GroupBy(grid_vars) .apply(data, self._transform, orient, grouping_vars) ) # Normalize, potentially within groups if not grouping_vars or self.common_norm is True: res = res.assign(group_weight=data["weight"].sum()) else: if self.common_norm is False: norm_vars = grouping_vars else: self._check_var_list_or_boolean("common_norm", grouping_vars) norm_vars = [v for v in self.common_norm if v in grouping_vars] res = res.join( data.groupby(norm_vars)["weight"].sum().rename("group_weight"), on=norm_vars, ) res["density"] *= res.eval("weight / group_weight") value = {"x": "y", "y": "x"}[orient] res[value] = res["density"] return res.drop(["weight", "group_weight"], axis=1)