Merge pull request #1504 from hvasbath/kdekwargs

kde plot added kwargs

Author: ColCarroll

pymc3/plots.py

@@ -150,10 +150,7 @@ def make_2d(a):
     return a
 
 
-def kde2plot_op(ax, x, y, grid=200, cmap=None):
-
-    if cmap is None:
-        cmap =plt.cm.gist_earth_r
+def kde2plot_op(ax, x, y, grid=200, **kwargs):
 
     xmin = x.min()
     xmax = x.max()
@@ -167,8 +164,7 @@ def kde2plot_op(ax, x, y, grid=200, cmap=None):
     kernel = kde.gaussian_kde(values)
     Z = np.reshape(kernel(positions).T, X.shape)
 
-    ax.imshow(np.rot90(Z), cmap=cmap,
-              extent=[xmin, xmax, ymin, ymax])
+    ax.imshow(np.rot90(Z), extent=[xmin, xmax, ymin, ymax], **kwargs)
 
 
 def kdeplot(data, ax=None):
@@ -178,10 +174,10 @@ def kdeplot(data, ax=None):
     return ax
 
 
-def kde2plot(x, y, grid=200, ax=None, cmap=None):
+def kde2plot(x, y, grid=200, ax=None, **kwargs):
     if ax is None:
         f, ax = plt.subplots(1, 1, squeeze=True)
-    kde2plot_op(ax, x, y, grid, cmap)
+    kde2plot_op(ax, x, y, grid, **kwargs)
     return ax