Fix: correctly pipe keyword arguments to plot_dendrogram

brian2tools/plotting/base.py

@@ -106,10 +106,7 @@ def brian_plot(brian_obj,
             kwds['rate_unit'] = _get_best_unit(smooth_rate)
         return plot_rate(brian_obj.t, smooth_rate, axes=axes, **kwds)
     elif isinstance(brian_obj, Morphology):
-        if kwds:
-            logger.warn('plot_dendrogram does not take any additional keyword '
-                        'arguments, ignoring them.')
-        return plot_dendrogram(brian_obj, axes=axes)
+        return plot_dendrogram(brian_obj, axes=axes, **kwds)
     elif isinstance(brian_obj, Synapses):
         if len(brian_obj) == 0:
             raise TypeError('Synapses object does not have any synapses.')

brian2tools/plotting/morphology.py

@@ -352,7 +352,7 @@ def plot_morphology(morphology, plot_3d=None, show_compartments=False,
     return axes
 
 
-def plot_dendrogram(morphology, axes=None):
+def plot_dendrogram(morphology, axes=None, **kwds):
     '''
     Plot a "dendrogram" of a morphology, i.e. an abstract representation which
     visualizes the branching structure and the length of each section.
@@ -365,6 +365,12 @@ def plot_dendrogram(morphology, axes=None):
         The `~matplotlib.axes.Axes` instance used for plotting. Defaults to
         ``None`` which means that a new `~matplotlib.axes.Axes` will be
         created for the plot.
+    **kwds
+        Any additional keyword arguments are passed to matplotlib's
+        `~matplotlib.axes.Axes.plot`, `~matplotlib.axes.Axes.vlines`, and
+        `~matplotlib.axes.Axes.hlines` calls. Only arguments accepted by all
+        three functions should be used (e.g. ``color``, ``alpha``,
+        ``linewidth``).
 
     Returns
     -------
@@ -424,18 +430,21 @@ def plot_dendrogram(morphology, axes=None):
 
     x_values = (np.array(min_index) + np.array(max_index)) / 2.0
 
+    clip_on = kwds.pop('clip_on', False)
+    lw = kwds.pop('lw', kwds.pop('linewidth', 2))
+
     # Plot the dendogram with lengths of the vertical lines representing the
     # total distance to the root
-    plt.plot(x_values[0], length_metric[0], 'ko', clip_on=False)
+    plt.plot(x_values[0], length_metric[0], marker='o', clip_on=clip_on, **kwds)
     for sec, (x, depth) in enumerate(zip(x_values, length_metric)):
         child_start_idx = (sec+1)*max_children
         num_children = flat_morpho.morph_children_num[sec+1]
         if num_children > 0:
             child_indices = children[child_start_idx:child_start_idx+num_children]
             child_depth = length_metric[child_indices-1]
             child_x = x_values[child_indices-1]
-            axes.vlines(child_x, depth, child_depth, clip_on=False, lw=2)
-            axes.hlines(depth, min(child_x), max(child_x), lw=2)
+            axes.vlines(child_x, depth, child_depth, clip_on=clip_on, lw=lw, **kwds)
+            axes.hlines(depth, min(child_x), max(child_x), lw=lw, **kwds)
     axes.set_xticks([])
     axes.set_ylabel('distance from root (um)')
     axes.set_xlim(-1, terminal_counter)

brian2tools/tests/test_plotting.py

@@ -122,6 +122,20 @@ def test_plot_morphology():
     ax = plot_dendrogram(morpho)
     assert isinstance(ax, matplotlib.axes.Axes)
     plt.close()
+    # Test that styling kwargs are accepted and forwarded without error
+    ax = plot_dendrogram(morpho, color='red')
+    assert isinstance(ax, matplotlib.axes.Axes)
+    plt.close()
+    ax = plot_dendrogram(morpho, color='blue', alpha=0.5)
+    assert isinstance(ax, matplotlib.axes.Axes)
+    plt.close()
+    ax = plot_dendrogram(morpho, color='green', linewidth=3)
+    assert isinstance(ax, matplotlib.axes.Axes)
+    plt.close()
+    # brian_plot routes Morphology to plot_dendrogram, so kwargs must work there too
+    ax = brian_plot(morpho, color='purple')
+    assert isinstance(ax, matplotlib.axes.Axes)
+    plt.close()
     ax = plot_morphology(morpho)
     assert isinstance(ax, matplotlib.axes.Axes)
     plt.close()

dev/doc_tools/plotting_examples/morphology_plots.py

@@ -13,6 +13,12 @@
 brian_plot(morpho)
 savefig(os.path.join(fig_dir, 'plot_dendrogram.svg'))
 close()
+
+figure()
+plot_dendrogram(morpho, color='red', alpha=0.7)
+savefig(os.path.join(fig_dir, 'plot_dendrogram_custom.svg'))
+close()
+
 plot_morphology(morpho)  # 3d plot
 mayavi.show()
 figure()