from __future__ import annotations
import typing
from ..doctools import document
from ..utils import to_rgba
from .geom import geom
from .geom_line import geom_line, geom_path
from .geom_ribbon import geom_ribbon
if typing.TYPE_CHECKING:
from typing import Any
import pandas as pd
from plotnine.iapi import panel_view
from plotnine.typing import Axes, Coord, DrawingArea, Layer
[docs]@document
class geom_smooth(geom):
"""
A smoothed conditional mean
{usage}
Parameters
----------
{common_parameters}
legend_fill_ratio : float (default: 0.5)
How much (vertically) of the legend box should be filled by
the color that indicates the confidence intervals. Should be
in the range [0, 1].
"""
DEFAULT_AES = {
"alpha": 0.4,
"color": "black",
"fill": "#999999",
"linetype": "solid",
"size": 1,
"ymin": None,
"ymax": None,
}
REQUIRED_AES = {"x", "y"}
DEFAULT_PARAMS = {
"stat": "smooth",
"position": "identity",
"na_rm": False,
"legend_fill_ratio": 0.5,
}
def use_defaults(
self, data: pd.DataFrame, aes_modifiers: dict[str, Any]
) -> pd.DataFrame:
has_ribbon = "ymin" in data and "ymax" in data
data = super().use_defaults(data, aes_modifiers)
# When there is no ribbon, the default values for 'ymin'
# and 'ymax' are None (not numeric). So we remove them
# prevent any computations that may use them without checking.
if not has_ribbon:
del data["ymin"]
del data["ymax"]
return data
def setup_data(self, data: pd.DataFrame) -> pd.DataFrame:
return data.sort_values(["PANEL", "group", "x"])
@staticmethod
def draw_group(
data: pd.DataFrame,
panel_params: panel_view,
coord: Coord,
ax: Axes,
**params: Any,
):
has_ribbon = "ymin" in data and "ymax" in data
if has_ribbon:
data2 = data.copy()
data2["color"] = "none"
params2 = params.copy()
params2["outline_type"] = "full"
geom_ribbon.draw_group(data2, panel_params, coord, ax, **params2)
data["alpha"] = 1
geom_line.draw_group(data, panel_params, coord, ax, **params)
@staticmethod
def draw_legend(
data: pd.Series[Any], da: DrawingArea, lyr: Layer
) -> DrawingArea:
"""
Draw letter 'a' in the box
Parameters
----------
data : dataframe
Legend data
da : DrawingArea
Canvas
lyr : layer
Layer
Returns
-------
out : DrawingArea
"""
from matplotlib.patches import Rectangle
try:
has_se = lyr.stat.params["se"]
except KeyError:
has_se = False
if has_se:
fill = to_rgba(data["fill"], data["alpha"])
r = lyr.geom.params["legend_fill_ratio"]
bg = Rectangle(
(0, (1 - r) * da.height / 2),
width=da.width,
height=r * da.height,
facecolor=fill,
linewidth=0,
)
da.add_artist(bg)
data["alpha"] = 1
return geom_path.draw_legend(data, da, lyr)