Source code for plotnine.stats.stat_quantile

from warnings import warn

import pandas as pd

from ..doctools import document
from ..exceptions import PlotnineWarning
from .stat import stat


# method_args are any of the keyword args (other than q) for
# statsmodels.regression.quantile_regression.QuantReg.fit
[docs]@document class stat_quantile(stat): """ Compute quantile regression lines {usage} Parameters ---------- {common_parameters} quantiles : tuple, optional (default: (0.25, 0.5, 0.75)) Quantiles of y to compute formula : str, optional (default: 'y ~ x') Formula relating y variables to x variables method_args : dict, optional Extra arguments passed on to the model fitting method, :meth:`statsmodels.regression.quantile_regression.QuantReg.fit`. See Also -------- statsmodels.regression.quantile_regression.QuantReg plotnine.geoms.geom_quantile """ _aesthetics_doc = """ {aesthetics_table} .. rubric:: Options for computed aesthetics :: 'quantile' # quantile 'group' # group identifier Calculated aesthetics are accessed using the `after_stat` function. e.g. :py:`after_stat('quantile')`. """ REQUIRED_AES = {"x", "y"} DEFAULT_PARAMS = { "geom": "quantile", "position": "identity", "na_rm": False, "quantiles": (0.25, 0.5, 0.75), "formula": "y ~ x", "method_args": {}, } CREATES = {"quantile", "group"} def setup_params(self, data): params = self.params.copy() if params["formula"] is None: params["formula"] = "y ~ x" warn("Formula not specified, using '{}'", PlotnineWarning) try: iter(params["quantiles"]) except TypeError: params["quantiles"] = (params["quantiles"],) return params @classmethod def compute_group(cls, data, scales, **params): res = [quant_pred(q, data, **params) for q in params["quantiles"]] return pd.concat(res, axis=0, ignore_index=True)
def quant_pred(q, data, **params): """ Quantile precitions """ import statsmodels.formula.api as smf mod = smf.quantreg(params["formula"], data) reg_res = mod.fit(q=q, **params["method_args"]) out = pd.DataFrame( { "x": [data["x"].min(), data["x"].max()], "quantile": q, "group": "{}-{}".format(data["group"].iloc[0], q), } ) out["y"] = reg_res.predict(out) return out