Merge pull request #163 from lukelbd/parametric-fixes

Parametric plot fixes

Author: lukelbd

docs/1dplots.py

@@ -410,38 +410,36 @@
 # %%
 import proplot as plot
 import numpy as np
-N = 50
-cmap = 'IceFire'
-values = np.linspace(-N / 2, N / 2, N)
 fig, axs = plot.subplots(
     share=0, ncols=2, wratios=(2, 1),
     axwidth='7cm', aspect=(2, 1)
 )
 axs.format(suptitle='Parametric plots demo')
+cmap = 'IceFire'
 
 # Parametric line with smooth gradations
 ax = axs[0]
 state = np.random.RandomState(51423)
+N = 50
+x = (state.rand(N) - 0.52).cumsum()
+y = state.rand(N)
+c = np.linspace(-N / 2, N / 2, N)  # color values
 m = ax.parametric(
-    (state.rand(N) - 0.5).cumsum(), state.rand(N),
-    cmap=cmap, values=values, lw=7, extend='both'
-)
-ax.format(
-    xlabel='xlabel', ylabel='ylabel',
-    title='Line with smooth gradations'
+    x, y, c, cmap=cmap, lw=7, interp=5, capstyle='round', joinstyle='round'
 )
-ax.format(xlim=(-1, 5), ylim=(-0.2, 1.2))
+ax.format(xlabel='xlabel', ylabel='ylabel', title='Line with smooth gradations')
+ax.format(xmargin=0.05, ymargin=0.05)
 ax.colorbar(m, loc='b', label='parametric coordinate', locator=5)
 
 # Parametric line with stepped gradations
 N = 12
 ax = axs[1]
-values = np.linspace(-N / 2, N / 2, N + 1)
 radii = np.linspace(1, 0.2, N + 1)
 angles = np.linspace(0, 4 * np.pi, N + 1)
 x = radii * np.cos(1.4 * angles)
 y = radii * np.sin(1.4 * angles)
-m = ax.parametric(x, y, cmap=cmap, values=values, linewidth=15, interp=False)
+c = np.linspace(-N / 2, N / 2, N + 1)
+m = ax.parametric(x, y, c, cmap=cmap, lw=15)
 ax.format(
     xlim=(-1, 1), ylim=(-1, 1), title='Step gradations',
     xlabel='cosine angle', ylabel='sine angle'

proplot/axes/base.py

@@ -14,7 +14,7 @@
     _add_errorbars, _bar_wrapper, _barh_wrapper, _boxplot_wrapper,
     _cmap_changer, _cycle_changer,
     _fill_between_wrapper, _fill_betweenx_wrapper,
-    _hist_wrapper, _plot_wrapper, _scatter_wrapper,
+    _hist_wrapper, _parametric_wrapper, _plot_wrapper, _scatter_wrapper,
     _standardize_1d, _standardize_2d,
     _text_wrapper, _violinplot_wrapper,
     colorbar_wrapper, legend_wrapper,
@@ -1506,6 +1506,7 @@ def panel_axes(self, side, **kwargs):
         side = self._loc_translate(side, 'panel')
         return self.figure._add_axes_panel(self, side, **kwargs)
 
+    @_parametric_wrapper
     @_standardize_1d
     @_cmap_changer
     def parametric(
@@ -1522,15 +1523,19 @@ def parametric(
 
         Parameters
         ----------
-        *args : (y,) or (x,y)
+        *args : (y,), (x, y), or (x, y, values)
             The coordinates. If `x` is not provided, it is inferred from `y`.
-        cmap : colormap spec, optional
-            The colormap specifier, passed to `~proplot.constructor.Colormap`.
         values : list of float
             The parametric values used to map points on the line to colors
-            in the colormap.
+            in the colormap. This can also be passed as a third positional argument.
+        cmap : colormap spec, optional
+            The colormap specifier, passed to `~proplot.constructor.Colormap`.
+        cmap_kw : dict, optional
+            Keyword arguments passed to `~proplot.constructor.Colormap`.
         norm : normalizer spec, optional
             The normalizer, passed to `~proplot.constructor.Norm`.
+        norm_kw : dict, optional
+            Keyword arguments passed to `~proplot.constructor.Norm`.
         interp : int, optional
             If greater than ``0``, we interpolate to additional points
             between the `values` coordinates. The number corresponds to the
@@ -1552,74 +1557,35 @@ def parametric(
             The parametric line. See `this matplotlib example \
 <https://matplotlib.org/gallery/lines_bars_and_markers/multicolored_line>`__.
         """
-        # First error check
-        # WARNING: So far this only works for 1D *x* and *y* coordinates.
-        # Cannot draw multiple colormap lines at once
-        if values is None:
-            raise ValueError('Requires a "values" keyword arg.')
-        if len(args) not in (1, 2):
-            raise ValueError(f'Requires 1-2 arguments, got {len(args)}.')
-        y = np.array(args[-1]).squeeze()
-        x = np.arange(
-            y.shape[-1]) if len(args) == 1 else np.array(args[0]).squeeze()
-        values = np.array(values).squeeze()
-        if x.ndim != 1 or y.ndim != 1 or values.ndim != 1:
-            raise ValueError(
-                f'x ({x.ndim}d), y ({y.ndim}d), and values ({values.ndim}d)'
-                ' must be 1-dimensional.'
-            )
-        if len(x) != len(y) or len(x) != len(values) or len(y) != len(values):
-            raise ValueError(
-                f'{len(x)} xs, {len(y)} ys, but {len(values)} '
-                ' colormap values.'
-            )
-
-        # Interpolate values to allow for smooth gradations between values
-        # (interp=False) or color switchover halfway between points
-        # (interp=True). Then optionally interpolate the colormap values.
-        if interp > 0:
-            x, y, values = [], [], []
-            xorig, yorig, vorig = x, y, values
-            for j in range(xorig.shape[0] - 1):
-                idx = slice(None)
-                if j + 1 < xorig.shape[0] - 1:
-                    idx = slice(None, -1)
-                x.extend(
-                    np.linspace(xorig[j], xorig[j + 1], interp + 2)[idx].flat
-                )
-                y.extend(
-                    np.linspace(yorig[j], yorig[j + 1], interp + 2)[idx].flat
-                )
-                values.extend(
-                    np.linspace(vorig[j], vorig[j + 1], interp + 2)[idx].flat
-                )
-            x, y, values = np.array(x), np.array(y), np.array(values)
-
         # Get x/y coordinates and values for points to the 'left' and 'right'
         # of each joint
+        x, y = args  # standardized by parametric wrapper
+        interp  # avoid U100 unused argument error (arg is handled by wrapper)
         coords = []
         levels = edges(values)
         for j in range(y.shape[0]):
             if j == 0:
                 xleft, yleft = [], []
             else:
-                xleft = [(x[j - 1] + x[j]) / 2, x[j]]
-                yleft = [(y[j - 1] + y[j]) / 2, y[j]]
+                xleft = [0.5 * (x[j - 1] + x[j]), x[j]]
+                yleft = [0.5 * (y[j - 1] + y[j]), y[j]]
             if j + 1 == y.shape[0]:
                 xright, yright = [], []
             else:
                 xleft = xleft[:-1]  # prevent repetition when joined with right
                 yleft = yleft[:-1]
-                xright = [x[j], (x[j + 1] + x[j]) / 2]
-                yright = [y[j], (y[j + 1] + y[j]) / 2]
+                xright = [x[j], 0.5 * (x[j + 1] + x[j])]
+                yright = [y[j], 0.5 * (y[j + 1] + y[j])]
             pleft = np.stack((xleft, yleft), axis=1)
             pright = np.stack((xright, yright), axis=1)
             coords.append(np.concatenate((pleft, pright), axis=0))
+        coords = np.array(coords)
 
         # Create LineCollection and update with values
+        # NOTE: Default capstyle is butt but this may look weird with vector graphics
         hs = mcollections.LineCollection(
-            np.array(coords), cmap=cmap, norm=norm,
-            linestyles='-', capstyle='butt', joinstyle='miter'
+            coords, cmap=cmap, norm=norm,
+            linestyles='-', capstyle='butt', joinstyle='miter',
         )
         hs.set_array(np.array(values))
         hs.update({

proplot/axes/plot.py

@@ -1062,6 +1062,56 @@ def add_errorbars(
     return obj
 
 
+def _parametric_wrapper_temporary(self, func, *args, interp=0, **kwargs):
+    """
+    Calls `~proplot.axes.Axes.parametric` and optionally interpolates values before
+    they get passed to `cmap_changer` and the colormap boundaries are drawn.
+    """
+    # Parse input arguments
+    # WARNING: So far this only works for 1D *x* and *y* coordinates.
+    # Cannot draw multiple colormap lines at once
+    if len(args) == 3:
+        x, y, values = args
+    elif 'values' in kwargs:
+        values = kwargs.pop('values')
+        if len(args) == 1:
+            y = np.asarray(args[0])
+            x = np.arange(y.shape[-1])
+        elif len(args) == 2:
+            x, y = args
+        else:
+            raise ValueError(f'1 to 3 positional arguments required, got {len(args)}.')
+    else:
+        raise ValueError('Missing required keyword argument "values".')
+    x, y, values = np.atleast_1d(x), np.atleast_1d(y), np.atleast_1d(values)
+    if (
+        any(_.ndim != 1 for _ in (x, y, values))
+        or len({x.size, y.size, values.size}) > 1
+    ):
+        raise ValueError(
+            f'x {x.shape}, y {y.shape}, and values {values.shape} '
+            'must be 1-dimensional and have the same size.'
+        )
+
+    # Interpolate values to allow for smooth gradations between values
+    # (interp=False) or color switchover halfway between points
+    # (interp=True). Then optionally interpolate the colormap values.
+    if interp > 0:
+        xorig, yorig, vorig = x, y, values
+        x, y, values = [], [], []
+        for j in range(xorig.shape[0] - 1):
+            idx = slice(None)
+            if j + 1 < xorig.shape[0] - 1:
+                idx = slice(None, -1)
+            x.extend(np.linspace(xorig[j], xorig[j + 1], interp + 2)[idx].flat)
+            y.extend(np.linspace(yorig[j], yorig[j + 1], interp + 2)[idx].flat)
+            values.extend(np.linspace(vorig[j], vorig[j + 1], interp + 2)[idx].flat)
+        x, y, values = np.array(x), np.array(y), np.array(values)
+
+    # Call main function
+    return func(self, x, y, values=values, **kwargs)
+
+
 def _plot_wrapper_deprecated(
     self, func, *args, cmap=None, values=None, **kwargs
 ):
@@ -1350,6 +1400,7 @@ def barh_wrapper(
     # NOTE: You *must* do juggling of barh keyword order --> bar keyword order
     # --> barh keyword order, because horizontal hist passes arguments to bar
     # directly and will not use a 'barh' method with overridden argument order!
+    func  # avoid U100 error
     kwargs.setdefault('orientation', 'horizontal')
     if y is None and width is None:
         raise ValueError(
@@ -3439,6 +3490,7 @@ def _wrapper(self, *args, **kwargs):
 _fill_between_wrapper = _wrapper_decorator(fill_between_wrapper)
 _fill_betweenx_wrapper = _wrapper_decorator(fill_betweenx_wrapper)
 _hist_wrapper = _wrapper_decorator(hist_wrapper)
+_parametric_wrapper = _wrapper_decorator(_parametric_wrapper_temporary)
 _plot_wrapper = _wrapper_decorator(_plot_wrapper_deprecated)
 _scatter_wrapper = _wrapper_decorator(scatter_wrapper)
 _standardize_1d = _wrapper_decorator(standardize_1d)