simplifying colorbat code and updating arrows kwarg to address networ…

…kx bug

wntr/graphics/network.py

@@ -212,7 +212,7 @@ def plot_network(wn, node_attribute=None, link_attribute=None, title=None,
             nodelist=nodelist, node_color=nodecolor, node_size=node_size, 
             alpha=node_alpha, cmap=node_cmap, vmin=node_range[0], vmax = node_range[1], 
             linewidths=0, ax=ax)
-    edges = nx.draw_networkx_edges(G, pos, edgelist=linklist, 
+    edges = nx.draw_networkx_edges(G, pos, edgelist=linklist, arrows=directed,
             edge_color=linkcolor, width=link_width, alpha=link_alpha, edge_cmap=link_cmap, 
             edge_vmin=link_range[0], edge_vmax=link_range[1], ax=ax)
     if node_labels:
@@ -228,14 +228,11 @@ def plot_network(wn, node_attribute=None, link_attribute=None, title=None,
         clb = plt.colorbar(nodes, shrink=0.5, pad=0, ax=ax)
         clb.ax.set_title(node_colorbar_label, fontsize=10)
     if add_link_colorbar and link_attribute:
-        if directed:
-            vmin = min(map(abs,link_attribute.values()))
-            vmax = max(map(abs,link_attribute.values())) 
-            sm = plt.cm.ScalarMappable(cmap=link_cmap, norm=plt.Normalize(vmin=vmin, vmax=vmax))
-            sm.set_array([])
-            clb = plt.colorbar(sm, shrink=0.5, pad=0.05, ax=ax)
-        else:
-            clb = plt.colorbar(edges, shrink=0.5, pad=0.05, ax=ax)
+        vmin = min(map(abs,link_attribute.values()))
+        vmax = max(map(abs,link_attribute.values())) 
+        sm = plt.cm.ScalarMappable(cmap=link_cmap, norm=plt.Normalize(vmin=vmin, vmax=vmax))
+        sm.set_array([])
+        clb = plt.colorbar(sm, shrink=0.5, pad=0.05, ax=ax)
         clb.ax.set_title(link_colorbar_label, fontsize=10)
 
     ax.axis('off')