Merge pull request #7454 from anntzer/step-converters

tacaswell · web-flow · commit 901562f47adc · 2016-11-14T22:52:30.000-05:00
PERF: Avoid temporaries when preparing step plots.
diff --git a/doc/api/api_changes/lines_removed_api.rst b/doc/api/api_changes/lines_removed_api.rst
@@ -0,0 +1,6 @@
+Functions removed from the `lines` module
+`````````````````````````````````````````
+
+The `matplotlib.lines` module no longer imports the `pts_to_prestep`,
+`pts_to_midstep` and `pts_to_poststep` functions from the `matplotlib.cbook`
+module.
diff --git a/lib/matplotlib/cbook.py b/lib/matplotlib/cbook.py
@@ -2302,147 +2302,112 @@ def get_instancemethod(self):
         return getattr(self.parent_obj, self.instancemethod_name)
 
 
-def _step_validation(x, *args):
-    """
-    Helper function of `pts_to_*step` functions
-
-    This function does all of the normalization required to the
-    input and generate the template for output
-
-
-    """
-    args = tuple(np.asanyarray(y) for y in args)
-    x = np.asanyarray(x)
-    if x.ndim != 1:
-        raise ValueError("x must be 1 dimensional")
-    if len(args) == 0:
-        raise ValueError("At least one Y value must be passed")
-
-    return np.vstack((x, ) + args)
-
-
 def pts_to_prestep(x, *args):
     """
-    Covert continuous line to pre-steps
+    Convert continuous line to pre-steps.
 
-    Given a set of N points convert to 2 N -1 points
-    which when connected linearly give a step function
-    which changes values at the beginning of the intervals.
+    Given a set of ``N`` points, convert to ``2N - 1`` points, which when
+    connected linearly give a step function which changes values at the
+    beginning of the intervals.
 
     Parameters
     ----------
     x : array
-        The x location of the steps
+        The x location of the steps.
 
-    y1, y2, ... : array
-        Any number of y arrays to be turned into steps.
-        All must be the same length as ``x``
+    y1, ..., yp : array
+        y arrays to be turned into steps; all must be the same length as ``x``.
 
     Returns
     -------
-    x, y1, y2, .. : array
-        The x and y values converted to steps in the same order
-        as the input.  If the input is length ``N``, each of these arrays
-        will be length ``2N + 1``
-
+    out : array
+        The x and y values converted to steps in the same order as the input;
+        can be unpacked as ``x_out, y1_out, ..., yp_out``.  If the input is
+        length ``N``, each of these arrays will be length ``2N + 1``.
 
     Examples
     --------
     >> x_s, y1_s, y2_s = pts_to_prestep(x, y1, y2)
     """
-    # do normalization
-    vertices = _step_validation(x, *args)
-    # create the output array
-    steps = np.zeros((vertices.shape[0], 2 * len(x) - 1), float)
-    # do the to step conversion logic
-    steps[0, 0::2], steps[0, 1::2] = vertices[0, :], vertices[0, :-1]
-    steps[1:, 0::2], steps[1:, 1:-1:2] = vertices[1:, :], vertices[1:, 1:]
-    # convert 2D array back to tuple
-    return tuple(steps)
+    steps = np.zeros((1 + len(args), 2 * len(x) - 1))
+    # In all `pts_to_*step` functions, only assign *once* using `x` and `args`,
+    # as converting to an array may be expensive.
+    steps[0, 0::2] = x
+    steps[0, 1::2] = steps[0, 0:-2:2]
+    steps[1:, 0::2] = args
+    steps[1:, 1::2] = steps[1:, 2::2]
+    return steps
 
 
 def pts_to_poststep(x, *args):
     """
-    Covert continuous line to pre-steps
+    Convert continuous line to post-steps.
 
-    Given a set of N points convert to 2 N -1 points
-    which when connected linearly give a step function
-    which changes values at the end of the intervals.
+    Given a set of ``N`` points convert to ``2N + 1`` points, which when
+    connected linearly give a step function which changes values at the end of
+    the intervals.
 
     Parameters
     ----------
     x : array
-        The x location of the steps
+        The x location of the steps.
 
-    y1, y2, ... : array
-        Any number of y arrays to be turned into steps.
-        All must be the same length as ``x``
+    y1, ..., yp : array
+        y arrays to be turned into steps; all must be the same length as ``x``.
 
     Returns
     -------
-    x, y1, y2, .. : array
-        The x and y values converted to steps in the same order
-        as the input.  If the input is length ``N``, each of these arrays
-        will be length ``2N + 1``
-
+    out : array
+        The x and y values converted to steps in the same order as the input;
+        can be unpacked as ``x_out, y1_out, ..., yp_out``.  If the input is
+        length ``N``, each of these arrays will be length ``2N + 1``.
 
     Examples
     --------
     >> x_s, y1_s, y2_s = pts_to_poststep(x, y1, y2)
     """
-    # do normalization
-    vertices = _step_validation(x, *args)
-    # create the output array
-    steps = np.zeros((vertices.shape[0], 2 * len(x) - 1), float)
-    # do the to step conversion logic
-    steps[0, ::2], steps[0, 1:-1:2] = vertices[0, :], vertices[0, 1:]
-    steps[1:, 0::2], steps[1:, 1::2] = vertices[1:, :], vertices[1:, :-1]
-
-    # convert 2D array back to tuple
-    return tuple(steps)
+    steps = np.zeros((1 + len(args), 2 * len(x) - 1))
+    steps[0, 0::2] = x
+    steps[0, 1::2] = steps[0, 2::2]
+    steps[1:, 0::2] = args
+    steps[1:, 1::2] = steps[1:, 0:-2:2]
+    return steps
 
 
 def pts_to_midstep(x, *args):
     """
-    Covert continuous line to pre-steps
+    Convert continuous line to mid-steps.
 
-    Given a set of N points convert to 2 N -1 points
-    which when connected linearly give a step function
-    which changes values at the middle of the intervals.
+    Given a set of ``N`` points convert to ``2N`` points which when connected
+    linearly give a step function which changes values at the middle of the
+    intervals.
 
     Parameters
     ----------
     x : array
-        The x location of the steps
+        The x location of the steps.
 
-    y1, y2, ... : array
-        Any number of y arrays to be turned into steps.
-        All must be the same length as ``x``
+    y1, ..., yp : array
+        y arrays to be turned into steps; all must be the same length as ``x``.
 
     Returns
     -------
-    x, y1, y2, .. : array
-        The x and y values converted to steps in the same order
-        as the input.  If the input is length ``N``, each of these arrays
-        will be length ``2N + 1``
-
+    out : array
+        The x and y values converted to steps in the same order as the input;
+        can be unpacked as ``x_out, y1_out, ..., yp_out``.  If the input is
+        length ``N``, each of these arrays will be length ``2N``.
 
     Examples
     --------
     >> x_s, y1_s, y2_s = pts_to_midstep(x, y1, y2)
     """
-    # do normalization
-    vertices = _step_validation(x, *args)
-    # create the output array
-    steps = np.zeros((vertices.shape[0], 2 * len(x)), float)
-    steps[0, 1:-1:2] = 0.5 * (vertices[0, :-1] + vertices[0, 1:])
-    steps[0, 2::2] = 0.5 * (vertices[0, :-1] + vertices[0, 1:])
-    steps[0, 0] = vertices[0, 0]
-    steps[0, -1] = vertices[0, -1]
-    steps[1:, 0::2], steps[1:, 1::2] = vertices[1:, :], vertices[1:, :]
-
-    # convert 2D array back to tuple
-    return tuple(steps)
+    steps = np.zeros((1 + len(args), 2 * len(x)))
+    x = np.asanyarray(x)
+    steps[0, 1:-1:2] = steps[0, 2::2] = (x[:-1] + x[1:]) / 2
+    steps[0, 0], steps[0, -1] = x[0], x[-1]
+    steps[1:, 0::2] = args
+    steps[1:, 1::2] = steps[1:, 0::2]
+    return steps
 
 
 STEP_LOOKUP_MAP = {'default': lambda x, y: (x, y),
diff --git a/lib/matplotlib/lines.py b/lib/matplotlib/lines.py
@@ -12,27 +12,24 @@
 import warnings
 
 import numpy as np
-from numpy import ma
-from . import artist, colors as mcolors
-from .artist import Artist
-from .cbook import (iterable, is_string_like, is_numlike, ls_mapper_r,
-                    pts_to_prestep, pts_to_poststep, pts_to_midstep, ls_mapper,
-                    is_hashable, STEP_LOOKUP_MAP)
 
+from . import artist, colors as mcolors, docstring, rcParams
+from .artist import Artist, allow_rasterization
+from .cbook import (
+    iterable, is_string_like, is_numlike, ls_mapper, ls_mapper_r, is_hashable,
+    STEP_LOOKUP_MAP)
+from .markers import MarkerStyle
 from .path import Path
 from .transforms import Bbox, TransformedPath, IdentityTransform
 
-from matplotlib import rcParams
-from .artist import allow_rasterization
-from matplotlib import docstring
-from matplotlib.markers import MarkerStyle
 # Imported here for backward compatibility, even though they don't
 # really belong.
-from matplotlib.markers import TICKLEFT, TICKRIGHT, TICKUP, TICKDOWN
-from matplotlib.markers import (
+from numpy import ma
+from . import _path
+from .markers import (
     CARETLEFT, CARETRIGHT, CARETUP, CARETDOWN,
-    CARETLEFTBASE, CARETRIGHTBASE, CARETUPBASE, CARETDOWNBASE)
-from matplotlib import _path
+    CARETLEFTBASE, CARETRIGHTBASE, CARETUPBASE, CARETDOWNBASE,
+    TICKLEFT, TICKRIGHT, TICKUP, TICKDOWN)
 
 
 def _get_dash_pattern(style):
diff --git a/lib/matplotlib/tests/test_cbook.py b/lib/matplotlib/tests/test_cbook.py
@@ -460,18 +460,14 @@ def test_to_midstep():
     assert_array_equal(y1_target, y1s)
 
 
-def test_step_fails():
+@pytest.mark.parametrize(
+    "args",
+    [(np.arange(12).reshape(3, 4), 'a'),
+     (np.arange(12), 'a'),
+     (np.arange(12), np.arange(3))])
+def test_step_fails(args):
     with pytest.raises(ValueError):
-        cbook._step_validation(np.arange(12).reshape(3, 4), 'a')
-
-    with pytest.raises(ValueError):
-        cbook._step_validation(np.arange(12), 'a')
-
-    with pytest.raises(ValueError):
-        cbook._step_validation(np.arange(12))
-
-    with pytest.raises(ValueError):
-        cbook._step_validation(np.arange(12), np.arange(3))
+        cbook.pts_to_prestep(*args)
 
 
 def test_grouper():
