diff --git a/doc/.gitignore b/doc/.gitignore index 4e48027747..19459cf67a 100644 --- a/doc/.gitignore +++ b/doc/.gitignore @@ -13,3 +13,5 @@ objects.json objects.txt objects.inv gallery/thumbnails + +**/*.quarto_ipynb diff --git a/doc/_quartodoc.yml b/doc/_quartodoc.yml index 6414031a5d..59e8218070 100644 --- a/doc/_quartodoc.yml +++ b/doc/_quartodoc.yml @@ -535,6 +535,8 @@ quartodoc: - coord_equal - coord_fixed - coord_flip + - coord_polar + - coord_radial - coord_trans - title: Composing Plots diff --git a/plotnine/__init__.py b/plotnine/__init__.py index 69526927bd..a4b9bcce4c 100644 --- a/plotnine/__init__.py +++ b/plotnine/__init__.py @@ -20,6 +20,8 @@ coord_equal, coord_fixed, coord_flip, + coord_polar, + coord_radial, coord_trans, ) from .facets import ( @@ -289,6 +291,8 @@ "coord_equal", "coord_fixed", "coord_flip", + "coord_polar", + "coord_radial", "coord_trans", "element_blank", "element_line", diff --git a/plotnine/coords/__init__.py b/plotnine/coords/__init__.py index d49f9f8ff2..ae14967bef 100644 --- a/plotnine/coords/__init__.py +++ b/plotnine/coords/__init__.py @@ -5,6 +5,8 @@ from .coord_cartesian import coord_cartesian from .coord_fixed import coord_equal, coord_fixed from .coord_flip import coord_flip +from .coord_polar import coord_polar +from .coord_radial import coord_radial from .coord_trans import coord_trans __all__ = ( @@ -12,5 +14,7 @@ "coord_fixed", "coord_equal", "coord_flip", + "coord_polar", + "coord_radial", "coord_trans", ) diff --git a/plotnine/coords/coord.py b/plotnine/coords/coord.py index 8dcbc378fa..300c826772 100644 --- a/plotnine/coords/coord.py +++ b/plotnine/coords/coord.py @@ -104,6 +104,22 @@ def aspect(self, panel_params: panel_view) -> float | None: """ return None + def draw(self, axs: list) -> None: + """ + Draw coordinate-system decorations onto each panel axes. + + Called after all layers are drawn. Subclasses override this to + add elements such as polar grid lines. + """ + + def post_setup_ax(self, ax: Any) -> None: + """ + Hook called for each axes after set_limits_breaks_and_labels. + + Override in subclasses to apply per-axes settings that must + run after the facet has set tick positions and label padding. + """ + def labels(self, cur_labels: labels_view) -> labels_view: """ Modify labels diff --git a/plotnine/coords/coord_polar.py b/plotnine/coords/coord_polar.py new file mode 100644 index 0000000000..d9519b8f6f --- /dev/null +++ b/plotnine/coords/coord_polar.py @@ -0,0 +1,246 @@ +from __future__ import annotations + +from dataclasses import replace +from typing import TYPE_CHECKING, cast + +import numpy as np + +from ..iapi import panel_ranges +from .coord import coord, dist_euclidean + +if TYPE_CHECKING: + import pandas as pd + from matplotlib.axes import Axes + from matplotlib.projections.polar import PolarAxes + + from plotnine.iapi import panel_view + from plotnine.scales.scale import scale + + +class coord_polar(coord): + """ + Polar coordinate system + + `coord_polar` maps one position aesthetic to the angle and the other + to the radius. It is commonly used for pie charts, which are stacked + bar charts in polar coordinates. + + Parameters + ---------- + theta : + Which variable maps to the angle axis, ``"x"`` (default) or ``"y"``. + start : + Starting angle in radians, measured clockwise from 12 o'clock + (i.e. from the positive-y axis). Default 0. + direction : + ``1`` = clockwise (default), ``-1`` = counter-clockwise. + expand : + Add a small buffer around the data on the radius axis. + Default ``True``. + + Notes + ----- + Unlike ggplot2, plotnine coordinate systems do not currently expose a + ``clip`` argument. + + For partial arcs, donut charts, and theta/radius limits, use + ``coord_radial``. + + Examples + -------- + A pie chart is a stacked bar chart with the y position mapped to angle. + + ```python + import pandas as pd + from plotnine import aes, coord_polar, geom_col, ggplot + + df = pd.DataFrame({ + "x": [1, 1, 1], + "y": [2, 3, 5], + "group": ["a", "b", "c"], + }) + + ggplot(df, aes("x", "y", fill="group")) + geom_col() + coord_polar("y") + ``` + """ + + is_linear = False + + def __init__( + self, + theta: str = "x", + start: float = 0, + direction: int = 1, + expand: bool = True, + ) -> None: + self.theta = theta + self.start = start + self.direction = direction + self.expand = expand + self.params: dict = {} + + # ------------------------------------------------------------------ + # Panel params + # ------------------------------------------------------------------ + + def setup_panel_params(self, scale_x: scale, scale_y: scale) -> panel_view: + from .coord_cartesian import coord_cartesian + + # Theta fills exactly one full revolution — no expansion on that axis. + # R uses the caller-controlled expand flag. + pv_no_exp = coord_cartesian(expand=False).setup_panel_params( + scale_x, scale_y + ) + pv_exp = coord_cartesian(expand=self.expand).setup_panel_params( + scale_x, scale_y + ) + + if self.theta == "x": + theta_range = pv_no_exp.x.range + r_sv = pv_exp.y + else: + theta_range = pv_no_exp.y.range + r_sv = pv_exp.x + + self.params["theta_range"] = theta_range + self.params["r_range"] = r_sv.range + + empty = np.array([], dtype=float) + + # x → theta axis: data ticks are in original units (not radians), so + # suppress them. Limits span [start, start+2π] so that bars rotated + # by a non-zero start angle stay within the displayed theta range. + theta_start = float(self.start) + new_x = replace( + pv_exp.x, + limits=(theta_start, theta_start + 2 * np.pi), + range=(theta_start, theta_start + 2 * np.pi), + breaks=[], + minor_breaks=empty, + labels=[], + ) + + # y → r axis: use the scale for the r dimension with its natural + # breaks. + new_y = replace(r_sv) + + return replace(pv_exp, x=new_x, y=new_y) + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + def _to_radians(self, vals: np.ndarray) -> np.ndarray: + """Normalise data-space theta values to [start, start + 2π].""" + t_min, t_max = self.params["theta_range"] + denom = float(t_max) - float(t_min) + if denom == 0: + return np.zeros_like(vals, dtype=float) + norm = (np.asarray(vals, dtype=float) - float(t_min)) / denom + return self.start + self.direction * norm * 2.0 * np.pi + + # ------------------------------------------------------------------ + # Data transformation + # ------------------------------------------------------------------ + + def transform( + self, + data: pd.DataFrame, + panel_params: panel_view, + munch: bool = False, + ) -> pd.DataFrame: + # Munch first (in original data space) so curved edges get enough + # interpolation points before we convert theta → radians. + if munch: + data = self.munch(data, panel_params) + + if self.theta == "x": + theta_col, r_col = "x", "y" + theta_end_col, r_end_col = "xend", "yend" + else: + theta_col, r_col = "y", "x" + theta_end_col, r_end_col = "yend", "xend" + + if theta_col not in data.columns or r_col not in data.columns: + return data + + data = data.copy() + data[theta_col] = self._to_radians(data[theta_col].to_numpy()) + has_endpoints = ( + theta_end_col in data.columns and r_end_col in data.columns + ) + if has_endpoints: + data[theta_end_col] = self._to_radians( + data[theta_end_col].to_numpy() + ) + + # PolarAxes always expects x = theta (radians) and y = r. + # When theta = "y" we need to swap the columns. + if self.theta == "y": + data["x"], data["y"] = data["y"].copy(), data["x"].copy() + if has_endpoints: + data["xend"], data["yend"] = ( + data["yend"].copy(), + data["xend"].copy(), + ) + + return data + + # ------------------------------------------------------------------ + # Distance (used by munch, called before transform) + # ------------------------------------------------------------------ + + def distance( + self, + x: pd.Series, + y: pd.Series, + panel_params: panel_view, + ) -> np.ndarray: + # Normalise theta and r to [0, 1] then compute Euclidean distance. + t_min, t_max = self.params["theta_range"] + r_min, r_max = self.params["r_range"] + t_denom = float(t_max - t_min) or 1.0 + r_denom = float(r_max - r_min) or 1.0 + + if self.theta == "x": + theta_vals = np.asarray(x, dtype=float) + r_vals = np.asarray(y, dtype=float) + else: + theta_vals = np.asarray(y, dtype=float) + r_vals = np.asarray(x, dtype=float) + + theta_norm = (theta_vals - float(t_min)) / t_denom + r_norm = (r_vals - float(r_min)) / r_denom + return dist_euclidean(theta_norm, r_norm) + + def backtransform_range(self, panel_params: panel_view) -> panel_ranges: + t_range = tuple(self.params["theta_range"]) + r_range = tuple(self.params["r_range"]) + if self.theta == "x": + return panel_ranges(x=t_range, y=r_range) + return panel_ranges(x=r_range, y=t_range) + + # ------------------------------------------------------------------ + # Draw decorations on PolarAxes + # ------------------------------------------------------------------ + + def draw(self, axs: list[Axes]) -> None: + """Configure each PolarAxes: zero location, direction, r limits.""" + r_min, r_max = self.params.get("r_range", (0.0, 1.0)) + + # Matplotlib PolarAxes theta_direction: -1 = clockwise, 1 = counter-CW. + mpl_direction = -1 if self.direction == 1 else 1 + + for ax in axs: + polar_ax = cast("PolarAxes", ax) + polar_ax.set_theta_zero_location("N") # 12 o'clock = 0 + polar_ax.set_theta_direction(mpl_direction) + if np.isfinite(r_min) and np.isfinite(r_max) and r_min != r_max: + polar_ax.set_rlim(float(r_min), float(r_max)) + + # ------------------------------------------------------------------ + # Misc + # ------------------------------------------------------------------ + + def aspect(self, panel_params: panel_view) -> float: + return 1.0 diff --git a/plotnine/coords/coord_radial.py b/plotnine/coords/coord_radial.py new file mode 100644 index 0000000000..302d4624c8 --- /dev/null +++ b/plotnine/coords/coord_radial.py @@ -0,0 +1,336 @@ +from __future__ import annotations + +from dataclasses import replace +from typing import TYPE_CHECKING, Sequence, cast + +import numpy as np + +from .coord_polar import coord_polar + +if TYPE_CHECKING: + import pandas as pd + from matplotlib.axes import Axes + from matplotlib.projections.polar import PolarAxes + + from plotnine.iapi import panel_view + from plotnine.scales.scale import scale + + +class coord_radial(coord_polar): + """ + Radial coordinate system + + `coord_radial` maps one position aesthetic to the angle and the other + to the radius. Compared with ``coord_polar``, it adds support for + partial arcs, inner radius holes, theta/radius limits, radial-axis + placement, and rotation of the ``angle`` aesthetic. + + Parameters + ---------- + theta : + Which variable maps to the angle axis, ``"x"`` (default) or ``"y"``. + start : + Starting angle in radians, measured clockwise from 12 o'clock. + Default 0. + end : + Ending angle in radians, measured clockwise from 12 o'clock. + ``None`` (default) gives a full circle (``start + 2π * direction``). + direction : + ``1`` = clockwise (default), ``-1`` = counter-clockwise. + Only used when *end* is ``None``. + expand : + Add a small buffer around the data on the radius axis. + Default ``True``. + inner_radius : + Size of the inner hole as a fraction of the outer radius, in + ``[0, 1)``. ``0`` (default) means no hole; ``0.3`` creates a 30 % + donut hole, useful for gauge and donut charts. + r_axis_inside : + Where to place the radial (r) axis tick labels. + + * ``None`` (default) — let Matplotlib decide (usually outside). + * ``True`` — force inside, aligned just past the *start* angle. + * ``False`` — force outside (Matplotlib default). + * *float* — place at this theta angle in radians (clockwise from + North). + + Unlike ggplot2's ``r.axis.inside``, a length-2 value for separate + primary and secondary axis placement is not supported. + rotate_angle : + If ``True``, automatically add the local theta angle (in degrees) to + the ``angle`` aesthetic so that text or other rotated marks align with + the spoke direction. Default ``False``. + thetalim : + Data-space limits for the theta axis as ``(lo, hi)``. Only data + within this range is mapped to the arc; equivalent to zooming on the + angular axis. ``None`` (default) uses the full data range. + rlim : + Data-space limits for the r axis as ``(lo, hi)``. Only data within + this range is shown; equivalent to zooming on the radial axis. + ``None`` (default) uses the full data range. + theta_labels : + If ``True``, show theta axis tick labels on the outer edge of the + circle for full-circle plots, using the breaks and labels from the + theta scale. Default ``False``. Partial-arc plots always show + theta labels (filtered to the visible arc) regardless of this flag. + theta_label_pad : + Distance in points between the outer circle spine and the theta tick + labels. Default ``8``. Only applied when theta labels are shown. + + Notes + ----- + The Python API uses snake_case names for arguments that are dotted in + ggplot2: ``inner_radius``, ``r_axis_inside``, and ``rotate_angle``. + + Unlike ggplot2, plotnine coordinate systems do not currently expose a + ``clip`` argument. The ggplot2 ``reverse`` argument is not currently + implemented. + + Examples + -------- + A donut chart is a stacked bar chart with an inner radius. + + ```python + import pandas as pd + from plotnine import aes, coord_radial, geom_col, ggplot + + df = pd.DataFrame({ + "x": [1, 1, 1], + "y": [2, 3, 5], + "group": ["a", "b", "c"], + }) + + ( + ggplot(df, aes("x", "y", fill="group")) + + geom_col() + + coord_radial(theta="y", inner_radius=0.4) + ) + ``` + + Partial arcs can be used for gauge-like displays. + + ```python + import numpy as np + import pandas as pd + from plotnine import aes, coord_radial, geom_point, ggplot + + df = pd.DataFrame({"x": [1, 2, 3], "y": [2, 4, 3]}) + + ggplot(df, aes("x", "y")) + geom_point() + coord_radial( + start=-0.4 * np.pi, + end=0.4 * np.pi, + inner_radius=0.3, + ) + ``` + """ + + def __init__( + self, + theta: str = "x", + start: float = 0, + end: float | None = None, + direction: int = 1, + expand: bool = True, + inner_radius: float = 0, + r_axis_inside: bool | float | None = None, + rotate_angle: bool = False, + thetalim: tuple[float, float] | None = None, + rlim: tuple[float, float] | None = None, + theta_labels: bool = False, + theta_label_pad: float = 8, + ) -> None: + super().__init__( + theta=theta, + start=start, + direction=direction, + expand=expand, + ) + self.end = end + self.inner_radius = inner_radius + self.r_axis_inside = r_axis_inside + self.rotate_angle = rotate_angle + self.thetalim = thetalim + self.rlim = rlim + self.theta_labels = theta_labels + self.theta_label_pad = theta_label_pad + + # ------------------------------------------------------------------ + # Panel params + # ------------------------------------------------------------------ + + def setup_panel_params(self, scale_x: scale, scale_y: scale) -> panel_view: + from .coord_cartesian import coord_cartesian + + # Capture data-space theta breaks before super() clears them. + pv_data = coord_cartesian(expand=False).setup_panel_params( + scale_x, scale_y + ) + if self.theta == "x": + theta_breaks = list(pv_data.x.breaks) + theta_labels = list(pv_data.x.labels) + else: + theta_breaks = list(pv_data.y.breaks) + theta_labels = list(pv_data.y.labels) + + pv = super().setup_panel_params(scale_x, scale_y) + + # thetalim: zoom the theta data range — only this slice maps to the + # arc. + if self.thetalim is not None: + self.params["theta_range"] = tuple(self.thetalim) + + # rlim: zoom the r data range — update params, panel view y axis, and + # filter breaks/labels to within rlim so set_yticks doesn't force the + # PolarAxes r-axis to expand beyond the requested limits. + if self.rlim is not None: + self.params["r_range"] = tuple(self.rlim) + rlo, rhi = self.rlim + breaks = cast("Sequence[float]", pv.y.breaks) + labels = pv.y.labels + mask = [rlo <= b <= rhi for b in breaks] + new_y = replace( + pv.y, + limits=tuple(self.rlim), + range=tuple(self.rlim), + breaks=[b for b, m in zip(breaks, mask) if m], + labels=[l for l, m in zip(labels, mask) if m], + ) + pv = replace(pv, y=new_y) + + # Compute arc bounds for partial-arc plots (None means full circle). + arc_lo = arc_hi = None + if self.end is not None: + arc = self._arc + arc_lo = min(self.start, self.start + arc) + arc_hi = max(self.start, self.start + arc) + + # Convert data-space theta breaks to radian positions and restore them + # as theta axis tick labels on the outer edge. Always done for partial + # arcs; for full circles only when theta_labels=True (opt-in, so that + # pac-man / coxcomb charts keep breaks=[] as set by super()). + x_updates: dict = {} + if theta_breaks and (arc_lo is not None or self.theta_labels): + radian_pos = list( + self._to_radians(np.asarray(theta_breaks, dtype=float)) + ) + if arc_lo is not None: + keep = [arc_lo <= r <= arc_hi for r in radian_pos] + radian_pos = [r for r, k in zip(radian_pos, keep) if k] + theta_labels = [l for l, k in zip(theta_labels, keep) if k] + x_updates["breaks"] = radian_pos + x_updates["labels"] = theta_labels + + # Partial arc: x panel range must match [arc_lo, arc_hi] so that + # set_limits_breaks_and_labels calls ax.set_xlim(arc_lo, arc_hi) rather + # than ax.set_xlim(0, 2π), which would override set_thetalim. + if arc_lo is not None: + x_updates["limits"] = (arc_lo, arc_hi) + x_updates["range"] = (arc_lo, arc_hi) + + if x_updates: + pv = replace(pv, x=replace(pv.x, **x_updates)) + + return pv + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + @property + def _arc(self) -> float: + """ + Total arc in radians. + + A positive value represents clockwise movement when ``direction=1``. + """ + if self.end is not None: + return self.end - self.start + return self.direction * 2.0 * np.pi + + def _to_radians(self, vals: np.ndarray) -> np.ndarray: + """Normalize theta values to [start, start + arc].""" + t_min, t_max = self.params["theta_range"] + denom = float(t_max) - float(t_min) + if denom == 0: + return np.zeros_like(vals, dtype=float) + norm = (np.asarray(vals, dtype=float) - float(t_min)) / denom + return self.start + norm * self._arc + + # ------------------------------------------------------------------ + # Data transformation + # ------------------------------------------------------------------ + + def transform( + self, + data: pd.DataFrame, + panel_params: panel_view, + munch: bool = False, + ) -> pd.DataFrame: + data = super().transform(data, panel_params, munch=munch) + # After super().transform(), data["x"] is always theta in radians. + if ( + self.rotate_angle + and "angle" in data.columns + and "x" in data.columns + ): + data = data.copy() + data["angle"] = data["angle"] + np.degrees(data["x"]) + return data + + # ------------------------------------------------------------------ + # Draw decorations on PolarAxes + # ------------------------------------------------------------------ + + def draw(self, axs: list[Axes]) -> None: + """Configure PolarAxes: arc limits, inner radius, axis placement.""" + super().draw(axs) + + r_min, r_max = self.params.get("r_range", (0.0, 1.0)) + arc = self._arc + + for ax in axs: + polar_ax = cast("PolarAxes", ax) + # Restrict visible theta range for partial arcs. + if self.end is not None: + theta_lo = min(self.start, self.start + arc) + theta_hi = max(self.start, self.start + arc) + polar_ax.set_thetalim(theta_lo, theta_hi) + + # Inner radius: push the data away from the centre by setting a + # virtual r-origin below r_min. Formula: solve + # inner_radius = (r_min - r_origin) / (r_max - r_origin) + if ( + self.inner_radius > 0 + and np.isfinite(r_min) + and np.isfinite(r_max) + and r_max > r_min + and self.inner_radius < 1.0 + ): + r_origin = (r_min - self.inner_radius * r_max) / ( + 1.0 - self.inner_radius + ) + polar_ax.set_rorigin(r_origin) + + # Radial axis label placement. + if self.r_axis_inside is not None: + if isinstance(self.r_axis_inside, bool): + if self.r_axis_inside: + # Just inside the start angle keeps it out of the data. + polar_ax.set_rlabel_position( + np.degrees(self.start) + 10 + ) + else: + polar_ax.set_rlabel_position( + np.degrees(float(self.r_axis_inside)) + ) + + def post_setup_ax(self, ax: Axes) -> None: + """ + Apply theta label pad after facet has set tick positions and padding. + """ + if self.theta_labels or self.end is not None: + ax.tick_params(axis="x", pad=self.theta_label_pad) + # Allow geom_text labels to extend past the polar axes bounding box + # (e.g. spoke labels placed just beyond the outermost bar tip). + for text in ax.texts: + text.set_clip_on(False) diff --git a/plotnine/facets/facet.py b/plotnine/facets/facet.py index ae784d8768..2904ae80d7 100644 --- a/plotnine/facets/facet.py +++ b/plotnine/facets/facet.py @@ -355,6 +355,7 @@ def _inf_to_none( ax.tick_params(axis="x", which="major", pad=pad_x) ax.tick_params(axis="y", which="major", pad=pad_y) + self.coordinates.post_setup_ax(ax) def __deepcopy__(self, memo: dict[Any, Any]) -> facet: """ @@ -397,9 +398,16 @@ def _make_axes(self) -> tuple[p9GridSpec, list[Axes]]: gs = self._make_gridspec() # Create axes + from ..coords.coord_polar import coord_polar + + subplot_kw = ( + {"projection": "polar"} + if isinstance(self.plot.coordinates, coord_polar) + else {} + ) it = itertools.product(range(self.nrow), range(self.ncol)) for i, (row, col) in enumerate(it): - axsarr[row, col] = self.figure.add_subplot(gs[i]) + axsarr[row, col] = self.figure.add_subplot(gs[i], **subplot_kw) # Rearrange axes # They are ordered to match the positions in the layout table diff --git a/plotnine/ggplot.py b/plotnine/ggplot.py index 08a00ca8e2..84ccb5e602 100755 --- a/plotnine/ggplot.py +++ b/plotnine/ggplot.py @@ -485,6 +485,11 @@ def _draw_panel_borders(self): # grid lines below the borders. We leave ax.patch for the # background only. if self.theme.T.is_blank("panel_border"): + # For PolarAxes the default circular spine is separate from the + # panel border Rectangle; hide it explicitly when blank. + for ax in self.axs: + if "polar" in ax.spines: + ax.spines["polar"].set_visible(False) return from matplotlib.patches import Rectangle @@ -511,6 +516,7 @@ def _draw_layers(self): """ # Draw the geoms self.layers.draw(self.layout, self.coordinates) + self.coordinates.draw(self.axs) def _draw_breaks_and_labels(self): """ diff --git a/plotnine/themes/themeable.py b/plotnine/themes/themeable.py index f3e8f0f806..656cdbe1bd 100644 --- a/plotnine/themes/themeable.py +++ b/plotnine/themes/themeable.py @@ -978,7 +978,7 @@ def apply_ax(self, ax: Axes): vinstalled = version.parse(mpl.__version__) v310 = version.parse("3.10.0") name = "labelbottom" if vinstalled >= v310 else "labelleft" - if not ax.xaxis.get_tick_params()[name]: + if not ax.xaxis.get_tick_params().get(name, True): return # if not ax.xaxis.get_tick_params()["labelbottom"]: @@ -1108,6 +1108,9 @@ class axis_line_x(themeable): def apply_ax(self, ax: Axes): super().apply_ax(ax) + # PolarAxes has no "top"/"bottom" spines — skip silently. + if "top" not in ax.spines: + return properties = self._get_properties(omit=("solid_capstyle",)) # MPL has a default zorder of 2.5 for spines # so layers 3+ would be drawn on top of the spines @@ -1118,6 +1121,8 @@ def apply_ax(self, ax: Axes): def blank_ax(self, ax: Axes): super().blank_ax(ax) + if "top" not in ax.spines: + return ax.spines["top"].set_visible(False) ax.spines["bottom"].set_visible(False) @@ -1135,6 +1140,9 @@ class axis_line_y(themeable): def apply_ax(self, ax: Axes): super().apply_ax(ax) + # PolarAxes has no "left"/"right" spines — skip silently. + if "left" not in ax.spines: + return properties = self._get_properties(omit=("solid_capstyle",)) # MPL has a default zorder of 2.5 for spines # so layers 3+ would be drawn on top of the spines @@ -1145,6 +1153,8 @@ def apply_ax(self, ax: Axes): def blank_ax(self, ax: Axes): super().blank_ax(ax) + if "left" not in ax.spines: + return ax.spines["left"].set_visible(False) ax.spines["right"].set_visible(False) diff --git a/tests/test_coord_polar.py b/tests/test_coord_polar.py new file mode 100644 index 0000000000..fd1d169625 --- /dev/null +++ b/tests/test_coord_polar.py @@ -0,0 +1,287 @@ +import numpy as np +import pandas as pd +from matplotlib import pyplot as plt +from numpy.testing import assert_allclose + +from plotnine import ( + aes, + coord_radial, + element_blank, + element_line, + geom_col, + ggplot, + theme, +) +from plotnine.coords.coord_polar import coord_polar +from plotnine.scales import scale_x_continuous, scale_y_continuous + + +def trained_scales( + x=(0, 10), + y=(0, 10), + x_breaks=(0, 5, 10), + y_breaks=(0, 5, 10), + x_labels=("0", "5", "10"), + y_labels=("0", "5", "10"), +): + scale_x = scale_x_continuous(breaks=x_breaks, labels=x_labels) + scale_y = scale_y_continuous(breaks=y_breaks, labels=y_labels) + scale_x.train(x) + scale_y.train(y) + return scale_x, scale_y + + +def test_coord_polar_setup_panel_params_theta_x(): + scale_x, scale_y = trained_scales( + y_breaks=(0, 2, 5, 10), + y_labels=("0", "2", "5", "10"), + ) + coord = coord_polar(theta="x", start=np.pi / 4, expand=False) + + panel_params = coord.setup_panel_params(scale_x, scale_y) + + assert coord.params["theta_range"] == (0, 10) + assert coord.params["r_range"] == (0, 10) + assert panel_params.x.range == (np.pi / 4, np.pi / 4 + 2 * np.pi) + assert panel_params.x.breaks == [] + assert panel_params.x.labels == [] + assert panel_params.y.breaks == [0, 2, 5, 10] + + +def test_coord_polar_setup_panel_params_theta_y(): + scale_x, scale_y = trained_scales( + x_breaks=(0, 2, 5, 10), + x_labels=("0", "2", "5", "10"), + ) + coord = coord_polar(theta="y", expand=False) + + panel_params = coord.setup_panel_params(scale_x, scale_y) + + assert coord.params["theta_range"] == (0, 10) + assert coord.params["r_range"] == (0, 10) + assert panel_params.y.breaks == [0, 2, 5, 10] + + +def test_coord_polar_to_radians_zero_width_range(): + coord = coord_polar() + coord.params = {"theta_range": (1, 1)} + + assert_allclose(coord._to_radians(np.array([1, 2, 3])), [0, 0, 0]) + + +def test_coord_polar_transforms_segment_endpoints_theta_x(): + coord = coord_polar(theta="x") + coord.params = {"theta_range": (0, 10), "r_range": (0, 10)} + data = pd.DataFrame({"x": [0], "y": [1], "xend": [10], "yend": [2]}) + + out = coord.transform(data, None) + + assert out.loc[0, "x"] == 0 + assert out.loc[0, "y"] == 1 + assert np.isclose(out.loc[0, "xend"], 2 * np.pi) + assert out.loc[0, "yend"] == 2 + + +def test_coord_polar_transforms_segment_endpoints_theta_y(): + coord = coord_polar(theta="y") + coord.params = {"theta_range": (0, 10), "r_range": (0, 10)} + data = pd.DataFrame({"x": [1], "y": [0], "xend": [2], "yend": [10]}) + + out = coord.transform(data, None) + + assert out.loc[0, "x"] == 0 + assert out.loc[0, "y"] == 1 + assert np.isclose(out.loc[0, "xend"], 2 * np.pi) + assert out.loc[0, "yend"] == 2 + + +def test_coord_polar_transforms_theta_y_without_endpoints(): + coord = coord_polar(theta="y") + coord.params = {"theta_range": (0, 10), "r_range": (0, 10)} + data = pd.DataFrame({"x": [1], "y": [5]}) + + out = coord.transform(data, None) + + assert_allclose(out.loc[0, "x"], np.pi) + assert out.loc[0, "y"] == 1 + + +def test_coord_polar_munches_before_radian_transform(): + coord = coord_polar() + coord.params = {"theta_range": (0, 10), "r_range": (0, 10)} + data = pd.DataFrame({"x": [0, 10], "y": [1, 2], "group": [1, 1]}) + + out = coord.transform(data, None, munch=True) + + assert len(out) > len(data) + assert out["x"].between(0, 2 * np.pi).all() + + +def test_coord_polar_leaves_non_position_data_unchanged(): + coord = coord_polar() + data = pd.DataFrame({"label": ["A"]}) + + assert coord.transform(data, None) is data + + +def test_coord_polar_distance_and_backtransform_theta_x(): + coord = coord_polar() + coord.params = {"theta_range": (0, 10), "r_range": (0, 20)} + + distance = coord.distance(pd.Series([0, 10]), pd.Series([0, 10]), None) + + assert_allclose(distance, [np.sqrt(1.25)]) + assert coord.backtransform_range(None).x == (0, 10) + assert coord.backtransform_range(None).y == (0, 20) + + +def test_coord_polar_distance_and_backtransform_theta_y(): + coord = coord_polar(theta="y") + coord.params = {"theta_range": (0, 10), "r_range": (0, 20)} + + distance = coord.distance(pd.Series([0, 10]), pd.Series([0, 10]), None) + + assert_allclose(distance, [np.sqrt(1.25)]) + assert coord.backtransform_range(None).x == (0, 20) + assert coord.backtransform_range(None).y == (0, 10) + + +def test_coord_polar_draw_sets_polar_axis(): + coord = coord_polar(direction=-1) + coord.params = {"r_range": (2, 8)} + fig, ax = plt.subplots(subplot_kw={"projection": "polar"}) + + try: + coord.draw([ax]) + assert ax.get_theta_direction() == 1 + assert_allclose(ax.get_ylim(), (2, 8)) + finally: + plt.close(fig) + + +def test_coord_polar_aspect_is_square(): + assert coord_polar().aspect(None) == 1 + + +def test_coord_polar_draw_uses_polar_axes_and_hides_blank_border(): + data = pd.DataFrame({"x": ["a", "b"], "y": [1, 2]}) + p = ( + ggplot(data, aes("x", "y")) + + geom_col() + + coord_polar() + + theme(panel_border=element_blank(), axis_line=element_line()) + ) + + fig = p.draw() + + try: + ax = fig.axes[0] + assert ax.name == "polar" + assert not ax.spines["polar"].get_visible() + finally: + plt.close(fig) + + +def test_coord_radial_arc_uses_end_or_direction(): + assert coord_radial(start=1, end=4)._arc == 3 + assert coord_radial(direction=-1)._arc == -2 * np.pi + + +def test_coord_radial_setup_panel_params_for_partial_arc(): + scale_x, scale_y = trained_scales( + y_breaks=(0, 2, 4, 8, 10), + y_labels=("0", "2", "4", "8", "10"), + ) + coord = coord_radial( + start=0, + end=np.pi, + thetalim=(0, 10), + rlim=(2, 8), + expand=False, + ) + + panel_params = coord.setup_panel_params(scale_x, scale_y) + + assert coord.params["theta_range"] == (0, 10) + assert coord.params["r_range"] == (2, 8) + assert_allclose(panel_params.x.breaks, [0, np.pi / 2, np.pi]) + assert panel_params.x.labels == ["0", "5", "10"] + assert panel_params.x.range == (0, np.pi) + assert panel_params.y.range == (2, 8) + assert panel_params.y.breaks == [2, 4, 8] + assert panel_params.y.labels == ["2", "4", "8"] + + +def test_coord_radial_setup_panel_params_theta_y_with_labels(): + scale_x, scale_y = trained_scales( + y_breaks=(0, 5, 10), + y_labels=("low", "mid", "high"), + ) + coord = coord_radial(theta="y", theta_labels=True, expand=False) + + panel_params = coord.setup_panel_params(scale_x, scale_y) + + assert_allclose(panel_params.x.breaks, [0, np.pi, 2 * np.pi]) + assert panel_params.x.labels == ["low", "mid", "high"] + + +def test_coord_radial_to_radians_zero_width_range(): + coord = coord_radial() + coord.params = {"theta_range": (1, 1)} + + assert_allclose(coord._to_radians(np.array([1, 2, 3])), [0, 0, 0]) + + +def test_coord_radial_transform_rotates_angle(): + coord = coord_radial(rotate_angle=True) + coord.params = {"theta_range": (0, 10), "r_range": (0, 10)} + data = pd.DataFrame({"x": [0, 5], "y": [1, 1], "angle": [10, 20]}) + + out = coord.transform(data, None) + + assert_allclose(out["x"], [0, np.pi]) + assert_allclose(out["angle"], [10, 200]) + + +def test_coord_radial_draw_sets_arc_inner_radius_and_axis_position(): + coord = coord_radial( + start=np.pi / 4, + end=3 * np.pi / 4, + inner_radius=0.5, + r_axis_inside=True, + ) + coord.params = {"r_range": (2, 10)} + fig, ax = plt.subplots(subplot_kw={"projection": "polar"}) + + try: + coord.draw([ax]) + assert_allclose(ax.get_xlim(), (np.pi / 4, 3 * np.pi / 4)) + assert_allclose(ax.get_rorigin(), -6) + assert ax.get_rlabel_position() == 55 + finally: + plt.close(fig) + + +def test_coord_radial_draw_float_r_axis_position(): + coord = coord_radial(r_axis_inside=np.pi / 2) + coord.params = {"r_range": (0, 10)} + fig, ax = plt.subplots(subplot_kw={"projection": "polar"}) + + try: + coord.draw([ax]) + assert ax.get_rlabel_position() == 90 + finally: + plt.close(fig) + + +def test_coord_radial_post_setup_ax_sets_pad_and_unclips_text(): + coord = coord_radial(theta_label_pad=17, theta_labels=True) + fig, ax = plt.subplots(subplot_kw={"projection": "polar"}) + text = ax.text(0, 1, "label", clip_on=True) + + try: + coord.post_setup_ax(ax) + assert ax.xaxis.get_tick_params()["pad"] == 17 + assert not text.get_clip_on() + finally: + plt.close(fig)