From 8d6974f67d63fe0543ed0af0ad576ea0c75c80c6 Mon Sep 17 00:00:00 2001
From: Thomas Breloff <tom@breloff.com>
Date: Tue, 26 Apr 2016 15:09:18 -0400
Subject: [PATCH] contours keyword for surface and wireframe plots

---
 src/args.jl               |  1 +
 src/backends/pyplot.jl    | 59 ++++++++++++++++++++++++---------------
 src/backends/supported.jl |  2 +-
 3 files changed, 39 insertions(+), 23 deletions(-)

diff --git a/src/args.jl b/src/args.jl
index 3f42f094..c9e0a039 100644
--- a/src/args.jl
+++ b/src/args.jl
@@ -147,6 +147,7 @@ _seriesDefaults[:ribbon]          = nothing
 _seriesDefaults[:quiver]          = nothing
 _seriesDefaults[:normalize]       = false     # do we want a normalized histogram?
 _seriesDefaults[:weights]         = nothing   # optional weights for histograms (1D and 2D)
+_seriesDefaults[:contours]        = false     # add contours to 3d surface and wireframe plots
 
 
 const _plotDefaults = KW()
diff --git a/src/backends/pyplot.jl b/src/backends/pyplot.jl
index 11bd781b..8466720e 100644
--- a/src/backends/pyplot.jl
+++ b/src/backends/pyplot.jl
@@ -327,6 +327,18 @@ function _add_series(pkg::PyPlotBackend, plt::Plot, d::KW)
     needs_colorbar = false
     discrete_colorbar_values = nothing
 
+
+    # pass in an integer value as an arg, but a levels list as a keyword arg
+    levels = d[:levels]
+    levelargs = if isscalar(levels)
+        (levels)
+    elseif isvector(levels)
+        extrakw[:levels] = levels
+        ()
+    else
+        error("Only numbers and vectors are supported with levels keyword")
+    end
+
     # for each plotting command, optionally build and add a series handle to the list
 
     # line plot
@@ -459,23 +471,13 @@ function _add_series(pkg::PyPlotBackend, plt::Plot, d::KW)
         z = z.surf'
         needs_colorbar = true
 
-        # pass in an integer value as an arg, but a levels list as a keyword arg
-        levels = d[:levels]
-        args = if isscalar(levels)
-            (levels)
-        elseif isvector(levels)
-            extrakw[:levels] = levels
-            ()
-        else
-            error("Only numbers and vectors are supported with levels keyword")
-        end
 
         if lt == :contour3d
             extrakw[:extend3d] = true
         end
 
         # contour lines
-        handle = ax[:contour](x, y, z, args...;
+        handle = ax[:contour](x, y, z, levelargs...;
             label = d[:label],
             zorder = plt.n,
             linewidths = d[:linewidth],
@@ -486,19 +488,19 @@ function _add_series(pkg::PyPlotBackend, plt::Plot, d::KW)
         push!(handles, handle)
 
         # contour fills
-        if lt == :contour
-            handle = ax[:contourf](x, y, z, args...;
-                label = d[:label],
-                zorder = plt.n + 0.5,
-                cmap = pyfillcolormap(d),
-                extrakw...
-            )
-            push!(handles, handle)
-        end
+        # if lt == :contour
+        handle = ax[:contourf](x, y, z, levelargs...;
+            label = d[:label],
+            zorder = plt.n + 0.5,
+            cmap = pyfillcolormap(d),
+            extrakw...
+        )
+        push!(handles, handle)
+        # end
     end
 
     if lt in (:surface, :wireframe)
-        if typeof(z) <: AbstractMatrix
+        if typeof(z) <: AbstractMatrix || typeof(z) <: Surface
             x, y, z = map(Array, (x,y,z))
             if !ismatrix(x) || !ismatrix(y)
                 x = repmat(x', length(y), 1)
@@ -520,6 +522,20 @@ function _add_series(pkg::PyPlotBackend, plt::Plot, d::KW)
             )
             push!(handles, handle)
 
+            # contours on the axis planes
+            if d[:contours]
+                for (zdir,mat) in (("x",x), ("y",y), ("z",z))
+                    offset = (zdir == "y" ? maximum : minimum)(mat)
+                    handle = ax[:contourf](x, y, z, levelargs...;
+                        zdir = zdir,
+                        cmap = pyfillcolormap(d),
+                        offset = (zdir == "y" ? maximum : minimum)(mat)  # where to draw the contour plane
+                    )
+                    push!(handles, handle)
+                    needs_colorbar = true
+                end
+            end
+
         elseif typeof(z) <: AbstractVector
             # tri-surface plot (http://matplotlib.org/mpl_toolkits/mplot3d/tutorial.html#tri-surface-plots)
             handle = ax[:plot_trisurf](x, y, z;
@@ -534,7 +550,6 @@ function _add_series(pkg::PyPlotBackend, plt::Plot, d::KW)
         else
             error("Unsupported z type $(typeof(z)) for linetype=$lt")
         end
-
     end
 
     if lt == :heatmap
diff --git a/src/backends/supported.jl b/src/backends/supported.jl
index b1c17bc4..fccef3e6 100644
--- a/src/backends/supported.jl
+++ b/src/backends/supported.jl
@@ -104,7 +104,7 @@ supportedArgs(::PyPlotBackend) = [
     :orientation,
     :overwrite_figure,
     :polar,
-    :normalize, :weights
+    :normalize, :weights, :contours
   ]
 supportedAxes(::PyPlotBackend) = _allAxes
 supportedTypes(::PyPlotBackend) = [