Source code for plotnine.geoms.geom_step

from __future__ import annotations

import typing

import numpy as np
import pandas as pd

from ..doctools import document
from ..exceptions import PlotnineError
from ..utils import copy_missing_columns
from .geom import geom
from .geom_path import geom_path

if typing.TYPE_CHECKING:
    from typing import Any

    from plotnine.iapi import panel_view
    from plotnine.typing import Axes, Coord


[docs]@document class geom_step(geom_path): """ Stepped connected points {usage} Parameters ---------- {common_parameters} direction : str, optional (default: hv) One of *hv*, *vh* or *mid*, for horizontal-vertical steps, vertical-horizontal steps or steps half-way between adjacent x values. See Also -------- plotnine.geoms.geom_path : For documentation of extra parameters. """ DEFAULT_PARAMS = { "stat": "identity", "position": "identity", "na_rm": False, "direction": "hv", } draw_panel = geom.draw_panel @staticmethod def draw_group( data: pd.DataFrame, panel_params: panel_view, coord: Coord, ax: Axes, **params: Any, ): direction = params["direction"] n = len(data) data = data.sort_values("x", kind="mergesort") x = data["x"].to_numpy() y = data["y"].to_numpy() if direction == "vh": # create stepped path -- interleave x with # itself and y with itself xidx = np.repeat(range(n), 2)[:-1] yidx = np.repeat(range(n), 2)[1:] new_x, new_y = x[xidx], y[yidx] elif direction == "hv": xidx = np.repeat(range(n), 2)[1:] yidx = np.repeat(range(n), 2)[:-1] new_x, new_y = x[xidx], y[yidx] elif direction == "mid": xidx = np.repeat(range(n - 1), 2) yidx = np.repeat(range(n), 2) diff = x[1::] - x[:-1:] mid_x = x[:-1:] + diff / 2 new_x = np.hstack([x[0], mid_x[xidx], x[-1]]) new_y = y[yidx] else: raise PlotnineError(f"Invalid direction `{direction}`") path_data = pd.DataFrame({"x": new_x, "y": new_y}) copy_missing_columns(path_data, data) geom_path.draw_group(path_data, panel_params, coord, ax, **params)