Source code for plotnine.stats.stat_sum

from ..doctools import document
from ..mapping.aes import ALL_AESTHETICS
from ..mapping.evaluation import after_stat
from ..utils import groupby_apply
from .stat import stat


[docs]@document class stat_sum(stat): """ Sum unique values Useful for overplotting on scatterplots. {usage} Parameters ---------- {common_parameters} """ _aesthetics_doc = """ {aesthetics_table} .. rubric:: Options for computed aesthetics :: 'n' # Number of observations at a position 'prop' # Ratio of points in that panel at a position """ REQUIRED_AES = {"x", "y"} DEFAULT_PARAMS = {"geom": "point", "position": "identity", "na_rm": False} DEFAULT_AES = {"size": after_stat("n"), "weight": 1} CREATES = {"n", "prop"} @classmethod def compute_panel(cls, data, scales, **params): if "weight" not in data: data["weight"] = 1 def count(df): """ Do a weighted count """ df["n"] = df["weight"].sum() return df.iloc[0:1] def ave(df): """ Calculate proportion values """ df["prop"] = df["n"] / df["n"].sum() return df # group by all present aesthetics other than the weight, # then sum them (i.e no. of uniques) to get the raw count # 'n', and the proportions 'prop' per group s: set[str] = set(data.columns) & ALL_AESTHETICS by = list(s.difference(["weight"])) counts = groupby_apply(data, by, count) counts = groupby_apply(counts, "group", ave) return counts