Merge pull request #308 from edmundus/master

Added functionality for plotting step plots

Author: nschloe

test/test_steps.py

@@ -0,0 +1,33 @@
+from helpers import assert_equality
+
+
+def plot():
+    import matplotlib.pyplot as plt
+    import numpy as np
+
+    fig = plt.figure()
+    x = np.arange(5)
+    y1 = np.array([1, 2, 1, 4, 2])
+    y2 = np.array([1, 2, 1, 4, 2])
+    y3 = np.array([1, 2, 1, 4, 2])
+    y4 = np.array([1, 2, 1, 4, 2])
+    y5 = np.array([2, 3, 2, 5, 3])
+
+    plt.step(x, y1, "r-")
+    plt.step(x, y2, "b--", where="pre")
+    plt.step(x, y3, "g-.", where="post")
+    plt.step(x, y4, "y:", where="mid")
+    plt.plot(x, y5, "c.")
+    plt.legend(["default", "pre", "post", "mid", "default"])
+
+    return fig
+
+
+def test():
+    assert_equality(plot, __file__[:-3] + "_reference.tex")
+
+
+if __name__ == "__main__":
+    import helpers
+
+    helpers.compare_mpl_latex(plot)

test/test_steps_reference.tex

@@ -0,0 +1,65 @@
+\begin{tikzpicture}
+
+\definecolor{color0}{rgb}{0.75,0.75,0}
+\definecolor{color1}{rgb}{0,0.75,0.75}
+
+\begin{axis}[
+legend cell align={left},
+legend style={at={(0.03,0.97)}, anchor=north west, draw=white!80.0!black},
+tick align=outside,
+tick pos=left,
+x grid style={white!69.01960784313725!black},
+xmin=-0.2, xmax=4.2,
+xtick style={color=black},
+y grid style={white!69.01960784313725!black},
+ymin=0.8, ymax=5.2,
+ytick style={color=black}
+]
+\addplot [semithick, red, const plot mark right]
+table {%
+0 1
+1 2
+2 1
+3 4
+4 2
+};
+\addlegendentry{default}
+\addplot [semithick, blue, const plot mark right, dashed]
+table {%
+0 1
+1 2
+2 1
+3 4
+4 2
+};
+\addlegendentry{pre}
+\addplot [semithick, green!50.0!black, const plot mark left, dash pattern=on 1pt off 3pt on 3pt off 3pt]
+table {%
+0 1
+1 2
+2 1
+3 4
+4 2
+};
+\addlegendentry{post}
+\addplot [semithick, color0, const plot mark mid, dotted]
+table {%
+0 1
+1 2
+2 1
+3 4
+4 2
+};
+\addlegendentry{mid}
+\addplot [semithick, color1, mark=*, mark size=3, mark options={solid}, only marks]
+table {%
+0 2
+1 3
+2 2
+3 5
+4 3
+};
+\addlegendentry{default}
+\end{axis}
+
+\end{tikzpicture}

tikzplotlib/line2d.py

@@ -31,6 +31,19 @@ def draw_line2d(data, obj):
     data, line_xcolor, _ = mycol.mpl_color2xcolor(data, color)
     addplot_options.append(line_xcolor)
 
+    # get draw style
+    drawstyle = obj.get_drawstyle()
+    if drawstyle in [None, "default"]:
+        pass
+    else:
+        if drawstyle == "steps-mid":
+            style = "const plot mark mid"
+        elif drawstyle in ["steps-pre", "steps"]:
+            style = "const plot mark right"
+        elif drawstyle == "steps-post":
+            style = "const plot mark left"
+        addplot_options.append("{}".format(style))
+
     alpha = obj.get_alpha()
     if alpha is not None:
         addplot_options.append("opacity={}".format(alpha))